@hexabot-ai/api 3.2.2-beta.2 → 3.2.3

package/src/websocket/utils/gateway-options.ts

@@ -4,40 +4,154 @@
  * Full terms: see LICENSE.md.
  */
 
+import type { IncomingMessage } from 'http';
 import util from 'util';
 
+import type { Source } from '@hexabot-ai/types';
+import { NotFoundException } from '@nestjs/common';
+import { isUUID } from 'class-validator';
 import type { ServerOptions } from 'socket.io';
 
 import { AppInstance } from '@/app.instance';
+import { SourceService } from '@/channel/services/source.service';
 import { config } from '@/config';
+import { CONSOLE_CHANNEL_NAME } from '@/extensions/channels/console/settings.schema';
+import { WEB_CHANNEL_NAME } from '@/extensions/channels/web/settings.schema';
 import { SettingService } from '@/setting/services/setting.service';
+import { getAllowedDomains, isAllowedOrigin } from '@/utils/helpers/origin';
 
-export const buildWebSocketGatewayOptions = (): Partial<ServerOptions> => {
+type SocketCorsSource = Pick<Source, 'channel' | 'settings'>;
+type AllowRequestCallback = (
+  err: string | null | undefined,
+  success: boolean,
+) => void;
+
+export type SocketOriginDependencies = {
+  getAllowedOrigins: () => Promise<string[]>;
+  findActiveSourceById: (sourceId: string) => Promise<SocketCorsSource | null>;
+};
+
+const getDefaultSocketOriginDependencies = (): SocketOriginDependencies => {
+  const app = AppInstance.getApp();
+  const settingService = app.get<SettingService>(SettingService);
+  const sourceService = app.get<SourceService>(SourceService);
+
+  return {
+    getAllowedOrigins: () => settingService.getAllowedOrigins(),
+    findActiveSourceById: (sourceId) => sourceService.findActiveById(sourceId),
+  };
+};
+
+export const isSocketRequestOriginAllowed = async (
+  handshake: IncomingMessage,
+  dependencies: SocketOriginDependencies = getDefaultSocketOriginDependencies(),
+): Promise<boolean> => {
+  const { origin: rawOrigin } = handshake.headers;
+  const origin = Array.isArray(rawOrigin) ? rawOrigin[0] : rawOrigin;
+  const rawSourceId = handshake.url
+    ? new URL(handshake.url, 'ws://localhost').searchParams.get('source_id')
+    : null;
+
+  if (!rawSourceId) {
+    return isAllowedOrigin(origin, await dependencies.getAllowedOrigins());
+  }
+
+  if (!isUUID(rawSourceId)) {
+    return false;
+  }
+
+  let source: SocketCorsSource | null;
+  try {
+    source = await dependencies.findActiveSourceById(rawSourceId);
+  } catch (error) {
+    if (error instanceof NotFoundException) {
+      return false;
+    }
+    throw error;
+  }
+  if (!source) {
+    return false;
+  }
+
+  if (source.channel === WEB_CHANNEL_NAME) {
+    return isAllowedOrigin(origin, getAllowedDomains(source.settings));
+  }
+
+  if (source.channel === CONSOLE_CHANNEL_NAME) {
+    return isAllowedOrigin(origin, await dependencies.getAllowedOrigins());
+  }
+
+  return false;
+};
+
+const runBeforeConnect = (
+  handshake: IncomingMessage,
+  cb: AllowRequestCallback,
+): boolean => {
+  try {
+    const result = config.sockets.beforeConnect(handshake);
+
+    if (!result) {
+      cb(null, false);
+
+      return false;
+    }
+
+    return true;
+  } catch (e) {
+    // eslint-disable-next-line no-console
+    console.log(
+      `A socket was rejected via the config.sockets.beforeConnect function.\n` +
+        `It attempted to connect with headers:\n` +
+        `${util.inspect(handshake.headers, { depth: null })}\n` +
+        `Details: ${e}`,
+    );
+
+    cb(e instanceof Error ? e.message : `${e}`, false);
+
+    return false;
+  }
+};
+
+export const buildWebSocketGatewayOptions = (
+  dependencies?: SocketOriginDependencies,
+): Partial<ServerOptions> => {
   const opts: Partial<ServerOptions> = {
     allowEIO3: true, // Allows support for Engine.io v3 clients.
     path: config.sockets.path,
     ...(typeof config.sockets.serveClient !== 'undefined' && {
       serveClient: config.sockets.serveClient,
     }),
-    ...(config.sockets.beforeConnect && {
-      allowRequest: (handshake, cb) => {
-        try {
-          const result = config.sockets.beforeConnect(handshake);
+    allowRequest: (handshake, cb) => {
+      if (!runBeforeConnect(handshake, cb)) {
+        return;
+      }
+
+      if (config.env === 'test') {
+        cb(null, true);
+
+        return;
+      }
+
+      isSocketRequestOriginAllowed(handshake, dependencies)
+        .then((allowed) => {
+          if (allowed) {
+            cb(null, true);
+
+            return;
+          }
 
-          return cb(null, result);
-        } catch (e) {
           // eslint-disable-next-line no-console
           console.log(
-            `A socket was rejected via the config.sockets.beforeConnect function.\n` +
-              `It attempted to connect with headers:\n` +
-              `${util.inspect(handshake.headers, { depth: null })}\n` +
-              `Details: ${e}`,
+            `A socket was rejected by the Socket.IO origin policy.\n` +
+              `It attempted to connect with origin: ${handshake.headers.origin}`,
           );
-
-          return cb(e, false);
-        }
-      },
-    }),
+          cb('Origin not allowed', false);
+        })
+        .catch((error) => {
+          cb(error instanceof Error ? error.message : `${error}`, false);
+        });
+    },
     ...(config.sockets.pingTimeout && {
       pingTimeout: config.sockets.pingTimeout,
     }),
@@ -54,31 +168,7 @@ export const buildWebSocketGatewayOptions = (): Partial<ServerOptions> => {
     ...(config.sockets.cookie && { cookie: config.sockets.cookie }),
     cors: {
       credentials: true,
-      origin: async (origin, cb) => {
-        if (config.env === 'test') {
-          cb(null, true);
-        } else {
-          // Retrieve the allowed origins from the settings
-          const app = AppInstance.getApp();
-          const settingService = app.get<SettingService>(SettingService);
-
-          await settingService
-            .getAllowedOrigins()
-            .then((allowedOrigins) => {
-              if (origin && allowedOrigins.includes(origin)) {
-                cb(null, true);
-              } else {
-                // eslint-disable-next-line no-console
-                console.log(
-                  `A socket was rejected via the config.sockets.onlyAllowOrigins array.\n` +
-                    `It attempted to connect with origin: ${origin}`,
-                );
-                cb(new Error('Origin not allowed'), false);
-              }
-            })
-            .catch(cb);
-        }
-      },
+      origin: true,
     },
   };
 

package/src/websocket/websocket.module.ts

@@ -9,12 +9,17 @@ import { Global, Module } from '@nestjs/common';
 import { UserModule } from '@/user/user.module';
 
 import { SocketEventDispatcherService } from './services/socket-event-dispatcher.service';
+import { WEBSOCKET_GATEWAY } from './tokens';
 import { WebsocketGateway } from './websocket.gateway';
 
 @Global()
 @Module({
   imports: [UserModule],
-  providers: [WebsocketGateway, SocketEventDispatcherService],
-  exports: [WebsocketGateway],
+  providers: [
+    WebsocketGateway,
+    { provide: WEBSOCKET_GATEWAY, useExisting: WebsocketGateway },
+    SocketEventDispatcherService,
+  ],
+  exports: [WebsocketGateway, WEBSOCKET_GATEWAY],
 })
 export class WebsocketModule {}

package/src/workflow/services/workflow.service.ts

@@ -8,6 +8,7 @@ import type { WorkflowEventMap } from '@hexabot-ai/agentic';
 import { WorkflowFull, Workflow } from '@hexabot-ai/types';
 import {
   BadRequestException,
+  Inject,
   Injectable,
   InternalServerErrorException,
 } from '@nestjs/common';
@@ -19,15 +20,16 @@ import { UpdateOneOptions } from '@/utils/generics/base-orm.repository';
 import { BaseOrmService } from '@/utils/generics/base-orm.service';
 import { InferCreateDto } from '@/utils/types/dto.types';
 import {
-  Room,
   SocketGet,
   SocketPost,
-  SocketReq,
-  SocketRequest,
-  SocketRes,
-  SocketResponse,
-  WebsocketGateway,
-} from '@/websocket';
+} from '@/websocket/decorators/socket-method.decorator';
+import { SocketReq } from '@/websocket/decorators/socket-req.decorator';
+import { SocketRes } from '@/websocket/decorators/socket-res.decorator';
+import { WEBSOCKET_GATEWAY } from '@/websocket/tokens';
+import { Room } from '@/websocket/types';
+import { SocketRequest } from '@/websocket/utils/socket-request';
+import { SocketResponse } from '@/websocket/utils/socket-response';
+import type { WebsocketGateway } from '@/websocket/websocket.gateway';
 
 import { WorkflowUpdateDto } from '../dto/workflow.dto';
 import { WorkflowOrmEntity } from '../entities/workflow.entity';
@@ -36,6 +38,11 @@ import { WorkflowType } from '../types';
 
 import { WorkflowRunService } from './workflow-run.service';
 
+type WorkflowSocketGateway = Pick<
+  WebsocketGateway,
+  'joinSockets' | 'broadcastWorkflowEvent'
+>;
+
 @Injectable()
 export class WorkflowService extends BaseOrmService<WorkflowOrmEntity> {
   private readonly MANUAL_INPUT_VALIDATION_ERROR = {
@@ -51,7 +58,8 @@ export class WorkflowService extends BaseOrmService<WorkflowOrmEntity> {
    */
   constructor(
     readonly repository: WorkflowRepository,
-    private readonly gateway: WebsocketGateway,
+    @Inject(WEBSOCKET_GATEWAY)
+    private readonly gateway: WorkflowSocketGateway,
     private readonly workflowRunService: WorkflowRunService,
   ) {
     super(repository);