From b24fc12e2b4caefef5082d4410793e5aa3c0d362 Mon Sep 17 00:00:00 2001 From: "glm5.2 agent 180" Date: Sat, 20 Jun 2026 04:45:22 +0300 Subject: [PATCH] test(ai-chat): cover safety-critical code paths (review follow-ups P1) Adds unit tests for five pieces of security-critical code that previously had zero coverage, closing the warning-level findings of the ai-chat multi-aspect review (docs/backlog/ai-chat-review-followups.md, priority 1). - crypto/secret-box.spec.ts (NEW): AES-256-GCM round-trip; non-determinism (two encrypts of the same input yield different blobs, both decrypt); tampered authTag / ciphertext bytes throw with the 'APP_SECRET may have changed' message; wrong APP_SECRET throws the same. Guards the only at-rest protection of provider API keys. - ai-chat/external-mcp/ssrf-guard.spec.ts (NEW): isIpAllowed blocks every forbidden class (loopback, link-local incl. metadata 169.254.169.254, private, CGNAT, ULA, unspecified, IPv4-mapped IPv6, unparseable) and allows a public IP; isUrlAllowed rejects bad scheme / invalid URL, blocks IP-literal private, and (with a mocked dns.lookup) blocks DNS-rebinding to a private address and an unresolvable host. - ai-chat/ai-chat.service.spec.ts (extended): assistantParts now covered - paired tool call -> output-available (compacted), unpaired call -> output-error with 'Tool call did not complete.' (regression guard for the MissingToolResultsError fix), broken calls skipped, step text and fallback text paths. Requires exporting assistantParts + StepLike (the only production change here, two export keywords). - ai-chat/tools/ai-chat-tools.service.spec.ts (extended): JSON-string coercion in patchNode / insertNode / updatePageJson - string parsed to object, invalid JSON throws the specific message, updatePageJson distinguishes undefined (title-only) / object / string. Guards the OpenAI tool-call compatibility fix. - database/repos/ai-chat/page-embedding.repo.spec.ts (NEW): searchByEmbedding with empty spaceIds returns [] without touching the DB (Proxy stub throws on any access). Guards the access-scoping early-return. 54 new tests, all green. No functional behaviour changed. --- .../src/core/ai-chat/ai-chat.service.spec.ts | 77 +++++++- .../src/core/ai-chat/ai-chat.service.ts | 6 +- .../ai-chat/external-mcp/ssrf-guard.spec.ts | 116 ++++++++++++ .../tools/ai-chat-tools.service.spec.ts | 168 ++++++++++++++++++ .../repos/ai-chat/page-embedding.repo.spec.ts | 30 ++++ .../integrations/crypto/secret-box.spec.ts | 66 +++++++ 6 files changed, 460 insertions(+), 3 deletions(-) create mode 100644 apps/server/src/core/ai-chat/external-mcp/ssrf-guard.spec.ts create mode 100644 apps/server/src/database/repos/ai-chat/page-embedding.repo.spec.ts create mode 100644 apps/server/src/integrations/crypto/secret-box.spec.ts diff --git a/apps/server/src/core/ai-chat/ai-chat.service.spec.ts b/apps/server/src/core/ai-chat/ai-chat.service.spec.ts index f1f3461a..cc733642 100644 --- a/apps/server/src/core/ai-chat/ai-chat.service.spec.ts +++ b/apps/server/src/core/ai-chat/ai-chat.service.spec.ts @@ -1,4 +1,4 @@ -import { compactToolOutput } from './ai-chat.service'; +import { compactToolOutput, assistantParts, type StepLike } from './ai-chat.service'; /** * Unit tests for compactToolOutput: the pure helper that shrinks LARGE tool @@ -66,3 +66,78 @@ describe('compactToolOutput', () => { expect(compactedBytes).toBeLessThan(originalBytes / 10); }); }); + +/** + * Unit tests for assistantParts: the pure helper that rebuilds UIMessage parts + * from SDK step objects so multi-turn history replays prior tool-calls and + * results to the model. The key invariant is that an UNPAIRED tool call (no + * matching result) is emitted as a synthetic `output-error` — without this, + * convertToModelMessages would throw MissingToolResultsError on the next turn. + */ +describe('assistantParts', () => { + it('emits output-available for a paired tool call with its (compacted) output', () => { + const steps: StepLike[] = [ + { + toolCalls: [{ toolCallId: 'c1', toolName: 'getPage', input: { pageId: 'p1' } }], + toolResults: [{ toolCallId: 'c1', output: { title: 'T' } }], + }, + ]; + const parts = assistantParts(steps, 'irrelevant') as Array< + Record + >; + const toolPart = parts.find((p) => p.type === 'tool-getPage'); + expect(toolPart).toBeDefined(); + expect(toolPart!.state).toBe('output-available'); + expect(toolPart!.output).toEqual({ title: 'T' }); + expect(toolPart!.input).toEqual({ pageId: 'p1' }); + }); + + it('emits output-error for an unpaired tool call (regression guard)', () => { + const steps: StepLike[] = [ + { + toolCalls: [{ toolCallId: 'c2', toolName: 'getPage', input: {} }], + toolResults: [], + }, + ]; + const parts = assistantParts(steps, '') as Array>; + const toolPart = parts.find((p) => p.type === 'tool-getPage'); + expect(toolPart).toBeDefined(); + expect(toolPart!.state).toBe('output-error'); + expect(toolPart!.errorText).toBe('Tool call did not complete.'); + }); + + it('skips tool calls missing toolName or toolCallId', () => { + const steps: StepLike[] = [ + { + toolCalls: [ + { input: { x: 1 } }, + { toolName: 'getPage' }, + { toolCallId: 'c4' }, + ], + toolResults: [], + }, + ]; + const parts = assistantParts(steps, '') as Array>; + expect(parts.filter((p) => String(p.type).startsWith('tool-'))).toHaveLength(0); + }); + + it('includes a text part for a step that carries text', () => { + const steps: StepLike[] = [{ text: 'hello' }]; + const parts = assistantParts(steps, '') as Array>; + expect(parts).toContainEqual({ type: 'text', text: 'hello' }); + }); + + it('appends the fallback text when no step carries text', () => { + const steps: StepLike[] = [ + { + toolCalls: [{ toolCallId: 'c5', toolName: 'getPage', input: {} }], + toolResults: [], + }, + ]; + const parts = assistantParts(steps, 'final answer') as Array< + Record + >; + const last = parts[parts.length - 1]; + expect(last).toEqual({ type: 'text', text: 'final answer' }); + }); +}); diff --git a/apps/server/src/core/ai-chat/ai-chat.service.ts b/apps/server/src/core/ai-chat/ai-chat.service.ts index 3119c3c4..045d5599 100644 --- a/apps/server/src/core/ai-chat/ai-chat.service.ts +++ b/apps/server/src/core/ai-chat/ai-chat.service.ts @@ -419,7 +419,8 @@ function textPart(text: string): Array<{ type: 'text'; text: string }> { * `toolResults` -> TypedToolResult). Typed loosely so this survives provider * variation; only the fields we persist are referenced. */ -type StepLike = { +// Exported for unit testing (ai-chat.service.spec.ts). +export type StepLike = { text?: string; toolCalls?: ReadonlyArray<{ toolCallId?: string; @@ -538,7 +539,8 @@ function compactValue(value: unknown, depth: number): unknown { * recovers the name. Falls back to a single `text` part built from * `fallbackText` when the steps carry no text. */ -function assistantParts( +// Exported for unit testing (ai-chat.service.spec.ts). +export function assistantParts( steps: ReadonlyArray | undefined, fallbackText: string, ): UIMessage['parts'] { diff --git a/apps/server/src/core/ai-chat/external-mcp/ssrf-guard.spec.ts b/apps/server/src/core/ai-chat/external-mcp/ssrf-guard.spec.ts new file mode 100644 index 00000000..214b8d61 --- /dev/null +++ b/apps/server/src/core/ai-chat/external-mcp/ssrf-guard.spec.ts @@ -0,0 +1,116 @@ +// Hoisted before imports: ssrf-guard.ts calls `promisify(dns.lookup)` at module +// load time, so the mock must be in place before the first import of the guard. +jest.mock('node:dns', () => ({ lookup: jest.fn() })); + +import { lookup as dnsLookup } from 'node:dns'; +import { isIpAllowed, isUrlAllowed } from './ssrf-guard'; + +const lookupMock = dnsLookup as unknown as jest.Mock; + +/** + * Unit tests for the SSRF guard that protects the admin-configured external MCP + * server URLs. `isIpAllowed` is pure and synchronous; `isUrlAllowed` resolves + * the hostname via dns.lookup (mocked here) and blocks if ANY resolved address + * is non-public. Every blocked range corresponds to a real exploit vector. + */ +describe('isIpAllowed', () => { + it.each([ + ['127.0.0.1', 'loopback'], + ['::1', 'loopback'], + ])('blocks IPv4/IPv6 loopback %s', (ip) => { + expect(isIpAllowed(ip)).toEqual({ ok: false, reason: expect.any(String) }); + expect(isIpAllowed(ip).ok).toBe(false); + }); + + it.each(['169.254.169.254', 'fe80::1'])( + 'blocks link-local / metadata endpoint %s', + (ip) => { + expect(isIpAllowed(ip).ok).toBe(false); + }, + ); + + it.each(['10.0.0.1', '172.16.0.1', '192.168.1.1'])( + 'blocks private range %s', + (ip) => { + expect(isIpAllowed(ip).ok).toBe(false); + }, + ); + + it('blocks CGNAT 100.64.0.1', () => { + expect(isIpAllowed('100.64.0.1').ok).toBe(false); + }); + + it('blocks ULA fd00::1', () => { + expect(isIpAllowed('fd00::1').ok).toBe(false); + }); + + it.each(['0.0.0.0', '::'])('blocks unspecified %s', (ip) => { + expect(isIpAllowed(ip).ok).toBe(false); + }); + + it('blocks IPv4-mapped IPv6 loopback ::ffff:127.0.0.1', () => { + expect(isIpAllowed('::ffff:127.0.0.1').ok).toBe(false); + }); + + it('returns an unparseable reason for garbage', () => { + expect(isIpAllowed('not-an-ip')).toEqual({ + ok: false, + reason: 'unparseable IP address', + }); + }); + + it.each(['8.8.8.8', '1.1.1.1'])('allows public IP %s', (ip) => { + expect(isIpAllowed(ip)).toEqual({ ok: true }); + }); +}); + +describe('isUrlAllowed', () => { + afterEach(() => { + lookupMock.mockReset(); + }); + + it('rejects a non-http scheme', async () => { + const res = await isUrlAllowed('ftp://example.com'); + expect(res).toEqual({ ok: false, reason: 'only http/https URLs are allowed' }); + }); + + it('rejects an invalid URL', async () => { + const res = await isUrlAllowed('not a url'); + expect(res).toEqual({ ok: false, reason: 'invalid URL' }); + }); + + it('blocks a private IP-literal host (no DNS needed)', async () => { + const res = await isUrlAllowed('http://127.0.0.1'); + expect(res.ok).toBe(false); + }); + + it('allows a public IP-literal host', async () => { + const res = await isUrlAllowed('http://8.8.8.8'); + expect(res.ok).toBe(true); + }); + + it('blocks DNS-rebinding: hostname resolves to a private address', async () => { + lookupMock.mockImplementation( + (host, opts, cb) => cb(null, [{ address: '10.0.0.5' }]), + ); + const res = await isUrlAllowed('http://rebind.test'); + expect(res.ok).toBe(false); + expect(res.reason).toContain('blocked address range'); + }); + + it('allows a hostname that resolves to a public address', async () => { + lookupMock.mockImplementation( + (host, opts, cb) => cb(null, [{ address: '93.184.216.34' }]), + ); + const res = await isUrlAllowed('http://example.com'); + expect(res.ok).toBe(true); + }); + + it('blocks an unresolvable hostname', async () => { + lookupMock.mockImplementation( + (host, opts, cb) => cb(new Error('ENOTFOUND')), + ); + const res = await isUrlAllowed('http://nope.test'); + expect(res).toEqual({ ok: false, reason: 'host could not be resolved' }); + }); +}); diff --git a/apps/server/src/core/ai-chat/tools/ai-chat-tools.service.spec.ts b/apps/server/src/core/ai-chat/tools/ai-chat-tools.service.spec.ts index 65218300..3ae0dd42 100644 --- a/apps/server/src/core/ai-chat/tools/ai-chat-tools.service.spec.ts +++ b/apps/server/src/core/ai-chat/tools/ai-chat-tools.service.spec.ts @@ -211,3 +211,171 @@ describe('AiChatToolsService expanded toolset guardrails', () => { expect(parsed).not.toHaveProperty('deleteComments'); }); }); + +/** + * Tests for the JSON-string node/content coercion in patchNode, insertNode, and + * updatePageJson. OpenAI tool-calls sometimes serialize object arguments as + * JSON strings; the execute body parses them before forwarding to the client. + * A regression that removes the parse would silently break insert/patch under + * OpenAI. The recording fake client captures forwarded args so we can assert + * the PARSED object (not the string) reaches the client. + */ +describe('AiChatToolsService JSON-string node/content coercion', () => { + const forwarded = { + patchNode: [] as unknown[][], + insertNode: [] as unknown[][], + updatePageJson: [] as unknown[][], + }; + + const fakeClient: Partial = { + patchNode: (...args: unknown[]) => { + forwarded.patchNode.push(args); + return Promise.resolve({}); + }, + insertNode: (...args: unknown[]) => { + forwarded.insertNode.push(args); + return Promise.resolve({}); + }, + updatePageJson: (...args: unknown[]) => { + forwarded.updatePageJson.push(args); + return Promise.resolve({}); + }, + }; + + const tokenServiceStub = { + generateAccessToken: jest.fn().mockResolvedValue('access-token'), + generateCollabToken: jest.fn().mockResolvedValue('collab-token'), + }; + + let service: AiChatToolsService; + + beforeEach(() => { + forwarded.patchNode.length = 0; + forwarded.insertNode.length = 0; + forwarded.updatePageJson.length = 0; + jest.spyOn(loader, 'loadDocmostMcp').mockResolvedValue({ + DocmostClient: function () { + return fakeClient as DocmostClientLike; + } as unknown as loader.DocmostClientCtor, + }); + service = new AiChatToolsService( + tokenServiceStub as never, + {} as never, + {} as never, + {} as never, + {} as never, + ); + }); + + afterEach(() => { + jest.restoreAllMocks(); + }); + + async function buildTools() { + return service.forUser( + { id: 'user-1', email: 'u@example.com', workspaceId: 'ws-1' } as never, + 'session-1', + 'ws-1', + 'chat-1', + ); + } + + it('patchNode: parses a JSON-string node and forwards the object', async () => { + const tools = await buildTools(); + await tools.patchNode.execute( + { pageId: 'p1', nodeId: 'n1', node: '{"type":"paragraph"}' } as never, + {} as never, + ); + expect(forwarded.patchNode[0]).toEqual([ + 'p1', + 'n1', + { type: 'paragraph' }, + ]); + }); + + it('patchNode: throws on an invalid JSON string', async () => { + const tools = await buildTools(); + await expect( + tools.patchNode.execute( + { pageId: 'p1', nodeId: 'n1', node: '{not json' } as never, + {} as never, + ), + ).rejects.toThrow('node was a string but not valid JSON'); + }); + + it('insertNode: parses a JSON-string node and forwards the object', async () => { + const tools = await buildTools(); + await tools.insertNode.execute( + { + pageId: 'p1', + node: '{"type":"paragraph"}', + position: 'append', + } as never, + {} as never, + ); + expect(forwarded.insertNode[0]).toEqual([ + 'p1', + { type: 'paragraph' }, + { position: 'append', anchorNodeId: undefined, anchorText: undefined }, + ]); + }); + + it('insertNode: throws on an invalid JSON string', async () => { + const tools = await buildTools(); + await expect( + tools.insertNode.execute( + { + pageId: 'p1', + node: 'not json', + position: 'append', + } as never, + {} as never, + ), + ).rejects.toThrow('node was a string but not valid JSON'); + }); + + it('updatePageJson: content undefined forwards doc undefined (title-only)', async () => { + const tools = await buildTools(); + await tools.updatePageJson.execute( + { pageId: 'p1', title: 'T' } as never, + {} as never, + ); + expect(forwarded.updatePageJson[0]).toEqual(['p1', undefined, 'T']); + }); + + it('updatePageJson: content as object is forwarded as-is', async () => { + const tools = await buildTools(); + const doc = { type: 'doc', content: [] }; + await tools.updatePageJson.execute( + { pageId: 'p1', content: doc } as never, + {} as never, + ); + expect(forwarded.updatePageJson[0]).toEqual(['p1', doc, undefined]); + }); + + it('updatePageJson: content as JSON string is parsed and forwarded as object', async () => { + const tools = await buildTools(); + await tools.updatePageJson.execute( + { + pageId: 'p1', + content: '{"type":"doc","content":[]}', + } as never, + {} as never, + ); + expect(forwarded.updatePageJson[0]).toEqual([ + 'p1', + { type: 'doc', content: [] }, + undefined, + ]); + }); + + it('updatePageJson: invalid JSON string throws', async () => { + const tools = await buildTools(); + await expect( + tools.updatePageJson.execute( + { pageId: 'p1', content: '{bad' } as never, + {} as never, + ), + ).rejects.toThrow('content was a string but not valid JSON'); + }); +}); diff --git a/apps/server/src/database/repos/ai-chat/page-embedding.repo.spec.ts b/apps/server/src/database/repos/ai-chat/page-embedding.repo.spec.ts new file mode 100644 index 00000000..0af8b750 --- /dev/null +++ b/apps/server/src/database/repos/ai-chat/page-embedding.repo.spec.ts @@ -0,0 +1,30 @@ +import { PageEmbeddingRepo } from './page-embedding.repo'; +import type { KyselyDB } from '../../types/kysely.types'; + +/** + * Unit test for the access-scoping early-return in searchByEmbedding: when + * spaceIds is empty the repo MUST return [] without touching the DB. This is a + * pure access-scoping guard — an empty space list means the caller has no + * readable spaces, so the result set is always empty. + * + * The dimension-mismatch case (filtering rows by model_dimensions) requires a + * real pgvector test database and is out of scope here. + */ +describe('PageEmbeddingRepo.searchByEmbedding early-return', () => { + it('returns [] without accessing the DB when spaceIds is empty', async () => { + const explosiveDb = new Proxy( + {}, + { + get() { + throw new Error('DB must not be accessed when spaceIds is empty'); + }, + }, + ) as unknown as KyselyDB; + + const repo = new PageEmbeddingRepo(explosiveDb); + + await expect( + repo.searchByEmbedding('ws-1', [1, 2, 3], [], 10), + ).resolves.toEqual([]); + }); +}); diff --git a/apps/server/src/integrations/crypto/secret-box.spec.ts b/apps/server/src/integrations/crypto/secret-box.spec.ts new file mode 100644 index 00000000..5db470cd --- /dev/null +++ b/apps/server/src/integrations/crypto/secret-box.spec.ts @@ -0,0 +1,66 @@ +import { SecretBoxService } from './secret-box'; + +/** + * Unit tests for SecretBoxService: the AES-256-GCM symmetric encryption that + * protects provider API keys at rest. The output layout is + * base64(salt[16] | iv[12] | authTag[16] | ciphertext). These tests guard the + * round-trip correctness, the per-record randomness (salt + iv), and the tamper + * / wrong-key detection (the error message the UI relies on). + */ +describe('SecretBoxService', () => { + function makeService(secret = 'test-secret'): SecretBoxService { + return new SecretBoxService({ + getAppSecret: () => secret, + } as never); + } + + const PLAIN = 'sk-test-1234567890'; + + it('round-trips: decryptSecret(encryptSecret(plain)) === plain', () => { + const svc = makeService(); + const blob = svc.encryptSecret(PLAIN); + expect(svc.decryptSecret(blob)).toBe(PLAIN); + }); + + it('produces different blobs for the same plaintext (random salt+iv), both decrypt back', () => { + const svc = makeService(); + const blob1 = svc.encryptSecret(PLAIN); + const blob2 = svc.encryptSecret(PLAIN); + expect(blob1).not.toBe(blob2); + expect(svc.decryptSecret(blob1)).toBe(PLAIN); + expect(svc.decryptSecret(blob2)).toBe(PLAIN); + }); + + it('detects a flipped authTag byte and throws the APP_SECRET message', () => { + const svc = makeService(); + const blob = svc.encryptSecret(PLAIN); + const data = Buffer.from(blob, 'base64'); + // authTag region: SALT[16] + IV[12] = offset 28, AUTH_TAG[16] = 28..43 + data[28] ^= 0xff; + const tampered = data.toString('base64'); + expect(() => svc.decryptSecret(tampered)).toThrow( + 'APP_SECRET may have changed', + ); + }); + + it('detects a flipped ciphertext byte and throws the APP_SECRET message', () => { + const svc = makeService(); + const blob = svc.encryptSecret(PLAIN); + const data = Buffer.from(blob, 'base64'); + // ciphertext region: SALT[16] + IV[12] + AUTH_TAG[16] = offset 44 + data[44] ^= 0xff; + const tampered = data.toString('base64'); + expect(() => svc.decryptSecret(tampered)).toThrow( + 'APP_SECRET may have changed', + ); + }); + + it('throws the APP_SECRET message when decrypted with a different APP_SECRET', () => { + const encrypter = makeService('secret-a'); + const blob = encrypter.encryptSecret(PLAIN); + const decrypter = makeService('secret-b'); + expect(() => decrypter.decryptSecret(blob)).toThrow( + 'APP_SECRET may have changed', + ); + }); +});