import { describe, expect, it } from "vitest";
import { generateKeyPair } from "@stablelib/x25519";
import { noiseXXInitiator, noiseXXResponder } from "./noise-xx";

function memoryPair() {
  const a: Uint8Array[] = [];
  const b: Uint8Array[] = [];
  const aR: Array<(v: Uint8Array) => void> = [];
  const bR: Array<(v: Uint8Array) => void> = [];
  const tA = {
    async send(m: Uint8Array) {
      const r = bR.shift();
      if (r) r(m); else b.push(m);
    },
    receive(): Promise<Uint8Array> {
      const v = a.shift();
      if (v) return Promise.resolve(v);
      return new Promise<Uint8Array>((resolve) => aR.push(resolve));
    },
  };
  const tB = {
    async send(m: Uint8Array) {
      const r = aR.shift();
      if (r) r(m); else a.push(m);
    },
    receive(): Promise<Uint8Array> {
      const v = b.shift();
      if (v) return Promise.resolve(v);
      return new Promise<Uint8Array>((resolve) => bR.push(resolve));
    },
  };
  return [tA, tB] as const;
}

describe("Noise XX mutual authentication", () => {
  it("completes the three-message handshake and exchanges static keys", async () => {
    const alice = generateKeyPair();
    const bob = generateKeyPair();
    const [tA, tB] = memoryPair();

    const aPromise = noiseXXInitiator(alice.secretKey, alice.publicKey, tA.send, tA.receive);
    const bPromise = noiseXXResponder(bob.secretKey, bob.publicKey, tB.send, tB.receive);

    const [a, b] = await Promise.all([aPromise, bPromise]);
    expect(a.remoteStaticPublicKey).toEqual(bob.publicKey);
    expect(b.remoteStaticPublicKey).toEqual(alice.publicKey);
    expect(a.handshakeHash).toEqual(b.handshakeHash);
  });

  it("transport frames decrypt in both directions and are bound to the handshake", async () => {
    const alice = generateKeyPair();
    const bob = generateKeyPair();
    const [tA, tB] = memoryPair();

    const [a, b] = await Promise.all([
      noiseXXInitiator(alice.secretKey, alice.publicKey, tA.send, tA.receive),
      noiseXXResponder(bob.secretKey, bob.publicKey, tB.send, tB.receive),
    ]);

    const ct1 = a.encrypt(new TextEncoder().encode("hello bob"));
    expect(new TextDecoder().decode(b.decrypt(ct1))).toBe("hello bob");

    const ct2 = b.encrypt(new TextEncoder().encode("hello alice"));
    expect(new TextDecoder().decode(a.decrypt(ct2))).toBe("hello alice");

    // Tampering with one byte breaks the MAC.
    const ct3 = a.encrypt(new TextEncoder().encode("not delivered"));
    ct3[0] ^= 1;
    expect(() => b.decrypt(ct3)).toThrow();
  });

  it("an attacker who swaps the responder static cannot complete msg2 auth", async () => {
    const alice = generateKeyPair();
    const bobReal = generateKeyPair();
    const bobImpostor = generateKeyPair();
    const [tA, tB] = memoryPair();

    // Responder uses impostor private key but advertises real pubkey to initiator.
    // Real XX has no out-of-band pubkey check; this test verifies the registry
    // check that the caller MUST perform on `remoteStaticPublicKey`.
    const aPromise = noiseXXInitiator(alice.secretKey, alice.publicKey, tA.send, tA.receive);
    const bPromise = noiseXXResponder(bobImpostor.secretKey, bobImpostor.publicKey, tB.send, tB.receive);
    const [a] = await Promise.all([aPromise, bPromise]);
    // The "swap" is the caller seeing a remote static they did not expect.
    expect(a.remoteStaticPublicKey).not.toEqual(bobReal.publicKey);
    expect(a.remoteStaticPublicKey).toEqual(bobImpostor.publicKey);
  });
});