92 lines
6.7 KiB
TypeScript
92 lines
6.7 KiB
TypeScript
/** @jest-environment node */
|
|
import { EventEmitter } from 'node:events'
|
|
import { lookup } from 'node:dns/promises'
|
|
import { request as httpsRequest } from 'node:https'
|
|
import { requestAiModelOperation, requestAiModelOutput } from '@/lib/ai-model-http'
|
|
|
|
jest.mock('node:dns/promises', () => ({ lookup: jest.fn() }))
|
|
jest.mock('node:https', () => ({ request: jest.fn() }))
|
|
jest.mock('@/lib/ai-model-access', () => ({ AiModelError: class extends Error { constructor(message: string, public status = 400, public code = 'bad_request') { super(message) } } }))
|
|
let sent: { url: string; options: any }[]
|
|
beforeEach(() => {
|
|
jest.clearAllMocks(); sent = []
|
|
;(lookup as jest.Mock).mockResolvedValue([{ address: '93.184.216.34', family: 4 }])
|
|
;(httpsRequest as jest.Mock).mockImplementation((url, options, callback) => {
|
|
sent.push({ url: String(url), options })
|
|
const req = new EventEmitter() as any
|
|
req.write = () => true; req.destroy = jest.fn()
|
|
req.end = () => queueMicrotask(() => { const response = new EventEmitter() as any; response.statusCode = 200; response.headers = {}; response.destroy = jest.fn(); callback(response); response.emit('data', Buffer.from('{}')); response.emit('end') })
|
|
return req
|
|
})
|
|
})
|
|
|
|
test('retains configured /v1 prefix for standard operations and roots the separate extension path', async () => {
|
|
await requestAiModelOperation('https://models.example.com/v1', 'synthetic-secret', '/v1/messages', '{}', 'request-1', { 'Anthropic-Version': '2023-06-01', 'X-Relay-Route-Option': 'route-1' })
|
|
await requestAiModelOperation('https://models.example.com/v1', 'synthetic-secret', '/extensions/v1/tasks/task-1/result', undefined, undefined, {})
|
|
expect(sent.map(item => item.url)).toEqual(['https://models.example.com/v1/messages', 'https://models.example.com/extensions/v1/tasks/task-1/result'])
|
|
expect(sent[0].options.headers).toMatchObject({ Authorization: 'Bearer synthetic-secret', 'Idempotency-Key': 'request-1', 'Anthropic-Version': '2023-06-01', 'X-Relay-Route-Option': 'route-1' })
|
|
expect(sent[1].options.method).toBe('GET')
|
|
expect(sent[1].options.headers['Idempotency-Key']).toBeUndefined()
|
|
})
|
|
|
|
test('retained response endpoint only permits an authenticated GET on the configured public origin', async () => {
|
|
const path = '/api/external/requests/1c8c2fa5-5ad0-436b-977f-a047a64fa845/result'
|
|
await requestAiModelOperation('https://models.example.com/v1', 'secret', path, undefined, undefined, {})
|
|
expect(sent).toHaveLength(1)
|
|
expect(sent[0].url).toBe('https://models.example.com' + path)
|
|
expect(sent[0].options.method).toBe('GET')
|
|
expect(sent[0].options.headers).toEqual({ Authorization: 'Bearer secret' })
|
|
await expect(requestAiModelOperation('https://models.example.com/v1', 'secret', path, '{}', undefined, {})).rejects.toMatchObject({ code: 'protocol_invalid' })
|
|
await expect(requestAiModelOperation('https://models.example.com/v1', 'secret', path, undefined, 'new-generation', {})).rejects.toMatchObject({ code: 'protocol_invalid' })
|
|
expect(sent).toHaveLength(1)
|
|
})
|
|
|
|
test('accepted task results reject a body, a new idempotency key or route overrides before transport', async () => {
|
|
const path = '/extensions/v1/tasks/original-task/result'
|
|
for (const [body, key, headers] of [
|
|
['{}', undefined, {}], [undefined, 'new-generation', {}], [undefined, undefined, { 'X-Relay-Route-Option': 'different-route' }],
|
|
] as const) await expect(requestAiModelOperation('https://models.example.com/v1', 'secret', path, body, key, headers)).rejects.toMatchObject({ code: 'protocol_invalid' })
|
|
expect(lookup).not.toHaveBeenCalled(); expect(httpsRequest).not.toHaveBeenCalled()
|
|
})
|
|
|
|
test.each(['/v1/unknown', '//169.254.169.254/metadata', '/extensions/v1/tasks/../../secret/result', '/extensions/v1/tasks/task-1/result?url=private'])('rejects unsupported or injected operation %s before DNS', async path => {
|
|
await expect(requestAiModelOperation('https://models.example.com/v1', 'secret', path, '{}', 'request-1', {})).rejects.toMatchObject({ code: 'protocol_unsupported' })
|
|
expect(lookup).not.toHaveBeenCalled(); expect(httpsRequest).not.toHaveBeenCalled()
|
|
})
|
|
|
|
test.each([{ Authorization: 'Bearer attacker' }, { Host: 'localhost' }, { 'X-Relay-Route-Option': 'route\r\ninjected' }])('rejects arbitrary headers before transport', async headers => {
|
|
await expect(requestAiModelOperation('https://models.example.com/v1', 'secret', '/v1/messages', '{}', 'request-1', headers as any)).rejects.toMatchObject({ code: 'protocol_invalid' })
|
|
expect(httpsRequest).not.toHaveBeenCalled()
|
|
})
|
|
|
|
test.each(['https://attacker.example/api/external/assets/file', 'http://models.example.com/api/external/assets/file', 'https://models.example.com/admin', 'https://models.example.com/api/external/assets/../admin', 'https://user:pass@models.example.com/api/external/assets/file'])('rejects generated asset target %s without disclosing credentials', async url => {
|
|
await expect(requestAiModelOutput('https://models.example.com/v1', 'secret', url)).rejects.toMatchObject({ code: 'output_invalid' })
|
|
expect(httpsRequest).not.toHaveBeenCalled()
|
|
})
|
|
|
|
test('same-origin assets keep fixed DNS and fresh authorization checks', async () => {
|
|
const guard = jest.fn()
|
|
await requestAiModelOutput('https://models.example.com/v1', 'secret', `https://models.example.com/api/external/assets/${'a'.repeat(216)}`, guard)
|
|
expect(guard).toHaveBeenCalledTimes(1)
|
|
expect(sent[0].options.headers.Authorization).toBe('Bearer secret')
|
|
const callback = jest.fn(); sent[0].options.lookup('models.example.com', {}, callback)
|
|
expect(callback).toHaveBeenCalledWith(null, '93.184.216.34', 4)
|
|
})
|
|
|
|
test('revoked user creates zero generation or asset sockets', async () => {
|
|
const guard = async () => { throw new Error('revoked') }
|
|
await expect(requestAiModelOperation('https://models.example.com/v1', 'secret', '/extensions/v1/tasks', '{}', 'one', {}, guard)).rejects.toThrow('revoked')
|
|
await expect(requestAiModelOutput('https://models.example.com/v1', 'secret', `https://models.example.com/api/external/assets/${'a'.repeat(32)}`, guard)).rejects.toThrow('revoked')
|
|
expect(httpsRequest).not.toHaveBeenCalled()
|
|
})
|
|
|
|
test.each([32, 216, 8192])('accepts a same-origin encrypted relay asset token of length %i', async length => {
|
|
await requestAiModelOutput('https://models.example.com/v1', 'secret', `https://models.example.com/api/external/assets/${'a'.repeat(length)}`)
|
|
expect(httpsRequest).toHaveBeenCalledTimes(1)
|
|
})
|
|
|
|
test.each([31, 8193])('rejects relay asset token length %i outside its published contract', async length => {
|
|
await expect(requestAiModelOutput('https://models.example.com/v1', 'secret', `https://models.example.com/api/external/assets/${'a'.repeat(length)}`)).rejects.toMatchObject({ code: 'output_invalid' })
|
|
expect(lookup).not.toHaveBeenCalled(); expect(httpsRequest).not.toHaveBeenCalled()
|
|
})
|