stack/apps/api/src/terminal/terminal.gateway.ts

/**
 * TerminalGateway
 *
 * WebSocket gateway for real-time PTY terminal sessions.
 * Uses the `/terminal` namespace to keep terminal traffic separate
 * from the main WebSocket gateway.
 *
 * Protocol:
 * 1. Client connects with auth token in handshake
 * 2. Client emits `terminal:create` to spawn a new PTY session
 * 3. Server emits `terminal:created` with { sessionId }
 * 4. Client emits `terminal:input` with { sessionId, data } to send keystrokes
 * 5. Server emits `terminal:output` with { sessionId, data } for stdout/stderr
 * 6. Client emits `terminal:resize` with { sessionId, cols, rows } on window resize
 * 7. Client emits `terminal:close` with { sessionId } to terminate the PTY
 * 8. Server emits `terminal:exit` with { sessionId, exitCode, signal } on PTY exit
 *
 * Authentication:
 * - Same pattern as websocket.gateway.ts and speech.gateway.ts
 * - Token extracted from handshake.auth.token / query.token / Authorization header
 *
 * Workspace isolation:
 * - Clients join room `terminal:{workspaceId}` on connect
 * - Sessions are scoped to workspace; cross-workspace access is denied
 */

import {
  WebSocketGateway as WSGateway,
  WebSocketServer,
  SubscribeMessage,
  OnGatewayConnection,
  OnGatewayDisconnect,
} from "@nestjs/websockets";
import { Logger } from "@nestjs/common";
import { Server, Socket } from "socket.io";
import { AuthService } from "../auth/auth.service";
import { PrismaService } from "../prisma/prisma.service";
import { TerminalService } from "./terminal.service";
import {
  CreateTerminalDto,
  TerminalInputDto,
  TerminalResizeDto,
  CloseTerminalDto,
} from "./terminal.dto";
import { validate } from "class-validator";
import { plainToInstance } from "class-transformer";

// ==========================================
// Types
// ==========================================

interface AuthenticatedSocket extends Socket {
  data: {
    userId?: string;
    workspaceId?: string;
  };
}

// ==========================================
// Gateway
// ==========================================

@WSGateway({
  namespace: "/terminal",
  cors: {
    origin: (process.env.TRUSTED_ORIGINS ?? process.env.WEB_URL ?? "http://localhost:3000")
      .split(",")
      .map((s) => s.trim()),
    credentials: true,
  },
})
export class TerminalGateway implements OnGatewayConnection, OnGatewayDisconnect {
  @WebSocketServer()
  server!: Server;

  private readonly logger = new Logger(TerminalGateway.name);
  private readonly CONNECTION_TIMEOUT_MS = 5000;

  constructor(
    private readonly authService: AuthService,
    private readonly prisma: PrismaService,
    private readonly terminalService: TerminalService
  ) {}

  // ==========================================
  // Connection lifecycle
  // ==========================================

  /**
   * Authenticate client on connection using handshake token.
   * Validates workspace membership and joins the workspace-scoped room.
   */
  async handleConnection(client: Socket): Promise<void> {
    const authenticatedClient = client as AuthenticatedSocket;

    const timeoutId = setTimeout(() => {
      if (!authenticatedClient.data.userId) {
        this.logger.warn(
          `Terminal client ${authenticatedClient.id} timed out during authentication`
        );
        authenticatedClient.emit("terminal:error", {
          message: "Authentication timed out.",
        });
        authenticatedClient.disconnect();
      }
    }, this.CONNECTION_TIMEOUT_MS);

    try {
      const token = this.extractTokenFromHandshake(authenticatedClient);

      if (!token) {
        this.logger.warn(`Terminal client ${authenticatedClient.id} connected without token`);
        authenticatedClient.emit("terminal:error", {
          message: "Authentication failed: no token provided.",
        });
        authenticatedClient.disconnect();
        clearTimeout(timeoutId);
        return;
      }

      const sessionData = await this.authService.verifySession(token);

      if (!sessionData) {
        this.logger.warn(`Terminal client ${authenticatedClient.id} has invalid token`);
        authenticatedClient.emit("terminal:error", {
          message: "Authentication failed: invalid or expired token.",
        });
        authenticatedClient.disconnect();
        clearTimeout(timeoutId);
        return;
      }

      const user = sessionData.user as { id: string };
      const userId = user.id;

      const workspaceMembership = await this.prisma.workspaceMember.findFirst({
        where: { userId },
        select: { workspaceId: true, userId: true, role: true },
      });

      if (!workspaceMembership) {
        this.logger.warn(`Terminal user ${userId} has no workspace access`);
        authenticatedClient.emit("terminal:error", {
          message: "Authentication failed: no workspace access.",
        });
        authenticatedClient.disconnect();
        clearTimeout(timeoutId);
        return;
      }

      authenticatedClient.data.userId = userId;
      authenticatedClient.data.workspaceId = workspaceMembership.workspaceId;

      // Join workspace-scoped terminal room
      const room = this.getWorkspaceRoom(workspaceMembership.workspaceId);
      await authenticatedClient.join(room);

      clearTimeout(timeoutId);
      this.logger.log(
        `Terminal client ${authenticatedClient.id} connected (user: ${userId}, workspace: ${workspaceMembership.workspaceId})`
      );
    } catch (error) {
      clearTimeout(timeoutId);
      this.logger.error(
        `Authentication failed for terminal client ${authenticatedClient.id}:`,
        error instanceof Error ? error.message : "Unknown error"
      );
      authenticatedClient.emit("terminal:error", {
        message: "Authentication failed: an unexpected error occurred.",
      });
      authenticatedClient.disconnect();
    }
  }

  /**
   * Clean up all PTY sessions for this client's workspace on disconnect.
   */
  handleDisconnect(client: Socket): void {
    const authenticatedClient = client as AuthenticatedSocket;
    const { workspaceId, userId } = authenticatedClient.data;

    if (workspaceId) {
      this.terminalService.closeWorkspaceSessions(workspaceId);

      const room = this.getWorkspaceRoom(workspaceId);
      void authenticatedClient.leave(room);
      this.logger.log(
        `Terminal client ${authenticatedClient.id} disconnected (user: ${userId ?? "unknown"}, workspace: ${workspaceId})`
      );
    } else {
      this.logger.debug(`Terminal client ${authenticatedClient.id} disconnected (unauthenticated)`);
    }
  }

  // ==========================================
  // Terminal events
  // ==========================================

  /**
   * Spawn a new PTY session for the connected client.
   *
   * Emits `terminal:created` with { sessionId, name, cols, rows } on success.
   * Emits `terminal:error` on failure.
   */
  @SubscribeMessage("terminal:create")
  async handleCreate(client: Socket, payload: unknown): Promise<void> {
    const authenticatedClient = client as AuthenticatedSocket;
    const { userId, workspaceId } = authenticatedClient.data;

    if (!userId || !workspaceId) {
      authenticatedClient.emit("terminal:error", {
        message: "Not authenticated. Connect with a valid token.",
      });
      return;
    }

    // Validate DTO
    const dto = plainToInstance(CreateTerminalDto, payload ?? {});
    const errors = await validate(dto);
    if (errors.length > 0) {
      const messages = errors.map((e) => Object.values(e.constraints ?? {}).join(", ")).join("; ");
      authenticatedClient.emit("terminal:error", {
        message: `Invalid payload: ${messages}`,
      });
      return;
    }

    try {
      const result = this.terminalService.createSession(authenticatedClient, {
        workspaceId,
        socketId: authenticatedClient.id,
        ...(dto.name !== undefined ? { name: dto.name } : {}),
        ...(dto.cols !== undefined ? { cols: dto.cols } : {}),
        ...(dto.rows !== undefined ? { rows: dto.rows } : {}),
        ...(dto.cwd !== undefined ? { cwd: dto.cwd } : {}),
      });

      authenticatedClient.emit("terminal:created", {
        sessionId: result.sessionId,
        name: result.name,
        cols: result.cols,
        rows: result.rows,
      });

      this.logger.log(
        `Terminal session ${result.sessionId} created for client ${authenticatedClient.id} (workspace: ${workspaceId})`
      );
    } catch (error) {
      const message = error instanceof Error ? error.message : String(error);
      this.logger.error(
        `Failed to create terminal session for client ${authenticatedClient.id}: ${message}`
      );
      authenticatedClient.emit("terminal:error", { message });
    }
  }

  /**
   * Write input data to an existing PTY session.
   *
   * Emits `terminal:error` if the session is not found or unauthorized.
   */
  @SubscribeMessage("terminal:input")
  async handleInput(client: Socket, payload: unknown): Promise<void> {
    const authenticatedClient = client as AuthenticatedSocket;
    const { userId, workspaceId } = authenticatedClient.data;

    if (!userId || !workspaceId) {
      authenticatedClient.emit("terminal:error", {
        message: "Not authenticated. Connect with a valid token.",
      });
      return;
    }

    const dto = plainToInstance(TerminalInputDto, payload ?? {});
    const errors = await validate(dto);
    if (errors.length > 0) {
      const messages = errors.map((e) => Object.values(e.constraints ?? {}).join(", ")).join("; ");
      authenticatedClient.emit("terminal:error", {
        message: `Invalid payload: ${messages}`,
      });
      return;
    }

    if (!this.terminalService.sessionBelongsToWorkspace(dto.sessionId, workspaceId)) {
      authenticatedClient.emit("terminal:error", {
        message: `Terminal session ${dto.sessionId} not found or unauthorized.`,
      });
      return;
    }

    try {
      this.terminalService.writeToSession(dto.sessionId, dto.data);
    } catch (error) {
      const message = error instanceof Error ? error.message : String(error);
      this.logger.warn(`Failed to write to terminal session ${dto.sessionId}: ${message}`);
      authenticatedClient.emit("terminal:error", { message });
    }
  }

  /**
   * Resize an existing PTY session.
   *
   * Emits `terminal:error` if the session is not found or unauthorized.
   */
  @SubscribeMessage("terminal:resize")
  async handleResize(client: Socket, payload: unknown): Promise<void> {
    const authenticatedClient = client as AuthenticatedSocket;
    const { userId, workspaceId } = authenticatedClient.data;

    if (!userId || !workspaceId) {
      authenticatedClient.emit("terminal:error", {
        message: "Not authenticated. Connect with a valid token.",
      });
      return;
    }

    const dto = plainToInstance(TerminalResizeDto, payload ?? {});
    const errors = await validate(dto);
    if (errors.length > 0) {
      const messages = errors.map((e) => Object.values(e.constraints ?? {}).join(", ")).join("; ");
      authenticatedClient.emit("terminal:error", {
        message: `Invalid payload: ${messages}`,
      });
      return;
    }

    if (!this.terminalService.sessionBelongsToWorkspace(dto.sessionId, workspaceId)) {
      authenticatedClient.emit("terminal:error", {
        message: `Terminal session ${dto.sessionId} not found or unauthorized.`,
      });
      return;
    }

    try {
      this.terminalService.resizeSession(dto.sessionId, dto.cols, dto.rows);
    } catch (error) {
      const message = error instanceof Error ? error.message : String(error);
      this.logger.warn(`Failed to resize terminal session ${dto.sessionId}: ${message}`);
      authenticatedClient.emit("terminal:error", { message });
    }
  }

  /**
   * Kill and close an existing PTY session.
   *
   * Emits `terminal:error` if the session is not found or unauthorized.
   */
  @SubscribeMessage("terminal:close")
  async handleClose(client: Socket, payload: unknown): Promise<void> {
    const authenticatedClient = client as AuthenticatedSocket;
    const { userId, workspaceId } = authenticatedClient.data;

    if (!userId || !workspaceId) {
      authenticatedClient.emit("terminal:error", {
        message: "Not authenticated. Connect with a valid token.",
      });
      return;
    }

    const dto = plainToInstance(CloseTerminalDto, payload ?? {});
    const errors = await validate(dto);
    if (errors.length > 0) {
      const messages = errors.map((e) => Object.values(e.constraints ?? {}).join(", ")).join("; ");
      authenticatedClient.emit("terminal:error", {
        message: `Invalid payload: ${messages}`,
      });
      return;
    }

    if (!this.terminalService.sessionBelongsToWorkspace(dto.sessionId, workspaceId)) {
      authenticatedClient.emit("terminal:error", {
        message: `Terminal session ${dto.sessionId} not found or unauthorized.`,
      });
      return;
    }

    const closed = this.terminalService.closeSession(dto.sessionId);
    if (!closed) {
      authenticatedClient.emit("terminal:error", {
        message: `Terminal session ${dto.sessionId} not found.`,
      });
      return;
    }

    this.logger.log(`Terminal session ${dto.sessionId} closed by client ${authenticatedClient.id}`);
  }

  // ==========================================
  // Private helpers
  // ==========================================

  /**
   * Extract authentication token from Socket.IO handshake.
   * Checks auth.token, query.token, and Authorization header (in that order).
   */
  private extractTokenFromHandshake(client: Socket): string | undefined {
    const authToken = client.handshake.auth.token as unknown;
    if (typeof authToken === "string" && authToken.length > 0) {
      return authToken;
    }

    const queryToken = client.handshake.query.token as unknown;
    if (typeof queryToken === "string" && queryToken.length > 0) {
      return queryToken;
    }

    const authHeader = client.handshake.headers.authorization as unknown;
    if (typeof authHeader === "string") {
      const parts = authHeader.split(" ");
      const [type, token] = parts;
      if (type === "Bearer" && token) {
        return token;
      }
    }

    return undefined;
  }

  /**
   * Get the workspace-scoped room name for the terminal namespace.
   */
  private getWorkspaceRoom(workspaceId: string): string {
    return `terminal:${workspaceId}`;
  }
}