CloudSSH/tests/ssh/kex.test.ts

import { describe, expect, it } from 'vitest';
import {
  SUPPORTED_ENCRYPTION_ALGORITHMS,
  SUPPORTED_KEX_ALGORITHMS,
  SUPPORTED_MAC_ALGORITHMS,
} from '../../src/ssh/algorithms';
import {
  filterExtInfo,
  KEXInitBuilder,
  negotiate,
  parseKEXInit,
  parseServerSigAlgs,
} from '../../src/ssh/kex';
import { SSH_MSG_EXT_INFO, SSH_MSG_KEXINIT } from '../../src/types';

// =====================================================================
// kex.test.ts
// ---------------------------------------------------------------
// KEXINIT 报文的构造/解析 + 算法协商函数。
// SSH 握手第一阶段就是 KEXINIT，长度前缀偏移、cookie 长度、列表数量
// 任何一个错位都会让整个握手失败。这里用 build → parse 往返测试固化
// 报文格式的稳定性，再用 negotiate 验证"客户端优先序"语义。
// =====================================================================

describe('kex — KEXInitBuilder.build()', () => {
  it('返回非空 Uint8Array', () => {
    const packet = KEXInitBuilder.build();
    expect(packet).toBeInstanceOf(Uint8Array);
    expect(packet.length).toBeGreaterThan(0);
  });

  it('首字节为 SSH_MSG_KEXINIT (=20)', () => {
    const packet = KEXInitBuilder.build();
    expect(packet[0]).toBe(SSH_MSG_KEXINIT);
    expect(SSH_MSG_KEXINIT).toBe(20);
  });

  it('cookie 占 16 字节（紧跟在消息类型后）', () => {
    const packet = KEXInitBuilder.build();
    const cookie = packet.subarray(1, 17);
    expect(cookie.length).toBe(16);
  });

  it('两次 build 的 cookie 应不同（随机生成）', () => {
    // cookie 是防重放保护，每次构造都应随机生成
    const p1 = KEXInitBuilder.build();
    const p2 = KEXInitBuilder.build();
    const c1 = p1.subarray(1, 17);
    const c2 = p2.subarray(1, 17);
    // 极小概率相等（16 字节随机），实际不会触发
    let same = true;
    for (let i = 0; i < 16; i++) {
      if (c1[i] !== c2[i]) {
        same = false;
        break;
      }
    }
    expect(same).toBe(false);
  });

  it('总长度可被解析且不抛错（结构完整）', () => {
    const packet = KEXInitBuilder.build();
    expect(() => parseKEXInit(packet)).not.toThrow();
  });
});

describe('kex — build → parse 往返', () => {
  // 这一组是本文件最关键的测试：用 build 构造，再 parse 回来，
  // 各算法列表应该一一对应于 algorithms.ts 中的常量声明。
  it('kexAlgorithms 解析后等于 SUPPORTED_KEX_ALGORITHMS（含 ext-info-c 前缀）', () => {
    const packet = KEXInitBuilder.build();
    const msg = parseKEXInit(packet);
    // RFC 8301: ext-info-c 作为伪算法名插入到 kex 算法列表最前面
    expect(msg.kexAlgorithms).toEqual(['ext-info-c', ...SUPPORTED_KEX_ALGORITHMS]);
  });
  it('hostKeyAlgorithms 包含 Ed25519 / ECDSA(P-256/P-384/P-521) / RSA-SHA2', () => {
    const packet = KEXInitBuilder.build();
    const msg = parseKEXInit(packet);
    expect(msg.hostKeyAlgorithms).toContain('ssh-ed25519');
    expect(msg.hostKeyAlgorithms).toContain('ecdsa-sha2-nistp256');
    expect(msg.hostKeyAlgorithms).toContain('ecdsa-sha2-nistp384');
    expect(msg.hostKeyAlgorithms).toContain('ecdsa-sha2-nistp521');
    expect(msg.hostKeyAlgorithms).toContain('rsa-sha2-512');
    expect(msg.hostKeyAlgorithms).toContain('rsa-sha2-256');
    expect(msg.hostKeyAlgorithms).toContain('ssh-rsa');
  });
  it('encryptionC2S / S2C 解析后等于 SUPPORTED_ENCRYPTION_ALGORITHMS', () => {
    const packet = KEXInitBuilder.build();
    const msg = parseKEXInit(packet);
    expect(msg.encryptionC2S).toEqual(SUPPORTED_ENCRYPTION_ALGORITHMS);
    expect(msg.encryptionS2C).toEqual(SUPPORTED_ENCRYPTION_ALGORITHMS);
  });
  it('macC2S / S2C 解析后等于 SUPPORTED_MAC_ALGORITHMS', () => {
    const packet = KEXInitBuilder.build();
    const msg = parseKEXInit(packet);
    expect(msg.macC2S).toEqual(SUPPORTED_MAC_ALGORITHMS);
    expect(msg.macS2C).toEqual(SUPPORTED_MAC_ALGORITHMS);
  });
  it('compressionC2S / S2C 均为 ["none"]', () => {
    const packet = KEXInitBuilder.build();
    const msg = parseKEXInit(packet);
    expect(msg.compressionC2S).toEqual(['none']);
    expect(msg.compressionS2C).toEqual(['none']);
  });
});

describe('kex — parseKEXInit 手工构造报文', () => {
  // 直接 build 一个随机 cookie 的报文只能验证"我们的 build 我们能 parse"，
  // 这里手工构造一个简化报文，验证 parse 能正确处理来自对方的 KEXINIT，
  // 包括长度前缀偏移、空列表等边界。
  function buildCustomPacket(opts: {
    kex: string[];
    hostKey: string[];
    encC2S: string[];
    encS2C: string[];
    macC2S: string[];
    macS2C: string[];
    compC2S: string[];
    compS2C: string[];
  }): Uint8Array {
    const parts: Uint8Array[] = [];
    parts.push(new Uint8Array([SSH_MSG_KEXINIT]));
    // 固定 cookie 16 字节为 0xAA（便于断言对照）
    parts.push(new Uint8Array(16).fill(0xaa));
    const lists = [
      opts.kex,
      opts.hostKey,
      opts.encC2S,
      opts.encS2C,
      opts.macC2S,
      opts.macS2C,
      opts.compC2S,
      opts.compS2C,
      // 第一语言、第二语言（项目里写死为空字符串，用 [] 模拟空列表）
      [],
      [],
    ];
    for (const list of lists) {
      const s = list.join(',');
      const enc = new TextEncoder().encode(s);
      const len = new Uint8Array(4);
      new DataView(len.buffer).setUint32(0, enc.length, false);
      parts.push(len);
      parts.push(enc);
    }
    parts.push(new Uint8Array([0])); // first_kex_packet_follows
    parts.push(new Uint8Array(4)); // reserved
    // concat
    let total = 0;
    for (const p of parts) total += p.length;
    const out = new Uint8Array(total);
    let off = 0;
    for (const p of parts) {
      out.set(p, off);
      off += p.length;
    }
    return out;
  }

  it('正确解析含 3 个元素的 kex 列表', () => {
    const pkt = buildCustomPacket({
      kex: ['curve25519-sha256', 'ecdh-sha2-nistp256', 'diffie-hellman-group14-sha256'],
      hostKey: ['ssh-ed25519'],
      encC2S: ['aes256-gcm@openssh.com', 'aes128-ctr'],
      encS2C: ['aes256-gcm@openssh.com'],
      macC2S: ['hmac-sha2-256'],
      macS2C: ['hmac-sha2-512'],
      compC2S: ['none'],
      compS2C: ['none'],
    });
    const msg = parseKEXInit(pkt);
    expect(msg.kexAlgorithms).toEqual([
      'curve25519-sha256',
      'ecdh-sha2-nistp256',
      'diffie-hellman-group14-sha256',
    ]);
    expect(msg.encryptionC2S).toEqual(['aes256-gcm@openssh.com', 'aes128-ctr']);
    expect(msg.encryptionS2C).toEqual(['aes256-gcm@openssh.com']);
    expect(msg.macC2S).toEqual(['hmac-sha2-256']);
    expect(msg.macS2C).toEqual(['hmac-sha2-512']);
  });

  it('正确解析空列表（应为 [""] 而非 []）', () => {
    // 注意 parseKEXInit 用 split(',') 解析空字符串会得到 ['']，这是项目当前行为；
    // 这个用例固化该行为，避免未来"修复"成 [] 而误伤协商逻辑
    const pkt = buildCustomPacket({
      kex: ['curve25519-sha256'],
      hostKey: ['ssh-ed25519'],
      encC2S: ['aes256-gcm@openssh.com'],
      encS2C: ['aes256-gcm@openssh.com'],
      macC2S: ['hmac-sha2-256'],
      macS2C: ['hmac-sha2-256'],
      compC2S: [],
      compS2C: [],
    });
    const msg = parseKEXInit(pkt);
    // 空列表序列化为空字符串，反序列化得到 [''] —— 这是当前实现行为，固化为不变性
    expect(msg.compressionC2S).toEqual(['']);
    expect(msg.compressionS2C).toEqual(['']);
  });
});

describe('kex — negotiate', () => {
  it('返回客户端列表中第一个服务端也支持的算法', () => {
    const client = ['curve25519-sha256', 'ecdh-sha2-nistp256'];
    const server = ['ecdh-sha2-nistp256', 'diffie-hellman-group14-sha256'];
    // 客户端第一项服务端没有，第二项有
    expect(negotiate(client, server)).toBe('ecdh-sha2-nistp256');
  });

  it('客户端优先序生效：优先返回客户端靠前的项', () => {
    const client = ['curve25519-sha256', 'ecdh-sha2-nistp256'];
    const server = ['ecdh-sha2-nistp256', 'curve25519-sha256'];
    // 两个都支持，但 curve25519 在客户端排第一，必须返回它
    expect(negotiate(client, server)).toBe('curve25519-sha256');
  });

  it('仅一个共同算法时返回该算法', () => {
    expect(negotiate(['a', 'b'], ['b'])).toBe('b');
  });

  it('无共同算法时抛错', () => {
    expect(() => negotiate(['a', 'b'], ['c', 'd'])).toThrow(/No common/);
  });

  it('抛错信息包含双方算法列表（便于排查握手失败）', () => {
    try {
      negotiate(['a'], ['b', 'c']);
      throw new Error('should have thrown');
    } catch (e) {
      const msg = (e as Error).message;
      expect(msg).toContain('a');
      expect(msg).toContain('b');
      expect(msg).toContain('c');
    }
  });

  it('默认 category 为 algorithm，抛错信息体现该词', () => {
    expect(() => negotiate(['a'], ['b'])).toThrow(/No common algorithm/);
  });

  it('自定义 category（如 cipher）时抛错信息体现该词', () => {
    expect(() => negotiate(['a'], ['b'], 'cipher')).toThrow(/No common cipher/);
  });

  it('客户端为空列表时抛错（边界）', () => {
    expect(() => negotiate([], ['a'])).toThrow(/No common/);
  });

  it('服务端为空列表时抛错（边界）', () => {
    expect(() => negotiate(['a'], [])).toThrow(/No common/);
  });

  it('双方都为空时抛错', () => {
    expect(() => negotiate([], [])).toThrow(/No common/);
  });

  it('相同列表时返回首项', () => {
    expect(negotiate(['a', 'b', 'c'], ['a', 'b', 'c'])).toBe('a');
  });

  it('真实场景：客户端 [curve25519, ecdh] vs 服务端 [diffie-hellman, curve25519]', () => {
    // 服务端虽把 diffie-hellman 排在首位，但客户端优先 curve25519，
    // 两者都有 curve25519，按 RFC 8732 应优先客户端序
    expect(
      negotiate(
        ['curve25519-sha256', 'ecdh-sha2-nistp256'],
        ['diffie-hellman-group14-sha256', 'curve25519-sha256']
      )
    ).toBe('curve25519-sha256');
  });
});

describe('kex — parseKEXInit 边界与错误', () => {
  it('对最小合法报文（仅类型 + cookie + 10 个长度为 0 的列表 + 尾部）不抛错', () => {
    // 构造一个所有列表都是空字符串的最小 KEXINIT
    const parts: Uint8Array[] = [];
    parts.push(new Uint8Array([SSH_MSG_KEXINIT]));
    parts.push(new Uint8Array(16)); // cookie 全 0
    for (let i = 0; i < 10; i++) {
      parts.push(new Uint8Array(4)); // 4 字节 0 = 空列表
    }
    parts.push(new Uint8Array([0]));
    parts.push(new Uint8Array(4));
    let total = 0;
    for (const p of parts) total += p.length;
    const out = new Uint8Array(total);
    let off = 0;
    for (const p of parts) {
      out.set(p, off);
      off += p.length;
    }
    expect(() => parseKEXInit(out)).not.toThrow();
    const msg = parseKEXInit(out);
    expect(msg.kexAlgorithms).toEqual(['']); // 空列表 → ['']
  });

  it('对过短报文（cookie 未读完）会越界读取 —— 固化当前行为', () => {
    // parseKEXInit 不做边界检查；固化"不做边界检查"这条不变性，
    // 提醒未来若加边界检查需更新此用例。
    const short = new Uint8Array([SSH_MSG_KEXINIT, 1, 2, 3]);
    // 当前实现会读 data[1..16] 但数组只有 4 字节，返回 undefined 而不抛
    // —— 不强制断言具体行为，只确认不会静默返回合理值
    let threw = false;
    try {
      parseKEXInit(short);
    } catch {
      threw = true;
    }
    // 接受两种行为：抛错 或 不抛错但返回畸形数据。重点是此用例提醒：边界检查缺失。
    expect(typeof threw).toBe('boolean');
  });
});

// ─── EXT_INFO parsing (RFC 8301) ───────────────────────────────────────
// 构造 SSH_MSG_EXT_INFO payload 的手工辅助
function buildExtInfoPayload(exts: { name: string; value: string }[]): Uint8Array {
  const enc = new TextEncoder();
  const parts: Uint8Array[] = [];
  parts.push(new Uint8Array([SSH_MSG_EXT_INFO])); // 7
  const nr = new Uint8Array(4);
  new DataView(nr.buffer).setUint32(0, exts.length, false);
  parts.push(nr);
  for (const { name, value } of exts) {
    const nb = enc.encode(name);
    const nLen = new Uint8Array(4);
    new DataView(nLen.buffer).setUint32(0, nb.length, false);
    parts.push(nLen, nb);
    const vb = enc.encode(value);
    const vLen = new Uint8Array(4);
    new DataView(vLen.buffer).setUint32(0, vb.length, false);
    parts.push(vLen, vb);
  }
  let total = 0;
  for (const p of parts) total += p.length;
  const out = new Uint8Array(total);
  let off = 0;
  for (const p of parts) {
    out.set(p, off);
    off += p.length;
  }
  return out;
}

describe('kex — parseServerSigAlgs (RFC 8301)', () => {
  it('提取 server-sig-algs 扩展', () => {
    const payload = buildExtInfoPayload([
      { name: 'server-sig-algs', value: 'rsa-sha2-512,rsa-sha2-256,ssh-rsa,ecdsa-sha2-nistp256' },
    ]);
    expect(parseServerSigAlgs(payload)).toEqual([
      'rsa-sha2-512',
      'rsa-sha2-256',
      'ssh-rsa',
      'ecdsa-sha2-nistp256',
    ]);
  });

  it('多个扩展时仍提取 server-sig-algs', () => {
    const payload = buildExtInfoPayload([
      { name: 'other-ext', value: 'some-value' },
      { name: 'server-sig-algs', value: 'rsa-sha2-256,ssh-ed25519' },
      { name: 'third', value: '' },
    ]);
    expect(parseServerSigAlgs(payload)).toEqual(['rsa-sha2-256', 'ssh-ed25519']);
  });

  it('未包含 server-sig-algs 时返回空数组', () => {
    const payload = buildExtInfoPayload([{ name: 'other-ext', value: 'val' }]);
    expect(parseServerSigAlgs(payload)).toEqual([]);
  });

  it('空扩展列表返回空数组', () => {
    const payload = buildExtInfoPayload([]);
    expect(parseServerSigAlgs(payload)).toEqual([]);
  });

  it('trim 与空字符串过滤', () => {
    const payload = buildExtInfoPayload([
      { name: 'server-sig-algs', value: '  rsa-sha2-256 , , ssh-ed25519  ' },
    ]);
    expect(parseServerSigAlgs(payload)).toEqual(['rsa-sha2-256', 'ssh-ed25519']);
  });

  it('nr-extensions 超过 1024 上限时抛错', () => {
    const payload = new Uint8Array(5);
    payload[0] = SSH_MSG_EXT_INFO;
    new DataView(payload.buffer).setUint32(1, 5000, false); // nr=5000
    expect(() => parseServerSigAlgs(payload)).toThrow(/nr-extensions 过大/);
  });

  it('畸形 payload（截断）抛错', () => {
    const payload = new Uint8Array([SSH_MSG_EXT_INFO, 0, 0, 0, 1, 0x00 /* 缺 name-len */]);
    expect(() => parseServerSigAlgs(payload)).toThrow();
  });
});

describe('kex — filterExtInfo (RFC 8301)', () => {
  it('过滤 ext-info-c 与 ext-info-s', () => {
    expect(filterExtInfo(['ext-info-c', 'curve25519-sha256', 'ecdh-sha2-nistp256'])).toEqual([
      'curve25519-sha256',
      'ecdh-sha2-nistp256',
    ]);
    expect(filterExtInfo(['ext-info-s', 'curve25519-sha256'])).toEqual(['curve25519-sha256']);
  });

  it('不过滤同名前缀但不同后缀的算法（事实不存在但恶意构造）', () => {
    expect(filterExtInfo(['ext-info-x', 'curve25519-sha256'])).toEqual(['curve25519-sha256']);
  });

  it('保留普通算法', () => {
    expect(filterExtInfo(['curve25519-sha256', 'ecdh-sha2-nistp256'])).toEqual([
      'curve25519-sha256',
      'ecdh-sha2-nistp256',
    ]);
  });
});