Phase2

apps/api/src/nodes/nodes.gateway.ts

@@ -0,0 +1,84 @@
+import {
+  Injectable,
+  Logger,
+  OnModuleInit,
+  UnauthorizedException,
+} from '@nestjs/common';
+import { HttpAdapterHost } from '@nestjs/core';
+import type { FastifyRequest } from 'fastify';
+import type { WebSocket } from 'ws';
+
+import { NodesService } from './nodes.service';
+
+@Injectable()
+export class NodesGateway implements OnModuleInit {
+  private readonly logger = new Logger(NodesGateway.name);
+
+  constructor(
+    private readonly adapterHost: HttpAdapterHost,
+    private readonly nodesService: NodesService,
+  ) {}
+
+  onModuleInit(): void {
+    const fastify = this.adapterHost.httpAdapter.getInstance();
+
+    fastify.get(
+      '/api/v1/nodes/ws',
+      { websocket: true },
+      (socket: WebSocket, request: FastifyRequest) => {
+        void this.handleConnection(socket, request);
+      },
+    );
+  }
+
+  private async handleConnection(
+    socket: WebSocket,
+    request: FastifyRequest,
+  ): Promise<void> {
+    const query = request.query as Record<string, string | string[] | undefined>;
+    const nodeId = this.readQueryParam(query['nodeId']);
+    const token = this.readQueryParam(query['token']);
+
+    if (!nodeId || !token) {
+      socket.close(4400, 'nodeId and token are required');
+      return;
+    }
+
+    try {
+      await this.nodesService.validateEnrollment(nodeId, token);
+    } catch (error) {
+      const reason =
+        error instanceof UnauthorizedException
+          ? 'Invalid enrollment credentials'
+          : 'Enrollment validation failed';
+      socket.close(4401, reason);
+      return;
+    }
+
+    this.nodesService.registerConnection(nodeId, socket);
+
+    socket.on('close', () => {
+      this.nodesService.unregisterConnection(nodeId, socket);
+    });
+
+    socket.on('message', (data) => {
+      void this.nodesService.handleAgentMessage(nodeId, data.toString());
+    });
+
+    socket.on('error', (error) => {
+      this.logger.error(`WebSocket error for node ${nodeId}: ${error.message}`);
+    });
+
+    socket.send(JSON.stringify(this.nodesService.createHelloAck(nodeId)));
+  }
+
+  private readQueryParam(
+    value: string | string[] | undefined,
+  ): string | undefined {
+    if (Array.isArray(value)) {
+      return value[0];
+    }
+
+    return value;
+  }
+}