smartproxy/ts/protocols/common/fragment-handler.ts

/**
 * Shared Fragment Handler for Protocol Detection
 * 
 * Provides unified fragment buffering and reassembly for protocols
 * that may span multiple TCP packets.
 */

import { Buffer } from 'buffer';

/**
 * Fragment tracking information
 */
export interface IFragmentInfo {
  buffer: Buffer;
  timestamp: number;
  connectionId: string;
}

/**
 * Options for fragment handling
 */
export interface IFragmentOptions {
  maxBufferSize?: number;
  timeout?: number;
  cleanupInterval?: number;
}

/**
 * Result of fragment processing
 */
export interface IFragmentResult {
  isComplete: boolean;
  buffer?: Buffer;
  needsMoreData: boolean;
  error?: string;
}

/**
 * Shared fragment handler for protocol detection
 */
export class FragmentHandler {
  private fragments = new Map<string, IFragmentInfo>();
  private cleanupTimer?: NodeJS.Timeout;
  
  constructor(private options: IFragmentOptions = {}) {
    // Start cleanup timer if not already running
    if (options.cleanupInterval && !this.cleanupTimer) {
      this.cleanupTimer = setInterval(
        () => this.cleanup(),
        options.cleanupInterval
      );
    }
  }
  
  /**
   * Add a fragment for a connection
   */
  addFragment(connectionId: string, fragment: Buffer): IFragmentResult {
    const existing = this.fragments.get(connectionId);
    
    if (existing) {
      // Append to existing buffer
      const newBuffer = Buffer.concat([existing.buffer, fragment]);
      
      // Check size limit
      const maxSize = this.options.maxBufferSize || 65536;
      if (newBuffer.length > maxSize) {
        this.fragments.delete(connectionId);
        return {
          isComplete: false,
          needsMoreData: false,
          error: 'Buffer size exceeded maximum allowed'
        };
      }
      
      // Update fragment info
      this.fragments.set(connectionId, {
        buffer: newBuffer,
        timestamp: Date.now(),
        connectionId
      });
      
      return {
        isComplete: false,
        buffer: newBuffer,
        needsMoreData: true
      };
    } else {
      // New fragment
      this.fragments.set(connectionId, {
        buffer: fragment,
        timestamp: Date.now(),
        connectionId
      });
      
      return {
        isComplete: false,
        buffer: fragment,
        needsMoreData: true
      };
    }
  }
  
  /**
   * Get the current buffer for a connection
   */
  getBuffer(connectionId: string): Buffer | undefined {
    return this.fragments.get(connectionId)?.buffer;
  }
  
  /**
   * Mark a connection as complete and clean up
   */
  complete(connectionId: string): void {
    this.fragments.delete(connectionId);
  }
  
  /**
   * Check if we're tracking a connection
   */
  hasConnection(connectionId: string): boolean {
    return this.fragments.has(connectionId);
  }
  
  /**
   * Clean up expired fragments
   */
  cleanup(): void {
    const now = Date.now();
    const timeout = this.options.timeout || 5000;
    
    for (const [connectionId, info] of this.fragments.entries()) {
      if (now - info.timestamp > timeout) {
        this.fragments.delete(connectionId);
      }
    }
  }
  
  /**
   * Clear all fragments
   */
  clear(): void {
    this.fragments.clear();
  }
  
  /**
   * Destroy the handler and clean up resources
   */
  destroy(): void {
    if (this.cleanupTimer) {
      clearInterval(this.cleanupTimer);
      this.cleanupTimer = undefined;
    }
    this.clear();
  }
  
  /**
   * Get the number of tracked connections
   */
  get size(): number {
    return this.fragments.size;
  }
}