diff --git a/greenlight/src/greenlight/drci_poke.py b/greenlight/src/greenlight/drci_poke.py index 3b277119ee..eadc6116f8 100644 --- a/greenlight/src/greenlight/drci_poke.py +++ b/greenlight/src/greenlight/drci_poke.py @@ -124,8 +124,9 @@ def poke( if status in _SUCCESS_STATUSES: logger.info("poked Dr. CI for %s#%d (HTTP %d)", repo, pr_number, status) else: - # An auth failure answers 500, not 403 -- the endpoint's auth branch sits outside its - # try/catch -- so the status code cannot classify the failure. Log it and move on. + # The status classifies the failure: 401 a rejected credential, 403 a request the + # endpoint refuses, 429 the per-user limit, 503 a dependency it could not reach. + # Nothing here acts on the distinction, so log it and move on. logger.error("Dr. CI poke for %s#%d returned HTTP %d", repo, pr_number, status) diff --git a/greenlight/tests/test_drci_poke.py b/greenlight/tests/test_drci_poke.py index 0072999643..b141f6ea67 100644 --- a/greenlight/tests/test_drci_poke.py +++ b/greenlight/tests/test_drci_poke.py @@ -170,14 +170,14 @@ def test_poke_logs_success_for_2xx(status, poke_config, caplog): assert any(f"poked Dr. CI for pytorch/pytorch#5 (HTTP {status})" in r.getMessage() for r in caplog.records) -@pytest.mark.parametrize("status", [300, 400, 403, 500, 502]) +@pytest.mark.parametrize("status", [300, 400, 401, 403, 429, 500, 502, 503]) def test_poke_swallows_non_2xx_and_logs_the_code(status, poke_config, caplog): rec = _Recorder() with caplog.at_level(logging.ERROR, logger="greenlight"): drci_poke.poke("pytorch/pytorch", 5, poke_config(), sleep=_FakeSleep(rec), post=_FakePost(rec, status)) - # An auth failure answers 500, so a non-2xx cannot be classified -- but it must never raise. + # A poke is best-effort: whatever the endpoint answers, log the code and never raise. assert any(f"returned HTTP {status}" in record.getMessage() for record in caplog.records) diff --git a/torchci/pages/api/drci/drci.ts b/torchci/pages/api/drci/drci.ts index 3140ecb699..175232db06 100644 --- a/torchci/pages/api/drci/drci.ts +++ b/torchci/pages/api/drci/drci.ts @@ -109,8 +109,12 @@ export default async function handler( }> ) { const authorization = req.headers.authorization; + const botKey = process.env.DRCI_BOT_KEY; - if (authorization == process.env.DRCI_BOT_KEY) { + // `botKey &&` matters: without it an unset DRCI_BOT_KEY makes the comparison + // undefined == undefined for a caller that sent no header, and every + // anonymous request is admitted as the bot. + if (botKey && authorization === botKey) { // Dr. CI bot key is used to update the comment, probably called from the // update Dr. CI workflow } else if (authorization) { @@ -121,14 +125,47 @@ export default async function handler( res.status(403).end(); return; } - // Check if they exceed the rate limit - const userOctokit = await getOctokitWithUserToken(authorization as string); - const user = await userOctokit.rest.users.getAuthenticated(); - if (await drCIRateLimitExceeded(user.data.login)) { - res.status(429).end(); + // Resolving the caller talks to GitHub and to the rate limiter, either of + // which can throw. Uncaught, that leaves the function with no response and + // the platform answers 500, which names no cause -- and trymerge reads any + // failure of this endpoint as "no classifications", so it degrades quietly. + // Answer a status that says which half failed. + try { + // Check if they exceed the rate limit + const userOctokit = await getOctokitWithUserToken( + authorization as string + ); + const user = await userOctokit.rest.users.getAuthenticated(); + if (await drCIRateLimitExceeded(user.data.login)) { + res.status(429).end(); + return; + } + // This insert IS the per-user limit, so a failed one must refuse rather + // than serve: reads and writes use separate ClickHouse credentials, so a + // lost INSERT permission can persist while the read above keeps + // answering, and every authenticated user would drive unbounded work + // under the service's own bot credentials without consuming quota. The + // bot-key path never reaches this limiter, so refusing here does not + // touch the scheduled comment updates. + await incrementDrCIRateLimit(user.data.login); + } catch (error) { + // Only GitHub saying the credential is bad is an auth failure. A + // ClickHouse or network fault means this endpoint could not reach its + // dependencies, which is ours, not the caller's. Messages are fixed -- + // dependency exception text can name backends, hosts and queries -- and + // the detail stays in the server log. + if ((error as { status?: number }).status === 401) { + console.error("Dr.CI rejected the caller's credential:", error); + res.setHeader("WWW-Authenticate", "Bearer"); + res.status(401).json({ error: "Invalid credentials" } as any); + } else { + console.error("Dr.CI could not resolve the caller:", error); + res + .status(503) + .json({ error: "Authentication service unavailable" } as any); + } return; } - incrementDrCIRateLimit(user.data.login); } else { // No authorization provided, return 403 res.status(403).end(); diff --git a/torchci/test/drciAuth.test.ts b/torchci/test/drciAuth.test.ts new file mode 100644 index 0000000000..70904ded54 --- /dev/null +++ b/torchci/test/drciAuth.test.ts @@ -0,0 +1,262 @@ +/** + * The Dr. CI endpoint's authorization block, which runs before the handler's + * own try/catch. + * + * Every fault here must still produce a response. One that escapes leaves the + * function with none and the platform answers 500, and trymerge reads any + * failure of this call as "no classifications" -- so an escaping fault + * degrades merges quietly instead of reporting itself. + */ + +import { NextApiRequest } from "next"; +import * as githubModule from "../lib/github"; +import * as rateLimitModule from "../lib/rateLimit"; +import handler from "../pages/api/drci/drci"; +import { mockRes } from "./nextApiMocks"; + +jest.mock("../lib/github", () => { + const actual = jest.requireActual("../lib/github"); + return { + ...actual, + getOctokitWithUserToken: jest.fn(), + getOctokit: jest.fn(), + }; +}); + +jest.mock("../lib/rateLimit", () => ({ + drCIRateLimitExceeded: jest.fn(), + incrementDrCIRateLimit: jest.fn(), +})); + +const mockGetOctokitWithUserToken = + githubModule.getOctokitWithUserToken as jest.Mock; +const mockGetOctokit = githubModule.getOctokit as jest.Mock; +const mockRateLimitExceeded = + rateLimitModule.drCIRateLimitExceeded as jest.Mock; +const mockIncrementRateLimit = + rateLimitModule.incrementDrCIRateLimit as jest.Mock; + +/** Rejecting getOctokit is the cheapest proof the handler body was reached. */ +const REACHED_BODY = "reached the handler body"; + +function mockReq( + authorization: string | undefined, + query: Record = {} +): NextApiRequest { + return { + method: "POST", + headers: authorization === undefined ? {} : { authorization }, + query, + body: { repo: "pytorch", org: "pytorch" }, + } as unknown as NextApiRequest; +} + +function octokitThatRejects(error: Error) { + return { + rest: { + users: { getAuthenticated: jest.fn().mockRejectedValue(error) }, + }, + }; +} + +function octokitForLogin(login: string) { + return { + rest: { + users: { + getAuthenticated: jest.fn().mockResolvedValue({ data: { login } }), + }, + }, + }; +} + +/** What Octokit raises when GitHub refuses the token. */ +function badCredentials() { + return Object.assign(new Error("Bad credentials"), { status: 401 }); +} + +describe("Dr. CI authorization", () => { + const savedBotKey = process.env.DRCI_BOT_KEY; + + beforeEach(() => { + jest.clearAllMocks(); + process.env.DRCI_BOT_KEY = "the-bot-key"; + mockRateLimitExceeded.mockResolvedValue(false); + mockIncrementRateLimit.mockResolvedValue(undefined); + mockGetOctokit.mockRejectedValue(new Error(REACHED_BODY)); + }); + + afterAll(() => { + if (savedBotKey === undefined) { + delete process.env.DRCI_BOT_KEY; + } else { + process.env.DRCI_BOT_KEY = savedBotKey; + } + }); + + test("a credential GitHub refuses answers 401, not an unhandled throw", async () => { + mockGetOctokitWithUserToken.mockResolvedValue( + octokitThatRejects(badCredentials()) + ); + const res = mockRes(); + + await expect( + handler( + mockReq("a-key-that-no-longer-matches", { prNumber: "1000" }), + res + ) + ).resolves.not.toThrow(); + + expect(res._status).toBe(401); + expect(res._json.error).toBe("Invalid credentials"); + expect(res._headers["WWW-Authenticate"]).toBe("Bearer"); + }); + + test("a GitHub fault that is not a refusal answers 503, not 401", async () => { + // Reaching GitHub at all is this endpoint's problem, not the caller's, so + // it must not be reported as bad credentials. + mockGetOctokitWithUserToken.mockResolvedValue( + octokitThatRejects(new Error("getaddrinfo ENOTFOUND api.github.com")) + ); + const res = mockRes(); + + await handler(mockReq("some-user-token", { prNumber: "1000" }), res); + + expect(res._status).toBe(503); + expect(res._json.error).toBe("Authentication service unavailable"); + }); + + test("a rate-limit read failure answers 503 and leaks no backend detail", async () => { + mockGetOctokitWithUserToken.mockResolvedValue(octokitForLogin("someone")); + mockRateLimitExceeded.mockRejectedValue( + new Error("clickhouse://secret-host:8123 refused the connection") + ); + const res = mockRes(); + + await handler(mockReq("some-user-token", { prNumber: "1000" }), res); + + expect(res._status).toBe(503); + expect(res._json.error).toBe("Authentication service unavailable"); + expect(JSON.stringify(res._json)).not.toContain("secret-host"); + }); + + test("a rate-limit WRITE failure refuses rather than serving unmetered", async () => { + // The increment used to be neither awaited nor caught, so its rejection + // escaped as an unhandled one. It is the limit itself, not bookkeeping: + // serving on a failed insert lets any authenticated user drive unbounded + // work under the service's bot credentials without consuming quota. + mockGetOctokitWithUserToken.mockResolvedValue(octokitForLogin("someone")); + mockIncrementRateLimit.mockRejectedValue(new Error("insert failed")); + const res = mockRes(); + + await handler(mockReq("some-user-token", { prNumber: "1000" }), res); + + expect(mockIncrementRateLimit).toHaveBeenCalledWith("someone"); + expect(res._status).toBe(503); + expect(mockGetOctokit).not.toHaveBeenCalled(); + }); + + test("a factory failure is handled like any other resolution failure", async () => { + mockGetOctokitWithUserToken.mockRejectedValue( + new Error("cannot construct") + ); + const res = mockRes(); + + await handler(mockReq("some-user-token", { prNumber: "1000" }), res); + + expect(res._status).toBe(503); + expect(res._json.error).toBe("Authentication service unavailable"); + }); + + test("an upstream 403 is reported as unavailable, not as bad credentials", async () => { + // GitHub answers 403 for secondary rate limiting and missing scope, which + // are not "your token is wrong" -- reporting them as 401 is the + // misdirection this split exists to avoid. + mockGetOctokitWithUserToken.mockResolvedValue( + octokitThatRejects( + Object.assign(new Error("API rate limit exceeded"), { status: 403 }) + ) + ); + const res = mockRes(); + + await handler(mockReq("some-user-token", { prNumber: "1000" }), res); + + expect(res._status).toBe(503); + }); + + test("an authenticated user under the limit reaches the handler body", async () => { + mockGetOctokitWithUserToken.mockResolvedValue(octokitForLogin("someone")); + const res = mockRes(); + + await handler(mockReq("some-user-token", { prNumber: "1000" }), res); + + expect(mockIncrementRateLimit).toHaveBeenCalledWith("someone"); + expect(res._status).toBe(400); + expect(res._json.error).toContain(REACHED_BODY); + }); + + test("an over-limit user still gets 429", async () => { + mockGetOctokitWithUserToken.mockResolvedValue(octokitForLogin("someone")); + mockRateLimitExceeded.mockResolvedValue(true); + const res = mockRes(); + + await handler(mockReq("some-user-token", { prNumber: "1000" }), res); + + expect(res._status).toBe(429); + // Asserted because 429 answers with .end(): a response that sets a status + // and never ends leaves the request hanging, which a status-only + // assertion cannot see. + expect(res._ended).toBe(true); + expect(mockIncrementRateLimit).not.toHaveBeenCalled(); + }); + + test("a user token without prNumber still gets 403", async () => { + const res = mockRes(); + + await handler(mockReq("some-user-token"), res); + + expect(res._status).toBe(403); + expect(res._ended).toBe(true); + expect(mockGetOctokitWithUserToken).not.toHaveBeenCalled(); + }); + + test("no Authorization header gets 403", async () => { + const res = mockRes(); + + await handler(mockReq(undefined), res); + + expect(res._status).toBe(403); + expect(res._ended).toBe(true); + }); + + test("an unset DRCI_BOT_KEY does not admit an anonymous caller as the bot", async () => { + // undefined == undefined is true, so the bare comparison used to treat + // every anonymous request as the bot whenever the deployment lost the var. + delete process.env.DRCI_BOT_KEY; + const res = mockRes(); + + await handler(mockReq(undefined), res); + + expect(res._status).toBe(403); + }); + + test("an empty DRCI_BOT_KEY does not admit an empty Authorization header", async () => { + process.env.DRCI_BOT_KEY = ""; + const res = mockRes(); + + await handler(mockReq(""), res); + + expect(res._status).toBe(403); + expect(mockGetOctokitWithUserToken).not.toHaveBeenCalled(); + }); + + test("the bot key skips the user path and reaches the handler body", async () => { + const res = mockRes(); + + await handler(mockReq("the-bot-key"), res); + + expect(mockGetOctokitWithUserToken).not.toHaveBeenCalled(); + expect(mockRateLimitExceeded).not.toHaveBeenCalled(); + expect(res._status).toBe(400); + expect(res._json.error).toContain(REACHED_BODY); + }); +}); diff --git a/torchci/test/nextApiMocks.ts b/torchci/test/nextApiMocks.ts index eba8e81d53..a9111394ca 100644 --- a/torchci/test/nextApiMocks.ts +++ b/torchci/test/nextApiMocks.ts @@ -4,6 +4,7 @@ export type MockApiResponse = NextApiResponse & { _status: number; _json: any; _headers: Record; + _ended: boolean; }; export function mockRes(): MockApiResponse { @@ -11,6 +12,7 @@ export function mockRes(): MockApiResponse { _status: 0, _json: null, _headers: {}, + _ended: false, setHeader(name: string, value: any) { res._headers[name] = value; return res; @@ -23,6 +25,13 @@ export function mockRes(): MockApiResponse { res._json = data; return res; }, + end(data?: any) { + res._ended = true; + if (data !== undefined) { + res._json = data; + } + return res; + }, }; return res; }