smartproxy/ts/core/models/wrapped-socket.ts

import { EventEmitter } from 'events';
import * as plugins from '../../plugins.js';

/**
 * WrappedSocket wraps a regular net.Socket to provide transparent access
 * to the real client IP and port when behind a proxy using PROXY protocol.
 * 
 * This is the FOUNDATION for all PROXY protocol support and must be implemented
 * before any protocol parsing can occur.
 */
export class WrappedSocket extends EventEmitter {
  private realClientIP?: string;
  private realClientPort?: number;
  
  constructor(
    public readonly socket: plugins.net.Socket,
    realClientIP?: string,
    realClientPort?: number
  ) {
    super();
    this.realClientIP = realClientIP;
    this.realClientPort = realClientPort;
    
    // Forward all socket events
    this.forwardSocketEvents();
  }
  
  /**
   * Returns the real client IP if available, otherwise the socket's remote address
   */
  get remoteAddress(): string | undefined {
    return this.realClientIP || this.socket.remoteAddress;
  }
  
  /**
   * Returns the real client port if available, otherwise the socket's remote port
   */
  get remotePort(): number | undefined {
    return this.realClientPort || this.socket.remotePort;
  }
  
  /**
   * Returns the remote family (IPv4 or IPv6)
   */
  get remoteFamily(): string | undefined {
    // If we have a real client IP, determine the family
    if (this.realClientIP) {
      if (this.realClientIP.includes(':')) {
        return 'IPv6';
      } else {
        return 'IPv4';
      }
    }
    return this.socket.remoteFamily;
  }
  
  /**
   * Returns the local address of the socket
   */
  get localAddress(): string | undefined {
    return this.socket.localAddress;
  }
  
  /**
   * Returns the local port of the socket
   */
  get localPort(): number | undefined {
    return this.socket.localPort;
  }
  
  /**
   * Indicates if this connection came through a trusted proxy
   */
  get isFromTrustedProxy(): boolean {
    return !!this.realClientIP;
  }
  
  /**
   * Updates the real client information (called after parsing PROXY protocol)
   */
  setProxyInfo(ip: string, port: number): void {
    this.realClientIP = ip;
    this.realClientPort = port;
  }
  
  // Pass-through all socket methods
  write(data: any, encoding?: any, callback?: any): boolean {
    return this.socket.write(data, encoding, callback);
  }
  
  end(data?: any, encoding?: any, callback?: any): this {
    this.socket.end(data, encoding, callback);
    return this;
  }
  
  destroy(error?: Error): this {
    this.socket.destroy(error);
    return this;
  }
  
  pause(): this {
    this.socket.pause();
    return this;
  }
  
  resume(): this {
    this.socket.resume();
    return this;
  }
  
  setTimeout(timeout: number, callback?: () => void): this {
    this.socket.setTimeout(timeout, callback);
    return this;
  }
  
  setNoDelay(noDelay?: boolean): this {
    this.socket.setNoDelay(noDelay);
    return this;
  }
  
  setKeepAlive(enable?: boolean, initialDelay?: number): this {
    this.socket.setKeepAlive(enable, initialDelay);
    return this;
  }
  
  ref(): this {
    this.socket.ref();
    return this;
  }
  
  unref(): this {
    this.socket.unref();
    return this;
  }
  
  /**
   * Pipe to another stream
   */
  pipe<T extends NodeJS.WritableStream>(destination: T, options?: {
    end?: boolean;
  }): T {
    return this.socket.pipe(destination, options);
  }
  
  /**
   * Cork the stream
   */
  cork(): void {
    if ('cork' in this.socket && typeof this.socket.cork === 'function') {
      this.socket.cork();
    }
  }
  
  /**
   * Uncork the stream
   */
  uncork(): void {
    if ('uncork' in this.socket && typeof this.socket.uncork === 'function') {
      this.socket.uncork();
    }
  }
  
  /**
   * Get the number of bytes read
   */
  get bytesRead(): number {
    return this.socket.bytesRead;
  }
  
  /**
   * Get the number of bytes written
   */
  get bytesWritten(): number {
    return this.socket.bytesWritten;
  }
  
  /**
   * Check if the socket is connecting
   */
  get connecting(): boolean {
    return this.socket.connecting;
  }
  
  /**
   * Check if the socket is destroyed
   */
  get destroyed(): boolean {
    return this.socket.destroyed;
  }
  
  /**
   * Check if the socket is readable
   */
  get readable(): boolean {
    return this.socket.readable;
  }
  
  /**
   * Check if the socket is writable
   */
  get writable(): boolean {
    return this.socket.writable;
  }
  
  /**
   * Get pending status
   */
  get pending(): boolean {
    return this.socket.pending;
  }
  
  /**
   * Get ready state
   */
  get readyState(): string {
    return this.socket.readyState;
  }
  
  /**
   * Address info
   */
  address(): plugins.net.AddressInfo | {} | null {
    const addr = this.socket.address();
    if (addr === null) return null;
    if (typeof addr === 'string') return addr as any;
    return addr;
  }
  
  /**
   * Set socket encoding
   */
  setEncoding(encoding?: BufferEncoding): this {
    this.socket.setEncoding(encoding);
    return this;
  }
  
  /**
   * Connect method (for client sockets)
   */
  connect(options: plugins.net.SocketConnectOpts, connectionListener?: () => void): this;
  connect(port: number, host?: string, connectionListener?: () => void): this;
  connect(path: string, connectionListener?: () => void): this;
  connect(...args: any[]): this {
    (this.socket as any).connect(...args);
    return this;
  }
  
  /**
   * Forward all events from the underlying socket
   */
  private forwardSocketEvents(): void {
    const events = ['data', 'end', 'close', 'error', 'drain', 'timeout', 'connect', 'ready', 'lookup'];
    events.forEach(event => {
      this.socket.on(event, (...args) => {
        this.emit(event, ...args);
      });
    });
  }
}