85 lines
2.2 KiB
TypeScript
85 lines
2.2 KiB
TypeScript
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;
|
|
}
|
|
}
|