paseo/packages/relay/src/encrypted-channel.test.ts

import { describe, it, expect, vi } from "vitest";
import {
  createClientChannel,
  createDaemonChannel,
  Transport,
} from "./encrypted-channel.js";
import { generateKeyPair, exportPublicKey } from "./crypto.js";

/**
 * Creates a pair of connected mock transports.
 * Messages sent on one are received on the other.
 */
function createMockTransportPair(): [Transport, Transport] {
  const transportA: Transport = {
    send: vi.fn(),
    close: vi.fn(),
    onmessage: null,
    onclose: null,
    onerror: null,
  };

  const transportB: Transport = {
    send: vi.fn(),
    close: vi.fn(),
    onmessage: null,
    onclose: null,
    onerror: null,
  };

  // Wire them together
  (transportA.send as ReturnType<typeof vi.fn>).mockImplementation(
    (data: string | ArrayBuffer) => {
      setTimeout(() => transportB.onmessage?.(data), 0);
    }
  );

  (transportB.send as ReturnType<typeof vi.fn>).mockImplementation(
    (data: string | ArrayBuffer) => {
      setTimeout(() => transportA.onmessage?.(data), 0);
    }
  );

  return [transportA, transportB];
}

describe("EncryptedChannel", () => {
  it("establishes encrypted channel between daemon and client", async () => {
    const [daemonTransport, clientTransport] = createMockTransportPair();

    // Daemon generates keypair (public key goes in QR)
    const daemonKeyPair = await generateKeyPair();
    const daemonPubKeyB64 = await exportPublicKey(daemonKeyPair.publicKey);

    // Start daemon waiting for client
    const daemonChannelPromise = createDaemonChannel(daemonTransport, daemonKeyPair);

    // Client connects (scanned QR, got daemon's public key)
    const clientChannel = await createClientChannel(
      clientTransport,
      daemonPubKeyB64
    );

    // Daemon receives hello and completes handshake
    const daemonChannel = await daemonChannelPromise;

    expect(clientChannel.isOpen()).toBe(true);
    expect(daemonChannel.isOpen()).toBe(true);
  });

  it("exchanges encrypted messages bidirectionally", async () => {
    const [daemonTransport, clientTransport] = createMockTransportPair();

    const daemonKeyPair = await generateKeyPair();
    const daemonPubKeyB64 = await exportPublicKey(daemonKeyPair.publicKey);

    const daemonMessages: (string | ArrayBuffer)[] = [];
    const clientMessages: (string | ArrayBuffer)[] = [];

    const daemonChannelPromise = createDaemonChannel(
      daemonTransport,
      daemonKeyPair,
      { onmessage: (data) => daemonMessages.push(data) }
    );

    const clientChannel = await createClientChannel(
      clientTransport,
      daemonPubKeyB64,
      { onmessage: (data) => clientMessages.push(data) }
    );

    const daemonChannel = await daemonChannelPromise;

    // Send messages both directions
    await clientChannel.send("Hello from client");
    await daemonChannel.send("Hello from daemon");
    await clientChannel.send("Second message from client");

    // Wait for async delivery
    await new Promise((r) => setTimeout(r, 50));

    expect(daemonMessages).toEqual([
      "Hello from client",
      "Second message from client",
    ]);
    expect(clientMessages).toEqual(["Hello from daemon"]);
  });

  it("encrypted messages are opaque to transport", async () => {
    const [daemonTransport, clientTransport] = createMockTransportPair();

    const daemonKeyPair = await generateKeyPair();
    const daemonPubKeyB64 = await exportPublicKey(daemonKeyPair.publicKey);

    const daemonChannelPromise = createDaemonChannel(daemonTransport, daemonKeyPair);
    const clientChannel = await createClientChannel(
      clientTransport,
      daemonPubKeyB64
    );
    await daemonChannelPromise;

    // Clear mock call history
    (clientTransport.send as ReturnType<typeof vi.fn>).mockClear();

    // Send a plaintext message
    const plaintext = "Secret message";
    await clientChannel.send(plaintext);

    // Check what was actually sent over the transport
    expect(clientTransport.send).toHaveBeenCalledTimes(1);
    const sentData = (clientTransport.send as ReturnType<typeof vi.fn>).mock
      .calls[0][0];

    // Should be base64 string (encrypted)
    expect(typeof sentData).toBe("string");
    // Should NOT contain the plaintext
    expect(sentData).not.toContain(plaintext);
    // Should be significantly longer than plaintext (IV + auth tag overhead)
    expect(sentData.length).toBeGreaterThan(plaintext.length + 20);
  });

  it("fails handshake on invalid hello", async () => {
    const [daemonTransport] = createMockTransportPair();

    const daemonKeyPair = await generateKeyPair();

    const daemonChannelPromise = createDaemonChannel(daemonTransport, daemonKeyPair);

    // Send invalid hello
    setTimeout(() => {
      daemonTransport.onmessage?.('{"type":"invalid"}');
    }, 0);

    await expect(daemonChannelPromise).rejects.toThrow("Invalid hello message");
  });
});