outposts: add support for provider-specific websocket messages

Signed-off-by: Jens Langhammer <jens@goauthentik.io>

authentik/outposts/channels.py

@@ -27,6 +27,9 @@ class WebsocketMessageInstruction(IntEnum):
     # Message sent by us to trigger an Update
     TRIGGER_UPDATE = 2
 
+    # Provider specific message
+    PROVIDER_SPECIFIC = 3
+
 
 @dataclass(slots=True)
 class WebsocketMessage:
@@ -131,3 +134,14 @@ def event_update(self, event):  # pragma: no cover
         self.send_json(
             asdict(WebsocketMessage(instruction=WebsocketMessageInstruction.TRIGGER_UPDATE))
         )
+
+    def event_provider_specific(self, event):
+        """Event handler which can be called by provider-specific
+        implementations to send specific messages to the outpost"""
+        self.send_json(
+            asdict(
+                WebsocketMessage(
+                    instruction=WebsocketMessageInstruction.PROVIDER_SPECIFIC, args=event
+                )
+            )
+        )

authentik/outposts/tasks.py

@@ -5,7 +5,6 @@
 from typing import Any, Optional
 from urllib.parse import urlparse
 
-import yaml
 from asgiref.sync import async_to_sync
 from channels.layers import get_channel_layer
 from django.core.cache import cache
@@ -16,6 +15,7 @@
 from kubernetes.config.incluster_config import SERVICE_TOKEN_FILENAME
 from kubernetes.config.kube_config import KUBE_CONFIG_DEFAULT_LOCATION
 from structlog.stdlib import get_logger
+from yaml import safe_load
 
 from authentik.events.monitored_tasks import (
     MonitoredTask,
@@ -279,7 +279,7 @@ def outpost_connection_discovery(self: MonitoredTask):
             with kubeconfig_path.open("r", encoding="utf8") as _kubeconfig:
                 KubernetesServiceConnection.objects.create(
                     name=kubeconfig_local_name,
-                    kubeconfig=yaml.safe_load(_kubeconfig),
+                    kubeconfig=safe_load(_kubeconfig),
                 )
     unix_socket_path = urlparse(DEFAULT_UNIX_SOCKET).path
     socket = Path(unix_socket_path)

internal/outpost/ak/api.go

@@ -22,6 +22,8 @@ import (
 	log "github.com/sirupsen/logrus"
 )
 
+type WSHandler func(ctx context.Context, args map[string]interface{})
+
 const ConfigLogLevel = "log_level"
 
 // APIController main controller which connects to the authentik api via http and ws
@@ -42,6 +44,7 @@ type APIController struct {
 	lastWsReconnect     time.Time
 	wsIsReconnecting    bool
 	wsBackoffMultiplier int
+	wsHandlers          []WSHandler
 	refreshHandlers     []func()
 
 	instanceUUID uuid.UUID
@@ -106,6 +109,7 @@ func NewAPIController(akURL url.URL, token string) *APIController {
 		reloadOffset:        time.Duration(rand.Intn(10)) * time.Second,
 		instanceUUID:        uuid.New(),
 		Outpost:             outpost,
+		wsHandlers:          []WSHandler{},
 		wsBackoffMultiplier: 1,
 		refreshHandlers:     make([]func(), 0),
 	}
@@ -156,6 +160,10 @@ func (a *APIController) AddRefreshHandler(handler func()) {
 	a.refreshHandlers = append(a.refreshHandlers, handler)
 }
 
+func (a *APIController) AddWSHandler(handler WSHandler) {
+	a.wsHandlers = append(a.wsHandlers, handler)
+}
+
 func (a *APIController) OnRefresh() error {
 	// Because we don't know the outpost UUID, we simply do a list and pick the first
 	// The service account this token belongs to should only have access to a single outpost

internal/outpost/ak/api_ws.go

@@ -1,6 +1,7 @@
 package ak
 
 import (
+	"context"
 	"crypto/tls"
 	"fmt"
 	"net/http"
@@ -145,6 +146,10 @@ func (ac *APIController) startWSHandler() {
 					"build":        constants.BUILD("tagged"),
 				}).SetToCurrentTime()
 			}
+		} else if wsMsg.Instruction == WebsocketInstructionProviderSpecific {
+			for _, h := range ac.wsHandlers {
+				h(context.Background(), wsMsg.Args)
+			}
 		}
 	}
 }

internal/outpost/ak/api_ws_msg.go

@@ -9,6 +9,8 @@ const (
 	WebsocketInstructionHello websocketInstruction = 1
 	// WebsocketInstructionTriggerUpdate Code received to trigger a config update
 	WebsocketInstructionTriggerUpdate websocketInstruction = 2
+	// WebsocketInstructionProviderSpecific Code received to trigger some provider specific function
+	WebsocketInstructionProviderSpecific websocketInstruction = 3
 )
 
 type websocketMessage struct {