smartproxy/ts/core/utils/proxy-protocol.ts

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

/**
 * Interface representing parsed PROXY protocol information
 */
export interface IProxyInfo {
  protocol: 'TCP4' | 'TCP6' | 'UNKNOWN';
  sourceIP: string;
  sourcePort: number;
  destinationIP: string;
  destinationPort: number;
}

/**
 * Interface for parse result including remaining data
 */
export interface IProxyParseResult {
  proxyInfo: IProxyInfo | null;
  remainingData: Buffer;
}

/**
 * Parser for PROXY protocol v1 (text format)
 * Spec: https://www.haproxy.org/download/1.8/doc/proxy-protocol.txt
 */
export class ProxyProtocolParser {
  static readonly PROXY_V1_SIGNATURE = 'PROXY ';
  static readonly MAX_HEADER_LENGTH = 107; // Max length for v1 header
  static readonly HEADER_TERMINATOR = '\r\n';
  
  /**
   * Parse PROXY protocol v1 header from buffer
   * Returns proxy info and remaining data after header
   */
  static parse(data: Buffer): IProxyParseResult {
    // Check if buffer starts with PROXY signature
    if (!data.toString('ascii', 0, 6).startsWith(this.PROXY_V1_SIGNATURE)) {
      return {
        proxyInfo: null,
        remainingData: data
      };
    }
    
    // Find header terminator
    const headerEndIndex = data.indexOf(this.HEADER_TERMINATOR);
    if (headerEndIndex === -1) {
      // Header incomplete, need more data
      if (data.length > this.MAX_HEADER_LENGTH) {
        // Header too long, invalid
        throw new Error('PROXY protocol header exceeds maximum length');
      }
      return {
        proxyInfo: null,
        remainingData: data
      };
    }
    
    // Extract header line
    const headerLine = data.toString('ascii', 0, headerEndIndex);
    const remainingData = data.slice(headerEndIndex + 2); // Skip \r\n
    
    // Parse header
    const parts = headerLine.split(' ');
    
    if (parts.length < 2) {
      throw new Error(`Invalid PROXY protocol header format: ${headerLine}`);
    }
    
    const [signature, protocol] = parts;
    
    // Validate protocol
    if (!['TCP4', 'TCP6', 'UNKNOWN'].includes(protocol)) {
      throw new Error(`Invalid PROXY protocol: ${protocol}`);
    }
    
    // For UNKNOWN protocol, ignore addresses
    if (protocol === 'UNKNOWN') {
      return {
        proxyInfo: {
          protocol: 'UNKNOWN',
          sourceIP: '',
          sourcePort: 0,
          destinationIP: '',
          destinationPort: 0
        },
        remainingData
      };
    }
    
    // For TCP4/TCP6, we need all 6 parts
    if (parts.length !== 6) {
      throw new Error(`Invalid PROXY protocol header format: ${headerLine}`);
    }
    
    const [, , srcIP, dstIP, srcPort, dstPort] = parts;
    
    // Validate and parse ports
    const sourcePort = parseInt(srcPort, 10);
    const destinationPort = parseInt(dstPort, 10);
    
    if (isNaN(sourcePort) || sourcePort < 0 || sourcePort > 65535) {
      throw new Error(`Invalid source port: ${srcPort}`);
    }
    
    if (isNaN(destinationPort) || destinationPort < 0 || destinationPort > 65535) {
      throw new Error(`Invalid destination port: ${dstPort}`);
    }
    
    // Validate IP addresses
    const protocolType = protocol as 'TCP4' | 'TCP6' | 'UNKNOWN';
    if (!this.isValidIP(srcIP, protocolType)) {
      throw new Error(`Invalid source IP for ${protocol}: ${srcIP}`);
    }
    
    if (!this.isValidIP(dstIP, protocolType)) {
      throw new Error(`Invalid destination IP for ${protocol}: ${dstIP}`);
    }
    
    return {
      proxyInfo: {
        protocol: protocol as 'TCP4' | 'TCP6',
        sourceIP: srcIP,
        sourcePort,
        destinationIP: dstIP,
        destinationPort
      },
      remainingData
    };
  }
  
  /**
   * Generate PROXY protocol v1 header
   */
  static generate(info: IProxyInfo): Buffer {
    if (info.protocol === 'UNKNOWN') {
      return Buffer.from(`PROXY UNKNOWN\r\n`, 'ascii');
    }
    
    const header = `PROXY ${info.protocol} ${info.sourceIP} ${info.destinationIP} ${info.sourcePort} ${info.destinationPort}\r\n`;
    
    if (header.length > this.MAX_HEADER_LENGTH) {
      throw new Error('Generated PROXY protocol header exceeds maximum length');
    }
    
    return Buffer.from(header, 'ascii');
  }
  
  /**
   * Validate IP address format
   */
  private static isValidIP(ip: string, protocol: 'TCP4' | 'TCP6' | 'UNKNOWN'): boolean {
    if (protocol === 'TCP4') {
      return plugins.net.isIPv4(ip);
    } else if (protocol === 'TCP6') {
      return plugins.net.isIPv6(ip);
    }
    return false;
  }
  
  /**
   * Attempt to read a complete PROXY protocol header from a socket
   * Returns null if no PROXY protocol detected or incomplete
   */
  static async readFromSocket(socket: plugins.net.Socket, timeout: number = 5000): Promise<IProxyParseResult | null> {
    return new Promise((resolve) => {
      let buffer = Buffer.alloc(0);
      let resolved = false;
      
      const cleanup = () => {
        socket.removeListener('data', onData);
        socket.removeListener('error', onError);
        clearTimeout(timer);
      };
      
      const timer = setTimeout(() => {
        if (!resolved) {
          resolved = true;
          cleanup();
          resolve({
            proxyInfo: null,
            remainingData: buffer
          });
        }
      }, timeout);
      
      const onData = (chunk: Buffer) => {
        buffer = Buffer.concat([buffer, chunk]);
        
        // Check if we have enough data
        if (!buffer.toString('ascii', 0, Math.min(6, buffer.length)).startsWith(this.PROXY_V1_SIGNATURE)) {
          // Not PROXY protocol
          resolved = true;
          cleanup();
          resolve({
            proxyInfo: null,
            remainingData: buffer
          });
          return;
        }
        
        // Try to parse
        try {
          const result = this.parse(buffer);
          if (result.proxyInfo) {
            // Successfully parsed
            resolved = true;
            cleanup();
            resolve(result);
          } else if (buffer.length > this.MAX_HEADER_LENGTH) {
            // Header too long
            resolved = true;
            cleanup();
            resolve({
              proxyInfo: null,
              remainingData: buffer
            });
          }
          // Otherwise continue reading
        } catch (error) {
          // Parse error
          logger.log('error', `PROXY protocol parse error: ${error.message}`);
          resolved = true;
          cleanup();
          resolve({
            proxyInfo: null,
            remainingData: buffer
          });
        }
      };
      
      const onError = (error: Error) => {
        logger.log('error', `Socket error while reading PROXY protocol: ${error.message}`);
        resolved = true;
        cleanup();
        resolve({
          proxyInfo: null,
          remainingData: buffer
        });
      };
      
      socket.on('data', onData);
      socket.on('error', onError);
    });
  }
}