// Cross-worker JWKS consistency gate (MAIN :3000).
//
// MAIN runs a multi-worker cluster, and each worker caches its OWN oidc-provider
// with the JWKS baked in (lib/oidc/provider.js). A key rotation must become visible
// in the JWKS served by EVERY worker -- not only the one that handled the rotate
// POST -- otherwise a relying party load-balanced to a stale worker cannot verify
// id_tokens signed with the new active key. This gate guards the fingerprint-based
// provider-cache invalidation that makes every worker converge after a rotation.
//
// It rotates once over a PINNED connection (one worker), then fetches /idp/jwks
// many times over FRESH connections so the cluster spreads them across workers, and
// asserts the new active kid (and the retained rotated-out kid) appears in EVERY
// response and that all workers serve an identical set.
//
// Run after a coherent boot:  node idp/test/jwksConsistencyGate.js  (MAIN up; admin@local).

const http = require("http");

const PORT            = 3000;
const ADMIN_EMAIL     = "admin@local";
const ADMIN_PW        = "AdminP@ss1";
const SPREAD_REQUESTS = 40;   // fresh connections -- enough to hit every cluster worker

const jar = {};
let pass = 0;
let fail = 0;


const ok = (cond, msg) => {
    if (cond) {
        pass++;
        console.log("  PASS  " + msg);
    } else {
        fail++;
        console.log("  FAIL  " + msg);
    }
};


// Auth + rotate are pinned to ONE worker (keep-alive) so the session and CSRF token
// stay consistent. The JWKS reads below deliberately use fresh connections to fan
// out across the cluster.
const pinAgent = new http.Agent({ keepAlive: true, maxSockets: 1 });


const storeCookies = (headers) => {
    const sc = headers["set-cookie"];
    if (!sc) {
        return;
    }
    for (const line of sc) {
        const pair = line.split(";")[0];
        const eq = pair.indexOf("=");
        if (eq > 0) {
            jar[pair.slice(0, eq).trim()] = pair.slice(eq + 1).trim();
        }
    }
};


const request = (method, path, opts) => {
    const options = opts || {};
    return new Promise((resolve, reject) => {
        const headers = Object.assign({}, options.headers || {});
        if (!options.noCookies && Object.keys(jar).length > 0) {
            headers["Cookie"] = Object.keys(jar).map((k) => k + "=" + jar[k]).join("; ");
        }
        let data = null;
        if (options.body) {
            data = new URLSearchParams(options.body).toString();
            headers["Content-Type"] = "application/x-www-form-urlencoded";
            headers["Content-Length"] = Buffer.byteLength(data);
        }
        // fresh -> a brand-new connection per request (the cluster round-robins it
        // to some worker); otherwise reuse the pinned keep-alive connection.
        const agent = options.fresh ? false : pinAgent;
        const r = http.request({ host: "localhost", port: PORT, method: method, path: path, headers: headers, agent: agent }, (resp) => {
            storeCookies(resp.headers);
            let body = "";
            resp.on("data", (c) => { body += c; });
            resp.on("end", () => resolve({ status: resp.statusCode, headers: resp.headers, body: body }));
        });
        r.on("error", reject);
        if (data !== null) {
            r.write(data);
        }
        r.end();
    });
};


const csrfOf = (h) => (h.match(/name="_csrf" value="([^"]+)"/) || [])[1] || "";
const activeKidFrom = (html) => (html.match(/Active signing key \(kid\)<\/th><td><code>([^<]+)<\/code>/) || [])[1] || "";


const jwksKids = async () => {
    const r = await request("GET", "/idp/jwks", { fresh: true });
    try {
        return (JSON.parse(r.body).keys || []).map((k) => k.kid).filter(Boolean);
    } catch (e) {
        return [];
    }
};


const run = async () => {
    const lp = await request("GET", "/auth/login");
    await request("POST", "/auth/login", { body: { email: ADMIN_EMAIL, password: ADMIN_PW, _csrf: csrfOf(lp.body) } });
    ok(/true/.test((await request("GET", "/auth/authenticated")).body), "admin session authenticated");

    const adminBefore  = await request("GET", "/admin/idp");
    const beforeActive = activeKidFrom(adminBefore.body);
    ok(!!beforeActive, "read pre-rotation active kid (" + beforeActive + ")");

    const rot = await request("POST", "/admin/idp/rotate-key", { body: { _csrf: csrfOf(adminBefore.body) } });
    ok(rot.status === 302, "admin rotate-key returned 302 (HTTP " + rot.status + ")");

    const afterActive = activeKidFrom((await request("GET", "/admin/idp")).body);
    ok(!!afterActive && afterActive !== beforeActive, "rotation minted a new active kid (" + beforeActive + " -> " + afterActive + ")");

    // Fan out JWKS reads across workers. EVERY response must carry the new active
    // kid (so no worker serves a stale JWKS) and the rotated-out kid (grace window),
    // and all responses must be the identical set.
    let withNew = 0;
    let withOld = 0;
    const seenSets = new Set();
    for (let i = 0; i < SPREAD_REQUESTS; i++) {
        const kids = await jwksKids();
        seenSets.add(kids.slice().sort().join(","));
        if (kids.includes(afterActive)) {
            withNew++;
        }
        if (kids.includes(beforeActive)) {
            withOld++;
        }
    }
    ok(withNew === SPREAD_REQUESTS, "new active kid in ALL " + SPREAD_REQUESTS + " JWKS responses across workers (" + withNew + "/" + SPREAD_REQUESTS + ")");
    ok(withOld === SPREAD_REQUESTS, "rotated-out kid retained in ALL JWKS responses for the grace window (" + withOld + "/" + SPREAD_REQUESTS + ")");
    ok(seenSets.size === 1, "every worker served an IDENTICAL JWKS set (distinct sets seen: " + seenSets.size + ")");

    console.log("\nRESULT: " + pass + " passed, " + fail + " failed");
    process.exit(fail ? 1 : 0);
};


const main = async () => {
    try {
        await run();
    } catch (e) {
        console.error("JWKS CONSISTENCY GATE ERROR:", e);
        process.exit(2);
    }
};


main();