// Tests for the Discord WS + REST reverse proxy.
// Spins up a mock Discord server (HTTP + WS) and verifies the proxy
// correctly forwards REST requests, rewrites /gateway/bot URLs,
// pipes WebSocket frames bidirectionally, and propagates close events.
import { describe, test, expect, beforeAll, afterAll } from 'vitest';
import http from 'node:http';
import WebSocket, { WebSocketServer } from 'ws';
import { createDiscordRestProxyHandler, createGatewayUpgradeHandler, deriveGatewayBaseFromRest, isLoopbackAddress, } from './discord-ws-proxy.js';
// ── Test infrastructure ──────────────────────────────────────────────────
const TEST_GATEWAY_TOKEN = 'test-client-id:test-secret';
let mockDiscordServer;
let mockDiscordPort;
let mockDiscordWss;
// Track active WS connections for cleanup
const mockWsConnections = [];
beforeAll(async () => {
    // Start mock Discord server (simulates the gateway-proxy REST API)
    mockDiscordWss = new WebSocketServer({ noServer: true });
    mockDiscordServer = http.createServer((req, res) => {
        const url = new URL(req.url || '/', 'http://localhost');
        // Mock GET /api/v10/gateway/bot
        if (url.pathname === '/api/v10/gateway/bot') {
            const body = JSON.stringify({
                url: `ws://127.0.0.1:${mockDiscordPort}/gateway-ws`,
                shards: 1,
                session_start_limit: {
                    total: 1000,
                    remaining: 999,
                    reset_after: 14400000,
                    max_concurrency: 1,
                },
            });
            res.writeHead(200, {
                'content-type': 'application/json',
                'x-ratelimit-remaining': '39',
            });
            res.end(body);
            return;
        }
        // Mock generic REST endpoint
        if (url.pathname.startsWith('/api/')) {
            const chunks = [];
            req.on('data', (c) => {
                chunks.push(c);
            });
            req.on('end', () => {
                const requestBody = chunks.length > 0 ? Buffer.concat(chunks).toString() : undefined;
                res.writeHead(200, { 'content-type': 'application/json' });
                res.end(JSON.stringify({
                    echo: true,
                    method: req.method,
                    path: url.pathname,
                    body: requestBody,
                    authorization: req.headers['authorization'],
                }));
            });
            return;
        }
        res.writeHead(404);
        res.end();
    });
    // Handle WebSocket upgrades on mock Discord server
    mockDiscordServer.on('upgrade', (req, socket, head) => {
        mockDiscordWss.handleUpgrade(req, socket, head, (ws) => {
            mockWsConnections.push(ws);
            mockDiscordWss.emit('connection', ws, req);
        });
    });
    await new Promise((resolve) => {
        mockDiscordServer.listen(0, '127.0.0.1', () => {
            const addr = mockDiscordServer.address();
            mockDiscordPort = typeof addr === 'object' && addr ? addr.port : 0;
            resolve();
        });
    });
});
afterAll(async () => {
    for (const ws of mockWsConnections) {
        if (ws.readyState === WebSocket.OPEN) {
            ws.close();
        }
    }
    mockWsConnections.length = 0;
    await new Promise((r) => {
        mockDiscordWss.close(() => {
            mockDiscordServer.close(() => {
                r();
            });
        });
    });
});
// ── REST proxy tests ─────────────────────────────────────────────────────
describe('REST proxy', () => {
    test('forwards GET request to upstream and returns response', async () => {
        const { server, port } = await createTestProxy();
        const res = await fetch(`http://127.0.0.1:${port}/api/v10/channels/123/messages`, { headers: { authorization: 'Bot test-token' } });
        const body = await res.json();
        expect(res.status).toBe(200);
        expect(body.echo).toBe(true);
        expect(body.method).toBe('GET');
        expect(body.path).toBe('/api/v10/channels/123/messages');
        expect(body.authorization).toBe('Bot test-token');
        await closeServer(server);
    });
    test('forwards POST request with body', async () => {
        const { server, port } = await createTestProxy();
        const res = await fetch(`http://127.0.0.1:${port}/api/v10/channels/123/messages`, {
            method: 'POST',
            headers: {
                'content-type': 'application/json',
                authorization: 'Bot test-token',
            },
            body: JSON.stringify({ content: 'hello' }),
        });
        const body = await res.json();
        expect(body.method).toBe('POST');
        expect(body.body).toBe('{"content":"hello"}');
        await closeServer(server);
    });
    test('returns 404 for non-/api/ paths', async () => {
        const { server, port } = await createTestProxy();
        const res = await fetch(`http://127.0.0.1:${port}/something-else`);
        expect(res.status).toBe(404);
        await closeServer(server);
    });
    test('rewrites /gateway/bot URL to local WS proxy', async () => {
        const { server, port } = await createTestProxy();
        const res = await fetch(`http://127.0.0.1:${port}/api/v10/gateway/bot`, { headers: { authorization: 'Bot test-token' } });
        const body = await res.json();
        // The url should be rewritten to point to local proxy
        expect(body.url).toBe(`ws://127.0.0.1:${port}/gateway`);
        // Other fields preserved
        expect(body.shards).toBe(1);
        // Rate limit headers forwarded
        expect(res.headers.get('x-ratelimit-remaining')).toBe('39');
        await closeServer(server);
    });
});
// ── WebSocket bridge tests ───────────────────────────────────────────────
describe('WebSocket bridge', () => {
    test('bidirectional frame forwarding between gateway and local', async () => {
        const { server, port } = await createTestProxyWithWs();
        // Gateway connects first
        const gateway = new WebSocket(`ws://127.0.0.1:${port}/gateway`, {
            headers: { 'x-kimaki-role': 'gateway', authorization: TEST_GATEWAY_TOKEN },
        });
        await waitForOpen(gateway);
        // Local (discord.js) connects
        const local = new WebSocket(`ws://127.0.0.1:${port}/gateway?token=${encodeURIComponent(TEST_GATEWAY_TOKEN)}`);
        await waitForOpen(local);
        const localMessages = [];
        const gatewayMessages = [];
        local.on('message', (data) => {
            localMessages.push(data.toString());
        });
        gateway.on('message', (data) => {
            gatewayMessages.push(data.toString());
        });
        // Gateway → Local
        gateway.send('event-from-gateway');
        await waitFor(() => localMessages.length >= 1, 2000);
        expect(localMessages).toEqual(['event-from-gateway']);
        // Local → Gateway
        local.send('identify-from-local');
        await waitFor(() => gatewayMessages.length >= 1, 2000);
        expect(gatewayMessages).toEqual(['identify-from-local']);
        gateway.close();
        local.close();
        await closeServer(server);
    });
    test('local connects before gateway', async () => {
        const { server, port } = await createTestProxyWithWs();
        // Local connects first
        const local = new WebSocket(`ws://127.0.0.1:${port}/gateway?token=${encodeURIComponent(TEST_GATEWAY_TOKEN)}`);
        await waitForOpen(local);
        const localMessages = [];
        local.on('message', (data) => {
            localMessages.push(data.toString());
        });
        // Gateway connects after
        const gateway = new WebSocket(`ws://127.0.0.1:${port}/gateway`, {
            headers: { 'x-kimaki-role': 'gateway', authorization: TEST_GATEWAY_TOKEN },
        });
        await waitForOpen(gateway);
        gateway.send('hello-from-gateway');
        await waitFor(() => localMessages.length >= 1, 2000);
        expect(localMessages).toEqual(['hello-from-gateway']);
        gateway.close();
        local.close();
        await closeServer(server);
    });
    test('gateway close propagates to local', async () => {
        const { server, port } = await createTestProxyWithWs();
        const gateway = new WebSocket(`ws://127.0.0.1:${port}/gateway`, {
            headers: { 'x-kimaki-role': 'gateway', authorization: TEST_GATEWAY_TOKEN },
        });
        const local = new WebSocket(`ws://127.0.0.1:${port}/gateway?token=${encodeURIComponent(TEST_GATEWAY_TOKEN)}`);
        await waitForOpen(gateway);
        await waitForOpen(local);
        const localClosed = new Promise((resolve) => {
            local.on('close', (code, reason) => {
                resolve({ code, reason: reason.toString() });
            });
        });
        gateway.close(4001, 'gateway shutdown');
        const result = await localClosed;
        expect(result.code).toBe(4001);
        expect(result.reason).toBe('gateway shutdown');
        await closeServer(server);
    });
    test('local close propagates to gateway', async () => {
        const { server, port } = await createTestProxyWithWs();
        const gateway = new WebSocket(`ws://127.0.0.1:${port}/gateway`, {
            headers: { 'x-kimaki-role': 'gateway', authorization: TEST_GATEWAY_TOKEN },
        });
        const local = new WebSocket(`ws://127.0.0.1:${port}/gateway?token=${encodeURIComponent(TEST_GATEWAY_TOKEN)}`);
        await waitForOpen(gateway);
        await waitForOpen(local);
        const gatewayClosed = new Promise((resolve) => {
            gateway.on('close', (code, reason) => {
                resolve({ code, reason: reason.toString() });
            });
        });
        local.close(4000, 'client done');
        const result = await gatewayClosed;
        expect(result.code).toBe(4000);
        expect(result.reason).toBe('client done');
        await closeServer(server);
    });
    test('rejects upgrade on non-/gateway path', async () => {
        const { server, port } = await createTestProxyWithWs();
        const client = new WebSocket(`ws://127.0.0.1:${port}/other-path`);
        const error = await new Promise((resolve) => {
            client.on('error', (err) => {
                resolve(err);
            });
        });
        expect(error).toBeDefined();
        await closeServer(server);
    });
    test('rejects gateway connection with wrong authorization', async () => {
        const { server, port } = await createTestProxyWithWs();
        const client = new WebSocket(`ws://127.0.0.1:${port}/gateway`, {
            headers: { 'x-kimaki-role': 'gateway', authorization: 'wrong-token' },
        });
        const error = await new Promise((resolve) => {
            client.on('error', (err) => {
                resolve(err);
            });
        });
        expect(error).toBeDefined();
        expect(error.message).toContain('401');
        await closeServer(server);
    });
    test('new gateway connection replaces waiting one', async () => {
        const { server, port } = await createTestProxyWithWs();
        // First gateway connects and waits
        const gateway1 = new WebSocket(`ws://127.0.0.1:${port}/gateway`, {
            headers: { 'x-kimaki-role': 'gateway', authorization: TEST_GATEWAY_TOKEN },
        });
        await waitForOpen(gateway1);
        const gateway1Closed = new Promise((resolve) => {
            gateway1.on('close', (code) => {
                resolve(code);
            });
        });
        // Second gateway replaces it
        const gateway2 = new WebSocket(`ws://127.0.0.1:${port}/gateway`, {
            headers: { 'x-kimaki-role': 'gateway', authorization: TEST_GATEWAY_TOKEN },
        });
        await waitForOpen(gateway2);
        // First should be closed
        const code = await gateway1Closed;
        expect(code).toBe(1000);
        // Now local connects and pairs with gateway2
        const local = new WebSocket(`ws://127.0.0.1:${port}/gateway?token=${encodeURIComponent(TEST_GATEWAY_TOKEN)}`);
        await waitForOpen(local);
        const localMessages = [];
        local.on('message', (data) => {
            localMessages.push(data.toString());
        });
        gateway2.send('from-gateway2');
        await waitFor(() => localMessages.length >= 1, 2000);
        expect(localMessages).toEqual(['from-gateway2']);
        gateway2.close();
        local.close();
        await closeServer(server);
    });
    test('multiple sequential bridge sessions', async () => {
        const { server, port } = await createTestProxyWithWs();
        // First session
        const g1 = new WebSocket(`ws://127.0.0.1:${port}/gateway`, {
            headers: { 'x-kimaki-role': 'gateway', authorization: TEST_GATEWAY_TOKEN },
        });
        const l1 = new WebSocket(`ws://127.0.0.1:${port}/gateway?token=${encodeURIComponent(TEST_GATEWAY_TOKEN)}`);
        await waitForOpen(g1);
        await waitForOpen(l1);
        const l1Messages = [];
        l1.on('message', (data) => {
            l1Messages.push(data.toString());
        });
        g1.send('session1');
        await waitFor(() => l1Messages.length >= 1, 2000);
        expect(l1Messages).toEqual(['session1']);
        // Close first session
        g1.close();
        l1.close();
        await waitFor(() => g1.readyState === WebSocket.CLOSED, 2000);
        await waitFor(() => l1.readyState === WebSocket.CLOSED, 2000);
        // Second session should work fine
        const g2 = new WebSocket(`ws://127.0.0.1:${port}/gateway`, {
            headers: { 'x-kimaki-role': 'gateway', authorization: TEST_GATEWAY_TOKEN },
        });
        const l2 = new WebSocket(`ws://127.0.0.1:${port}/gateway?token=${encodeURIComponent(TEST_GATEWAY_TOKEN)}`);
        await waitForOpen(g2);
        await waitForOpen(l2);
        const l2Messages = [];
        l2.on('message', (data) => {
            l2Messages.push(data.toString());
        });
        g2.send('session2');
        await waitFor(() => l2Messages.length >= 1, 2000);
        expect(l2Messages).toEqual(['session2']);
        g2.close();
        l2.close();
        await closeServer(server);
    });
});
// ── Gateway URL derivation tests ─────────────────────────────────────────
describe('deriveGatewayBaseFromRest', () => {
    test('discord.com returns gateway.discord.gg', () => {
        expect(deriveGatewayBaseFromRest('https://discord.com')).toBe('wss://gateway.discord.gg');
    });
    test('gateway proxy swaps https to wss', () => {
        expect(deriveGatewayBaseFromRest('https://discord-gateway.kimaki.xyz')).toBe('wss://discord-gateway.kimaki.xyz');
    });
    test('http proxy swaps to ws', () => {
        expect(deriveGatewayBaseFromRest('http://127.0.0.1:8080')).toBe('ws://127.0.0.1:8080');
    });
    test('strips trailing slash', () => {
        expect(deriveGatewayBaseFromRest('https://proxy.example.com/')).toBe('wss://proxy.example.com');
    });
    test('invalid URL falls back to default', () => {
        expect(deriveGatewayBaseFromRest('not-a-url')).toBe('wss://gateway.discord.gg');
    });
});
describe('isLoopbackAddress', () => {
    test('returns true for loopback variants', () => {
        expect(isLoopbackAddress('::1')).toBe(true);
        expect(isLoopbackAddress('127.0.0.1')).toBe(true);
        expect(isLoopbackAddress('127.12.34.56')).toBe(true);
        expect(isLoopbackAddress('::ffff:127.0.0.1')).toBe(true);
    });
    test('returns false for non-loopback addresses', () => {
        expect(isLoopbackAddress('192.168.1.5')).toBe(false);
        expect(isLoopbackAddress('10.0.0.1')).toBe(false);
        expect(isLoopbackAddress('::ffff:192.168.1.5')).toBe(false);
        expect(isLoopbackAddress('')).toBe(false);
    });
});
// ── Helpers ──────────────────────────────────────────────────────────────
/**
 * Create a test proxy server with REST handler only (pointing to mock Discord).
 */
async function createTestProxy() {
    const handler = createMockRestProxyHandler();
    const server = http.createServer(handler);
    const port = await new Promise((resolve) => {
        server.listen(0, '127.0.0.1', () => {
            const addr = server.address();
            resolve(typeof addr === 'object' && addr ? addr.port : 0);
        });
    });
    return { server, port };
}
/**
 * Create a test proxy server with REST and WS bridge handling.
 */
async function createTestProxyWithWs() {
    const handler = createMockRestProxyHandler();
    const server = http.createServer(handler);
    // Mount the real bridge handler
    const upgradeHandler = createGatewayUpgradeHandler({
        getGatewayToken: () => TEST_GATEWAY_TOKEN,
    });
    server.on('upgrade', (req, socket, head) => {
        upgradeHandler(req, socket, head);
    });
    const port = await new Promise((resolve) => {
        server.listen(0, '127.0.0.1', () => {
            const addr = server.address();
            resolve(typeof addr === 'object' && addr ? addr.port : 0);
        });
    });
    return { server, port };
}
/**
 * REST proxy handler for tests — forwards to mock Discord server.
 */
function createMockRestProxyHandler() {
    return async (req, res) => {
        const url = new URL(req.url || '/', 'http://localhost');
        if (!url.pathname.startsWith('/api/')) {
            res.writeHead(404);
            res.end();
            return;
        }
        const upstreamUrl = `http://127.0.0.1:${mockDiscordPort}${url.pathname}${url.search}`;
        const headers = {};
        for (const [key, value] of Object.entries(req.headers)) {
            if (key === 'host' || key === 'connection') {
                continue;
            }
            if (typeof value === 'string') {
                headers[key] = value;
            }
        }
        const chunks = [];
        for await (const chunk of req) {
            chunks.push(chunk);
        }
        const body = chunks.length > 0 ? Buffer.concat(chunks) : undefined;
        const upstreamRes = await fetch(upstreamUrl, {
            method: req.method || 'GET',
            headers,
            body,
        }).catch((err) => {
            res.writeHead(502);
            res.end(err.message);
            return null;
        });
        if (!upstreamRes) {
            return;
        }
        // Intercept /gateway/bot for URL rewrite
        const isGatewayBot = url.pathname.endsWith('/gateway/bot');
        if (isGatewayBot && upstreamRes.ok) {
            const responseBody = await upstreamRes.json();
            const addr = res.socket?.localPort;
            responseBody.url = `ws://127.0.0.1:${addr}/gateway`;
            const rewritten = JSON.stringify(responseBody);
            const responseHeaders = {
                'content-type': 'application/json',
            };
            upstreamRes.headers.forEach((value, key) => {
                if (key.startsWith('x-ratelimit') || key === 'retry-after') {
                    responseHeaders[key] = value;
                }
            });
            res.writeHead(upstreamRes.status, responseHeaders);
            res.end(rewritten);
            return;
        }
        const responseBody = Buffer.from(await upstreamRes.arrayBuffer());
        const responseHeaders = {};
        upstreamRes.headers.forEach((value, key) => {
            if (key !== 'transfer-encoding') {
                responseHeaders[key] = value;
            }
        });
        responseHeaders['content-length'] = responseBody.length.toString();
        res.writeHead(upstreamRes.status, responseHeaders);
        res.end(responseBody);
    };
}
function closeServer(server) {
    return new Promise((resolve) => {
        server.close(() => {
            resolve();
        });
    });
}
function sleep(ms) {
    return new Promise((resolve) => {
        setTimeout(resolve, ms);
    });
}
function waitForOpen(ws) {
    return new Promise((resolve, reject) => {
        if (ws.readyState === WebSocket.OPEN) {
            resolve();
            return;
        }
        ws.on('open', () => {
            resolve();
        });
        ws.on('error', (err) => {
            reject(err);
        });
    });
}
async function waitFor(condition, timeoutMs) {
    const start = Date.now();
    while (!condition()) {
        if (Date.now() - start > timeoutMs) {
            throw new Error(`waitFor timed out after ${timeoutMs}ms`);
        }
        await sleep(20);
    }
}
