Shade/packages/shade-transport/src/ws-adapter.ts

import type { ShadeSessionManager } from '@shade/core';
import { encodeEnvelope, decodeEnvelope } from '@shade/proto';

/**
 * ShadeWebSocket — wraps an existing WebSocket with automatic E2EE.
 *
 * All outgoing messages are encrypted via the Double Ratchet.
 * All incoming messages are decrypted transparently.
 *
 * Usage:
 * ```ts
 * const ws = new WebSocket('wss://example.com/sync');
 * const shade = new ShadeWebSocket(ws, sessionManager, 'server');
 *
 * shade.onMessage((plaintext) => {
 *   console.log('Received:', plaintext);
 * });
 *
 * await shade.send('Hello encrypted world!');
 * ```
 */
export class ShadeWebSocket {
  private messageHandlers: Array<(plaintext: string) => void> = [];
  private errorHandlers: Array<(error: Error) => void> = [];

  constructor(
    private readonly ws: WebSocket,
    private readonly manager: ShadeSessionManager,
    private readonly peerAddress: string,
  ) {
    this.ws.addEventListener('message', (event) => {
      this.handleIncoming(event.data).catch((err) => {
        for (const handler of this.errorHandlers) handler(err);
      });
    });
  }

  /** Send an encrypted message to the peer */
  async send(plaintext: string): Promise<void> {
    const envelope = await this.manager.encrypt(this.peerAddress, plaintext);
    const bytes = encodeEnvelope(envelope);

    // Send as binary. Cast through BufferSource to satisfy
    // Uint8Array<ArrayBufferLike> vs Uint8Array<ArrayBuffer> mismatch under
    // strict TS configurations (5.7+).
    this.ws.send(bytes as unknown as ArrayBuffer);
  }

  /** Register a handler for decrypted incoming messages */
  onMessage(handler: (plaintext: string) => void): void {
    this.messageHandlers.push(handler);
  }

  /** Register a handler for decryption errors */
  onError(handler: (error: Error) => void): void {
    this.errorHandlers.push(handler);
  }

  /** Close the underlying WebSocket */
  close(): void {
    this.ws.close();
  }

  private async handleIncoming(data: any): Promise<void> {
    let bytes: Uint8Array;

    if (data instanceof ArrayBuffer) {
      bytes = new Uint8Array(data);
    } else if (data instanceof Uint8Array) {
      bytes = data;
    } else if (typeof data === 'string') {
      // Base64-encoded fallback for environments that don't support binary WS
      bytes = new Uint8Array(Buffer.from(data, 'base64'));
    } else {
      throw new Error('Unexpected WebSocket message type');
    }

    const envelope = decodeEnvelope(bytes);
    const plaintext = await this.manager.decrypt(this.peerAddress, envelope);

    for (const handler of this.messageHandlers) {
      handler(plaintext);
    }
  }
}