Answer Dr. CI auth failures with a status instead of an unhandled 500

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)
 
 

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)
 
 

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();

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<string, string> = {}
+): 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);
+  });
+});

torchci/test/nextApiMocks.ts

@@ -4,13 +4,15 @@ export type MockApiResponse = NextApiResponse & {
   _status: number;
   _json: any;
   _headers: Record<string, any>;
+  _ended: boolean;
 };
 
 export function mockRes(): MockApiResponse {
   const res: any = {
     _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;
 }