KSP-MissionControl/packages/krpc-client/tests/service-client.test.ts

/**
 * End-to-end test of KrpcServices with a tiny mock kRPC server.
 *
 * We start a real TCP server that:
 *   1. Accepts the RPC + stream handshakes
 *   2. Handles a small whitelist of procedure calls by responding
 *      with hand-encoded values
 *
 * The test then drives a real KRPCClient + KrpcServices pair through
 * the same wire format a real kRPC server uses, and verifies the
 * decoded values match the hand-encoded responses.
 *
 * For the server side we use the same SocketReader-based pattern that
 * the client uses, so the two sides share framing semantics.
 */
import { describe, it, expect, beforeAll, afterAll } from 'vitest';
import * as net from 'node:net';
import { Buffer } from 'node:buffer';
import { KRPC, encodeMessage } from '../src/schema.js';
import { KRPCClient } from '../src/client.js';
import { loadServices, KrpcServices } from '../src/service-client.js';
import { encodeVarint, recvMessage, sendMessage } from '../src/connection.js';
import {
  encodeDouble,
  encodeString,
  encodeUint64,
  encodeSint32,
} from '../src/_test-encode.js';

/** Mock services message that the server will hand back on GetServices. */
const MOCK_SERVICES_RAW = {
  services: [
    {
      name: 'KRPC',
      procedures: [
        {
          name: 'GetStatus',
          parameters: [],
          returnType: { code: 203, service: '', name: '', types: [] },
          returnIsNullable: false,
        },
        {
          name: 'GetServices',
          parameters: [],
          returnType: { code: 204, service: '', name: '', types: [] },
          returnIsNullable: false,
        },
      ],
      classes: [],
      enumerations: [],
    },
    {
      name: 'SpaceCenter',
      procedures: [
        {
          name: 'GetUT',
          parameters: [],
          returnType: { code: 1, service: '', name: '', types: [] },
          returnIsNullable: false,
        },
        {
          name: 'GetActiveVessel',
          parameters: [],
          returnType: { code: 100, service: 'SpaceCenter', name: 'Vessel', types: [] },
          returnIsNullable: false,
        },
        {
          name: 'GetBodies',
          parameters: [],
          returnType: {
            code: 301,
            service: '',
            name: '',
            types: [{ code: 100, service: 'SpaceCenter', name: 'CelestialBody', types: [] }],
          },
          returnIsNullable: false,
        },
        {
          name: 'CelestialBody.GetName',
          parameters: [
            {
              name: 'self',
              type: { code: 100, service: 'SpaceCenter', name: 'CelestialBody', types: [] },
              nullable: false,
            },
          ],
          returnType: { code: 8, service: '', name: '', types: [] },
          returnIsNullable: false,
        },
        {
          name: 'Vessel.GetName',
          parameters: [
            {
              name: 'self',
              type: { code: 100, service: 'SpaceCenter', name: 'Vessel', types: [] },
              nullable: false,
            },
          ],
          returnType: { code: 8, service: '', name: '', types: [] },
          returnIsNullable: false,
        },
        {
          name: 'Vessel.GetType',
          parameters: [
            {
              name: 'self',
              type: { code: 100, service: 'SpaceCenter', name: 'Vessel', types: [] },
              nullable: false,
            },
          ],
          returnType: { code: 101, service: 'SpaceCenter', name: 'VesselType', types: [] },
          returnIsNullable: false,
        },
      ],
      classes: [{ name: 'CelestialBody' }, { name: 'Vessel' }, { name: 'Orbit' }],
      enumerations: [
        {
          name: 'VesselType',
          values: [
            { name: 'Ship', value: 0 },
            { name: 'Probe', value: 3 },
          ],
        },
      ],
    },
  ],
};

async function freePort(): Promise<number> {
  return new Promise<number>((resolve, reject) => {
    const srv = net.createServer();
    srv.unref();
    srv.on('error', reject);
    srv.listen(0, '127.0.0.1', () => {
      const addr = srv.address();
      if (addr && typeof addr === 'object') {
        const p = addr.port;
        srv.close(() => resolve(p));
      } else {
        srv.close(() => reject(new Error('no addr')));
      }
    });
  });
}

function replyWithValue(socket: net.Socket, value: Uint8Array): void {
  const resultMsg = KRPC.ProcedureResult.create({
    value: Buffer.from(value),
  });
  const respMsg = KRPC.Response.create({ results: [resultMsg] });
  const payload = Buffer.from(KRPC.Response.encode(respMsg).finish());
  sendMessage(socket, KRPC.Response, { results: [resultMsg] });
  void payload;
}

function replyWithError(socket: net.Socket, name: string, description: string): void {
  const errMsg = KRPC.Error.create({ service: 'Test', name, description });
  const resultMsg = KRPC.ProcedureResult.create({ error: errMsg, value: Buffer.from([]) });
  sendMessage(socket, KRPC.Response, { results: [resultMsg] });
}

interface MockServer {
  rpcPort: number;
  streamPort: number;
  close: () => Promise<void>;
  stub: (
    service: string,
    procedure: string,
    handler: (args: Uint8Array[]) => Uint8Array,
  ) => void;
  /** Counts the number of calls received for (service, procedure). */
  callCount: (service: string, procedure: string) => number;
  /** Captured arguments of the last call to (service, procedure). */
  lastArgs: (service: string, procedure: string) => Uint8Array[];
}

async function startMockServer(): Promise<MockServer> {
  const rpcPort = await freePort();
  const streamPort = await freePort();
  const stubs = new Map<string, (args: Uint8Array[]) => Uint8Array>();
  const counts = new Map<string, number>();
  const lastArgsMap = new Map<string, Uint8Array[]>();

  const stub = (svc: string, proc: string, handler: (args: Uint8Array[]) => Uint8Array) => {
    stubs.set(`${svc}.${proc}`, handler);
  };

  // Default stubs
  stub('KRPC', 'GetServices', () => {
    const servicesMsg = KRPC.Services.create(MOCK_SERVICES_RAW);
    return new Uint8Array(KRPC.Services.encode(servicesMsg).finish());
  });
  stub('KRPC', 'GetStatus', () => {
    const statusMsg = KRPC.Status.create({ version: 'test' });
    return new Uint8Array(KRPC.Status.encode(statusMsg).finish());
  });
  stub('SpaceCenter', 'GetUT', () => encodeDouble(4_700_000.5));

  // ── RPC server ─────────────────────────────────────────────────────
  const rpcServer = net.createServer((socket) => {
    void (async () => {
      try {
        // 1. Handshake
        const req = await recvMessage<{ type: number; clientName: string }>(
          socket,
          KRPC.ConnectionRequest,
        );
        if (req.type !== 0) {
          socket.destroy();
          return;
        }
        sendMessage(socket, KRPC.ConnectionResponse, {
          status: 0,
          message: '',
          clientIdentifier: Buffer.from('test-client'),
        });

        // 2. Procedure loop
        while (!socket.destroyed) {
          let callReq: {
            calls: { service: string; procedure: string; arguments?: { value: Uint8Array }[] }[];
          };
          try {
            callReq = await recvMessage(socket, KRPC.Request);
          } catch {
            return;
          }
          const call = callReq.calls[0];
          if (!call) {
            replyWithError(socket, 'Malformed', 'no call');
            continue;
          }
          const key = `${call.service}.${call.procedure}`;
          counts.set(key, (counts.get(key) ?? 0) + 1);
          const args = (call.arguments ?? []).map((a) => new Uint8Array(a.value));
          lastArgsMap.set(key, args);
          const handler = stubs.get(key);
          if (!handler) {
            replyWithError(socket, 'UnknownProcedure', `${key} not stubbed`);
            continue;
          }
          try {
            const value = handler(args);
            replyWithValue(socket, value);
          } catch (e) {
            replyWithError(socket, 'HandlerError', String(e));
          }
        }
      } catch {
        if (!socket.destroyed) socket.destroy();
      }
    })();
  });
  await new Promise<void>((resolve) => rpcServer.listen(rpcPort, '127.0.0.1', resolve));

  // ── Stream server (handshake only, then idle) ─────────────────────
  const streamServer = net.createServer((socket) => {
    void (async () => {
      try {
        const req = await recvMessage<{ type: number; clientIdentifier: Uint8Array }>(
          socket,
          KRPC.ConnectionRequest,
        );
        if (req.type !== 1) {
          socket.destroy();
          return;
        }
        sendMessage(socket, KRPC.ConnectionResponse, { status: 0, message: '' });
        // Keep alive
        await new Promise(() => undefined);
      } catch {
        socket.destroy();
      }
    })();
  });
  await new Promise<void>((resolve) => streamServer.listen(streamPort, '127.0.0.1', resolve));

  return {
    rpcPort,
    streamPort,
    stub,
    callCount: (svc, proc) => counts.get(`${svc}.${proc}`) ?? 0,
    lastArgs: (svc, proc) => lastArgsMap.get(`${svc}.${proc}`) ?? [],
    close: () =>
      new Promise<void>((resolve) => {
        rpcServer.close(() => {
          streamServer.close(() => resolve());
        });
      }),
  };
}

describe('KrpcServices against a mock kRPC server', () => {
  let server: MockServer;
  let client: KRPCClient;
  let services: KrpcServices;

  beforeAll(async () => {
    server = await startMockServer();
    client = new KRPCClient({
      host: '127.0.0.1',
      rpcPort: server.rpcPort,
      streamPort: server.streamPort,
      clientName: 'test',
    });
    await client.connect();
    const loaded = await loadServices(client);
    services = loaded.services;
  }, 10_000);

  afterAll(async () => {
    await client.close();
    await server.close();
  });

  it('lists services', () => {
    const names = services.getCache().serviceNames();
    expect(names).toContain('KRPC');
    expect(names).toContain('SpaceCenter');
  });

  it('looks up the GetUT procedure', () => {
    const r = services.getCache().lookup('SpaceCenter', 'GetUT');
    expect(r.found).toBe(true);
  });

  it('invokes SpaceCenter.GetUT and decodes as double', async () => {
    const ut = await services.invoke<number>('SpaceCenter', 'GetUT');
    expect(ut).toBe(4_700_000.5);
  });

  it('invokes SpaceCenter.GetBodies and decodes a list of class ids', async () => {
    server.stub('SpaceCenter', 'GetBodies', () => {
      const items = [encodeUint64(1n), encodeUint64(2n), encodeUint64(3n)];
      const listMsg = KRPC.List.create({ items });
      return new Uint8Array(KRPC.List.encode(listMsg).finish());
    });
    const ids = await services.invoke<bigint[]>('SpaceCenter', 'GetBodies');
    expect(ids).toEqual([1n, 2n, 3n]);
  });

  it('invokes a class method and decodes the string return', async () => {
    server.stub('SpaceCenter', 'CelestialBody.GetName', () => encodeString('Kerbin'));
    const name = await services.invoke<string>('SpaceCenter', 'CelestialBody.GetName', 42n);
    expect(name).toBe('Kerbin');
  });

  it('rejects an unknown procedure', async () => {
    await expect(
      services.invoke('SpaceCenter', 'GetNothing'),
    ).rejects.toThrow(/procedure not found/);
  });

  it('rejects a wrong number of arguments', async () => {
    await expect(
      services.invoke('SpaceCenter', 'GetBodies', 1n, 2n),
    ).rejects.toThrow(/wrong number of arguments/);
  });

  it('decodes an enum return value', async () => {
    server.stub('SpaceCenter', 'Vessel.GetType', () => encodeSint32(3));
    const t = await services.invoke<number>('SpaceCenter', 'Vessel.GetType', 7n);
    expect(t).toBe(3);
  });

  it('passes BigInt class ids through to class methods', async () => {
    let receivedId: bigint | null = null;
    server.stub('SpaceCenter', 'CelestialBody.GetName', (args) => {
      const argBytes = args[0] ?? new Uint8Array();
      let v = 0n;
      let shift = 0n;
      for (const b of argBytes) {
        v |= BigInt(b & 0x7f) << shift;
        if ((b & 0x80) === 0) break;
        shift += 7n;
      }
      receivedId = v;
      return encodeString('Mun');
    });
    await services.invoke<string>('SpaceCenter', 'CelestialBody.GetName', 99n);
    expect(receivedId).toBe(99n);
  });
});