chore: move connection upgrade to router, add tests for manager

Signed-off-by: Heikki Hellgren <heikki.hellgren@op.fi>

plugins/signals-backend/src/service/SignalManager.test.ts

@@ -0,0 +1,178 @@
+/*
+ * Copyright 2024 The Backstage Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+import { WebSocket } from 'ws';
+import { EventSubscriber } from '@backstage/plugin-events-node';
+import { SignalManager } from './SignalManager';
+import { getVoidLogger } from '@backstage/backend-common';
+
+class MockWebSocket {
+  closed: boolean = false;
+  readyState: number = WebSocket.OPEN;
+  callbacks: Map<string | symbol, (this: WebSocket, ...args: any[]) => void> =
+    new Map();
+  data: any[] = [];
+
+  close(_: number, __: string | Buffer): void {
+    this.readyState = WebSocket.CLOSED;
+    this.closed = true;
+  }
+
+  on(
+    event: string | symbol,
+    listener: (this: WebSocket, ...args: any[]) => void,
+  ) {
+    this.callbacks.set(event, listener);
+    return this;
+  }
+
+  // @ts-ignore
+  send(data: any, _?: (err?: Error) => void): void {
+    this.data.push(data);
+  }
+
+  trigger(event: string | symbol, ...args: any[]): void {
+    const cb = this.callbacks.get(event);
+    if (!cb) {
+      throw new Error(`No callback for ${event.toString()}`);
+    }
+    // @ts-ignore
+    cb(...args);
+  }
+}
+
+describe('SignalManager', () => {
+  let onEvent: Function;
+
+  const mockEventBroker = {
+    publish: async () => {},
+    subscribe: (subscriber: EventSubscriber) => {
+      onEvent = subscriber.onEvent;
+    },
+  };
+
+  const manager = SignalManager.create({
+    eventBroker: mockEventBroker,
+    logger: getVoidLogger(),
+  });
+
+  it('should close connection on error', () => {
+    const ws = new MockWebSocket();
+    manager.addConnection(ws as unknown as WebSocket);
+
+    ws.trigger('error', new Error('error'));
+    expect(ws.closed).toBeTruthy();
+  });
+
+  it('should allow subscribing and unsubscribing to events', async () => {
+    const ws = new MockWebSocket();
+    manager.addConnection(ws as unknown as WebSocket);
+
+    ws.trigger(
+      'message',
+      JSON.stringify({ action: 'subscribe', channel: 'test' }),
+      false,
+    );
+
+    await onEvent({
+      topic: 'signals',
+      eventPayload: {
+        recipients: null,
+        channel: 'test',
+        message: { msg: 'test' },
+      },
+    });
+
+    expect(ws.data.length).toEqual(1);
+    expect(ws.data[0]).toEqual(
+      JSON.stringify({ channel: 'test', message: { msg: 'test' } }),
+    );
+
+    ws.trigger(
+      'message',
+      JSON.stringify({ action: 'unsubscribe', channel: 'test' }),
+      false,
+    );
+
+    await onEvent({
+      topic: 'signals',
+      eventPayload: {
+        recipients: null,
+        channel: 'test',
+        message: { msg: 'test' },
+      },
+    });
+
+    expect(ws.data.length).toEqual(1);
+  });
+
+  it('should only send to users from identity', async () => {
+    // Connection without identity
+    const ws1 = new MockWebSocket();
+    manager.addConnection(ws1 as unknown as WebSocket);
+
+    // Connection with identity and subscription
+    const ws2 = new MockWebSocket();
+    manager.addConnection(ws2 as unknown as WebSocket, {
+      identity: {
+        type: 'user',
+        ownershipEntityRefs: ['user:default/john.doe'],
+        userEntityRef: 'user:default/john.doe',
+      },
+      expiresInSeconds: 3600,
+      token: '1234',
+    });
+
+    // Connection without subscription
+    const ws3 = new MockWebSocket();
+    manager.addConnection(ws3 as unknown as WebSocket, {
+      identity: {
+        type: 'user',
+        ownershipEntityRefs: ['user:default/john.doe'],
+        userEntityRef: 'user:default/john.doe',
+      },
+      expiresInSeconds: 3600,
+      token: '1234',
+    });
+
+    ws1.trigger(
+      'message',
+      JSON.stringify({ action: 'subscribe', channel: 'test' }),
+      false,
+    );
+
+    ws2.trigger(
+      'message',
+      JSON.stringify({ action: 'subscribe', channel: 'test' }),
+      false,
+    );
+
+    await onEvent({
+      topic: 'signals',
+      eventPayload: {
+        recipients: 'user:default/john.doe',
+        channel: 'test',
+        message: { msg: 'test' },
+      },
+    });
+
+    expect(ws1.data.length).toEqual(0);
+    expect(ws3.data.length).toEqual(0);
+    expect(ws2.data.length).toEqual(1);
+    expect(ws2.data[0]).toEqual(
+      JSON.stringify({ channel: 'test', message: { msg: 'test' } }),
+    );
+  });
+});

plugins/signals-backend/src/service/SignalManager.ts

@@ -15,24 +15,11 @@
  */
 import { EventBroker, EventParams } from '@backstage/plugin-events-node';
 import { SignalPayload } from '@backstage/plugin-signals-node';
-import { RawData, WebSocket, WebSocketServer } from 'ws';
-import { IncomingMessage } from 'http';
+import { RawData, WebSocket } from 'ws';
 import { v4 as uuid } from 'uuid';
 import { JsonObject } from '@backstage/types';
-import {
-  BackstageIdentityResponse,
-  IdentityApi,
-  IdentityApiGetIdentityRequest,
-} from '@backstage/plugin-auth-node';
+import { BackstageIdentityResponse } from '@backstage/plugin-auth-node';
 import { LoggerService } from '@backstage/backend-plugin-api';
-import { Duplex } from 'stream';
-
-/** @internal */
-export type ConnectionUpgradeOptions = {
-  request: IncomingMessage;
-  socket: Duplex;
-  head: Buffer;
-};
 
 /**
  * @internal
@@ -52,7 +39,6 @@ export type SignalManagerOptions = {
   // TODO: Remove optional when events-backend can offer this service
   eventBroker?: EventBroker;
   logger: LoggerService;
-  identity: IdentityApi;
 };
 
 /** @internal */
@@ -63,24 +49,13 @@ export class SignalManager {
   >();
   private eventBroker?: EventBroker;
   private logger: LoggerService;
-  private identity: IdentityApi;
-  private server: WebSocketServer;
 
   static create(options: SignalManagerOptions) {
     return new SignalManager(options);
   }
 
   private constructor(options: SignalManagerOptions) {
-    ({
-      eventBroker: this.eventBroker,
-      logger: this.logger,
-      identity: this.identity,
-    } = options);
-
-    this.server = new WebSocketServer({
-      noServer: true,
-      clientTracking: false,
-    });
+    ({ eventBroker: this.eventBroker, logger: this.logger } = options);
 
     this.eventBroker?.subscribe({
       supportsEventTopics: () => ['signals'],
@@ -89,37 +64,7 @@ export class SignalManager {
     });
   }
 
-  /**
-   * Handles request upgrade to websocket and adds the connection to internal
-   * list for publish/subscribe functionality
-   * @param req - Request
-   */
-  async handleUpgrade(options: ConnectionUpgradeOptions) {
-    const { request, socket, head } = options;
-    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 = request.headers['sec-websocket-protocol'];
-    if (token) {
-      identity = await this.identity.getIdentity({
-        request: {
-          headers: { authorization: token },
-        },
-      } as IdentityApiGetIdentityRequest);
-    }
-
-    this.server.handleUpgrade(
-      request,
-      socket,
-      head,
-      (ws: WebSocket, __: IncomingMessage) => {
-        this.addConnection(ws, identity);
-      },
-    );
-  }
-
-  private addConnection(ws: WebSocket, identity?: BackstageIdentityResponse) {
+  addConnection(ws: WebSocket, identity?: BackstageIdentityResponse) {
     const id = uuid();
 
     const conn = {
@@ -207,7 +152,11 @@ export class SignalManager {
         return;
       }
 
-      conn.ws.send(jsonMessage);
+      conn.ws.send(jsonMessage, err => {
+        if (err) {
+          this.logger.error(`Failed to send message to ${conn.id}: ${err}`);
+        }
+      });
     });
   }
 }

plugins/signals-backend/src/service/router.ts

@@ -18,10 +18,15 @@ import express, { NextFunction, Request, Response } from 'express';
 import Router from 'express-promise-router';
 import { LoggerService } from '@backstage/backend-plugin-api';
 import * as https from 'https';
-import http from 'http';
+import http, { IncomingMessage } from 'http';
 import { SignalManager } from './SignalManager';
-import { IdentityApi } from '@backstage/plugin-auth-node';
+import {
+  BackstageIdentityResponse,
+  IdentityApi,
+  IdentityApiGetIdentityRequest,
+} from '@backstage/plugin-auth-node';
 import { EventBroker } from '@backstage/plugin-events-node';
+import { WebSocket, WebSocketServer } from 'ws';
 
 /** @public */
 export interface RouterOptions {
@@ -34,13 +39,23 @@ export interface RouterOptions {
 export async function createRouter(
   options: RouterOptions,
 ): Promise<express.Router> {
-  const { logger } = options;
+  const { logger, identity } = options;
   const manager = SignalManager.create(options);
-  let subscribed = false;
+  let subscribedToUpgradeRequests = false;
 
-  const upgradeMiddleware = (req: Request, _: Response, next: NextFunction) => {
+  const webSocketServer = new WebSocketServer({
+    noServer: true,
+    clientTracking: false,
+  });
+
+  const upgradeMiddleware = async (
+    req: Request,
+    _: Response,
+    next: NextFunction,
+  ) => {
     const server: https.Server | http.Server = (req.socket as any)?.server;
     if (
+      subscribedToUpgradeRequests ||
       !server ||
       !req.headers ||
       req.headers.upgrade === undefined ||
@@ -50,15 +65,35 @@ export async function createRouter(
       return;
     }
 
-    if (!subscribed) {
-      subscribed = true;
-      server.on('upgrade', async (request, socket, head) => {
-        // TODO: Find a way to make this more generic
-        if (request.url === '/api/signals') {
-          await manager.handleUpgrade({ request, socket, head });
-        }
-      });
-    }
+    subscribedToUpgradeRequests = true;
+    server.on('upgrade', async (request, socket, head) => {
+      // TODO: Find a way to make this more generic
+      if (request.url !== '/api/signals') {
+        return;
+      }
+
+      let userIdentity: 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) {
+        userIdentity = await identity.getIdentity({
+          request: {
+            headers: { authorization: token },
+          },
+        } as IdentityApiGetIdentityRequest);
+      }
+
+      webSocketServer.handleUpgrade(
+        request,
+        socket,
+        head,
+        (ws: WebSocket, __: IncomingMessage) => {
+          manager.addConnection(ws, userIdentity);
+        },
+      );
+    });
   };
 
   const router = Router();