// Multi-tenant OIDC gate against the Postgres instance (:3002). Drives a full
// authorization-code + PKCE flow on tenant t1's issuer and asserts per-tenant
// isolation: an authorization code, access token, and id_token minted by t1 are
// REJECTED at t2 (per-tenant store + per-issuer Provider + per-tenant signing
// key). Tenants are addressed by Host header (tNN.localhost.localdomain:3002 ->
// tenant tNN; the token issuer is the 2-label base_url http://tNN.localhost:3002/idp).
// Each tenant uses a separate cookie jar. Run: node test/mtGate.js
// (PG :3002 up, tenants t1/t2 present with the plugin installed).

const http   = require("http");
const crypto = require("crypto");

const PG_PORT      = 3002;
const TENANTS      = ["t1", "t2"];
const CLIENT_ID    = "mt-rp";
const REDIRECT_URI = "http://localhost:9099/cb";
const ADMIN_PW     = "AdminP@ss1";

const jars = { t1: {}, t2: {} };
let pass = 0;
let fail = 0;


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


const HOST   = (t) => t + ".localhost.localdomain:" + PG_PORT;
const ISSUER = (t) => "http://" + t + ".localhost:" + PG_PORT + "/idp";
const ADMIN  = (t) => "admin@" + t + ".local";


const storeCookies = (t, 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) {
            jars[t][pair.slice(0, eq).trim()] = pair.slice(eq + 1).trim();
        }
    }
};


const request = (t, method, path, opts) => {
    const options = opts || {};
    return new Promise((resolve, reject) => {
        const headers = Object.assign({ Host: HOST(t) }, options.headers || {});
        const jar = jars[t];
        if (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);
        }
        const r = http.request({ host: "127.0.0.1", port: PG_PORT, method: method, path: path, headers: headers }, (resp) => {
            storeCookies(t, 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 = (html) => {
    const m = html.match(/name="_csrf" value="([^"]+)"/);
    return m ? m[1] : "";
};


const getJson = async (t, path) => {
    const r = await request(t, "GET", path);
    try {
        return JSON.parse(r.body);
    } catch (e) {
        return {};
    }
};


const authed = async (t) => {
    return /true/.test((await request(t, "GET", "/auth/authenticated")).body);
};


const resolvePath = (loc) => {
    const u = loc.indexOf("http") === 0 ? new URL(loc) : new URL(loc, "http://x");
    return u.pathname + u.search;
};


// Ensure a tenant admin exists and the jar holds an authenticated session.
const bootstrapTenant = async (t) => {
    const lp = await request(t, "GET", "/auth/login");
    await request(t, "POST", "/auth/login", { body: { email: ADMIN(t), password: ADMIN_PW, _csrf: csrfOf(lp.body) } });
    if (await authed(t)) {
        return true;
    }
    const cp = await request(t, "GET", "/auth/create_first_user");
    const cc = csrfOf(cp.body);
    if (cc) {
        await request(t, "POST", "/auth/create_first_user", { body: { email: ADMIN(t), password: ADMIN_PW, default_language: "en", _csrf: cc } });
    }
    if (await authed(t)) {
        return true;
    }
    const lp2 = await request(t, "GET", "/auth/login");
    await request(t, "POST", "/auth/login", { body: { email: ADMIN(t), password: ADMIN_PW, _csrf: csrfOf(lp2.body) } });
    return await authed(t);
};


const clientExists = async (t, clientId) => {
    const html = (await request(t, "GET", "/admin/idp/clients")).body;
    return html.indexOf("<code>" + clientId + "</code>") >= 0;
};


const registerRp = async (t) => {
    const c1 = csrfOf((await request(t, "GET", "/admin/idp/clients")).body);
    await request(t, "POST", "/admin/idp/clients/delete", { body: { client_id: CLIENT_ID, _csrf: c1 } });
    const c2 = csrfOf((await request(t, "GET", "/admin/idp/clients")).body);
    await request(t, "POST", "/admin/idp/clients/create", { body: {
        client_id: CLIENT_ID, label: "MT RP", redirect_uris: REDIRECT_URI,
        auth_method: "none", scope: "openid email profile groups", _csrf: c2
    } });
    return clientExists(t, CLIENT_ID);
};


// Run the auth-code + PKCE flow up to obtaining the authorization code (does NOT
// exchange it), so the caller can replay an unconsumed code cross-tenant.
const getAuthCode = async (t) => {
    const verifier  = crypto.randomBytes(32).toString("base64url");
    const challenge = crypto.createHash("sha256").update(verifier).digest().toString("base64url");
    const q = new URLSearchParams({
        client_id:             CLIENT_ID,
        response_type:         "code",
        redirect_uri:          REDIRECT_URI,
        scope:                 "openid email profile groups",
        state:                 crypto.randomBytes(8).toString("base64url"),
        nonce:                 crypto.randomBytes(8).toString("base64url"),
        code_challenge:        challenge,
        code_challenge_method: "S256"
    });
    let path = "/idp/auth?" + q.toString();
    let code = null;
    let hops = 0;
    while (hops < 15) {
        hops++;
        const r = await request(t, "GET", path);
        if (r.status >= 300 && r.status < 400 && r.headers.location) {
            if (r.headers.location.indexOf(REDIRECT_URI) === 0) {
                code = new URL(r.headers.location).searchParams.get("code");
                break;
            }
            path = resolvePath(r.headers.location);
        } else if (r.status === 200 && /\/confirm"/.test(r.body)) {
            const m = r.body.match(/action="([^"]*\/confirm)"/);
            if (!m) {
                throw new Error("consent page without a confirm action");
            }
            const pr = await request(t, "POST", resolvePath(m[1]), { body: { allow: "1" } });
            if (pr.status >= 300 && pr.status < 400 && pr.headers.location) {
                path = resolvePath(pr.headers.location);
            } else {
                throw new Error("consent confirm did not redirect: HTTP " + pr.status);
            }
        } else {
            throw new Error("auth flow stalled on " + t + ": HTTP " + r.status + " " + r.body.slice(0, 120));
        }
    }
    return { code: code, verifier: verifier };
};


const exchangeToken = async (t, code, verifier) => {
    const r = await request(t, "POST", "/idp/token", { body: {
        grant_type:    "authorization_code",
        code:          code,
        redirect_uri:  REDIRECT_URI,
        client_id:     CLIENT_ID,
        code_verifier: verifier
    } });
    let tokens = {};
    try {
        tokens = JSON.parse(r.body);
    } catch (e) {
        /* leave empty */
    }
    let header = {};
    let payload = {};
    if (tokens.id_token) {
        tokens.__parts = tokens.id_token.split(".");
        header  = JSON.parse(Buffer.from(tokens.__parts[0], "base64url").toString());
        payload = JSON.parse(Buffer.from(tokens.__parts[1], "base64url").toString());
    }
    return { status: r.status, body: r.body, tokens: tokens, header: header, payload: payload };
};


const verifyIdToken = async (t, tok) => {
    const jwks = await getJson(t, "/idp/jwks");
    const jwk = (jwks.keys || []).find((k) => k.kid === tok.header.kid) || (jwks.keys || [])[0];
    if (!jwk) {
        return false;
    }
    const pub = crypto.createPublicKey({ key: jwk, format: "jwk" });
    return crypto.verify("sha256", Buffer.from(tok.tokens.__parts[0] + "." + tok.tokens.__parts[1]), pub, Buffer.from(tok.tokens.__parts[2], "base64url"));
};


const run = async () => {
    // Phase 0: distinct issuers + distinct keys per tenant
    const d1 = await getJson("t1", "/idp/.well-known/openid-configuration");
    const d2 = await getJson("t2", "/idp/.well-known/openid-configuration");
    ok(d1.issuer === ISSUER("t1") && d2.issuer === ISSUER("t2"), "distinct per-tenant issuers (" + d1.issuer + " / " + d2.issuer + ")");
    const k1 = await getJson("t1", "/idp/jwks");
    const k2 = await getJson("t2", "/idp/jwks");
    const kid1 = k1.keys && k1.keys[0] && k1.keys[0].kid;
    const kid2 = k2.keys && k2.keys[0] && k2.keys[0].kid;
    ok(kid1 && kid2 && kid1 !== kid2, "distinct per-tenant signing keys (" + kid1 + " / " + kid2 + ")");

    // Bootstrap admins + register the RP in each tenant
    for (const t of TENANTS) {
        ok(await bootstrapTenant(t), t + " admin session established");
        ok(await registerRp(t), t + " RP registered (" + CLIENT_ID + ")");
    }

    // Positive: full flow on t1
    const c1  = await getAuthCode("t1");
    ok(!!c1.code, "t1 issued an authorization code");
    const tok = await exchangeToken("t1", c1.code, c1.verifier);
    ok(tok.status === 200 && !!tok.tokens.access_token && !!tok.tokens.id_token, "t1 token endpoint issued access_token + id_token");
    ok(await verifyIdToken("t1", tok), "t1 id_token verifies against t1 JWKS");
    ok(tok.payload.iss === ISSUER("t1"), "t1 id_token iss = " + tok.payload.iss);

    // Cross-tenant (A): an unconsumed t1 code is rejected at t2's token endpoint
    const c1b = await getAuthCode("t1");
    const xTok = await exchangeToken("t2", c1b.code, c1b.verifier);
    let xErr = {};
    try {
        xErr = JSON.parse(xTok.body);
    } catch (e) {
        /* leave empty */
    }
    ok(xTok.status >= 400 && !xTok.tokens.access_token, "t1 authorization code REJECTED at t2 token endpoint (HTTP " + xTok.status + " " + (xErr.error || "") + ")");

    // Cross-tenant (B): a t1 access token is rejected at t2 userinfo
    const meT2 = await request("t2", "GET", "/idp/me", { headers: { Authorization: "Bearer " + tok.tokens.access_token } });
    ok(meT2.status === 401, "t1 access_token REJECTED at t2 userinfo (HTTP " + meT2.status + ")");

    // Cross-tenant (C): a t1 id_token does NOT verify against t2's JWKS
    const sigVsT2 = await verifyIdToken("t2", tok);
    ok(sigVsT2 === false, "t1 id_token signature does NOT verify against t2 JWKS");
    ok(tok.payload.iss !== ISSUER("t2"), "t1 id_token iss is not t2's issuer");

    // Positive control: the same t1 access token still works at t1 userinfo
    const meT1 = await request("t1", "GET", "/idp/me", { headers: { Authorization: "Bearer " + tok.tokens.access_token } });
    ok(meT1.status === 200, "t1 access_token still valid at t1 userinfo (rejection is tenant-specific, not blanket)");

    // Cleanup
    for (const t of TENANTS) {
        const cc = csrfOf((await request(t, "GET", "/admin/idp/clients")).body);
        await request(t, "POST", "/admin/idp/clients/delete", { body: { client_id: CLIENT_ID, _csrf: cc } });
    }

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


run().catch((e) => {
    console.error("MT GATE ERROR:", e);
    process.exit(2);
});