From 7af9750ea8f9d7bf815ffe21565f50f993700368 Mon Sep 17 00:00:00 2001
From: Heikki Hellgren <heikki.hellgren@op.fi>
Date: Tue, 5 Dec 2023 09:19:50 +0200
Subject: [PATCH] feat: use websocket header to pass the auth token

Signed-off-by: Heikki Hellgren <heikki.hellgren@op.fi>
---
 plugins/signals-node/src/SignalsService.ts | 45 ++++++++--------------
 plugins/signals-node/src/types.ts          |  8 ++++
 plugins/signals/src/api/SignalsClient.ts   | 16 ++------
 3 files changed, 28 insertions(+), 41 deletions(-)

diff --git a/plugins/signals-node/src/SignalsService.ts b/plugins/signals-node/src/SignalsService.ts
index 9cbe4b1565..6a1547a6bb 100644
--- a/plugins/signals-node/src/SignalsService.ts
+++ b/plugins/signals-node/src/SignalsService.ts
@@ -19,7 +19,11 @@ import {
   EventSubscriber,
 } from '@backstage/plugin-events-node';
 import { Logger } from 'winston';
-import { ServiceOptions, SignalConnection } from './types';
+import {
+  ServiceOptions,
+  SignalConnection,
+  SignalsEventBrokerPayload,
+} from './types';
 import { RawData, WebSocket, WebSocketServer } from 'ws';
 import { IncomingMessage } from 'http';
 import { v4 as uuid } from 'uuid';
@@ -31,13 +35,6 @@ import {
   IdentityApiGetIdentityRequest,
 } from '@backstage/plugin-auth-node';
 
-/** @public */
-export type SignalsEventBrokerPayload = {
-  recipients?: string[];
-  topic?: string;
-  message?: JsonObject;
-};
-
 /** @public */
 export class SignalsService implements EventSubscriber {
   private readonly serverId: string;
@@ -83,9 +80,18 @@ export class SignalsService implements EventSubscriber {
    * @param req - Request
    */
   handleUpgrade = async (req: Request) => {
-    const identity = await this.identity.getIdentity({
-      request: req,
-    });
+    let identity: BackstageIdentityResponse | undefined = undefined;
+
+    // Authentication token is passed in Sec-WebSocket-Protocol header as there
+    // is no other way to pass the token with plain websockets
+    const token = req.headers['sec-websocket-protocol'];
+    if (token) {
+      identity = await this.identity.getIdentity({
+        request: {
+          headers: { authorization: token },
+        },
+      } as IdentityApiGetIdentityRequest);
+    }
 
     this.server.handleUpgrade(
       req,
@@ -161,23 +167,6 @@ export class SignalsService implements EventSubscriber {
       );
       connection.subscriptions.delete(message.topic as string);
     }
-
-    if (message.action === 'authenticate' && message.token) {
-      this.logger.info(`Connection ${connection.id} authenticated`);
-      this.identity
-        .getIdentity({
-          request: {
-            headers: { authorization: message.token },
-          },
-        } as IdentityApiGetIdentityRequest)
-        .then(identity => {
-          if (identity) {
-            connection.user = identity.identity.userEntityRef;
-            connection.ownershipEntityRefs =
-              identity.identity.ownershipEntityRefs;
-          }
-        });
-    }
   }
 
   /**
diff --git a/plugins/signals-node/src/types.ts b/plugins/signals-node/src/types.ts
index c53bfdb20e..090e4b2ff1 100644
--- a/plugins/signals-node/src/types.ts
+++ b/plugins/signals-node/src/types.ts
@@ -17,6 +17,7 @@ import { IdentityApi } from '@backstage/plugin-auth-node';
 import { EventBroker } from '@backstage/plugin-events-node';
 import { Logger } from 'winston';
 import { WebSocket } from 'ws';
+import { JsonObject } from '@backstage/types';
 
 /**
  * @public
@@ -27,6 +28,13 @@ export type ServiceOptions = {
   identity: IdentityApi;
 };
 
+/** @public */
+export type SignalsEventBrokerPayload = {
+  recipients?: string[];
+  topic?: string;
+  message?: JsonObject;
+};
+
 /**
  * @internal
  */
diff --git a/plugins/signals/src/api/SignalsClient.ts b/plugins/signals/src/api/SignalsClient.ts
index cc1254578f..d8d3328411 100644
--- a/plugins/signals/src/api/SignalsClient.ts
+++ b/plugins/signals/src/api/SignalsClient.ts
@@ -123,9 +123,11 @@ export class SignalsClient implements SignalsApi {
     }
 
     const apiUrl = `${await this.discoveryApi.getBaseUrl('signals')}`;
+    const { token } = await this.identity.getCredentials();
+
     const url = new URL(apiUrl);
     url.protocol = url.protocol === 'http:' ? 'ws' : 'wss';
-    this.ws = new WebSocket(url.toString());
+    this.ws = new WebSocket(url.toString(), token);
 
     this.ws.onmessage = (data: MessageEvent) => {
       this.handleMessage(data);
@@ -155,9 +157,6 @@ export class SignalsClient implements SignalsApi {
     if (!this.ws || this.ws.readyState !== WebSocket.OPEN) {
       throw new Error('Connect timeout');
     }
-
-    // Authenticate
-    await this.authenticate();
   }
 
   private handleMessage(data: MessageEvent) {
@@ -175,15 +174,6 @@ export class SignalsClient implements SignalsApi {
     }
   }
 
-  private async authenticate() {
-    const { token } = await this.identity.getCredentials();
-    if (token) {
-      // Authentication is done with websocket message to server as the plain
-      // websocket does not allow sending headers during connection upgrade
-      this.send({ action: 'authenticate', token: token });
-    }
-  }
-
   private reconnect() {
     if (this.reconnectTimeout) {
       clearTimeout(this.reconnectTimeout);