fix: implement device auth rate limit handling (coder#17079)

hugodutka · web-flow · commit d5557fcbf5c3 · 2025-03-25T11:32:03.000+01:00
The [OAuth2 specification](https://datatracker.ietf.org/doc/html/rfc8628) describes how clients in the device flow should handle retrying requests when they are rate limited. We didn't respect it, which sometimes prevented users from logging in or setting up external auth. They'd see a `slow_down` error in the UI and would be unable to complete the authentication flow. This PR implements rate limit handling according to the spec.
diff --git a/site/src/components/GitDeviceAuth/GitDeviceAuth.test.ts b/site/src/components/GitDeviceAuth/GitDeviceAuth.test.ts
@@ -0,0 +1,38 @@
+import { AxiosError, type AxiosResponse } from "axios";
+import { newRetryDelay } from "./GitDeviceAuth";
+
+test("device auth retry delay", async () => {
+	const slowDownError = new AxiosError(
+		"slow_down",
+		"500",
+		undefined,
+		undefined,
+		{
+			data: {
+				detail: "slow_down",
+			},
+		} as AxiosResponse,
+	);
+	const retryDelay = newRetryDelay(undefined);
+
+	// If no initial interval is provided, the default must be 5 seconds.
+	expect(retryDelay(0, undefined)).toBe(5000);
+	// If the error is a slow down error, the interval should increase by 5 seconds
+	// for this and all subsequent requests, and by 5 seconds extra delay for this
+	// request.
+	expect(retryDelay(1, slowDownError)).toBe(15000);
+	expect(retryDelay(1, slowDownError)).toBe(15000);
+	expect(retryDelay(2, undefined)).toBe(10000);
+
+	// Like previous request.
+	expect(retryDelay(3, slowDownError)).toBe(20000);
+	expect(retryDelay(3, undefined)).toBe(15000);
+	// If the error is not a slow down error, the interval should not increase.
+	expect(retryDelay(4, new AxiosError("other", "500"))).toBe(15000);
+
+	// If the initial interval is provided, it should be used.
+	const retryDelayWithInitialInterval = newRetryDelay(1);
+	expect(retryDelayWithInitialInterval(0, undefined)).toBe(1000);
+	expect(retryDelayWithInitialInterval(1, slowDownError)).toBe(11000);
+	expect(retryDelayWithInitialInterval(2, undefined)).toBe(6000);
+});
diff --git a/site/src/components/GitDeviceAuth/GitDeviceAuth.tsx b/site/src/components/GitDeviceAuth/GitDeviceAuth.tsx
@@ -5,6 +5,7 @@ import CircularProgress from "@mui/material/CircularProgress";
 import Link from "@mui/material/Link";
 import type { ApiErrorResponse } from "api/errors";
 import type { ExternalAuthDevice } from "api/typesGenerated";
+import { isAxiosError } from "axios";
 import { Alert, AlertDetail } from "components/Alert/Alert";
 import { CopyButton } from "components/CopyButton/CopyButton";
 import type { FC } from "react";
@@ -14,6 +15,59 @@ interface GitDeviceAuthProps {
 	deviceExchangeError?: ApiErrorResponse;
 }
 
+const DeviceExchangeError = {
+	AuthorizationPending: "authorization_pending",
+	SlowDown: "slow_down",
+	ExpiredToken: "expired_token",
+	AccessDenied: "access_denied",
+} as const;
+
+export const isExchangeErrorRetryable = (_: number, error: unknown) => {
+	if (!isAxiosError(error)) {
+		return false;
+	}
+	const detail = error.response?.data?.detail;
+	return (
+		detail === DeviceExchangeError.AuthorizationPending ||
+		detail === DeviceExchangeError.SlowDown
+	);
+};
+
+/**
+ * The OAuth2 specification (https://datatracker.ietf.org/doc/html/rfc8628)
+ * describes how the client should handle retries. This function returns a
+ * closure that implements the retry logic described in the specification.
+ * The closure should be memoized because it stores state.
+ */
+export const newRetryDelay = (initialInterval: number | undefined) => {
+	// "If no value is provided, clients MUST use 5 as the default."
+	// https://datatracker.ietf.org/doc/html/rfc8628#section-3.2
+	let interval = initialInterval ?? 5;
+	let lastFailureCountHandled = 0;
+	return (failureCount: number, error: unknown) => {
+		const isSlowDown =
+			isAxiosError(error) &&
+			error.response?.data.detail === DeviceExchangeError.SlowDown;
+		// We check the failure count to ensure we increase the interval
+		// at most once per failure.
+		if (isSlowDown && lastFailureCountHandled < failureCount) {
+			lastFailureCountHandled = failureCount;
+			// https://datatracker.ietf.org/doc/html/rfc8628#section-3.5
+			// "the interval MUST be increased by 5 seconds for this and all subsequent requests"
+			interval += 5;
+		}
+		let extraDelay = 0;
+		if (isSlowDown) {
+			// I found GitHub is very strict about their rate limits, and they'll block
+			// even if the request is 500ms earlier than they expect. This may happen due to
+			// e.g. network latency, so it's best to cool down for longer if GitHub just
+			// rejected our request.
+			extraDelay = 5;
+		}
+		return (interval + extraDelay) * 1000;
+	};
+};
+
 export const GitDeviceAuth: FC<GitDeviceAuthProps> = ({
 	externalAuthDevice,
 	deviceExchangeError,
@@ -27,16 +81,26 @@ export const GitDeviceAuth: FC<GitDeviceAuthProps> = ({
 	if (deviceExchangeError) {
 		// See https://datatracker.ietf.org/doc/html/rfc8628#section-3.5
 		switch (deviceExchangeError.detail) {
-			case "authorization_pending":
+			case DeviceExchangeError.AuthorizationPending:
+				break;
+			case DeviceExchangeError.SlowDown:
+				status = (
+					<div>
+						{status}
+						<Alert severity="warning">
+							Rate limit reached. Waiting a few seconds before retrying...
+						</Alert>
+					</div>
+				);
 				break;
-			case "expired_token":
+			case DeviceExchangeError.ExpiredToken:
 				status = (
 					<Alert severity="error">
 						The one-time code has expired. Refresh to get a new one!
 					</Alert>
 				);
 				break;
-			case "access_denied":
+			case DeviceExchangeError.AccessDenied:
 				status = (
 					<Alert severity="error">Access to the Git provider was denied.</Alert>
 				);
diff --git a/site/src/pages/ExternalAuthPage/ExternalAuthPage.tsx b/site/src/pages/ExternalAuthPage/ExternalAuthPage.tsx
@@ -6,10 +6,15 @@ import {
 	externalAuthProvider,
 } from "api/queries/externalAuth";
 import { isAxiosError } from "axios";
+import {
+	isExchangeErrorRetryable,
+	newRetryDelay,
+} from "components/GitDeviceAuth/GitDeviceAuth";
 import { SignInLayout } from "components/SignInLayout/SignInLayout";
 import { Welcome } from "components/Welcome/Welcome";
 import { useAuthenticated } from "contexts/auth/RequireAuth";
 import type { FC } from "react";
+import { useMemo } from "react";
 import { useQuery, useQueryClient } from "react-query";
 import { useParams, useSearchParams } from "react-router-dom";
 import ExternalAuthPageView from "./ExternalAuthPageView";
@@ -32,17 +37,22 @@ const ExternalAuthPage: FC = () => {
 			Boolean(externalAuthProviderQuery.data?.device),
 		refetchOnMount: false,
 	});
+	const retryDelay = useMemo(
+		() => newRetryDelay(externalAuthDeviceQuery.data?.interval),
+		[externalAuthDeviceQuery.data],
+	);
 	const exchangeExternalAuthDeviceQuery = useQuery({
 		...exchangeExternalAuthDevice(
 			provider,
 			externalAuthDeviceQuery.data?.device_code ?? "",
 			queryClient,
 		),
 		enabled: Boolean(externalAuthDeviceQuery.data),
-		retry: true,
-		retryDelay: (externalAuthDeviceQuery.data?.interval || 5) * 1000,
-		refetchOnWindowFocus: (query) =>
-			query.state.status === "success" ? false : "always",
+		retry: isExchangeErrorRetryable,
+		retryDelay,
+		// We don't want to refetch the query outside of the standard retry
+		// logic, because the device auth flow is very strict about rate limits.
+		refetchOnWindowFocus: false,
 	});
 
 	if (externalAuthProviderQuery.isLoading || !externalAuthProviderQuery.data) {
diff --git a/site/src/pages/LoginOAuthDevicePage/LoginOAuthDevicePage.tsx b/site/src/pages/LoginOAuthDevicePage/LoginOAuthDevicePage.tsx
@@ -4,21 +4,18 @@ import {
 	getGitHubDeviceFlowCallback,
 } from "api/queries/oauth2";
 import { isAxiosError } from "axios";
+import {
+	isExchangeErrorRetryable,
+	newRetryDelay,
+} from "components/GitDeviceAuth/GitDeviceAuth";
 import { SignInLayout } from "components/SignInLayout/SignInLayout";
 import { Welcome } from "components/Welcome/Welcome";
-import { useEffect } from "react";
+import { useEffect, useMemo } from "react";
 import type { FC } from "react";
 import { useQuery } from "react-query";
 import { useSearchParams } from "react-router-dom";
 import LoginOAuthDevicePageView from "./LoginOAuthDevicePageView";
 
-const isErrorRetryable = (error: unknown) => {
-	if (!isAxiosError(error)) {
-		return false;
-	}
-	return error.response?.data?.detail === "authorization_pending";
-};
-
 // The page is hardcoded to only use GitHub,
 // as that's the only OAuth2 login provider in our backend
 // that currently supports the device flow.
@@ -38,19 +35,23 @@ const LoginOAuthDevicePage: FC = () => {
 		...getGitHubDevice(),
 		refetchOnMount: false,
 	});
+
+	const retryDelay = useMemo(
+		() => newRetryDelay(externalAuthDeviceQuery.data?.interval),
+		[externalAuthDeviceQuery.data],
+	);
+
 	const exchangeExternalAuthDeviceQuery = useQuery({
 		...getGitHubDeviceFlowCallback(
 			externalAuthDeviceQuery.data?.device_code ?? "",
 			state,
 		),
 		enabled: Boolean(externalAuthDeviceQuery.data),
-		retry: (_, error) => isErrorRetryable(error),
-		retryDelay: (externalAuthDeviceQuery.data?.interval || 5) * 1000,
-		refetchOnWindowFocus: (query) =>
-			query.state.status === "success" ||
-			(query.state.error != null && !isErrorRetryable(query.state.error))
-				? false
-				: "always",
+		retry: isExchangeErrorRetryable,
+		retryDelay,
+		// We don't want to refetch the query outside of the standard retry
+		// logic, because the device auth flow is very strict about rate limits.
+		refetchOnWindowFocus: false,
 	});
 
 	useEffect(() => {
