trivial-pow-limiter

-- Proof-of-Work Rate Limiter for nginx
-- Pure Lua implementation with no external dependencies
-- Place this in your nginx config with: access_by_lua_file /path/to/pow_ratelimit.lua;
-- entirely generated by claude with relatively minimal prompting, probably has bugs

-- Configuration
local cookie_name = "pow_token"
local difficulty = 3  -- Number of leading zeros required in hex (adjust for difficulty)
local token_lifetime = 3600  -- Token valid for 1 hour (in seconds)
local challenge_secret = "PUT SOMETHING HERE"  -- Secret for challenge generation
local time_bucket_size = 300  -- Time bucket size in seconds (5 minutes) - allows some time drift during solving

-- ============================================================================
-- Pure Lua SHA-256 Implementation
-- ============================================================================

-- Detect bit operations library (LuaJIT vs Lua 5.3+)
local bit_module = bit or bit32
local band, bor, bxor, bnot, rshift, lshift

if bit_module then
    -- LuaJIT or bit32 library
    band = bit_module.band
    bor = bit_module.bor
    bxor = bit_module.bxor
    bnot = bit_module.bnot
    rshift = bit_module.rshift
    lshift = bit_module.lshift
end

local function sha256(msg)
    local MOD = 2^32

    local function rotr(n, b)
        n = n % MOD
        return bor(rshift(n, b), lshift(n, 32 - b))
    end

    local function shr(n, b)
        return rshift(n % MOD, b)
    end

    local k = {
        0x428a2f98, 0x71374491, 0xb5c0fbcf, 0xe9b5dba5, 0x3956c25b, 0x59f111f1, 0x923f82a4, 0xab1c5ed5,
        0xd807aa98, 0x12835b01, 0x243185be, 0x550c7dc3, 0x72be5d74, 0x80deb1fe, 0x9bdc06a7, 0xc19bf174,
        0xe49b69c1, 0xefbe4786, 0x0fc19dc6, 0x240ca1cc, 0x2de92c6f, 0x4a7484aa, 0x5cb0a9dc, 0x76f988da,
        0x983e5152, 0xa831c66d, 0xb00327c8, 0xbf597fc7, 0xc6e00bf3, 0xd5a79147, 0x06ca6351, 0x14292967,
        0x27b70a85, 0x2e1b2138, 0x4d2c6dfc, 0x53380d13, 0x650a7354, 0x766a0abb, 0x81c2c92e, 0x92722c85,
        0xa2bfe8a1, 0xa81a664b, 0xc24b8b70, 0xc76c51a3, 0xd192e819, 0xd6990624, 0xf40e3585, 0x106aa070,
        0x19a4c116, 0x1e376c08, 0x2748774c, 0x34b0bcb5, 0x391c0cb3, 0x4ed8aa4a, 0x5b9cca4f, 0x682e6ff3,
        0x748f82ee, 0x78a5636f, 0x84c87814, 0x8cc70208, 0x90befffa, 0xa4506ceb, 0xbef9a3f7, 0xc67178f2,
    }

    local h = {
        0x6a09e667, 0xbb67ae85, 0x3c6ef372, 0xa54ff53a,
        0x510e527f, 0x9b05688c, 0x1f83d9ab, 0x5be0cd19,
    }

    local len = #msg
    local msg_len_bits = len * 8

    msg = msg .. string.char(0x80)

    local padding_len = (56 - (#msg % 64)) % 64
    msg = msg .. string.rep(string.char(0), padding_len)

    for i = 7, 0, -1 do
        msg = msg .. string.char(math.floor(msg_len_bits / 2^(i * 8)) % 256)
    end

    for chunk_start = 1, #msg, 64 do
        local w = {}

        for i = 0, 15 do
            local offset = chunk_start + i * 4
            local b1, b2, b3, b4 = string.byte(msg, offset, offset + 3)
            w[i] = ((b1 * 256 + b2) * 256 + b3) * 256 + b4
        end

        for i = 16, 63 do
            local s0 = bxor(bxor(rotr(w[i-15], 7), rotr(w[i-15], 18)), shr(w[i-15], 3))
            local s1 = bxor(bxor(rotr(w[i-2], 17), rotr(w[i-2], 19)), shr(w[i-2], 10))
            w[i] = (w[i-16] + s0 + w[i-7] + s1) % MOD
        end

        local a, b, c, d, e, f, g, h_val = h[1], h[2], h[3], h[4], h[5], h[6], h[7], h[8]

        for i = 0, 63 do
            a, b, c, d, e, f, g, h_val = a % MOD, b % MOD, c % MOD, d % MOD, e % MOD, f % MOD, g % MOD, h_val % MOD
            local S1 = bxor(bxor(rotr(e, 6), rotr(e, 11)), rotr(e, 25))
            local ch = bxor(band(e, f), band(bnot(e), g))
            local temp1 = (h_val + S1 + ch + k[i+1] + w[i]) % MOD
            local S0 = bxor(bxor(rotr(a, 2), rotr(a, 13)), rotr(a, 22))
            local maj = bxor(bxor(band(a, b), band(a, c)), band(b, c))
            local temp2 = (S0 + maj) % MOD

            h_val = g
            g = f
            f = e
            e = (d + temp1) % MOD
            d = c
            c = b
            b = a
            a = (temp1 + temp2) % MOD
        end

        h[1] = (h[1] + a) % MOD
        h[2] = (h[2] + b) % MOD
        h[3] = (h[3] + c) % MOD
        h[4] = (h[4] + d) % MOD
        h[5] = (h[5] + e) % MOD
        h[6] = (h[6] + f) % MOD
        h[7] = (h[7] + g) % MOD
        h[8] = (h[8] + h_val) % MOD
    end

    local result = ""
    for i = 1, 8 do
        result = result .. string.format("%08x", h[i] % MOD)
    end

    return result
end


-- ============================================================================
-- Logging functions (compatible with nginx and standalone)
-- ============================================================================

local log_funcs = {}

if ngx then
    -- Running in nginx
    log_funcs.warn = function(msg) ngx.log(ngx.WARN, msg) end
    log_funcs.info = function(msg) ngx.log(ngx.INFO, msg) end
    log_funcs.err = function(msg) ngx.log(ngx.ERR, msg) end
else
    -- Running standalone
    log_funcs.warn = function(msg) print("[WARN] " .. msg) end
    log_funcs.info = function(msg) print("[INFO] " .. msg) end
    log_funcs.err = function(msg) print("[ERR] " .. msg) end
end

-- ============================================================================
-- Challenge Generation
-- ============================================================================

-- Generate a challenge based on secret, client IP, and time bucket
-- The time bucket allows for some clock drift between challenge generation and solution
local function generate_challenge(client_ip, timestamp)
    local time_bucket = math.floor(timestamp / time_bucket_size)
    local challenge_input = challenge_secret .. client_ip .. tostring(time_bucket)
    return sha256(challenge_input)
end

-- Validate that a challenge is valid for the given client IP within the time window
local function is_valid_challenge(challenge, client_ip, current_time)
    -- Check current time bucket
    local current_challenge = generate_challenge(client_ip, current_time)
    if challenge == current_challenge then
        return true
    end

    -- Check previous time bucket (to allow for time drift during solving)
    local previous_challenge = generate_challenge(client_ip, current_time - time_bucket_size)
    if challenge == previous_challenge then
        return true
    end

    -- Check next time bucket (for minor clock skew)
    local next_challenge = generate_challenge(client_ip, current_time + time_bucket_size)
    if challenge == next_challenge then
        return true
    end

    return false
end

-- ============================================================================
-- Token Validation Logic
-- ============================================================================

local function validate_token(token_string, current_time, client_ip)
    -- Parse the token format: challenge-nonce
    local challenge, nonce = string.match(token_string, "^([^-]+)-(.+)$")

    if not challenge or not nonce then
        log_funcs.warn("Invalid PoW token: invalid format (expected challenge-nonce)")
        return false
    end

    -- Convert nonce to number
    nonce = tonumber(nonce)
    if not nonce then
        log_funcs.warn("Invalid PoW token: nonce is not a valid number")
        return false
    end

    -- Validate that the challenge is valid for this client IP within the time window
    if not is_valid_challenge(challenge, client_ip, current_time) then
        log_funcs.warn("Invalid PoW token: challenge not valid for client IP or expired")
        return false
    end

    -- Verify the hash is actually correct for the challenge + nonce
    local input = tostring(challenge) .. tostring(nonce)
    local computed_hash = sha256(input)

    -- Verify hash meets difficulty requirement
    local target_prefix = string.rep("0", difficulty)
    if string.sub(computed_hash, 1, difficulty) ~= target_prefix then
        log_funcs.warn("Invalid PoW token: insufficient difficulty")
        return false
    end

    return true
end

-- ============================================================================
-- Main execution (only runs in nginx context)
-- ============================================================================

if ngx then
    -- Get current time
    local current_time = ngx.now()

    -- Get client IP from the specified header
    local client_ip = ngx.var.http_proxy_forwarded_for
    if not client_ip or client_ip == "" then
        -- Fall back to direct connection IP if header is not present
        client_ip = ngx.var.remote_addr
        log_funcs.warn("proxy_forwarded_for header not found, using remote_addr: " .. client_ip)
    end

    -- Check if request is from a legitimate search engine crawler and allow it through
    local user_agent = ngx.var.http_user_agent
    if user_agent then
        local user_agent_lower = string.lower(user_agent)
        -- Check for various search engine bot user agents
        if string.find(user_agent_lower, "googlebot") or
           string.find(user_agent_lower, "adsbot%-google") or
           string.find(user_agent_lower, "mediapartners%-google") or
           string.find(user_agent_lower, "google%-inspectiontool") or
           string.find(user_agent_lower, "bingbot") or
           string.find(user_agent_lower, "msnbot") or
           string.find(user_agent_lower, "slurp") or  -- Yahoo
           string.find(user_agent_lower, "duckduckbot") or
           string.find(user_agent_lower, "baiduspider") or
           string.find(user_agent_lower, "yandexbot") or
           string.find(user_agent_lower, "sogou") or  -- Sogou (Chinese search engine)
           string.find(user_agent_lower, "exabot") or  -- Exalead
           string.find(user_agent_lower, "facebot") or  -- Facebook
           string.find(user_agent_lower, "chatgpt-user") or -- ChatGPT live sessions
           string.find(user_agent_lower, "oai-searchbot") or -- ChatGPT search feature
           string.find(user_agent_lower, "ia_archiver") then  -- Alexa/Internet Archive
            log_funcs.info("Search engine crawler detected, allowing request: " .. user_agent)
            return  -- Allow the request to proceed without PoW challenge
        end
    end

    -- Check if the PoW cookie exists and validate it
    local cookies = ngx.var.http_cookie
    local has_valid_token = false

    if cookies then
        local token = ngx.var["cookie_" .. cookie_name]
        if token and token ~= "" then
            has_valid_token = validate_token(token, current_time, client_ip)
        end
    end

    -- If no valid token, serve the PoW challenge page
    if not has_valid_token then
        -- Generate the challenge for this client
        local challenge = generate_challenge(client_ip, current_time)
        ngx.header.content_type = "text/html"
        ngx.header["Cache-Control"] = "no-store, no-cache, must-revalidate, max-age=0"
        ngx.header["Pragma"] = "no-cache"
        ngx.header["Expires"] = "0"
        ngx.status = 403

        local html = string.format([[
<!DOCTYPE html>
<html>
<head>
    <title>Proof of Work Challenge</title>
    <style>
        body {
            font-family: Arial, sans-serif;
            display: flex;
            justify-content: center;
            align-items: center;
            height: 100vh;
            margin: 0;
            background: linear-gradient(135deg, #667eea 0%%, #764ba2 100%%);
            color: white;
        }
        .container {
            text-align: center;
            background: rgba(0, 0, 0, 0.3);
            padding: 40px;
            border-radius: 10px;
            box-shadow: 0 8px 32px rgba(0, 0, 0, 0.3);
        }
        .spinner {
            border: 4px solid rgba(255, 255, 255, 0.3);
            border-radius: 50%%;
            border-top: 4px solid white;
            width: 40px;
            height: 40px;
            animation: spin 1s linear infinite;
            margin: 20px auto;
        }
        @keyframes spin {
            0%% { transform: rotate(0deg); }
            100%% { transform: rotate(360deg); }
        }
        #status {
            margin-top: 20px;
            font-size: 14px;
            opacity: 0.8;
        }
    </style>
</head>
<body>
    <div class="container">
        <h1>Computing Proof of Work...</h1>
        <div class="spinner"></div>
        <div id="status">Please wait while we verify your request</div>
    </div>

    <script>
        const DIFFICULTY = %d;
        const CHALLENGE = '%s';

        // Precompute difficulty check parameters
        const fullZeroBytes = Math.floor(DIFFICULTY / 2);
        const hasPartialByte = DIFFICULTY %% 2 === 1;

        // Reuse TextEncoder instance (significant performance gain)
        const encoder = new TextEncoder();

        // Check if hash meets difficulty requirement by examining raw bytes
        // This is much faster than converting to hex string first
        function meetsdifficulty(hashBytes) {
            // Check full zero bytes
            for (let i = 0; i < fullZeroBytes; i++) {
                if (hashBytes[i] !== 0) return false;
            }

            // Check partial byte (high nibble must be 0)
            if (hasPartialByte) {
                if (hashBytes[fullZeroBytes] >= 16) return false;
            }

            return true;
        }

        async function solveProofOfWork() {
            const statusEl = document.getElementById('status');
            let nonce = 0;
            let attempts = 0;
            const startTime = Date.now();

            while (true) {
                // Build input string and encode it
                const input = CHALLENGE + nonce;
                const inputBytes = encoder.encode(input);

                // Compute SHA-256
                const hashBuffer = await crypto.subtle.digest('SHA-256', inputBytes);
                const hashBytes = new Uint8Array(hashBuffer);

                attempts++;

                // Update UI less frequently for better performance
                if (attempts %% 10000 === 0) {
                    const elapsed = ((Date.now() - startTime) / 1000).toFixed(1);
                    const rate = Math.round(attempts / elapsed);
                    statusEl.textContent = `Attempts: ${attempts.toLocaleString()} | Time: ${elapsed}s | Rate: ${rate.toLocaleString()}/s`;
                    // Allow UI to update
                    await new Promise(resolve => setTimeout(resolve, 0));
                }

                // Check difficulty on raw bytes (no hex conversion needed!)
                if (meetsdifficulty(hashBytes)) {
                    const elapsed = ((Date.now() - startTime) / 1000).toFixed(2);
                    const rate = Math.round(attempts / elapsed);
                    statusEl.textContent = `Solution found! Attempts: ${attempts.toLocaleString()} | Time: ${elapsed}s | Rate: ${rate.toLocaleString()}/s`;

                    // Create the proof token (simple format: challenge-nonce)
                    const token = CHALLENGE + '-' + nonce;

                    // Set cookie (expires in 1 hour)
                    const expires = new Date(Date.now() + 3600000).toUTCString();
                    document.cookie = `%s=${token}; expires=${expires}; path=/; SameSite=Lax`;

                    // Refresh the page
                    setTimeout(() => {
                        window.location.reload();
                    }, 500);

                    return;
                }

                nonce++;
            }
        }

        // Start solving immediately
        solveProofOfWork().catch(err => {
            document.getElementById('status').textContent = 'Error: ' + err.message;
        });
    </script>
</body>
</html>
]], difficulty, challenge, cookie_name)

        ngx.say(html)
        ngx.exit(403)
    end

    -- If we reach here, the request has a valid token and should proceed normally
end

-- Export functions for testing
return {
    sha256 = sha256,
    validate_token = validate_token,
    generate_challenge = generate_challenge,
    is_valid_challenge = is_valid_challenge,
    config = {
        cookie_name = cookie_name,
        difficulty = difficulty,
        token_lifetime = token_lifetime,
        challenge_secret = challenge_secret,
        time_bucket_size = time_bucket_size,
    }
}