From 560b3c44e4eb3e653a274d59b955260790f2f2b1 Mon Sep 17 00:00:00 2001 From: jkooe <62494645+jkooe@users.noreply.github.com> Date: Mon, 26 Oct 2026 00:14:17 +0800 Subject: [PATCH] =?UTF-8?q?fix(agent):=20=E9=99=90=E9=80=9F=E6=A1=B6=20key?= =?UTF-8?q?=20=E6=94=B9=E8=83=BD=E5=8A=9B=E5=90=8D=20+=20describeArgs=20?= =?UTF-8?q?=E6=88=AA=E6=96=AD=EF=BC=88A-P0-2=20/=20A-P0-4=EF=BC=89?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 架构审查两个 P0： A-P0-2 限速桶 key 用错值（server.ts:364）： 原实现 `allowByRate(clientId, limit, String(limit))` 把**限速值** 当桶名——ACL 支持按能力限速（精确 > glob > 全局），配 「screen.capture ≤5/min、app.install ≤5/min」两条同值规则时 两者 key 都是 `clientId|5` 共用一个桶：截图打满后装软件被误伤， 且两者总量仍按 5/min（漏过）。同时 rateBuckets 只增不删 + psk 模式 client_id 自报 → 外部可枚举任意 client_id 造键 → 长期运行内存无界增长。 - 抽到 util/rate-limit.ts：key = `${认证身份 ?? 远端地址}:${能力名}` （未认证时退化按远端地址，不用空串把所有人混为一桶） - 过期键清扫（滑动窗口边界与旧实现逐字节一致：now - t < 60_000 才算窗口内）+ 容量上限淘汰最旧（Map 插入序 = 时间序） - e2e limited 客户端改配两个同限速值能力（各 2/min）， 打满 system.info 后 system.audit.list 必须仍有完整配额 A-P0-4 describeArgs 大参数同步序列化阻塞事件循环（audit.ts:172）： 原实现每次 invoke 全量 JSON.stringify + sha256。fs.write 的 data 可达数 MB——同步序列化 + 哈希 ≈ 数十毫秒级 CPU 且阻塞事件循环， 期间 GUI 注入、事件推送、保活 ping 全部延迟。 - 超过 64KB 只哈希前 64KB（按**字节**切：多字节字符下按字符切 会得到非确定性边界），总长记 args_bytes、截断记 args_truncated - 口径从「全量哈希」变为「前 64KB 哈希」只影响审计摘要， 不影响任何安全属性（摘要用途是比对与防泄密，不是完整性证明） 测试（均故障注入实测会红）： - rate-limit 6 例：桶名是能力名 / 身份进 key / 不限速不建桶 / 过期键清扫（59.999s 窗口内、恰好 60s 出窗）/ 容量上限 / 同值多能力各占一桶 - audit-args-digest 6 例：小参数口径不回退 / 大参数截断口径 / 按字节切（字符切口径必须不同）/ 尾部不进摘要 / 端到端入链 校验通过 / logArgs 默认关不产预览 - e2e 注入验证：桶名退回 String(rateLimit) → v3 场景红；恢复 → 绿 单测 445/445、e2e 33/33、typecheck 全绿 --- apps/agent/src/audit.ts | 37 +++++++- apps/agent/src/server.ts | 30 +++--- apps/agent/src/util/rate-limit.ts | 83 ++++++++++++++++ tests/e2e/run.mjs | 23 ++++- tests/unit/audit-args-digest.test.mjs | 130 ++++++++++++++++++++++++++ tests/unit/rate-limit.test.mjs | 89 ++++++++++++++++++ 6 files changed, 368 insertions(+), 24 deletions(-) create mode 100644 apps/agent/src/util/rate-limit.ts create mode 100644 tests/unit/audit-args-digest.test.mjs create mode 100644 tests/unit/rate-limit.test.mjs diff --git a/apps/agent/src/audit.ts b/apps/agent/src/audit.ts index 3aa3c3e..af449bb 100644 --- a/apps/agent/src/audit.ts +++ b/apps/agent/src/audit.ts @@ -72,6 +72,13 @@ export interface AuditEntry { error?: string | null; /** 参数摘要（sha256 前缀），用于比对而不泄露内容 */ args_digest?: string; + /** + * v2.1.0（A-P0-4）：参数序列化后的**总字节数**。 + * 与 `args_digest_truncated` 同时出现即表示「摘要只覆盖前 64KB」。 + */ + args_bytes?: number; + /** v2.1.0（A-P0-4）：摘要是否为截断哈希（序列化超过 64KB 时只哈希前 64KB）。 */ + args_truncated?: boolean; /** 参数预览（默认不记录；开启后自动脱敏） */ args_preview?: string; reason?: string; @@ -168,12 +175,34 @@ function sanitize(value: unknown, depth = 0): unknown { return value; } +/** + * 参数摘要阈值（字节）：超过即只哈希前 64KB。 + * + * ⚠️ A-P0-4（2026-10-25）：原实现每次 invoke 都全量 `JSON.stringify` + sha256。 + * `fs.write` 的 data 可达数 MB（协议单帧，ws 默认 maxPayload 隐式放行）—— + * 同步序列化 + 哈希一个 8MB 字符串 ≈ 数十毫秒级 CPU，且**阻塞事件循环**： + * 期间 GUI 注入（`input.*`）、事件推送、保活 ping 全部延迟。 + * 截断后：摘要口径从「全量哈希」变为「前 64KB 哈希」，**只影响审计摘要， + * 不影响任何安全属性**（摘要用途是比对与去泄密，从来不是完整性证明）。 + * 总长记入 `args_bytes`，截断记入 `args_truncated`，语义不断代。 + */ +const ARGS_DIGEST_MAX_BYTES = 64 * 1024; + /** 生成参数摘要与（可选）脱敏预览 —— 字段名与 AuditEntry 对齐，便于展开。 */ -export function describeArgs(args: unknown): { args_digest: string; args_preview?: string } { +export function describeArgs(args: unknown): { + args_digest: string; + args_bytes?: number; + args_truncated?: boolean; + args_preview?: string; +} { const json = JSON.stringify(args ?? {}); - const digest = `sha256:${createHash('sha256').update(json).digest('hex').slice(0, 16)}`; - if (!opts.logArgs) return { args_digest: digest }; - return { args_digest: digest, args_preview: JSON.stringify(sanitize(args)) }; + const bytes = Buffer.byteLength(json, 'utf8'); + // 截断按**字节**切（多字节字符下按字符切会得到非确定性的 64KB 边界） + const payload = bytes > ARGS_DIGEST_MAX_BYTES ? Buffer.from(json, 'utf8').subarray(0, ARGS_DIGEST_MAX_BYTES) : json; + const digest = `sha256:${createHash('sha256').update(payload).digest('hex').slice(0, 16)}`; + const base = bytes > ARGS_DIGEST_MAX_BYTES ? { args_bytes: bytes, args_truncated: true } : {}; + if (!opts.logArgs) return { args_digest: digest, ...base }; + return { args_digest: digest, ...base, args_preview: JSON.stringify(sanitize(args)) }; } /** 写入一条审计记录。任何异常都被吞掉 —— 审计绝不阻断主流程。 */ diff --git a/apps/agent/src/server.ts b/apps/agent/src/server.ts index e954788..1515caf 100644 --- a/apps/agent/src/server.ts +++ b/apps/agent/src/server.ts @@ -32,6 +32,7 @@ import { createCapabilityRegistry } from './capabilities/index.js'; import { audit, describeArgs } from './audit.js'; import { disposeWatchesByOwner, watchCount } from './events.js'; import { runningTaskCount } from './capabilities/task.js'; +import { createRateLimiter } from './util/rate-limit.js'; import type { AgentConfig } from './config.js'; import type { TlsMaterial } from './certs.js'; @@ -96,21 +97,15 @@ export function createAgentCore(cfg: AgentConfig): AgentCore { */ const hasAcl = Boolean(cfg.acl); - /** 按调用方的滑动窗口限速（每分钟）。 */ - const rateBuckets = new Map(); - function allowByRate(clientId: string, limitPerMin?: number, bucket?: string): boolean { - if (!limitPerMin || limitPerMin <= 0) return true; - const key = bucket ? `${clientId}|${bucket}` : clientId; - const now = Date.now(); - const recent = (rateBuckets.get(key) ?? []).filter((t) => now - t < 60_000); - if (recent.length >= limitPerMin) { - rateBuckets.set(key, recent); - return false; - } - recent.push(now); - rateBuckets.set(key, recent); - return true; - } + /** + * 按调用方的滑动窗口限速（每分钟）。 + * + * ⚠️ A-P0-2（2026-10-25）：原实现 `allowByRate(clientId, limit, String(limit))` + * 把**限速值**当桶名 —— 同限速值的不同能力共用一个桶（互相挤占/漏过）， + * 且 `rateBuckets` 只增不删（psk 模式 client_id 自报 → 可枚举造键 → 内存无界增长）。 + * 现抽到 `util/rate-limit.ts`：桶名 = 能力名；过期键顺手删 + 容量上限淘汰最旧。 + */ + const rateLimiter = createRateLimiter(); const levelOrder: Record = { debug: 0, info: 1, warn: 2 }; const log = (level: LogLevel, msg: string): void => { if (levelOrder[level] >= levelOrder[cfg.log_level]) { @@ -360,8 +355,11 @@ export function createAgentCore(cfg: AgentConfig): AgentCore { } } + // 限速键 = `${认证身份 ?? 远端地址}:${能力名}`： + // ① 不同能力各占一桶（同限速值不再互挤/漏过）； + // ② 未认证身份（psk 自报缺失）时退化按远端地址，不用空串把所有人混为一桶。 const rateLimit = clientRule ? rateLimitFor(clientRule, name) : undefined; - if (!allowByRate(state.clientId ?? '', rateLimit, String(rateLimit ?? ''))) { + if (!rateLimiter.allow(state.clientId ?? state.remote, rateLimit, name)) { log('warn', `限速拒绝 client_id=${state.clientId} capability=${name}`); audit({ type: 'rate.limited', diff --git a/apps/agent/src/util/rate-limit.ts b/apps/agent/src/util/rate-limit.ts new file mode 100644 index 0000000..efa5722 --- /dev/null +++ b/apps/agent/src/util/rate-limit.ts @@ -0,0 +1,83 @@ +/** + * 按调用方的滑动窗口限速器（每分钟）。 + * + * ## 为什么从 server.ts 抽出来（A-P0-2，2026-10-25） + * 原实现有两处缺陷，且逻辑埋在 `createAgentCore` 闭包里无法单测： + * + * 1. **桶 key 用错值**：原 `allowByRate(clientId, limit, String(limit))` 第三参 + * （bucket）传的是**限速值本身**。于是「screen.capture ≤ 5/min」与 + * 「app.install ≤ 5/min」两条同值规则共用 `clientId|5` 一个桶： + * 截图 5 次后装软件也被限速（误伤），且两者总量仍按 5/min（漏过）。 + * 正确口径：key = `${身份}:${能力名}`。 + * 2. **桶只增不删**：`rateBuckets` Map 无任何清扫，key 永久留存。而 psk 模式下 + * `client_id` 由客户端**自报**（配置级约束）→ 外部可枚举任意 client_id 造键 + * → 长期运行内存无界增长。现在：每次允许判定顺手删过期键 + 容量上限淘汰最旧。 + * + * ## 语义 + * - 窗口 60s 滑动：只数「最近 60 秒内」的调用； + * - `limitPerMin` 为空/≤0 表示不限速（直接放行，**不建桶** —— 与旧行为一致）； + * - 身份用「认证身份 ?? 远端地址」：psk 自报 client_id 不可信，但限速键仍以它为纲 + * （配了限速就说明该 client_id 已在 ACL 里登记），无身份的调用方退化按远端地址。 + */ + +/** 桶容量上限：超出淘汰最旧（Map 保持插入顺序，首键即最旧）。 */ +const DEFAULT_MAX_KEYS = 10_000; + +const WINDOW_MS = 60_000; + +export interface RateLimiter { + /** + * 请求一次调用配额。 + * + * @param identity 调用方身份（认证身份 ?? 远端地址） + * @param limitPerMin 每分钟上限；空/≤0 = 不限速 + * @param bucket 桶名（能力名）—— 不同能力各占一桶，互不挤占 + */ + allow(identity: string, limitPerMin: number | undefined, bucket: string): boolean; + /** 当前桶数（供测试与观测）。 */ + size(): number; +} + +export interface RateLimiterOptions { + /** 桶容量上限（超出淘汰最旧）。默认 10_000。 */ + maxKeys?: number; + /** 时钟注入（单测用假时钟推进 60s 窗口，不必真等）。 */ + now?: () => number; +} + +export function createRateLimiter(opts: RateLimiterOptions = {}): RateLimiter { + const maxKeys = opts.maxKeys ?? DEFAULT_MAX_KEYS; + const now = opts.now ?? Date.now; + const buckets = new Map(); + + return { + allow(identity: string, limitPerMin: number | undefined, bucket: string): boolean { + if (!limitPerMin || limitPerMin <= 0) return true; + const key = `${identity}|${bucket}`; + const t = now(); + + // 清扫①：过期键直接删（只增不删 = 内存无界增长；自报 client_id 可枚举造键） + for (const [k, times] of buckets) { + if (times.length === 0 || t - times[times.length - 1]! >= WINDOW_MS) buckets.delete(k); + } + // 清扫②：容量上限，超出淘汰最旧（插入序 = 时间序，首键最旧） + while (buckets.size >= maxKeys) { + const oldest = buckets.keys().next(); + if (oldest.done) break; + buckets.delete(oldest.value); + } + + const recent = (buckets.get(key) ?? []).filter((ts) => t - ts < WINDOW_MS); + if (recent.length >= limitPerMin) { + buckets.set(key, recent); + return false; + } + recent.push(t); + buckets.set(key, recent); + return true; + }, + size(): number { + return buckets.size; + }, + }; +} diff --git a/tests/e2e/run.mjs b/tests/e2e/run.mjs index 952a11e..2d636c9 100644 --- a/tests/e2e/run.mjs +++ b/tests/e2e/run.mjs @@ -335,11 +335,13 @@ async function testEd25519Acl() { deny: [], }, { - // 用于验证限速：每分钟仅 2 次 + // 用于验证限速：每分钟仅 2 次。 + // A-P0-2：两个能力配**相同**限速值 2 —— 旧实现把限速值当桶名 + // （key=`limited|2`）会共桶互挤；修正后各能力独占一桶。 client_id: 'limited', pubkey: kp.publicKey, - allow: ['system.info'], - max_calls_per_min: 2, + allow: ['system.info', 'system.audit.list'], + rate_limits: { 'system.info': 2, 'system.audit.list': 2 }, }, ], }, @@ -407,7 +409,7 @@ async function testEd25519Acl() { assert.ok(withPreview, '开启 log_args 后应有参数预览'); c.close(); - // 4. 限速：limited 配额 2/分钟，第 3 次应被拒 + // 4. 限速：limited 每个能力 2/分钟，第 3 次应被拒 const lim = mk('limited', kp.privateKey); await lim.connect(); const r1 = await lim.invoke('system.info'); @@ -421,6 +423,19 @@ async function testEd25519Acl() { return true; }, ); + // A-P0-2 关键断言：system.info 已打满 2/分钟，但同限速值的 system.audit.list + // 必须仍有完整配额（旧实现共桶 → 此处第 1 次即被 -32407 误伤）。 + const a1 = await lim.invoke('system.audit.list', { limit: 5 }); + assert.equal(a1.status, 'ok', '同限速值的另一能力不得被挤占（A-P0-2）'); + const a2 = await lim.invoke('system.audit.list', { limit: 5 }); + assert.equal(a2.status, 'ok'); + await assert.rejects( + () => lim.invoke('system.audit.list', { limit: 5 }), + (e) => { + assert.equal(e.code, -32407, `期望 -32407（限速），实际 ${e.code}`); + return true; + }, + ); lim.close(); // 5. 未注册的调用方 → 握手失败 diff --git a/tests/unit/audit-args-digest.test.mjs b/tests/unit/audit-args-digest.test.mjs new file mode 100644 index 0000000..0fa28e0 --- /dev/null +++ b/tests/unit/audit-args-digest.test.mjs @@ -0,0 +1,130 @@ +/** + * 参数摘要守卫（`apps/agent/src/audit.ts` describeArgs，A-P0-4，2026-10-25）。 + * + * ⚠️ 背景：原实现每次 invoke 都全量 `JSON.stringify` + sha256。`fs.write` 的 data + * 可达数 MB——同步序列化 + 哈希 ≈ 数十毫秒级 CPU 且阻塞事件循环，期间 GUI + * 注入、事件推送、保活 ping 全部延迟。现改为：超过 64KB 只哈希**前 64KB** + * （按字节切），总长记 `args_bytes`、截断记 `args_truncated`。 + * 摘要口径从「全量」变为「前 64KB」，只影响审计摘要，不影响任何安全属性 + * （摘要用途是比对与防泄密，不是完整性证明）。 + * + * 本文件钉死五条（均可用故障注入打红）： + * ① 小参数口径不回退：digest 与独立重算的 sha256(json).slice(0,16) 一致； + * ② 大参数截断口径：digest = sha256(前 64KB 字节).slice(0,16)，args_bytes/args_truncated 正确； + * ③ 按**字节**切而非按字符切（多字节边界，字符切口径必须不同）； + * ④ 尾部不进摘要：前 64KB 相同、尾部不同的两个大参数 digest 相同； + * ⑤ 端到端：audit() 写入的记录带两字段，链校验仍通过；logArgs 预览行为不变。 + */ +import { test } from 'node:test'; +import assert from 'node:assert/strict'; +import { createHash } from 'node:crypto'; +import { mkdtempSync, readFileSync, rmSync } from 'node:fs'; +import { tmpdir } from 'node:os'; +import { join } from 'node:path'; +import { initAudit, audit, describeArgs, verifyAudit } from '../../apps/agent/dist/audit.js'; + +const MAX_BYTES = 64 * 1024; + +/** 独立重算摘要（不 import 实现内的私有常量路径，口径在此显式重写一遍）。 */ +function digestOf(payload) { + return `sha256:${createHash('sha256').update(payload).digest('hex').slice(0, 16)}`; +} + +function bigArgs(totalBytes, tailFill = '') { + // 构造 {capability, data}，使 JSON 序列化后的总字节数 ≥ totalBytes。 + // data 主体为 'x'（1 字节/字符，字节数可精确预估），tailFill 可接一段不同内容的尾部， + // 用于「前 64KB 相同、尾部不同」的对照。 + const skeleton = { capability: 'fs.write', data: '' }; + const baseBytes = Buffer.byteLength(JSON.stringify(skeleton), 'utf8'); + const need = Math.max(1, totalBytes - baseBytes - tailFill.length); + return { capability: 'fs.write', data: 'x'.repeat(need) + tailFill }; +} + +test('A-P0-4①：小参数口径不回退 —— digest 与旧实现逐字节一致，无截断字段', () => { + const args = { path: '/tmp/a.txt', data: 'hello' }; + const r = describeArgs(args); + assert.equal(r.args_digest, digestOf(JSON.stringify(args)), '小参数必须仍走全量哈希'); + assert.equal('args_bytes' in r, false, '小参数不得出现 args_bytes'); + assert.equal('args_truncated' in r, false, '小参数不得出现 args_truncated'); + assert.match(r.args_digest, /^sha256:[0-9a-f]{16}$/, 'digest 形状不变'); +}); + +test('A-P0-4②：大参数截断口径 —— 只哈希前 64KB，总长与截断标记正确', () => { + const total = 8 * 1024 * 1024; // 8MB，模拟 fs.write 大 data + const args = bigArgs(total); + const json = JSON.stringify(args); + const bytes = Buffer.byteLength(json, 'utf8'); + assert.ok(bytes > MAX_BYTES, '夹具必须真的超过 64KB'); + + const r = describeArgs(args); + assert.equal(r.args_bytes, bytes, 'args_bytes 必须是序列化总字节数'); + assert.equal(r.args_truncated, true); + // 口径：sha256(前 64KB 字节) 的前 16 位十六进制 + const expected = digestOf(Buffer.from(json, 'utf8').subarray(0, MAX_BYTES)); + assert.equal(r.args_digest, expected, '大参数 digest 必须等于前 64KB 的哈希'); + // 且必须与全量哈希不同（证明真的截断了） + assert.notEqual(r.args_digest, digestOf(json), '不得退化成全量哈希'); +}); + +test('A-P0-4③：按字节切而非按字符切 —— 多字节边界下字符切口径必须不同', () => { + // 用中文填充：'中' = 3 字节，64KB(65536) 不是 3 的整数倍 → 边界必然落字符中间 + const args = { capability: 'fs.write', data: '中'.repeat(40000) }; // ≈120KB + const json = JSON.stringify(args); + assert.ok(Buffer.byteLength(json, 'utf8') > MAX_BYTES); + + const r = describeArgs(args); + const byBytes = digestOf(Buffer.from(json, 'utf8').subarray(0, MAX_BYTES)); + assert.equal(r.args_digest, byBytes, '必须按字节切'); + + // 若有人改成按字符切（json.slice(0, 65536)），边界落在 65535 字节处 → 口径不同 → 本断言打红 + const byChars = digestOf(json.slice(0, MAX_BYTES)); + assert.notEqual(r.args_digest, byChars, '按字符切是错误口径，必须被此断言挡住'); +}); + +test('A-P0-4④：尾部不进摘要 —— 前 64KB 相同、尾部不同 → digest 相同但 args_bytes 不同', () => { + const a = bigArgs(2 * 1024 * 1024, ''); + const b = bigArgs(3 * 1024 * 1024, 'y'.repeat(1024 * 1024)); // 更长且尾部内容不同 + const ra = describeArgs(a); + const rb = describeArgs(b); + assert.equal(ra.args_digest, rb.args_digest, '前 64KB 一致则摘要必须一致'); + assert.notEqual(ra.args_bytes, rb.args_bytes, '总长仍如实记录'); + assert.equal(ra.args_truncated, true); + assert.equal(rb.args_truncated, true); +}); + +test('A-P0-4⑤：端到端 —— audit() 写入带两字段且链校验通过；logArgs 预览行为不变', () => { + const dir = mkdtempSync(join(tmpdir(), 'na-audit-args-')); + try { + initAudit({ enabled: true, logArgs: true }, dir); + const big = bigArgs(1024 * 1024); + audit({ type: 'invoke', capability: 'fs.write', status: 'ok', ...describeArgs(big) }); + audit({ type: 'invoke', capability: 'fs.write', status: 'ok', ...describeArgs({ path: '/tmp/x', data: 'y' }) }); + + const lines = readFileSync(join(dir, 'audit.log'), 'utf8').trim().split('\n'); + const rec1 = JSON.parse(lines[0]); + const rec2 = JSON.parse(lines[1]); + assert.equal(rec1.args_truncated, true, '大参数记录必须带截断标记'); + assert.equal(rec1.args_bytes, Buffer.byteLength(JSON.stringify(big), 'utf8')); + assert.equal(rec1.args_truncated, true); + assert.equal('args_bytes' in rec2, false, '小参数记录不得带截断字段'); + assert.equal(typeof rec1.args_preview, 'string', 'logArgs 开启时预览仍产出（大参数也一样）'); + + const v = verifyAudit(); + assert.equal(v.ok, true, `新字段不得破坏审计链，实际: ${JSON.stringify(v.broken_at)}`); + assert.equal(v.checked, 2); + } finally { + rmSync(dir, { recursive: true, force: true }); + } +}); + +test('A-P0-4⑤b：logArgs 默认关闭时不产出预览（行为不回退）', () => { + const dir = mkdtempSync(join(tmpdir(), 'na-audit-args2-')); + try { + initAudit({ enabled: true }, dir); // logArgs 默认 false + const r = describeArgs({ path: '/tmp/x', data: 'hello' }); + assert.equal('args_preview' in r, false, '默认不得记录预览'); + assert.equal(r.args_digest, digestOf(JSON.stringify({ path: '/tmp/x', data: 'hello' }))); + } finally { + rmSync(dir, { recursive: true, force: true }); + } +}); diff --git a/tests/unit/rate-limit.test.mjs b/tests/unit/rate-limit.test.mjs new file mode 100644 index 0000000..1fc690e --- /dev/null +++ b/tests/unit/rate-limit.test.mjs @@ -0,0 +1,89 @@ +/** + * 限速器守卫（`apps/agent/src/util/rate-limit.ts`，A-P0-2，2026-10-25）。 + * + * ⚠️ 背景：原实现 `allowByRate(clientId, limit, String(limit))` 把**限速值**当桶名。 + * ACL 支持按能力限速（`rateLimitFor`，精确 > glob > 全局），当某 client 配了 + * 「screen.capture ≤ 5/min、app.install ≤ 5/min」两条同值规则时，两者 key 都是 + * `clientId|5` → 共用一个桶：截图 5 次后装软件也被限速（误伤），且两者总量 + * 仍按 5/min（漏过）。 + * 同时 `rateBuckets` 只增不删 + psk 模式 client_id 自报 → 外部可枚举任意 + * client_id 造键 → 长期运行内存无界增长。 + * + * 本文件钉死四条（均可用故障注入打红）： + * ① 桶名是**能力名**：同限速值的不同能力互不挤占； + * ② 身份进 key：不同调用方互不挤占； + * ③ 过期键被清扫（假时钟推进 61s，桶数归零、重新放行）； + * ④ 容量上限：自报 client_id 枚举造键 → 桶数有界。 + */ +import { test } from 'node:test'; +import assert from 'node:assert/strict'; +import { createRateLimiter } from '../../apps/agent/dist/util/rate-limit.js'; + +test('A-P0-2①：桶名是能力名 —— 同限速值的不同能力互不挤占', () => { + const rl = createRateLimiter(); + // 同一 client、同一限速值、两个能力 + assert.equal(rl.allow('c1', 2, 'screen.capture'), true); + assert.equal(rl.allow('c1', 2, 'screen.capture'), true); + assert.equal(rl.allow('c1', 2, 'screen.capture'), false, '第 3 次截图应被限'); + // 关键：app.install 不受截图影响（旧实现此处共桶 → 误伤） + assert.equal(rl.allow('c1', 2, 'app.install'), true, '同限速值的另一能力不得被挤占'); + assert.equal(rl.allow('c1', 2, 'app.install'), true); + assert.equal(rl.allow('c1', 2, 'app.install'), false, 'app.install 自己也按 2/min 限'); + // 两个能力各占一桶 + assert.equal(rl.size(), 2); +}); + +test('A-P0-2②：身份进 key —— 不同调用方互不挤占', () => { + const rl = createRateLimiter(); + assert.equal(rl.allow('c1', 1, 'screen.capture'), true); + assert.equal(rl.allow('c1', 1, 'screen.capture'), false); + assert.equal(rl.allow('c2', 1, 'screen.capture'), true, '另一调用方有独立配额'); + assert.equal(rl.size(), 2); +}); + +test('A-P0-2③：不限速时不建桶（与旧行为一致：直接放行）', () => { + const rl = createRateLimiter(); + for (let i = 0; i < 100; i += 1) assert.equal(rl.allow('c1', undefined, 'screen.capture'), true); + assert.equal(rl.allow('c1', 0, 'screen.capture'), true, 'limit=0 同样视为不限速'); + assert.equal(rl.allow('c1', -1, 'screen.capture'), true); + assert.equal(rl.size(), 0, '不限速的调用不得留下桶'); +}); + +test('A-P0-2③：过期键被清扫（假时钟推进 61s → 桶清空且重新放行）', () => { + let t = 1_000_000; + const rl = createRateLimiter({ now: () => t }); + assert.equal(rl.allow('c1', 1, 'screen.capture'), true); + assert.equal(rl.allow('c1', 1, 'screen.capture'), false, '窗口内第 2 次被限'); + assert.equal(rl.size(), 1); + + // 滑动窗口边界，与旧实现逐字节同语义：`now - t < 60_000` 才算窗口内。 + t += 60_000 - 1; // 59.999s：仍在窗口内，第 2 次被限 + assert.equal(rl.allow('c1', 1, 'screen.capture'), false, '59.999s 仍算窗口内'); + + t += 1; // 恰好 60s：`60_000 < 60_000` 为 false → 旧时间戳出窗 → 重新放行 + assert.equal(rl.allow('c1', 1, 'screen.capture'), true, '恰好 60s 出窗（沿用旧实现边界）'); + assert.equal(rl.size(), 1, '自己的桶重建'); + // 再来一个不相关的旧桶，它应在下次 allow 时被顺手清掉 + rl.allow('c2', 1, 'other'); + t += 61_000; + rl.allow('c3', 1, 'x'); + assert.equal(rl.size(), 1, `过期键必须被清扫（只增不删 = 内存无界增长），实际 ${rl.size()}`); +}); + +test('A-P0-2④：容量上限 —— 自报 client_id 枚举造键，桶数有界', () => { + let t = 0; + const rl = createRateLimiter({ maxKeys: 100, now: () => (t += 1) }); + for (let i = 0; i < 1000; i += 1) { + rl.allow(`attacker-${i}`, 5, 'screen.capture'); + } + assert.ok(rl.size() <= 100, `桶数必须被容量上限约束，实际 ${rl.size()}`); + // 最早的键已被淘汰：它若再来，配额从零开始（可接受：枚举者本就拿不到稳定配额） + assert.equal(rl.allow('attacker-0', 5, 'screen.capture'), true); +}); + +test('A-P0-2：限速值与桶名解耦的结构证据（同值多能力各占一桶）', () => { + const rl = createRateLimiter(); + const caps = ['a.x', 'b.x', 'c.x', 'd.x']; + for (const c of caps) rl.allow('c1', 3, c); + assert.equal(rl.size(), caps.length, '每个能力一个桶（key 里没有限速值）'); +});