pangolin/server/private/lib/rateLimit.test.ts

/*
 * This file is part of a proprietary work.
 *
 * Copyright (c) 2025 Fossorial, Inc.
 * All rights reserved.
 *
 * This file is licensed under the Fossorial Commercial License.
 * You may not use this file except in compliance with the License.
 * Unauthorized use, copying, modification, or distribution is strictly prohibited.
 *
 * This file is not licensed under the AGPLv3.
 */

// Simple test file for the rate limit service with Redis
// Run with: npx ts-node rateLimitService.test.ts

import { RateLimitService } from "./rateLimit";

function generateClientId() {
    return "client-" + Math.random().toString(36).substring(2, 15);
}

async function runTests() {
    console.log("Starting Rate Limit Service Tests...\n");

    const rateLimitService = new RateLimitService();
    let testsPassed = 0;
    let testsTotal = 0;

    // Helper function to run a test
    async function test(name: string, testFn: () => Promise<void>) {
        testsTotal++;
        try {
            await testFn();
            console.log(`✅ ${name}`);
            testsPassed++;
        } catch (error) {
            console.log(`❌ ${name}: ${error}`);
        }
    }

    // Helper function for assertions
    function assert(condition: boolean, message: string) {
        if (!condition) {
            throw new Error(message);
        }
    }

    // Test 1: Basic rate limiting
    await test("Should allow requests under the limit", async () => {
        const clientId = generateClientId();
        const maxRequests = 5;

        for (let i = 0; i < maxRequests - 1; i++) {
            const result = await rateLimitService.checkRateLimit(
                clientId,
                undefined,
                maxRequests
            );
            assert(!result.isLimited, `Request ${i + 1} should be allowed`);
            assert(
                result.totalHits === i + 1,
                `Expected ${i + 1} hits, got ${result.totalHits}`
            );
        }
    });

    // Test 2: Rate limit blocking
    await test("Should block requests over the limit", async () => {
        const clientId = generateClientId();
        const maxRequests = 30;

        // Use up all allowed requests
        for (let i = 0; i < maxRequests - 1; i++) {
            const result = await rateLimitService.checkRateLimit(
                clientId,
                undefined,
                maxRequests
            );
            assert(!result.isLimited, `Request ${i + 1} should be allowed`);
        }

        // Next request should be blocked
        const blockedResult = await rateLimitService.checkRateLimit(
            clientId,
            undefined,
            maxRequests
        );
        assert(blockedResult.isLimited, "Request should be blocked");
        assert(
            blockedResult.reason === "global",
            "Should be blocked for global reason"
        );
    });

    // Test 3: Message type limits
    await test("Should handle message type limits", async () => {
        const clientId = generateClientId();
        const globalMax = 10;
        const messageTypeMax = 2;

        // Send messages of type 'ping' up to the limit
        for (let i = 0; i < messageTypeMax - 1; i++) {
            const result = await rateLimitService.checkRateLimit(
                clientId,
                "ping",
                globalMax,
                messageTypeMax
            );
            assert(
                !result.isLimited,
                `Ping message ${i + 1} should be allowed`
            );
        }

        // Next 'ping' should be blocked
        const blockedResult = await rateLimitService.checkRateLimit(
            clientId,
            "ping",
            globalMax,
            messageTypeMax
        );
        assert(blockedResult.isLimited, "Ping message should be blocked");
        assert(
            blockedResult.reason === "message_type:ping",
            "Should be blocked for message type"
        );

        // Other message types should still work
        const otherResult = await rateLimitService.checkRateLimit(
            clientId,
            "pong",
            globalMax,
            messageTypeMax
        );
        assert(!otherResult.isLimited, "Pong message should be allowed");
    });

    // Test 4: Reset functionality
    await test("Should reset client correctly", async () => {
        const clientId = generateClientId();
        const maxRequests = 3;

        // Use up some requests
        await rateLimitService.checkRateLimit(clientId, undefined, maxRequests);
        await rateLimitService.checkRateLimit(clientId, "test", maxRequests);

        // Reset the client
        await rateLimitService.resetKey(clientId);

        // Should be able to make fresh requests
        const result = await rateLimitService.checkRateLimit(
            clientId,
            undefined,
            maxRequests
        );
        assert(!result.isLimited, "Request after reset should be allowed");
        assert(result.totalHits === 1, "Should have 1 hit after reset");
    });

    // Test 5: Different clients are independent
    await test("Should handle different clients independently", async () => {
        const client1 = generateClientId();
        const client2 = generateClientId();
        const maxRequests = 2;

        // Client 1 uses up their limit
        await rateLimitService.checkRateLimit(client1, undefined, maxRequests);
        await rateLimitService.checkRateLimit(client1, undefined, maxRequests);
        const client1Blocked = await rateLimitService.checkRateLimit(
            client1,
            undefined,
            maxRequests
        );
        assert(client1Blocked.isLimited, "Client 1 should be blocked");

        // Client 2 should still be able to make requests
        const client2Result = await rateLimitService.checkRateLimit(
            client2,
            undefined,
            maxRequests
        );
        assert(!client2Result.isLimited, "Client 2 should not be blocked");
        assert(client2Result.totalHits === 1, "Client 2 should have 1 hit");
    });

    // Test 6: Decrement functionality
    await test("Should decrement correctly", async () => {
        const clientId = generateClientId();
        const maxRequests = 5;

        // Make some requests
        await rateLimitService.checkRateLimit(clientId, undefined, maxRequests);
        await rateLimitService.checkRateLimit(clientId, undefined, maxRequests);
        let result = await rateLimitService.checkRateLimit(
            clientId,
            undefined,
            maxRequests
        );
        assert(result.totalHits === 3, "Should have 3 hits before decrement");

        // Decrement
        await rateLimitService.decrementRateLimit(clientId);

        // Next request should reflect the decrement
        result = await rateLimitService.checkRateLimit(
            clientId,
            undefined,
            maxRequests
        );
        assert(
            result.totalHits === 3,
            "Should have 3 hits after decrement + increment"
        );
    });

    // Wait a moment for any pending Redis operations
    console.log("\nWaiting for Redis sync...");
    await new Promise((resolve) => setTimeout(resolve, 1000));

    // Force sync to test Redis integration
    await test("Should sync to Redis", async () => {
        await rateLimitService.forceSyncAllPendingData();
        // If this doesn't throw, Redis sync is working
        assert(true, "Redis sync completed");
    });

    // Cleanup
    await rateLimitService.cleanup();

    // Results
    console.log(`\n--- Test Results ---`);
    console.log(`✅ Passed: ${testsPassed}/${testsTotal}`);
    console.log(`❌ Failed: ${testsTotal - testsPassed}/${testsTotal}`);

    if (testsPassed === testsTotal) {
        console.log("\n🎉 All tests passed!");
        process.exit(0);
    } else {
        console.log("\n💥 Some tests failed!");
        process.exit(1);
    }
}

// Run the tests
runTests().catch((error) => {
    console.error("Test runner error:", error);
    process.exit(1);
});