smartsocket/ts/smartsocket.classes.socketserver.ts

import * as plugins from './smartsocket.plugins.js';
import * as pluginsTyped from './smartsocket.pluginstyped.js';

// used in case no other server is supplied
import { Smartsocket } from './smartsocket.classes.smartsocket.js';
import { logger } from './smartsocket.logging.js';

/**
 * class SocketServer
 * handles the WebSocket server, either standalone or integrated with smartserve
 */
export class SocketServer {
  private smartsocket: Smartsocket;
  private httpServer: pluginsTyped.http.Server | pluginsTyped.https.Server;
  private wsServer: pluginsTyped.ws.WebSocketServer;

  /**
   * whether we're using an external server (smartserve)
   */
  private externalServerMode = false;
  private externalServer: any = null;
  private externalWebSocketHooks: pluginsTyped.ISmartserveWebSocketHooks = null;

  /**
   * whether httpServer is standalone
   */
  private standaloneServer = false;

  constructor(smartSocketInstance: Smartsocket) {
    this.smartsocket = smartSocketInstance;
  }

  /**
   * Set an external server (smartserve) for WebSocket handling
   */
  public async setExternalServer(
    serverType: 'smartserve',
    serverArg: any,
    websocketHooks?: pluginsTyped.ISmartserveWebSocketHooks
  ) {
    if (serverType !== 'smartserve') {
      throw new Error(`Unsupported server type: ${serverType}. Only 'smartserve' is supported.`);
    }
    this.externalServerMode = true;
    this.externalServer = serverArg;
    this.externalWebSocketHooks = websocketHooks || null;
  }

  /**
   * starts listening to incoming websocket connections
   */
  public async start() {
    const done = plugins.smartpromise.defer();

    if (this.externalServerMode) {
      // Using external smartserve server
      // The smartserve server should be configured with websocket hooks
      // that call our handleNewConnection method
      logger.log('info', 'Using external smartserve server for WebSocket handling');

      // If smartserve provides a way to get the underlying http server for upgrade,
      // we could attach ws to it. For now, we expect smartserve to handle WS
      // and call us back via the hooks.
      done.resolve();
    } else {
      // Standalone mode - create our own HTTP server and WebSocket server
      const httpModule = await this.smartsocket.smartenv.getSafeNodeModule('http');
      const wsModule = await this.smartsocket.smartenv.getSafeNodeModule('ws');

      if (!this.smartsocket.options.port) {
        logger.log('error', 'there should be a port specified for smartsocket!');
        throw new Error('there should be a port specified for smartsocket');
      }

      this.httpServer = httpModule.createServer();
      this.standaloneServer = true;

      // Create WebSocket server attached to HTTP server
      this.wsServer = new wsModule.WebSocketServer({ server: this.httpServer });

      this.wsServer.on('connection', (ws: pluginsTyped.ws.WebSocket) => {
        this.smartsocket.handleNewConnection(ws);
      });

      this.httpServer.listen(this.smartsocket.options.port, () => {
        logger.log(
          'success',
          `Server started in standalone mode on port ${this.smartsocket.options.port}`
        );
        done.resolve();
      });
    }

    await done.promise;
  }

  /**
   * closes the server
   */
  public async stop() {
    const done = plugins.smartpromise.defer<void>();
    let resolved = false;

    if (this.wsServer) {
      // Close all WebSocket connections
      this.wsServer.clients.forEach((client) => {
        client.terminate();
      });
      this.wsServer.close();
      this.wsServer = null;
    }

    if (this.httpServer && this.standaloneServer) {
      const resolveOnce = () => {
        if (!resolved) {
          resolved = true;
          this.httpServer = null;
          this.standaloneServer = false;
          done.resolve();
        }
      };

      this.httpServer.close(() => {
        resolveOnce();
      });

      // Add a timeout in case close callback doesn't fire
      const timeoutId = setTimeout(() => {
        resolveOnce();
      }, 2000);

      // Ensure timeout doesn't keep process alive
      if (timeoutId.unref) {
        timeoutId.unref();
      }
    } else {
      done.resolve();
    }

    await done.promise;
  }

  /**
   * Returns WebSocket hooks for integration with smartserve
   * Call this to get hooks that you can pass to smartserve's websocket config
   */
  public getSmartserveWebSocketHooks(): pluginsTyped.ISmartserveWebSocketHooks {
    return {
      onOpen: async (peer: pluginsTyped.ISmartserveWebSocketPeer) => {
        // Create a wrapper that adapts ISmartserveWebSocketPeer to WebSocket-like interface
        const wsLikeSocket = this.createWsLikeFromPeer(peer);
        await this.smartsocket.handleNewConnection(wsLikeSocket as any);

        // Call external hooks if provided
        if (this.externalWebSocketHooks?.onOpen) {
          await this.externalWebSocketHooks.onOpen(peer);
        }
      },
      onMessage: async (peer: pluginsTyped.ISmartserveWebSocketPeer, message: pluginsTyped.ISmartserveWebSocketMessage) => {
        // Messages are handled by SocketConnection via the adapter
        // But we still call external hooks if provided
        if (this.externalWebSocketHooks?.onMessage) {
          await this.externalWebSocketHooks.onMessage(peer, message);
        }
      },
      onClose: async (peer: pluginsTyped.ISmartserveWebSocketPeer, code: number, reason: string) => {
        if (this.externalWebSocketHooks?.onClose) {
          await this.externalWebSocketHooks.onClose(peer, code, reason);
        }
      },
      onError: async (peer: pluginsTyped.ISmartserveWebSocketPeer, error: Error) => {
        if (this.externalWebSocketHooks?.onError) {
          await this.externalWebSocketHooks.onError(peer, error);
        }
      },
    };
  }

  /**
   * Creates a WebSocket-like object from a smartserve peer
   * This allows our SocketConnection to work with both native WebSocket and smartserve peers
   */
  private createWsLikeFromPeer(peer: pluginsTyped.ISmartserveWebSocketPeer): any {
    const messageListeners: Array<(event: any) => void> = [];
    const closeListeners: Array<() => void> = [];
    const errorListeners: Array<() => void> = [];

    // Store the adapter on the peer for message routing
    peer.data.set('smartsocket_adapter', {
      dispatchMessage: (data: string) => {
        messageListeners.forEach((listener) => {
          listener({ data });
        });
      },
      dispatchClose: () => {
        closeListeners.forEach((listener) => listener());
      },
      dispatchError: () => {
        errorListeners.forEach((listener) => listener());
      },
    });

    return {
      readyState: peer.readyState,
      send: (data: string) => peer.send(data),
      close: (code?: number, reason?: string) => peer.close(code, reason),
      addEventListener: (event: string, listener: any) => {
        if (event === 'message') {
          messageListeners.push(listener);
        } else if (event === 'close') {
          closeListeners.push(listener);
        } else if (event === 'error') {
          errorListeners.push(listener);
        }
      },
      removeEventListener: (event: string, listener: any) => {
        if (event === 'message') {
          const idx = messageListeners.indexOf(listener);
          if (idx >= 0) messageListeners.splice(idx, 1);
        } else if (event === 'close') {
          const idx = closeListeners.indexOf(listener);
          if (idx >= 0) closeListeners.splice(idx, 1);
        } else if (event === 'error') {
          const idx = errorListeners.indexOf(listener);
          if (idx >= 0) errorListeners.splice(idx, 1);
        }
      },
    };
  }
}