File size: 4,598 Bytes
4e23b01 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 | import { mkdtemp, readFile, rm, stat } from 'node:fs/promises';
import { tmpdir } from 'node:os';
import { join } from 'node:path';
import { afterAll, afterEach, beforeAll, describe, expect, it } from 'vitest';
import { WebSocket, type RawData } from 'ws';
import { type RunningServer, startServer } from '../src/start';
import { TEST_HOST_IDENTITY } from './helpers/hostIdentity';
function rawToString(data: RawData): string {
if (typeof data === 'string') return data;
if (Buffer.isBuffer(data)) return data.toString('utf8');
if (Array.isArray(data)) return Buffer.concat(data).toString('utf8');
return Buffer.from(data as ArrayBuffer).toString('utf8');
}
function openConn(url: string, protocols: string[]): Promise<{ ws: WebSocket; firstFrame: unknown }> {
return new Promise((resolve, reject) => {
const ws = new WebSocket(url, protocols);
ws.once('message', (data) => {
try {
resolve({ ws, firstFrame: JSON.parse(rawToString(data)) });
} catch {
resolve({ ws, firstFrame: null });
}
});
ws.once('error', reject);
});
}
function expectRejected(url: string): Promise<void> {
return new Promise((resolve, reject) => {
const ws = new WebSocket(url);
const done = (err?: Error): void => {
clearTimeout(t);
ws.removeAllListeners();
try {
ws.terminate();
} catch {
}
if (err === undefined) resolve();
else reject(err);
};
const t = setTimeout(
() => done(new Error('connection was not rejected within timeout')),
1500,
);
ws.once('open', () => done(new Error('connection unexpectedly opened')));
ws.once('error', () => done());
ws.once('close', () => done());
});
}
describe('production auth wiring', () => {
let server: RunningServer | undefined;
let home: string | undefined;
let base: string;
const sockets: WebSocket[] = [];
beforeAll(async () => {
home = await mkdtemp(join(tmpdir(), 'kimi-server-v2-auth-wiring-'));
await boot();
});
async function boot(): Promise<void> {
server = await startServer({ hostIdentity: TEST_HOST_IDENTITY, host: '127.0.0.1', port: 0, homeDir: home, logLevel: 'silent' });
base = `http://127.0.0.1:${server.port}`;
}
afterEach(() => {
for (const ws of sockets.splice(0)) {
try {
ws.close();
} catch {
}
}
});
afterAll(async () => {
if (server !== undefined) {
await server.close();
server = undefined;
}
if (home !== undefined) {
await rm(home, { recursive: true, force: true });
home = undefined;
}
});
it.skipIf(process.platform === 'win32')('writes a 0600 token file at boot and keeps it on close', async () => {
const p = join(home as string, 'server.token');
const info = await stat(p);
expect(info.mode & 0o777).toBe(0o600);
const token = (await readFile(p, 'utf8')).trim();
expect(token.length).toBeGreaterThan(0);
await (server as RunningServer).close();
server = undefined;
const after = await stat(p);
expect(after.mode & 0o777).toBe(0o600);
await boot();
});
it('gates HTTP: 200 with the token, 401 without', async () => {
const token = (await readFile(join(home as string, 'server.token'), 'utf8')).trim();
const ok = await fetch(`${base}/openapi.json`, {
headers: { Authorization: `Bearer ${token}` },
});
expect(ok.status).toBe(200);
const bad = await fetch(`${base}/openapi.json`);
expect(bad.status).toBe(401);
const body = (await bad.json()) as { code: number };
expect(body.code).toBe(40101);
});
it('gates /asyncapi.json: 200 with the token, 401 without', async () => {
const token = (await readFile(join(home as string, 'server.token'), 'utf8')).trim();
const ok = await fetch(`${base}/asyncapi.json`, {
headers: { Authorization: `Bearer ${token}` },
});
expect(ok.status).toBe(200);
const doc = (await ok.json()) as { asyncapi?: string };
expect(doc.asyncapi).toBeDefined();
const bad = await fetch(`${base}/asyncapi.json`);
expect(bad.status).toBe(401);
});
it('gates WS: server_hello with the token, rejected without', async () => {
const token = (await readFile(join(home as string, 'server.token'), 'utf8')).trim();
const wsUrl = `ws://127.0.0.1:${(server as RunningServer).port}/api/v1/ws`;
const { ws, firstFrame } = await openConn(wsUrl, [`kimi-code.bearer.${token}`]);
sockets.push(ws);
expect(firstFrame).toMatchObject({ type: 'server_hello' });
await expectRejected(wsUrl);
});
});
|