stack/packages/mosaic/src/lease-broker/broker-test-client.ts

import { createHash } from 'node:crypto';
import { createConnection, type Socket } from 'node:net';

const DEFAULT_TIMEOUT_MS = 3_000;
const MAX_REPLY_BYTES = 64 * 1024;
const MAX_DIAGNOSTIC_BYTES = 256;
const SENSITIVE_REPLY_FIELD = /"(?:promotion_token|session_id|token)"\s*:/;

export interface BrokerTestClientOptions {
  timeoutMs?: number;
}

export type BrokerTransportFailureKind =
  | 'early-close'
  | 'malformed-reply'
  | 'socket-error'
  | 'timeout';

export class BrokerTransportError extends Error {
  public readonly responseLength: number;
  public readonly responseBytes: string;
  public readonly responseHex: string;
  public readonly responsePreview: string;
  public readonly responseSha256: string;

  public constructor(
    public readonly kind: BrokerTransportFailureKind,
    description: string,
    response: Buffer,
    public readonly attempts = 1,
  ) {
    const diagnostics = responseDiagnostics(response);
    super(`${description}; attempts=${attempts}; ${responseContext(response, diagnostics)}`);
    this.name = 'BrokerTransportError';
    this.responseLength = response.length;
    this.responseBytes = diagnostics.preview;
    this.responseHex = diagnostics.hex;
    this.responsePreview = diagnostics.preview;
    this.responseSha256 = diagnostics.sha256;
  }
}

interface ResponseDiagnostics {
  preview: string;
  hex: string;
  sha256: string;
}

function responseDiagnostics(response: Buffer): ResponseDiagnostics {
  const fullText = response.toString('utf8');
  const sensitive = SENSITIVE_REPLY_FIELD.test(fullText);
  const bounded = response.subarray(0, MAX_DIAGNOSTIC_BYTES);
  const suffix = response.length > MAX_DIAGNOSTIC_BYTES ? '…' : '';
  return {
    preview:
      response.length === 0
        ? '<empty>'
        : sensitive
          ? '<redacted-sensitive-reply>'
          : `${bounded.toString('utf8')}${suffix}`,
    hex: sensitive ? '<redacted>' : `${bounded.toString('hex')}${suffix}`,
    sha256: createHash('sha256').update(response).digest('hex'),
  };
}

function responseContext(response: Buffer, diagnostics = responseDiagnostics(response)): string {
  const hex = diagnostics.hex.length === 0 ? '<empty>' : diagnostics.hex;
  return `length=${response.length}; bytes=${JSON.stringify(diagnostics.preview)}; hex=${hex}; sha256=${diagnostics.sha256}`;
}

function malformedReply(response: Buffer, reason: string): BrokerTransportError {
  return new BrokerTransportError('malformed-reply', `Malformed broker reply: ${reason}`, response);
}

function readBrokerReplyAttempt<T extends object>(
  socketPath: string,
  write: (socket: Socket) => void,
  timeoutMs: number,
): Promise<T> {
  return new Promise<T>((resolve, reject) => {
    const socket = createConnection(socketPath);
    let response = Buffer.alloc(0);
    let settled = false;

    const settle = (callback: () => void): void => {
      if (settled) return;
      settled = true;
      clearTimeout(timer);
      callback();
      socket.destroy();
    };
    const fail = (error: Error): void => settle(() => reject(error));
    const timer = setTimeout(
      () =>
        fail(
          new BrokerTransportError(
            'timeout',
            `Broker reply timed out before newline after ${timeoutMs}ms`,
            response,
          ),
        ),
      timeoutMs,
    );

    socket.once('error', (error) =>
      fail(
        new BrokerTransportError(
          'socket-error',
          `Broker reply socket error before newline: ${error.message}`,
          response,
        ),
      ),
    );
    socket.on('data', (chunk: Buffer) => {
      if (settled) return;
      response = Buffer.concat([response, chunk]);
      if (response.length > MAX_REPLY_BYTES) {
        fail(malformedReply(response, `exceeds ${MAX_REPLY_BYTES} bytes`));
      }
    });
    socket.once('end', () => {
      if (settled) return;
      const newline = response.indexOf(0x0a);
      if (response.length === 0 || newline < 0) {
        fail(
          new BrokerTransportError(
            'early-close',
            'Broker reply socket closed before newline',
            response,
          ),
        );
        return;
      }

      if (newline !== response.length - 1) {
        fail(malformedReply(response, 'contains bytes after the newline terminator'));
        return;
      }

      try {
        const parsed: unknown = JSON.parse(response.subarray(0, newline).toString('utf8'));
        if (typeof parsed !== 'object' || parsed === null || Array.isArray(parsed)) {
          fail(malformedReply(response, 'JSON value is not an object'));
          return;
        }
        settle(() => resolve(parsed as T));
      } catch (error: unknown) {
        fail(
          malformedReply(response, error instanceof Error ? error.message : 'JSON parsing failed'),
        );
      }
    });
    socket.once('connect', () => {
      try {
        write(socket);
      } catch (error: unknown) {
        fail(error instanceof Error ? error : new Error(String(error)));
      }
    });
  });
}

/** Read exactly one complete newline-framed JSON object; transport failures reject. */
export async function readBrokerReply<T extends object>(
  socketPath: string,
  write: (socket: Socket) => void,
  options: BrokerTestClientOptions = {},
): Promise<T> {
  return await readBrokerReplyAttempt<T>(
    socketPath,
    write,
    options.timeoutMs ?? DEFAULT_TIMEOUT_MS,
  );
}

/** Send one newline-framed request and read its complete broker reply. */
export async function requestBrokerReply<T extends object>(
  socketPath: string,
  requestValue: object,
  options?: BrokerTestClientOptions,
): Promise<T> {
  return await readBrokerReply<T>(
    socketPath,
    (socket) => socket.end(`${JSON.stringify(requestValue)}\n`),
    options,
  );
}