smartproxy/ts/protocols/websocket/utils.ts

/**
 * WebSocket Protocol Utilities
 */

import * as crypto from 'crypto';
import { WEBSOCKET_MAGIC_STRING } from './constants.js';
import type { RawData } from './types.js';

/**
 * Get the length of a WebSocket message regardless of its type
 * (handles all possible WebSocket message data types)
 */
export function getMessageSize(data: RawData): number {
  if (typeof data === 'string') {
    // For string data, get the byte length
    return Buffer.from(data, 'utf8').length;
  } else if (data instanceof Buffer) {
    // For Node.js Buffer
    return data.length;
  } else if (data instanceof ArrayBuffer) {
    // For ArrayBuffer
    return data.byteLength;
  } else if (Array.isArray(data)) {
    // For array of buffers, sum their lengths
    return data.reduce((sum, chunk) => {
      if (chunk instanceof Buffer) {
        return sum + chunk.length;
      } else if (chunk instanceof ArrayBuffer) {
        return sum + chunk.byteLength;
      }
      return sum;
    }, 0);
  } else {
    // For other types, try to determine the size or return 0
    try {
      return Buffer.from(data).length;
    } catch (e) {
      return 0;
    }
  }
}

/**
 * Convert any raw WebSocket data to Buffer for consistent handling
 */
export function toBuffer(data: RawData): Buffer {
  if (typeof data === 'string') {
    return Buffer.from(data, 'utf8');
  } else if (data instanceof Buffer) {
    return data;
  } else if (data instanceof ArrayBuffer) {
    return Buffer.from(data);
  } else if (Array.isArray(data)) {
    // For array of buffers, concatenate them
    return Buffer.concat(data.map(chunk => {
      if (chunk instanceof Buffer) {
        return chunk;
      } else if (chunk instanceof ArrayBuffer) {
        return Buffer.from(chunk);
      }
      return Buffer.from(chunk);
    }));
  } else {
    // For other types, try to convert to Buffer or return empty Buffer
    try {
      return Buffer.from(data);
    } catch (e) {
      return Buffer.alloc(0);
    }
  }
}

/**
 * Generate WebSocket accept key from client key
 */
export function generateAcceptKey(clientKey: string): string {
  const hash = crypto.createHash('sha1');
  hash.update(clientKey + WEBSOCKET_MAGIC_STRING);
  return hash.digest('base64');
}

/**
 * Validate WebSocket upgrade request
 */
export function isWebSocketUpgrade(headers: Record<string, string>): boolean {
  const upgrade = headers['upgrade'];
  const connection = headers['connection'];
  
  return upgrade?.toLowerCase() === 'websocket' && 
         connection?.toLowerCase().includes('upgrade');
}

/**
 * Generate random WebSocket key for client handshake
 */
export function generateWebSocketKey(): string {
  return crypto.randomBytes(16).toString('base64');
}