From 9ac1ceb4b68bb2f5b9b50522f4ce281d424ec401 Mon Sep 17 00:00:00 2001 From: Heikki Hellgren Date: Mon, 15 Jan 2024 09:16:28 +0200 Subject: [PATCH] chore: move connection upgrade to router, add tests for manager Signed-off-by: Heikki Hellgren --- .../src/service/SignalManager.test.ts | 178 ++++++++++++++++++ .../src/service/SignalManager.ts | 69 +------ plugins/signals-backend/src/service/router.ts | 63 +++++-- 3 files changed, 236 insertions(+), 74 deletions(-) create mode 100644 plugins/signals-backend/src/service/SignalManager.test.ts diff --git a/plugins/signals-backend/src/service/SignalManager.test.ts b/plugins/signals-backend/src/service/SignalManager.test.ts new file mode 100644 index 0000000000..5720d1ff52 --- /dev/null +++ b/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 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' } }), + ); + }); +}); diff --git a/plugins/signals-backend/src/service/SignalManager.ts b/plugins/signals-backend/src/service/SignalManager.ts index 988db47465..983496b96b 100644 --- a/plugins/signals-backend/src/service/SignalManager.ts +++ b/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}`); + } + }); }); } } diff --git a/plugins/signals-backend/src/service/router.ts b/plugins/signals-backend/src/service/router.ts index ec56860bf3..4394a9f480 100644 --- a/plugins/signals-backend/src/service/router.ts +++ b/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 { - 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();