KSP-MissionControl/packages/krpc-client/src/decoder.ts

/**
 * kRPC value decoder.
 *
 * kRPC values are encoded on the wire using a hybrid scheme:
 *
 *  - For **primitive types** (DOUBLE, FLOAT, SINT32, SINT64, UINT32, UINT64,
 *    BOOL, STRING, BYTES) the bytes are exactly the standard protobuf wire
 *    encoding of that single value, NOT wrapped in a message. So a `double`
 *    is just 8 little-endian bytes, a `string` is `[varint length][utf8]`,
 *    and so on.
 *
 *  - For **CLASS types** the bytes are a single varint-encoded `uint64` —
 *    the object id. An object id of 0 means `None` (the CLASS is nullable).
 *    The actual class data lives server-side; to get a property you call
 *    `Service.ClassName.GetX(id)`.
 *
 *  - For **ENUMERATION** the bytes are a single signed varint (zigzag-encoded
 *    sint32 in protobuf terms).
 *
 *  - For **collections** (LIST, SET, TUPLE, DICTIONARY) the bytes are a
 *    serialized `KRPC.List` / `KRPC.Set` / `KRPC.Tuple` / `KRPC.Dictionary`
 *    message. Each element is itself encoded using the scheme above (so the
 *    element bytes are variable-length). A null collection is a single byte
 *    `\x00` (NOT a length-prefixed empty list).
 *
 *  - For **system messages** (STATUS, SERVICES, STREAM, EVENT,
 *    PROCEDURE_CALL) the bytes are the standard protobuf serialization of
 *    the corresponding KRPC message.
 *
 * Reference: the Python client's `krpc/decoder.py` (Krpc 0.5.x).
 */
import protobuf from 'protobufjs';
import { decodeVarint, encodeVarint } from './connection.js';
import { TypeCode, type KrpcType, typeName } from './types.js';

/** Result of decoding a value. */
export type DecodedValue =
  | number
  | bigint
  | boolean
  | string
  | Uint8Array
  | null
  | DecodedValue[]
  | Set<DecodedValue>
  | Map<DecodedValue, DecodedValue>
  | { [k: string]: unknown };

/**
 * Decode a value from its wire bytes according to a KrpcType.
 *
 * @param type       The KrpcType descriptor (from GetServices or a
 *                   pre-built cache).
 * @param data       The raw bytes (response.value for returns, or
 *                   the value field of a KRPC.Argument for args).
 * @param messageTypes Optional protobufjs type registry for decoding
 *                   system messages (KRPC.Status, etc.) by name.
 *                   Only needed for MESSAGE-style TypeCodes.
 */
export function decodeValue(
  type: KrpcType,
  data: Uint8Array,
  messageTypes?: Record<string, protobuf.Type>,
): DecodedValue {
  switch (type.code) {
    case TypeCode.NONE:
      return null;

    case TypeCode.DOUBLE:
      return decodeDouble(data);
    case TypeCode.FLOAT:
      return decodeFloat(data);
    case TypeCode.SINT32:
      return decodeSint32(data);
    case TypeCode.SINT64:
      return decodeSint64(data);
    case TypeCode.UINT32:
      return decodeUint32(data);
    case TypeCode.UINT64:
      return decodeUint64(data);
    case TypeCode.BOOL:
      return decodeBool(data);
    case TypeCode.STRING:
      return decodeString(data);
    case TypeCode.BYTES:
      return decodeBytes(data);

    case TypeCode.CLASS: {
      // The wire form of a CLASS is a uint64 object id; 0 = None.
      // We decode as a BigInt so callers can pass the id back as an
      // argument to other class methods.
      const id = bigFromVarint(data);
      return id === 0n ? null : id;
    }

    case TypeCode.ENUMERATION: {
      // Wire form: a single signed varint (sint32 in protobuf terms).
      return decodeSint32(data);
    }

    case TypeCode.STATUS: {
      return decodeSystemMessage('Status', data, messageTypes);
    }
    case TypeCode.SERVICES: {
      return decodeSystemMessage('Services', data, messageTypes);
    }
    case TypeCode.STREAM: {
      return decodeSystemMessage('Stream', data, messageTypes);
    }
    case TypeCode.EVENT: {
      return decodeSystemMessage('Event', data, messageTypes);
    }
    case TypeCode.PROCEDURE_CALL: {
      return decodeSystemMessage('ProcedureCall', data, messageTypes);
    }

    case TypeCode.LIST:
      return decodeList(type, data, messageTypes);
    case TypeCode.SET:
      return decodeSet(type, data, messageTypes);
    case TypeCode.TUPLE:
      return decodeTuple(type, data, messageTypes);
    case TypeCode.DICTIONARY:
      return decodeDictionary(type, data, messageTypes);

    default:
      throw new Error(`unknown TypeCode ${type.code} (${typeName(type)})`);
  }
}

// ── primitive decoders ──────────────────────────────────────────────────────

/**
 * Read a fixed-width little-endian IEEE 754 number from a buffer.
 *
 * JavaScript's `DataView.getFloat64()` already does the right thing on
 * little-endian platforms, which Node always is. We still go through
 * `DataView` so the intent is explicit and the code is portable to
 * big-endian platforms (if we ever run on one).
 */
function readFloatLE(buf: Uint8Array, offset: number, bytes: 4 | 8): number {
  // Pad short reads with zero bytes. This shouldn't happen in practice,
  // but it's defensive against a truncated response.
  if (buf.length < offset + bytes) {
    const padded = new Uint8Array(bytes);
    padded.set(buf.subarray(offset, offset + bytes));
    buf = padded;
    offset = 0;
  }
  const view = new DataView(buf.buffer, buf.byteOffset + offset, bytes);
  return bytes === 8 ? view.getFloat64(0, true) : view.getFloat32(0, true);
}

export function decodeDouble(data: Uint8Array): number {
  return readFloatLE(data, 0, 8);
}

export function decodeFloat(data: Uint8Array): number {
  return readFloatLE(data, 0, 4);
}

export function decodeSint32(data: Uint8Array): number {
  // Protobuf sint32 is zigzag-encoded. `decodeVarint` returns a regular
  // varint, so we need the zigzag step too.
  const [raw] = decodeVarint(Buffer.from(data));
  return zigzagDecode(Number(raw));
}

export function decodeSint64(data: Uint8Array): number {
  // The Python client decodes sint64 to a plain int (BigInt in their
  // case is a separate branch). For our use case the values we care
  // about (enum cases) are well within int32 range, so we decode to
  // a JS number. If you really need full int64, decode to BigInt.
  return decodeSint32(data);
}

export function decodeUint32(data: Uint8Array): number {
  const [v] = decodeVarint(Buffer.from(data));
  return Number(v);
}

export function decodeUint64(data: Uint8Array): bigint {
  // Use BigInt for the 64-bit case. The kRPC ObjectId and StreamId are
  // uint64s and can exceed Number.MAX_SAFE_INTEGER in principle
  // (though kRPC never generates ids that large in practice).
  return bigFromVarint(data);
}

export function decodeBool(data: Uint8Array): boolean {
  if (data.length < 1) return false;
  return data[0] !== 0;
}

export function decodeString(data: Uint8Array): string {
  // Wire form: [varint length][utf8 bytes].
  const buf = Buffer.from(data);
  const [len, pos] = decodeVarint(buf);
  return buf.subarray(pos, pos + Number(len)).toString('utf-8');
}

export function decodeBytes(data: Uint8Array): Uint8Array {
  const buf = Buffer.from(data);
  const [len, pos] = decodeVarint(buf);
  return new Uint8Array(buf.subarray(pos, pos + Number(len)));
}

function zigzagDecode(n: number): number {
  return (n >>> 1) ^ -(n & 1);
}

/**
 * Decode a varint (assumed ≤ 64 bits) as a BigInt. Used for uint64 values
 * where the precision loss of Number() is unacceptable.
 */
function bigFromVarint(data: Uint8Array): bigint {
  let result = 0n;
  let shift = 0n;
  for (let i = 0; i < data.length; i++) {
    const b = data[i];
    if (b === undefined) break;
    result |= BigInt(b & 0x7f) << shift;
    if ((b & 0x80) === 0) return result;
    shift += 7n;
  }
  return result;
}

// ── collection decoders ─────────────────────────────────────────────────────

/**
 * Decode a KRPC.List message (when its serialized form is provided).
 * Used internally by `decodeList`. Also exported for tests.
 */
export function decodeKrpcList(
  data: Uint8Array,
  messageType: protobuf.Type,
): Uint8Array[] {
  if (data.length === 1 && data[0] === 0) {
    // A list may be serialized as a single 0x00 byte to indicate None.
    return [];
  }
  const msg = messageType.decode(data) as unknown as { items: Uint8Array[] };
  return msg.items ?? [];
}

export function decodeList(
  type: KrpcType,
  data: Uint8Array,
  messageTypes?: Record<string, protobuf.Type>,
): DecodedValue[] {
  if (data.length === 1 && data[0] === 0) return null as unknown as DecodedValue[];
  const listType = messageTypes?.['List'];
  if (!listType) {
    throw new Error('decodeList: no protobufjs type for KRPC.List registered');
  }
  const items = decodeKrpcList(data, listType);
  const elemType = type.types[0];
  if (!elemType) throw new Error('LIST type missing element type');
  return items.map((it) => decodeValue(elemType, it, messageTypes));
}

export function decodeSet(
  type: KrpcType,
  data: Uint8Array,
  messageTypes?: Record<string, protobuf.Type>,
): Set<DecodedValue> {
  if (data.length === 1 && data[0] === 0) return null as unknown as Set<DecodedValue>;
  const setType = messageTypes?.['Set'];
  if (!setType) {
    throw new Error('decodeSet: no protobufjs type for KRPC.Set registered');
  }
  const msg = setType.decode(data) as unknown as { items: Uint8Array[] };
  const elemType = type.types[0];
  if (!elemType) throw new Error('SET type missing element type');
  return new Set((msg.items ?? []).map((it) => decodeValue(elemType, it, messageTypes)));
}

export function decodeTuple(
  type: KrpcType,
  data: Uint8Array,
  messageTypes?: Record<string, protobuf.Type>,
): DecodedValue[] {
  if (data.length === 1 && data[0] === 0) return null as unknown as DecodedValue[];
  const tupleType = messageTypes?.['Tuple'];
  if (!tupleType) {
    throw new Error('decodeTuple: no protobufjs type for KRPC.Tuple registered');
  }
  const msg = tupleType.decode(data) as unknown as { items: Uint8Array[] };
  return type.types.map((inner, i) =>
    decodeValue(inner, msg.items[i] ?? new Uint8Array(0), messageTypes),
  );
}

export function decodeDictionary(
  type: KrpcType,
  data: Uint8Array,
  messageTypes?: Record<string, protobuf.Type>,
): Map<DecodedValue, DecodedValue> {
  if (data.length === 1 && data[0] === 0) {
    return null as unknown as Map<DecodedValue, DecodedValue>;
  }
  const dictType = messageTypes?.['Dictionary'];
  if (!dictType) {
    throw new Error('decodeDictionary: no protobufjs type for KRPC.Dictionary registered');
  }
  const msg = dictType.decode(data) as unknown as {
    entries: { key: Uint8Array; value: Uint8Array }[];
  };
  const keyType = type.types[0];
  const valType = type.types[1];
  if (!keyType || !valType) throw new Error('DICTIONARY type missing key/value types');
  const out = new Map<DecodedValue, DecodedValue>();
  for (const e of msg.entries ?? []) {
    out.set(
      decodeValue(keyType, e.key, messageTypes),
      decodeValue(valType, e.value, messageTypes),
    );
  }
  return out;
}

function decodeSystemMessage(
  name: string,
  data: Uint8Array,
  messageTypes?: Record<string, protobuf.Type>,
): { [k: string]: unknown } {
  const t = messageTypes?.[name];
  if (!t) throw new Error(`decodeSystemMessage: no type registered for ${name}`);
  return t.decode(data) as unknown as { [k: string]: unknown };
}

// ── encoders (mirrors, for sending arguments) ───────────────────────────────

/**
 * Encode a value into the kRPC wire format for `type`.
 *
 * This is the inverse of `decodeValue`. Used by the service client to
 * serialize procedure arguments. Collections and system messages use the
 * protobufjs registry the same way the decoder does.
 */
export function encodeValue(
  type: KrpcType,
  value: DecodedValue,
  messageTypes?: Record<string, protobuf.Type>,
): Uint8Array {
  switch (type.code) {
    case TypeCode.NONE:
      return new Uint8Array(0);
    case TypeCode.DOUBLE:
      return encodeDouble(value as number);
    case TypeCode.FLOAT:
      return encodeFloat(value as number);
    case TypeCode.SINT32:
      return encodeSint32(value as number);
    case TypeCode.SINT64:
      return encodeSint32(value as number);
    case TypeCode.UINT32:
      return encodeUint32(value as number);
    case TypeCode.UINT64:
      return encodeUint64(value as bigint);
    case TypeCode.BOOL:
      return encodeBool(value as boolean);
    case TypeCode.STRING:
      return encodeString(value as string);
    case TypeCode.BYTES:
      return encodeBytes(value as Uint8Array);

    case TypeCode.CLASS: {
      // CLASS args are encoded as uint64 object id; null = 0.
      if (value === null || value === undefined) return encodeUint64(0n);
      if (typeof value === 'bigint') return encodeUint64(value);
      if (typeof value === 'number') return encodeUint64(BigInt(value));
      throw new Error(`encodeValue: CLASS expects bigint object id, got ${typeof value}`);
    }

    case TypeCode.ENUMERATION:
      return encodeSint32(value as number);

    case TypeCode.LIST: {
      const listType = messageTypes?.['List'];
      if (!listType) throw new Error('encodeValue: no type for KRPC.List');
      const elemType = type.types[0];
      if (!elemType) throw new Error('LIST type missing element type');
      const items = (value as DecodedValue[]).map((v) => encodeValue(elemType, v, messageTypes));
      const msg = listType.create({ items });
      return listType.encode(msg).finish();
    }
    case TypeCode.TUPLE: {
      const tupleType = messageTypes?.['Tuple'];
      if (!tupleType) throw new Error('encodeValue: no type for KRPC.Tuple');
      const items = (value as DecodedValue[]).map((v, i) =>
        encodeValue(type.types[i] as KrpcType, v, messageTypes),
      );
      const msg = tupleType.create({ items });
      return tupleType.encode(msg).finish();
    }
    case TypeCode.DICTIONARY: {
      const dictType = messageTypes?.['Dictionary'];
      if (!dictType) throw new Error('encodeValue: no type for KRPC.Dictionary');
      const keyType = type.types[0];
      const valType = type.types[1];
      if (!keyType || !valType) throw new Error('DICTIONARY type missing key/value types');
      const entries: { key: Uint8Array; value: Uint8Array }[] = [];
      for (const [k, v] of (value as Map<DecodedValue, DecodedValue>).entries()) {
        entries.push({
          key: encodeValue(keyType, k, messageTypes),
          value: encodeValue(valType, v, messageTypes),
        });
      }
      const msg = dictType.create({ entries });
      return dictType.encode(msg).finish();
    }
    case TypeCode.SET: {
      // kRPC 0.5 doesn't really use SET much; if needed we'd encode as
      // KRPC.Set. For now, only LIST/TUPLE/DICT are exercised by our
      // SpaceCenter calls.
      throw new Error('encodeValue: SET not implemented (kRPC 0.5 does not use it)');
    }
    case TypeCode.STATUS:
    case TypeCode.SERVICES:
    case TypeCode.STREAM:
    case TypeCode.EVENT:
    case TypeCode.PROCEDURE_CALL:
      throw new Error(`encodeValue: system message ${typeName(type)} cannot be sent as argument`);

    default:
      throw new Error(`encodeValue: unknown TypeCode ${type.code}`);
  }
}

function encodeDouble(v: number): Uint8Array {
  const out = new Uint8Array(8);
  new DataView(out.buffer).setFloat64(0, v, true);
  return out;
}

function encodeFloat(v: number): Uint8Array {
  const out = new Uint8Array(4);
  new DataView(out.buffer).setFloat32(0, v, true);
  return out;
}

function encodeSint32(v: number): Uint8Array {
  return encodeVarint(zigzagEncode(v));
}

function encodeUint32(v: number): Uint8Array {
  return encodeVarint(v);
}

function encodeUint64(v: bigint): Uint8Array {
  // Manual varint encoding for BigInt to avoid Number truncation.
  const out: number[] = [];
  let x = v;
  while (x >= 0x80n) {
    out.push(Number((x & 0x7fn) | 0x80n));
    x >>= 7n;
  }
  out.push(Number(x));
  return new Uint8Array(out);
}

function encodeBool(v: boolean): Uint8Array {
  return new Uint8Array([v ? 1 : 0]);
}

function encodeString(v: string): Uint8Array {
  const utf8 = new TextEncoder().encode(v);
  return Buffer.concat([encodeVarint(utf8.length), Buffer.from(utf8)]);
}

function encodeBytes(v: Uint8Array): Uint8Array {
  return Buffer.concat([encodeVarint(v.length), Buffer.from(v)]);
}

function zigzagEncode(n: number): number {
  return (n << 1) ^ (n >> 31);
}