source-gamepanel/apps/api/src/plugins/socket.ts

import fp from 'fastify-plugin';
import type { FastifyInstance } from 'fastify';
import { and, eq } from 'drizzle-orm';
import { Server as SocketIOServer } from 'socket.io';
import { nodes, organizationMembers, servers } from '@source/database';
import { ROLES } from '@source/shared';
import type { Role } from '@source/shared';
import { getJwt } from '../lib/jwt.js';
import type { AccessTokenPayload } from '../lib/jwt.js';
import {
  daemonOpenConsoleStream,
  daemonSendCommand,
  type DaemonConsoleStreamHandle,
  type DaemonNodeConnection,
} from '../lib/daemon.js';

declare module 'fastify' {
  interface FastifyInstance {
    io: SocketIOServer;
  }
}

type ConsolePermission = 'console.read' | 'console.write';
type ConsoleCommandAck = {
  requestId: string | null;
  ok: boolean;
  error?: string;
};

interface SharedConsoleStream {
  handle: DaemonConsoleStreamHandle;
  subscribers: number;
}

function roomForServer(serverId: string): string {
  return `server:console:${serverId}`;
}

export default fp(async (app: FastifyInstance) => {
  const io = new SocketIOServer(app.server, {
    path: '/socket.io',
    cors: {
      origin: process.env.CORS_ORIGIN || 'http://localhost:5173',
      credentials: true,
    },
  });

  app.decorate('io', io);

  const serverStreams = new Map<string, SharedConsoleStream>();
  const socketSubscriptions = new Map<string, string>();

  const clearServerSubscriptions = (serverId: string) => {
    for (const [socketId, subscribedServerId] of socketSubscriptions.entries()) {
      if (subscribedServerId === serverId) {
        socketSubscriptions.delete(socketId);
      }
    }
  };

  io.use((socket, next) => {
    const token =
      typeof socket.handshake.auth?.token === 'string' ? socket.handshake.auth.token : null;

    if (!token) {
      next(new Error('Unauthorized'));
      return;
    }

    const verifier = getJwt(app)?.verify;
    if (typeof verifier !== 'function') {
      next(new Error('Authentication is not configured'));
      return;
    }

    try {
      const payload = verifier(token) as AccessTokenPayload;
      (socket.data as { user?: AccessTokenPayload }).user = payload;
      next();
    } catch {
      next(new Error('Unauthorized'));
    }
  });

  io.on('connection', (socket) => {
    const cleanupSocketStream = () => {
      const subscribedServerId = socketSubscriptions.get(socket.id);
      if (!subscribedServerId) return;

      socketSubscriptions.delete(socket.id);
      socket.leave(roomForServer(subscribedServerId));

      const shared = serverStreams.get(subscribedServerId);
      if (!shared) return;

      shared.subscribers = Math.max(0, shared.subscribers - 1);
      if (shared.subscribers === 0) {
        shared.handle.close();
        serverStreams.delete(subscribedServerId);
      }
    };

    socket.on('server:console:join', async (payload: unknown) => {
      const serverId =
        typeof (payload as { serverId?: unknown })?.serverId === 'string'
          ? (payload as { serverId: string }).serverId
          : '';
      if (!serverId) {
        socket.emit('server:console:output', { line: '[error] Invalid server id' });
        return;
      }

      const user = (socket.data as { user?: AccessTokenPayload }).user;
      if (!user) {
        socket.emit('server:console:output', { line: '[error] Unauthorized' });
        return;
      }

      const server = await getServerContext(app, serverId);
      if (!server) {
        socket.emit('server:console:output', { line: '[error] Server not found' });
        return;
      }

      const allowed = await hasConsolePermission(app, user, server.organizationId, 'console.read');
      if (!allowed) {
        socket.emit('server:console:output', { line: '[error] Missing permission: console.read' });
        return;
      }

      const previousSubscription = socketSubscriptions.get(socket.id);
      if (previousSubscription === serverId) {
        return;
      }
      cleanupSocketStream();
      socket.join(roomForServer(serverId));

      let shared = serverStreams.get(serverId);
      if (!shared) {
        try {
          const streamHandle = await daemonOpenConsoleStream(server.node, server.serverUuid);
          const room = roomForServer(serverId);

          streamHandle.stream.on('data', (output) => {
            io.to(room).emit('server:console:output', { line: output.line });
          });

          streamHandle.stream.on('end', () => {
            const current = serverStreams.get(serverId);
            if (current?.handle !== streamHandle) return;
            serverStreams.delete(serverId);
            clearServerSubscriptions(serverId);
            io.to(room).emit('server:console:output', { line: '[console] Stream ended' });
            io.in(room).socketsLeave(room);
          });

          streamHandle.stream.on('error', (error) => {
            const current = serverStreams.get(serverId);
            if (current?.handle !== streamHandle) return;
            serverStreams.delete(serverId);
            clearServerSubscriptions(serverId);
            app.log.warn(
              { error, serverId, serverUuid: server.serverUuid },
              'Console stream failed',
            );
            io.to(room).emit('server:console:output', { line: '[error] Console stream failed' });
            io.in(room).socketsLeave(room);
          });

          shared = {
            handle: streamHandle,
            subscribers: 0,
          };
          serverStreams.set(serverId, shared);
        } catch (error) {
          app.log.warn(
            { error, serverId, serverUuid: server.serverUuid, socketId: socket.id },
            'Failed to open console stream',
          );
          socket.leave(roomForServer(serverId));
          socket.emit('server:console:output', { line: '[error] Failed to open console stream' });
          return;
        }
      }

      shared.subscribers += 1;
      socketSubscriptions.set(socket.id, serverId);
    });

    socket.on('server:console:leave', () => {
      cleanupSocketStream();
    });

    socket.on('server:console:command', async (payload: unknown) => {
      const body = payload as {
        serverId?: unknown;
        orgId?: unknown;
        command?: unknown;
        requestId?: unknown;
      };

      const serverId = typeof body.serverId === 'string' ? body.serverId : '';
      const orgId = typeof body.orgId === 'string' ? body.orgId : '';
      const command = typeof body.command === 'string' ? body.command.trim() : '';
      const requestId =
        typeof body.requestId === 'string' && body.requestId.trim() ? body.requestId.trim() : null;

      if (!serverId || !orgId || !command) {
        socket.emit('server:console:output', { line: '[error] Invalid command payload' });
        const ack: ConsoleCommandAck = {
          requestId,
          ok: false,
          error: 'Invalid command payload',
        };
        socket.emit('server:console:command:ack', ack);
        return;
      }

      const user = (socket.data as { user?: AccessTokenPayload }).user;
      if (!user) {
        socket.emit('server:console:output', { line: '[error] Unauthorized' });
        const ack: ConsoleCommandAck = { requestId, ok: false, error: 'Unauthorized' };
        socket.emit('server:console:command:ack', ack);
        return;
      }

      const server = await getServerContext(app, serverId, orgId);
      if (!server) {
        socket.emit('server:console:output', { line: '[error] Server not found' });
        const ack: ConsoleCommandAck = { requestId, ok: false, error: 'Server not found' };
        socket.emit('server:console:command:ack', ack);
        return;
      }

      const allowed = await hasConsolePermission(app, user, orgId, 'console.write');
      if (!allowed) {
        socket.emit('server:console:output', { line: '[error] Missing permission: console.write' });
        const ack: ConsoleCommandAck = {
          requestId,
          ok: false,
          error: 'Missing permission: console.write',
        };
        socket.emit('server:console:command:ack', ack);
        return;
      }

      try {
        await daemonSendCommand(server.node, server.serverUuid, command);
        const ack: ConsoleCommandAck = { requestId, ok: true };
        socket.emit('server:console:command:ack', ack);
      } catch (error) {
        app.log.warn(
          { error, serverId, serverUuid: server.serverUuid, socketId: socket.id },
          'Failed to send console command',
        );
        // The daemon explains *why* (server not running, no RCON password, …) —
        // showing that beats a generic failure the user cannot act on.
        const reason = daemonErrorReason(error);
        socket.emit('server:console:output', { line: `[error] ${reason}` });
        const ack: ConsoleCommandAck = { requestId, ok: false, error: reason };
        socket.emit('server:console:command:ack', ack);
      }
    });

    socket.on('disconnect', () => {
      cleanupSocketStream();
    });
  });

  app.addHook('onClose', async () => {
    for (const stream of serverStreams.values()) {
      stream.handle.close();
    }
    serverStreams.clear();
    socketSubscriptions.clear();

    await new Promise<void>((resolve) => {
      io.close(() => resolve());
    });
  });
});

/** Strip the gRPC status prefix so the console shows the daemon's own wording. */
function daemonErrorReason(error: unknown): string {
  const raw = error instanceof Error ? error.message.trim() : '';
  if (!raw) return 'Failed to send command';

  const withoutStatus = raw.replace(/^\d+\s+[A-Z_]+:\s*/, '').trim();
  return withoutStatus || 'Failed to send command';
}

async function hasConsolePermission(
  app: FastifyInstance,
  user: AccessTokenPayload,
  orgId: string,
  permission: ConsolePermission,
): Promise<boolean> {
  if (user.isSuperAdmin) return true;

  const member = await app.db.query.organizationMembers.findFirst({
    where: and(
      eq(organizationMembers.organizationId, orgId),
      eq(organizationMembers.userId, user.sub),
    ),
    columns: {
      role: true,
      customPermissions: true,
    },
  });

  if (!member) return false;

  const custom = (member.customPermissions ?? {}) as Record<string, boolean>;
  if (permission in custom) {
    return Boolean(custom[permission]);
  }

  const rolePerms = ROLES[member.role as Role]?.permissions ?? [];
  return (rolePerms as readonly string[]).includes(permission);
}

async function getServerContext(
  app: FastifyInstance,
  serverId: string,
  orgId?: string,
): Promise<{
  organizationId: string;
  serverUuid: string;
  node: DaemonNodeConnection;
} | null> {
  const whereClause = orgId
    ? and(eq(servers.id, serverId), eq(servers.organizationId, orgId))
    : eq(servers.id, serverId);

  const [row] = await app.db
    .select({
      organizationId: servers.organizationId,
      serverUuid: servers.uuid,
      nodeFqdn: nodes.fqdn,
      nodeGrpcPort: nodes.grpcPort,
      nodeDaemonToken: nodes.daemonToken,
    })
    .from(servers)
    .innerJoin(nodes, eq(servers.nodeId, nodes.id))
    .where(whereClause);

  if (!row) return null;

  return {
    organizationId: row.organizationId,
    serverUuid: row.serverUuid,
    node: {
      fqdn: row.nodeFqdn,
      grpcPort: row.nodeGrpcPort,
      daemonToken: row.nodeDaemonToken,
    },
  };
}