Merge pull request #2938 from headlamp-k8s/handle-token-multiplexer

backend: frontend: Pass authentication token to websocket multiplexer

Author: joaquimrocha

backend/cmd/multiplexer.go

@@ -72,6 +72,8 @@ type Connection struct {
 	writeMu sync.Mutex
 	// closed is a flag to indicate if the connection is closed.
 	closed bool
+	// Authentication token.
+	Token *string
 }
 
 // Message represents a WebSocket message structure.
@@ -90,6 +92,8 @@ type Message struct {
 	Binary bool `json:"binary,omitempty"`
 	// Type is the type of the message.
 	Type string `json:"type"`
+	// Authentication token.
+	Token *string `json:"token"`
 }
 
 // Multiplexer manages multiple WebSocket connections.
@@ -246,14 +250,15 @@ func (m *Multiplexer) establishClusterConnection(
 	path,
 	query string,
 	clientConn *WSConnLock,
+	token *string,
 ) (*Connection, error) {
 	config, err := m.getClusterConfigWithFallback(clusterID, userID)
 	if err != nil {
 		logger.Log(logger.LevelError, map[string]string{"clusterID": clusterID}, err, "getting cluster config")
 		return nil, err
 	}
 
-	connection := m.createConnection(clusterID, userID, path, query, clientConn)
+	connection := m.createConnection(clusterID, userID, path, query, clientConn, token)
 
 	wsURL := createWebSocketURL(config.Host, path, query)
 
@@ -264,7 +269,7 @@ func (m *Multiplexer) establishClusterConnection(
 		return nil, fmt.Errorf("failed to get TLS config: %v", err)
 	}
 
-	conn, err := m.dialWebSocket(wsURL, tlsConfig, config.Host)
+	conn, err := m.dialWebSocket(wsURL, tlsConfig, config.Host, token)
 	if err != nil {
 		connection.updateStatus(StateError, err)
 
@@ -309,6 +314,7 @@ func (m *Multiplexer) createConnection(
 	path,
 	query string,
 	clientConn *WSConnLock,
+	token *string,
 ) *Connection {
 	return &Connection{
 		ClusterID: clusterID,
@@ -321,16 +327,29 @@ func (m *Multiplexer) createConnection(
 			State:   StateConnecting,
 			LastMsg: time.Now(),
 		},
+		Token: token,
 	}
 }
 
 // dialWebSocket establishes a WebSocket connection.
-func (m *Multiplexer) dialWebSocket(wsURL string, tlsConfig *tls.Config, host string) (*websocket.Conn, error) {
+func (m *Multiplexer) dialWebSocket(
+	wsURL string,
+	tlsConfig *tls.Config,
+	host string,
+	token *string,
+) (*websocket.Conn, error) {
 	dialer := websocket.Dialer{
 		TLSClientConfig:  tlsConfig,
 		HandshakeTimeout: HandshakeTimeout,
 	}
 
+	if token != nil {
+		dialer.Subprotocols = []string{
+			"base64.binary.k8s.io",
+			"base64url.bearer.authorization.k8s.io." + base64.RawStdEncoding.EncodeToString([]byte(*token)),
+		}
+	}
+
 	conn, resp, err := dialer.Dial(
 		wsURL,
 		http.Header{
@@ -339,6 +358,7 @@ func (m *Multiplexer) dialWebSocket(wsURL string, tlsConfig *tls.Config, host st
 	)
 	if err != nil {
 		logger.Log(logger.LevelError, nil, err, "dialing WebSocket")
+		logger.Log(logger.LevelError, nil, resp, "WebSocket response")
 		// We only attempt to close the response body if there was an error and resp is not nil.
 		// In the successful case (when err is nil), the resp will actually be nil for WebSocket connections,
 		// so we don't need to close anything.
@@ -393,6 +413,7 @@ func (m *Multiplexer) reconnect(conn *Connection) (*Connection, error) {
 		conn.Path,
 		conn.Query,
 		conn.Client,
+		conn.Token,
 	)
 	if err != nil {
 		logger.Log(logger.LevelError, map[string]string{"clusterID": conn.ClusterID}, err, "reconnecting to cluster")
@@ -482,7 +503,7 @@ func (m *Multiplexer) getOrCreateConnection(msg Message, clientConn *WSConnLock)
 	if !exists {
 		var err error
 
-		conn, err = m.establishClusterConnection(msg.ClusterID, msg.UserID, msg.Path, msg.Query, clientConn)
+		conn, err = m.establishClusterConnection(msg.ClusterID, msg.UserID, msg.Path, msg.Query, clientConn, msg.Token)
 		if err != nil {
 			logger.Log(
 				logger.LevelError,

backend/cmd/multiplexer_test.go

@@ -112,7 +112,7 @@ func TestCreateConnection(t *testing.T) {
 	clientConn, _ := createTestWebSocketConnection()
 
 	// Add RequestID to the createConnection call
-	conn := m.createConnection("test-cluster", "test-user", "/api/v1/pods", "", clientConn)
+	conn := m.createConnection("test-cluster", "test-user", "/api/v1/pods", "", clientConn, nil)
 	assert.NotNil(t, conn)
 	assert.Equal(t, "test-cluster", conn.ClusterID)
 	assert.Equal(t, "test-user", conn.UserID)
@@ -153,7 +153,7 @@ func TestDialWebSocket(t *testing.T) {
 	defer server.Close()
 
 	wsURL := "ws" + strings.TrimPrefix(server.URL, "http")
-	conn, err := m.dialWebSocket(wsURL, &tls.Config{InsecureSkipVerify: true}, server.URL) //nolint:gosec
+	conn, err := m.dialWebSocket(wsURL, &tls.Config{InsecureSkipVerify: true}, server.URL, nil) //nolint:gosec
 
 	assert.NoError(t, err)
 	assert.NotNil(t, conn)
@@ -170,12 +170,12 @@ func TestDialWebSocket_Errors(t *testing.T) {
 	// Test invalid URL
 	tlsConfig := &tls.Config{InsecureSkipVerify: true} //nolint:gosec
 
-	ws, err := m.dialWebSocket("invalid-url", tlsConfig, "")
+	ws, err := m.dialWebSocket("invalid-url", tlsConfig, "", nil)
 	assert.Error(t, err)
 	assert.Nil(t, ws)
 
 	// Test unreachable URL
-	ws, err = m.dialWebSocket("ws://localhost:12345", tlsConfig, "")
+	ws, err = m.dialWebSocket("ws://localhost:12345", tlsConfig, "", nil)
 	assert.Error(t, err)
 	assert.Nil(t, ws)
 }
@@ -535,7 +535,7 @@ func TestEstablishClusterConnection(t *testing.T) {
 	defer clientServer.Close()
 
 	// Test successful connection establishment
-	conn, err := m.establishClusterConnection("test-cluster", "test-user", "/api/v1/pods", "watch=true", clientConn)
+	conn, err := m.establishClusterConnection("test-cluster", "test-user", "/api/v1/pods", "watch=true", clientConn, nil)
 	assert.NoError(t, err)
 	assert.NotNil(t, conn)
 	assert.Equal(t, "test-cluster", conn.ClusterID)
@@ -544,7 +544,7 @@ func TestEstablishClusterConnection(t *testing.T) {
 	assert.Equal(t, "watch=true", conn.Query)
 
 	// Test with invalid cluster
-	conn, err = m.establishClusterConnection("non-existent", "test-user", "/api/v1/pods", "watch=true", clientConn)
+	conn, err = m.establishClusterConnection("non-existent", "test-user", "/api/v1/pods", "watch=true", clientConn, nil)
 	assert.Error(t, err)
 	assert.Nil(t, conn)
 }
@@ -572,7 +572,7 @@ func TestReconnect(t *testing.T) {
 	defer clientServer.Close()
 
 	// Create initial connection
-	conn := m.createConnection("test-cluster", "test-user", "/api/v1/services", "watch=true", clientConn)
+	conn := m.createConnection("test-cluster", "test-user", "/api/v1/services", "watch=true", clientConn, nil)
 	wsConn, wsServer := createTestWebSocketConnection()
 
 	defer wsServer.Close()
@@ -598,7 +598,7 @@ func TestReconnect(t *testing.T) {
 	assert.Contains(t, err.Error(), "getting context: key not found")
 
 	// Test reconnection with closed connection
-	conn = m.createConnection("test-cluster", "test-user", "/api/v1/pods", "watch=true", clientConn)
+	conn = m.createConnection("test-cluster", "test-user", "/api/v1/pods", "watch=true", clientConn, nil)
 	wsConn2, wsServer2 := createTestWebSocketConnection()
 
 	defer wsServer2.Close()
@@ -829,7 +829,7 @@ func TestMonitorConnection_ReconnectFailure(t *testing.T) {
 	clientConn, clientServer := createTestWebSocketConnection()
 	defer clientServer.Close()
 
-	conn := m.createConnection("test-cluster", "test-user", "/api/v1/pods", "", clientConn)
+	conn := m.createConnection("test-cluster", "test-user", "/api/v1/pods", "", clientConn, nil)
 	wsConn, wsServer := createTestWebSocketConn()
 
 	defer wsServer.Close()
@@ -1097,7 +1097,7 @@ func TestMonitorConnection_Reconnect(t *testing.T) {
 	wsURL := "ws" + strings.TrimPrefix(server.URL, "http")
 	tlsConfig := &tls.Config{InsecureSkipVerify: true} //nolint:gosec
 
-	ws, err := m.dialWebSocket(wsURL, tlsConfig, "")
+	ws, err := m.dialWebSocket(wsURL, tlsConfig, "", nil)
 	require.NoError(t, err)
 
 	conn.WSConn = ws

frontend/src/lib/k8s/api/v2/webSocket.test.ts

@@ -88,6 +88,7 @@ describe('WebSocket Tests', () => {
         path,
         query,
         userId,
+        token: 'test-token',
         type: 'REQUEST',
       });
 
@@ -130,6 +131,7 @@ describe('WebSocket Tests', () => {
           path: sub.path,
           query: sub.query,
           userId,
+          token: 'test-token',
           type: 'REQUEST',
         });
 

frontend/src/lib/k8s/api/v2/webSocket.ts

@@ -50,6 +50,9 @@ interface WebSocketMessage {
    * - COMPLETE: Server indicates the watch request has completed (e.g., due to timeout or error)
    */
   type: 'REQUEST' | 'CLOSE' | 'COMPLETE';
+
+  /** Authentication token */
+  token?: string;
 }
 
 /**
@@ -210,6 +213,7 @@ export const WebSocketManager = {
       query,
       userId: userId || '',
       type: 'REQUEST',
+      token: getToken(clusterId),
     };
     socket.send(JSON.stringify(requestMsg));