Merge pull request #1048 from snyk/chore/ACC-3063-add-typings-for-primus

chore: ACC-3064 add typings for Primus

Author: gclapperton

lib/broker-workload/requestFiltering.ts

@@ -1,4 +1,5 @@
 import { getFilterConfig } from '../hybrid-sdk/client/config/filters';
+import { WebSocketConnection } from '../hybrid-sdk/client/types/client';
 import { LOADEDFILTERSET } from '../hybrid-sdk/common/types/filter';
 
 export const filterRequest = (payload, options, websocketConnectionHandler) => {
@@ -28,7 +29,7 @@ export const filterRequest = (payload, options, websocketConnectionHandler) => {
 export const filterClientRequest = (
   payload,
   options,
-  websocketConnectionHandler,
+  websocketConnectionHandler: WebSocketConnection,
 ) => {
   let filterResponse;
   if (

lib/broker-workload/websocketRequests.ts

@@ -24,15 +24,17 @@ import {
   incrementHttpRequestsTotal,
 } from '../hybrid-sdk/common/utils/metrics';
 import { maskToken, hashToken } from '../hybrid-sdk/common/utils/token';
+import { WebSocketServer } from '../hybrid-sdk/server/types/socket';
+import { WebSocketConnection } from '../hybrid-sdk/client/types/client';
 
 export class BrokerWorkload extends Workload<WorkloadType.remoteServer> {
   options;
   connectionIdentifier: string;
-  websocketConnectionHandler;
+  websocketConnectionHandler: WebSocketServer | WebSocketConnection;
   constructor(
     connectionIdentifier: string,
     options,
-    websocketConnectionHandler,
+    websocketConnectionHandler: WebSocketServer | WebSocketConnection,
   ) {
     super('broker', WorkloadType['remote-server']);
     this.options = options;

lib/hybrid-sdk/client/socket.ts

@@ -158,10 +158,10 @@ export const createWebSocket = (
       },
     };
   }
-  const websocket: WebSocketConnection = new Socket(
+  const websocket = new Socket(
     localClientOps.config.brokerServerUrlForSocket,
     socketSettings,
-  );
+  ) as unknown as WebSocketConnection;
   websocket.socketVersion = 1;
   websocket.socketType = 'client';
   if (localClientOps.config.universalBrokerEnabled) {

lib/hybrid-sdk/client/socketHandlers/identifyHandler.ts

@@ -1,3 +1,5 @@
-export const identifyHandler = (serverData, websocket) => {
+import { WebSocketConnection } from '../types/client';
+
+export const identifyHandler = (serverData, websocket: WebSocketConnection) => {
   websocket.capabilities = serverData.capabilities;
 };

lib/hybrid-sdk/client/socketHandlers/requestHandler.ts

@@ -1,12 +1,17 @@
 import { forwardWebSocketRequest } from '../../common/connectionToWorkloadInterface/forwardWebsocketRequest';
 import { RequestPayload } from '../../common/types/http';
 import { LoadedClientOpts } from '../../common/types/options';
+import { WebSocketConnection } from '../types/client';
+import { WebSocketServer } from '../../server/types/socket';
 
 let initializedReqHandler: (
   webSocketIdentifier: string,
 ) => (payload: RequestPayload, emit: any) => void;
 
-export const initRequestHandler = (websocket, clientOps: LoadedClientOpts) => {
+export const initRequestHandler = (
+  websocket: WebSocketServer | WebSocketConnection,
+  clientOps: LoadedClientOpts,
+) => {
   initializedReqHandler = forwardWebSocketRequest(clientOps, websocket);
 };
 

lib/hybrid-sdk/client/types/client.ts

@@ -1,5 +1,6 @@
 import { FiltersType } from '../../common/types/filter';
 import { CheckResult } from '../checks/types';
+import { Primus } from 'primus';
 
 export interface HookResults {
   preflightCheckResults?: CheckResult[];
@@ -49,42 +50,30 @@ export interface ConnectionMetadata {
   serverId?: string;
 }
 
-export interface WebSocketConnection {
-  options: {
-    reconnect: any;
-    ping: number;
-    pong: number;
-    timeout: number;
-    transport: any;
-    queueSize: any;
-    stategy: any;
-  };
-  transport: any;
-  socketVersion?: any;
-  socketType?: string;
-  identifier?: string;
+export interface WebSocketConnection
+  extends Pick<Primus, 'destroy' | 'emit' | 'end' | 'on'> {
+  capabilities?: string[];
   clientConfig?: any;
-  role: Role;
   friendlyName?: string;
-  supportedIntegrationType: string;
+  identifier?: string;
+  role: Role;
   serverId: string;
-  url: any;
-  latency: any;
+  socketType: 'client';
+  socketVersion?: number;
+  supportedIntegrationType: string;
+  timeoutHandlerId?: NodeJS.Timeout;
+
+  // Added by primus, but specific to the client
   socket: any;
-  destroy: any;
-  send: any;
-  end: any;
-  open: any;
-  emit: any;
-  capabilities?: any;
-  on: (string, any) => any;
-  readyState: any;
-  timeoutHandlerId?: any;
-}
-// export interface WebSocketConnection {
-//   websocket: Connection;
-// }
+  readyState: number;
+  transport: {
+    extraHeaders?: Record<string, string>;
+  };
+  url: URL;
 
+  // Added by primus-emitter plugin
+  send: (event: string, ...args: any[]) => void;
+}
 export interface ValidationResult {
   connectionName: string;
   validated: boolean;

lib/hybrid-sdk/client/utils/socketHelpers.ts

@@ -1,5 +1,6 @@
 import primus from 'primus';
+import { WebSocketConnection } from '../types/client';
 
-export const isWebsocketConnOpen = (io) => {
+export const isWebsocketConnOpen = (io: WebSocketConnection) => {
   return io.readyState === primus.Spark.OPEN;
 };

lib/hybrid-sdk/common/connectionToWorkloadInterface/forwardWebsocketRequest.ts

@@ -1,37 +1,43 @@
 import { RequestPayload } from '../types/http';
-import { WebSocketConnection } from '../../client/types/client';
 import { LoadedClientOpts, LoadedServerOpts } from '../types/options';
 import { BrokerWorkload } from '../../../broker-workload/websocketRequests';
 import {
   Workload,
   WorkloadType,
   RemoteServerWorkloadRuntimeParams,
 } from '../../workloadFactory';
+import { WebSocketServer } from '../../server/types/socket';
+import { WebSocketConnection } from '../../client/types/client';
+import { HybridResponse } from '../../responseSenders';
 
 export const forwardWebSocketRequest = (
   options: LoadedClientOpts | LoadedServerOpts,
-  websocketConnectionHandler: WebSocketConnection,
+  websocketConnectionHandler: WebSocketServer | WebSocketConnection,
 ) => {
   // 1. Request coming in over websocket conn (logged)
   // 2. Filter for rule match (log and block if no match)
   // 3. Relay over HTTP conn (logged)
   // 4. Get response over HTTP conn (logged)
   // 5. Send response over websocket conn
 
-  return (connectionIdentifier) => async (payload: RequestPayload, emit) => {
-    const workloadName = options.config.remoteWorkloadName;
-    const workloadModulePath = options.config.remoteWorkloadModulePath;
-    const workload = (await Workload.instantiate(
-      workloadName,
-      workloadModulePath,
-      WorkloadType.remoteServer,
-      { connectionIdentifier, options, websocketConnectionHandler },
-    )) as BrokerWorkload;
+  return (connectionIdentifier) =>
+    async (
+      payload: RequestPayload,
+      emit: (response: HybridResponse) => void,
+    ) => {
+      const workloadName = options.config.remoteWorkloadName;
+      const workloadModulePath = options.config.remoteWorkloadModulePath;
+      const workload = (await Workload.instantiate(
+        workloadName,
+        workloadModulePath,
+        WorkloadType.remoteServer,
+        { connectionIdentifier, options, websocketConnectionHandler },
+      )) as BrokerWorkload;
 
-    const data: RemoteServerWorkloadRuntimeParams = {
-      payload,
-      websocketHandler: emit,
+      const data: RemoteServerWorkloadRuntimeParams = {
+        payload,
+        websocketHandler: emit,
+      };
+      await workload.handler(data);
     };
-    await workload.handler(data);
-  };
 };

lib/hybrid-sdk/responseSenders.ts

@@ -6,6 +6,8 @@ import { logError, logResponse } from '../logs/log';
 import { isJson } from './common/utils/json';
 import { replaceUrlPartialChunk } from './common/utils/replace-vars';
 import { RequestMetadata } from './types';
+import { WebSocketConnection } from './client/types/client';
+import { WebSocketServer } from './server/types/socket';
 
 export interface HybridResponse {
   status: number;
@@ -16,16 +18,16 @@ export interface HybridResponse {
 }
 export class HybridResponseHandler {
   connectionIdentifier;
-  websocketConnectionHandler;
+  websocketConnectionHandler: WebSocketServer | WebSocketConnection;
   logContext;
   config;
   requestMetadata: RequestMetadata;
-  websocketResponseHandler;
+  websocketResponseHandler: (response: HybridResponse) => void;
   responseHandler;
   constructor(
     requestMetadata: RequestMetadata,
-    websocketConnectionHandler,
-    websocketResponseHandler,
+    websocketConnectionHandler: WebSocketServer | WebSocketConnection,
+    websocketResponseHandler: (response: HybridResponse) => void,
     config,
     logContext,
   ) {
@@ -38,9 +40,9 @@ export class HybridResponseHandler {
     // WebsocketResponseHandler provided means WS response expected
     // header x-broker-ws-response:true used on server side
     if (
-      this.websocketConnectionHandler?.capabilities?.includes(
-        'receive-post-streams',
-      ) &&
+      (
+        this.websocketConnectionHandler as { capabilities?: string[] }
+      )?.capabilities?.includes('receive-post-streams') &&
       !this.websocketResponseHandler
     ) {
       // Response Traffic over HTTP Post
@@ -107,9 +109,10 @@ export class HybridResponseHandler {
         this.logContext,
         this.config,
         this.connectionIdentifier,
-        this.websocketConnectionHandler?.serverId ?? '',
+        (this.websocketConnectionHandler as { serverId?: string })?.serverId ??
+          '',
         this.requestMetadata.requestId,
-        this.websocketConnectionHandler.role,
+        (this.websocketConnectionHandler as { role?: string })?.role ?? '',
       );
       if (streamingRequestData) {
         // POST Streaming

lib/hybrid-sdk/server/socket.ts

@@ -1,7 +1,7 @@
 import Primus from 'primus';
 import Emitter from 'primus-emitter';
 import { LoadedServerOpts } from '../common/types/options';
-import { SocketHandler } from './types/socket';
+import { WebSocketServer, SocketHandler } from './types/socket';
 import { handleIoError } from './socketHandlers/errorHandler';
 import { handleSocketConnection } from './socketHandlers/connectionHandler';
 import { initConnectionHandler } from './socketHandlers/initHandlers';
@@ -47,17 +47,20 @@ const socket = ({ server, loadedServerOpts }): SocketHandler => {
     compression: Boolean(loadedServerOpts.config.socketUseCompression) || false,
   };
 
-  const websocket = new Primus(server, ioConfig);
+  const websocket = new Primus(server, ioConfig) as WebSocketServer;
   if (loadedServerOpts.config.BROKER_SERVER_MANDATORY_AUTH_ENABLED) {
     websocket.authorize(async (req, done) => {
       const connectionIdentifier = req.uri.pathname
         .replaceAll(/^\/primus\/([^/]+)\//g, '$1')
         .toLowerCase();
       const maskedToken = maskToken(connectionIdentifier);
-      const authHeader =
-        req.headers['Authorization'] ?? req.headers['authorization'];
-      const brokerClientId = req.headers['x-snyk-broker-client-id'] ?? null;
-      const role = req.headers['x-snyk-broker-client-role'] ?? null;
+      const authHeader = (req.headers['Authorization'] ??
+        req.headers['authorization']) as string | undefined;
+      const brokerClientId =
+        (req.headers['x-snyk-broker-client-id'] as string | undefined) ?? null;
+      const role =
+        (req.headers['x-snyk-broker-client-role'] as string | undefined) ??
+        null;
       if (
         (!authHeader ||
           !authHeader.toLowerCase().startsWith('bearer') ||
@@ -93,8 +96,8 @@ const socket = ({ server, loadedServerOpts }): SocketHandler => {
         );
         // deepcode ignore Ssrf: request URL comes from the filter response, with the origin url being injected by the filtered version
         const credsCheckResponse = await validateBrokerClientCredentials(
-          authHeader,
-          brokerClientId,
+          authHeader!,
+          brokerClientId!,
           connectionIdentifier,
         );
         if (!credsCheckResponse) {
@@ -121,9 +124,9 @@ const socket = ({ server, loadedServerOpts }): SocketHandler => {
         const currentClient: ClientSocket = {
           socketType: 'server',
           socketVersion: 1,
-          brokerClientId: brokerClientId,
+          brokerClientId: brokerClientId!,
           brokerAppClientId: brokerAppClientId,
-          role: role ?? Role.primary,
+          role: (role ?? Role.primary) as Role,
           credsValidationTime: nowDate,
         };
         const connections = getSocketConnections();