stack/packages/mosaic/src/lease-broker/lease-broker.acceptance.spec.ts

import { randomUUID } from 'node:crypto';
import { chmod, mkdtemp, readFile, stat, symlink, writeFile } from 'node:fs/promises';
import { createConnection, type Socket } from 'node:net';
import { tmpdir } from 'node:os';
import { join } from 'node:path';
import { spawn, spawnSync, type ChildProcess } from 'node:child_process';

import { afterEach, describe, expect, test, vi } from 'vitest';

interface BrokerReply {
  ok: boolean;
  code?: string;
  session_id?: string;
  peer?: { pid: number; uid: number; gid: number; starttime: string };
  token?: string;
}

const daemonPath = new URL('../../framework/tools/lease-broker/daemon.py', import.meta.url)
  .pathname;
const children: ChildProcess[] = [];

async function withTimeout<T>(
  promise: Promise<T>,
  label: string,
  milliseconds = 3_000,
): Promise<T> {
  let timer: NodeJS.Timeout | undefined;
  try {
    return await Promise.race([
      promise,
      new Promise<never>((_, reject) => {
        timer = setTimeout(() => reject(new Error(`${label} timed out`)), milliseconds);
      }),
    ]);
  } finally {
    if (timer !== undefined) clearTimeout(timer);
  }
}

async function rawRequest(
  socketPath: string,
  write: (socket: Socket) => void,
): Promise<BrokerReply> {
  return await withTimeout(
    new Promise<BrokerReply>((resolve, reject) => {
      const socket = createConnection(socketPath);
      let response = '';
      socket.setEncoding('utf8');
      socket.once('error', reject);
      socket.on('data', (chunk: string) => {
        response += chunk;
      });
      socket.once('end', () => resolve(JSON.parse(response) as BrokerReply));
      socket.once('connect', () => write(socket));
    }),
    'raw broker request',
  );
}

async function request(socketPath: string, requestValue: object): Promise<BrokerReply> {
  return await new Promise<BrokerReply>((resolve, reject) => {
    const socket = createConnection(socketPath);
    let response = '';
    socket.setEncoding('utf8');
    socket.once('error', reject);
    socket.on('data', (chunk: string) => {
      response += chunk;
    });
    socket.once('end', () => resolve(JSON.parse(response) as BrokerReply));
    socket.once('connect', () => socket.end(`${JSON.stringify(requestValue)}\n`));
  });
}

async function startBroker(
  parentMode = 0o700,
): Promise<{ root: string; socket: string; state: string }> {
  const root = await mkdtemp(join(tmpdir(), 'mosaic-lease-broker-'));
  await chmod(root, parentMode);
  const socket = join(root, 'broker.sock');
  const state = join(root, 'state.json');
  const child = spawn('python3', [daemonPath, '--socket', socket, '--state', state], {
    stdio: ['ignore', 'pipe', 'pipe'],
  });
  children.push(child);
  await new Promise<void>((resolve, reject) => {
    let stderr = '';
    child.stderr?.setEncoding('utf8');
    child.stderr?.on('data', (chunk: string) => (stderr += chunk));
    child.once('error', reject);
    child.once('exit', (code: number | null) =>
      reject(new Error(`broker exited ${code}: ${stderr}`)),
    );
    child.stdout?.once('data', () => resolve());
  });
  return { root, socket, state };
}

async function startBrokerWithState(stateValue: string): Promise<string> {
  const root = await mkdtemp(join(tmpdir(), 'mosaic-lease-broker-'));
  await chmod(root, 0o700);
  const state = join(root, 'state.json');
  await writeFile(state, stateValue, { mode: 0o600 });
  const child = spawn('python3', [
    daemonPath,
    '--socket',
    join(root, 'broker.sock'),
    '--state',
    state,
  ]);
  children.push(child);
  return await new Promise<string>((resolve) => {
    let raw = '';
    child.stderr?.on('data', (chunk: Buffer) => (raw += chunk.toString()));
    child.once('exit', () => resolve(raw));
  });
}

afterEach(() => {
  for (const child of children.splice(0)) child.kill('SIGTERM');
  vi.restoreAllMocks();
});

describe('authenticated external lease broker', () => {
  test('peercred returns true kernel (pid,starttime)', async () => {
    const getuid = process.getuid;
    const getgid = process.getgid;
    if (getuid === undefined || getgid === undefined) {
      throw new Error('Linux peer credentials require process.getuid() and process.getgid()');
    }

    const { socket } = await startBroker();
    const reply = await request(socket, { action: 'register_anchor', runtime_generation: 1 });
    const statText = await readFile(`/proc/${process.pid}/stat`, 'utf8');
    const fields = statText.slice(statText.lastIndexOf(')') + 2).split(' ');
    expect(reply).toMatchObject({
      ok: true,
      peer: {
        pid: process.pid,
        uid: getuid(),
        gid: getgid(),
        starttime: fields[19],
      },
    });
  });

  test.each([null, '', 'chosen'])('caller-asserted session_id refused (%j)', async (session_id) => {
    const { socket } = await startBroker();
    const reply = await request(socket, {
      action: 'register_anchor',
      runtime_generation: 1,
      session_id,
    });
    expect(reply).toMatchObject({ ok: false, code: 'CALLER_SESSION_ID_REFUSED' });
  });

  test('sibling-substitution rejected', async () => {
    const { socket } = await startBroker();
    const launcher = spawn(
      process.execPath,
      [
        '-e',
        `const n=require('net');const s=n.connect(${JSON.stringify(socket)},()=>s.end(JSON.stringify({action:'register_anchor',runtime_generation:1})+'\\n'));s.on('data',d=>{process.send(JSON.parse(d));setInterval(()=>{},1000)})`,
      ],
      { stdio: ['ignore', 'ignore', 'ignore', 'ipc'] },
    );
    children.push(launcher);
    const registration = await new Promise<BrokerReply>((resolve) =>
      launcher.once('message', (message) => resolve(message as BrokerReply)),
    );
    const attacker = spawn(
      process.execPath,
      [
        '-e',
        `const n=require('net');const s=n.connect(${JSON.stringify(socket)},()=>s.end(JSON.stringify({action:'authenticate',session_id:${JSON.stringify(registration.session_id)},runtime_generation:1})+'\\n'));s.pipe(process.stdout)`,
      ],
      { stdio: ['ignore', 'pipe', 'ignore'] },
    );
    children.push(attacker);
    let raw = '';
    attacker.stdout?.setEncoding('utf8');
    attacker.stdout?.on('data', (chunk: string) => (raw += chunk));
    await new Promise<void>((resolve) => attacker.once('exit', () => resolve()));
    expect(JSON.parse(raw)).toMatchObject({ ok: false, code: 'ANCESTRY_MISMATCH' });
  });

  test('generation bump revokes prior incarnation', async () => {
    const { socket } = await startBroker();
    const registered = await request(socket, { action: 'register_anchor', runtime_generation: 1 });
    expect(
      await request(socket, {
        action: 'authenticate',
        session_id: registered.session_id,
        runtime_generation: 2,
      }),
    ).toMatchObject({ ok: true });
    expect(
      await request(socket, {
        action: 'authenticate',
        session_id: registered.session_id,
        runtime_generation: 1,
      }),
    ).toMatchObject({ ok: false, code: 'STALE_GENERATION' });
  });

  test('same anchor re-registration reuses its session and revokes the prior incarnation', async () => {
    const { socket } = await startBroker();
    const first = await request(socket, { action: 'register_anchor', runtime_generation: 1 });
    const binding = {
      compaction_epoch: 2,
      request_epoch: 3,
      h_source: 'a'.repeat(64),
      h_payload: 'b'.repeat(64),
      schema_version: 1,
    };
    const minted = await request(socket, {
      action: 'mint_token',
      session_id: first.session_id,
      runtime_generation: 1,
      binding,
    });

    const bumped = await request(socket, { action: 'register_anchor', runtime_generation: 2 });
    const repeated = await request(socket, { action: 'register_anchor', runtime_generation: 2 });
    const lower = await request(socket, { action: 'register_anchor', runtime_generation: 1 });

    expect(bumped).toMatchObject({ ok: true, session_id: first.session_id });
    expect(repeated).toMatchObject({ ok: true, session_id: first.session_id });
    expect(lower).toMatchObject({ ok: false, code: 'STALE_GENERATION' });
    expect(
      await request(socket, {
        action: 'authenticate',
        session_id: first.session_id,
        runtime_generation: 1,
      }),
    ).toMatchObject({ ok: false, code: 'STALE_GENERATION' });
    expect(
      await request(socket, {
        action: 'consume_token',
        session_id: first.session_id,
        runtime_generation: 2,
        token: minted.token,
      }),
    ).toMatchObject({ ok: false, code: 'TOKEN_REPLAY' });
  });

  test('crypto token path works when Math.random is poisoned', async () => {
    const { socket } = await startBroker();
    vi.spyOn(Math, 'random').mockImplementation(() => {
      throw new Error('Math.random forbidden');
    });
    const registered = await request(socket, { action: 'register_anchor', runtime_generation: 1 });
    const binding = {
      compaction_epoch: 2,
      request_epoch: 3,
      h_source: 'a'.repeat(64),
      h_payload: 'b'.repeat(64),
      schema_version: 1,
    };
    const first = await request(socket, {
      action: 'mint_token',
      session_id: registered.session_id,
      runtime_generation: 1,
      binding,
    });
    const second = await request(socket, {
      action: 'mint_token',
      session_id: registered.session_id,
      runtime_generation: 1,
      binding,
    });
    expect(first.token).toMatch(/^[a-f0-9]{64}$/);
    expect(second.token).not.toBe(first.token);
    expect(
      await request(socket, {
        action: 'consume_token',
        session_id: registered.session_id,
        runtime_generation: 1,
        token: first.token,
      }),
    ).toMatchObject({ ok: true });
    expect(
      await request(socket, {
        action: 'consume_token',
        session_id: registered.session_id,
        runtime_generation: 1,
        token: first.token,
      }),
    ).toMatchObject({ ok: false, code: 'TOKEN_REPLAY' });
  });

  test('socket parent 0700 and socket 0600 enforced', async () => {
    const { root, socket, state } = await startBroker();
    await request(socket, { action: 'register_anchor', runtime_generation: 1 });
    expect((await stat(root)).mode & 0o777).toBe(0o700);
    expect((await stat(socket)).mode & 0o777).toBe(0o600);
    expect((await stat(state)).mode & 0o777).toBe(0o600);
  });

  test('insecure existing posture refused', async () => {
    await expect(startBroker(0o755)).rejects.toThrow(/INSECURE_PARENT_MODE/);
  });

  test('malformed and oversized frames fail closed without killing broker', async () => {
    const { socket } = await startBroker();
    const malformed = await new Promise<string>((resolve, reject) => {
      const connection = createConnection(socket, () => connection.end('{nope}\n'));
      let raw = '';
      connection.on('data', (chunk: Buffer) => (raw += chunk.toString()));
      connection.once('end', () => resolve(raw));
      connection.once('error', reject);
    });
    expect(JSON.parse(malformed)).toMatchObject({ ok: false, code: 'MALFORMED_REQUEST' });
    const registered = await request(socket, {
      action: 'register_anchor',
      runtime_generation: 1,
      nonce: randomUUID(),
    });
    expect(registered.ok).toBe(true);
  });

  test('silent connection deadline cannot prevent the next valid registration', async () => {
    const { socket } = await startBroker();
    const silent = createConnection(socket);
    await new Promise<void>((resolve, reject) => {
      silent.once('connect', resolve);
      silent.once('error', reject);
    });
    const registered = await withTimeout(
      request(socket, { action: 'register_anchor', runtime_generation: 1 }),
      'registration behind silent connection',
    );
    expect(registered.ok).toBe(true);
    silent.destroy();
  });

  test('queued silent peers cannot serialize the next valid registration', async () => {
    const { socket } = await startBroker();
    const silentConnections = await Promise.all(
      Array.from(
        { length: 4 },
        () =>
          new Promise<Socket>((resolve, reject) => {
            const connection = createConnection(socket);
            connection.once('connect', () => resolve(connection));
            connection.once('error', reject);
          }),
      ),
    );

    try {
      await new Promise((resolve) => setTimeout(resolve, 100));
      const started = performance.now();
      const registered = await withTimeout(
        request(socket, { action: 'register_anchor', runtime_generation: 1 }),
        'registration behind queued silent connections',
        6_000,
      );
      const elapsed = performance.now() - started;

      expect(registered.ok).toBe(true);
      expect(elapsed).toBeLessThan(1_500);
    } finally {
      for (const connection of silentConnections) connection.destroy();
    }
  });

  test('silent peers are reaped at the concurrency bound and their slots are reclaimed', async () => {
    const { socket } = await startBroker();
    const concurrencyCap = 16;
    const peers = Array.from({ length: concurrencyCap }, () => {
      const connection = createConnection(socket);
      return {
        connection,
        connected: new Promise<void>((resolve, reject) => {
          connection.once('connect', resolve);
          connection.once('error', reject);
        }),
        closed: new Promise<void>((resolve) => connection.once('close', () => resolve())),
      };
    });

    try {
      await Promise.all(peers.map(({ connected }) => connected));
      await new Promise((resolve) => setTimeout(resolve, 100));
      const started = performance.now();
      const registration = withTimeout(
        request(socket, { action: 'register_anchor', runtime_generation: 1 }),
        'registration while silent peers hold the concurrency bound',
        2_500,
      );
      const reaping = withTimeout(
        Promise.all(peers.map(({ closed }) => closed)),
        'silent peer deadline reaping',
        2_500,
      );
      const [registered] = await Promise.all([registration, reaping]);
      const elapsed = performance.now() - started;

      expect(registered.ok).toBe(true);
      expect(elapsed).toBeGreaterThan(500);
      expect(elapsed).toBeLessThan(2_500);
    } finally {
      for (const { connection } of peers) connection.destroy();
    }
  });

  test('newline-only client without half-close gets no success and cannot block next request', async () => {
    const { socket } = await startBroker();
    const incomplete = createConnection(socket);
    let raw = '';
    incomplete.setEncoding('utf8');
    incomplete.on('data', (chunk: string) => (raw += chunk));
    await new Promise<void>((resolve, reject) => {
      incomplete.once('connect', () => {
        incomplete.write(
          `${JSON.stringify({ action: 'register_anchor', runtime_generation: 1 })}\n`,
        );
        resolve();
      });
      incomplete.once('error', reject);
    });
    await new Promise((resolve) => setTimeout(resolve, 1_100));
    expect(raw).not.toContain('"ok":true');
    const registered = await withTimeout(
      request(socket, { action: 'register_anchor', runtime_generation: 1 }),
      'registration after non-half-closed client',
    );
    expect(registered.ok).toBe(true);
    incomplete.destroy();
  });

  test('client disconnect cannot prevent the next valid authentication', async () => {
    const { socket } = await startBroker();
    const registered = await request(socket, { action: 'register_anchor', runtime_generation: 1 });
    const reset = createConnection(socket);
    await new Promise<void>((resolve, reject) => {
      reset.once('connect', () => {
        reset.write(`${JSON.stringify({ action: 'register_anchor', runtime_generation: 1 })}\n`);
        reset.destroy();
        resolve();
      });
      reset.once('error', reject);
    });
    const authenticated = await withTimeout(
      request(socket, {
        action: 'authenticate',
        session_id: registered.session_id,
        runtime_generation: 1,
      }),
      'authentication after client disconnect',
    );
    expect(authenticated.ok).toBe(true);
  });

  test('delayed second frame is rejected and the next request succeeds', async () => {
    const { socket } = await startBroker();
    const reply = await rawRequest(socket, (connection) => {
      connection.write(`${JSON.stringify({ action: 'register_anchor', runtime_generation: 1 })}\n`);
      setTimeout(() => connection.end('{}\n'), 50);
    });
    expect(reply).toMatchObject({ ok: false, code: 'MALFORMED_REQUEST' });
    expect(
      await request(socket, { action: 'register_anchor', runtime_generation: 1 }),
    ).toMatchObject({
      ok: true,
    });
  });

  test('unterminated frame is rejected and the next request succeeds', async () => {
    const { socket } = await startBroker();
    const reply = await rawRequest(socket, (connection) => connection.end('{}'));
    expect(reply).toMatchObject({ ok: false, code: 'MALFORMED_REQUEST' });
    expect(
      await request(socket, { action: 'register_anchor', runtime_generation: 1 }),
    ).toMatchObject({
      ok: true,
    });
  });

  test('genuinely oversized frame is rejected and the next request succeeds', async () => {
    const { socket } = await startBroker();
    const reply = await rawRequest(socket, (connection) =>
      connection.end(`${JSON.stringify({ padding: 'x'.repeat(64 * 1024) })}\n`),
    );
    expect(reply).toMatchObject({ ok: false, code: 'MALFORMED_REQUEST' });
    expect(
      await request(socket, { action: 'register_anchor', runtime_generation: 1 }),
    ).toMatchObject({
      ok: true,
    });
  });

  test('boolean runtime generations fail closed', async () => {
    const { socket } = await startBroker();
    expect(
      await request(socket, { action: 'register_anchor', runtime_generation: true }),
    ).toMatchObject({ ok: false, code: 'INVALID_GENERATION' });
    const registered = await request(socket, { action: 'register_anchor', runtime_generation: 1 });
    expect(
      await request(socket, {
        action: 'authenticate',
        session_id: registered.session_id,
        runtime_generation: false,
      }),
    ).toMatchObject({ ok: false, code: 'INVALID_IDENTITY' });
  });

  test.each([
    { compaction_epoch: true, request_epoch: 0, schema_version: 1 },
    { compaction_epoch: 0, request_epoch: -1, schema_version: 1 },
    { compaction_epoch: 0, request_epoch: 0, schema_version: false },
    { compaction_epoch: 0, request_epoch: 0, schema_version: -1 },
    { compaction_epoch: 0, request_epoch: 0, schema_version: 1, h_source: 'A'.repeat(64) },
    { compaction_epoch: 0, request_epoch: 0, schema_version: 1, h_payload: 'a'.repeat(63) },
  ])('invalid cycle binding fails closed without persisting a token (%j)', async (override) => {
    const { socket, state } = await startBroker();
    const registered = await request(socket, { action: 'register_anchor', runtime_generation: 1 });
    const binding = Object.assign(
      {
        compaction_epoch: 0,
        request_epoch: 0,
        h_source: 'a'.repeat(64),
        h_payload: 'b'.repeat(64),
        schema_version: 1,
      },
      override,
    );
    expect(
      await request(socket, {
        action: 'mint_token',
        session_id: registered.session_id,
        runtime_generation: 1,
        binding,
      }),
    ).toMatchObject({ ok: false, code: 'INVALID_BINDING' });
    const persisted = JSON.parse(await readFile(state, 'utf8')) as { tokens: object };
    expect(persisted.tokens).toEqual({});
  });

  test('StateStore write-all unit path handles partial writes and cleans failed temp files', () => {
    const result = spawnSync('python3', [join(import.meta.dirname, 'state_store_unittest.py')], {
      encoding: 'utf8',
    });
    expect(result.status, result.stderr).toBe(0);
  });

  test('persistence integrity failure refuses startup', async () => {
    expect(await startBrokerWithState('{corrupt')).toContain('STATE_INTEGRITY');
  });

  test.each([
    { version: 1, sessions: {}, tokens: {}, unexpected: true },
    { version: 1, sessions: { bad: {} }, tokens: {} },
    {
      version: 1,
      sessions: {
        ['a'.repeat(64)]: { anchor_pid: true, anchor_starttime: '1', runtime_generation: 0 },
      },
      tokens: {},
    },
    {
      version: 1,
      sessions: {
        ['a'.repeat(64)]: { anchor_pid: 1, anchor_starttime: '01', runtime_generation: 0 },
      },
      tokens: {},
    },
    {
      version: 1,
      sessions: {
        ['a'.repeat(64)]: { anchor_pid: 1, anchor_starttime: '1', runtime_generation: 0 },
        ['b'.repeat(64)]: { anchor_pid: 1, anchor_starttime: '1', runtime_generation: 1 },
      },
      tokens: {},
    },
    {
      version: 1,
      sessions: {
        ['a'.repeat(64)]: { anchor_pid: 1, anchor_starttime: '1', runtime_generation: 0 },
      },
      tokens: {
        ['b'.repeat(64)]: {
          session_id: 'c'.repeat(64),
          runtime_generation: 0,
          binding: {
            compaction_epoch: 0,
            request_epoch: 0,
            h_source: 'd'.repeat(64),
            h_payload: 'e'.repeat(64),
            schema_version: 1,
          },
          consumed: false,
        },
      },
    },
    {
      version: 1,
      sessions: {
        ['a'.repeat(64)]: { anchor_pid: 1, anchor_starttime: '1', runtime_generation: 1 },
      },
      tokens: {
        ['b'.repeat(64)]: {
          session_id: 'a'.repeat(64),
          runtime_generation: 2,
          binding: {
            compaction_epoch: 0,
            request_epoch: 0,
            h_source: 'd'.repeat(64),
            h_payload: 'e'.repeat(64),
            schema_version: 1,
          },
          consumed: false,
        },
      },
    },
  ])('nested corrupt state refuses startup (%#)', async (stateValue) => {
    expect(await startBrokerWithState(JSON.stringify(stateValue))).toContain('STATE_INTEGRITY');
  });

  test('symlink state refuses startup', async () => {
    const root = await mkdtemp(join(tmpdir(), 'mosaic-lease-broker-'));
    await chmod(root, 0o700);
    const target = join(root, 'target.json');
    const state = join(root, 'state.json');
    await writeFile(target, JSON.stringify({ version: 1, sessions: {}, tokens: {} }), {
      mode: 0o600,
    });
    await symlink(target, state);
    const child = spawn('python3', [
      daemonPath,
      '--socket',
      join(root, 'broker.sock'),
      '--state',
      state,
    ]);
    children.push(child);
    const stderr = await new Promise<string>((resolve) => {
      let raw = '';
      child.stderr?.on('data', (chunk: Buffer) => (raw += chunk.toString()));
      child.once('exit', () => resolve(raw));
    });
    expect(stderr).toContain('STATE_INTEGRITY');
  });

  test('oversized state refuses startup', async () => {
    expect(await startBrokerWithState(' '.repeat(4 * 1024 * 1024 + 1))).toContain(
      'STATE_INTEGRITY',
    );
  });
});