From 92842a65befda2b5bbcd8ea3c08c2c3e3bbdc0d9 Mon Sep 17 00:00:00 2001
From: Cole MacKenzie <cole@cloudflare.com>
Date: Mon, 6 Jan 2025 10:19:11 -0800
Subject: [PATCH] fix: clear timeout if token retrieved successfully

This uses the promise based version of `setTimeout` from NodeJS and
registers the AbortController to handle cancellation signal. The http
server `.close()` method is also registered to the abort controller for
cleanup as `controller.abort()` is always called before returning the
result.
---
 packages/wrangler/src/pipelines/client.ts | 38 +++++++++++++----------
 1 file changed, 22 insertions(+), 16 deletions(-)

diff --git a/packages/wrangler/src/pipelines/client.ts b/packages/wrangler/src/pipelines/client.ts
index 23c21c33f09f..e8fbdfbad89c 100644
--- a/packages/wrangler/src/pipelines/client.ts
+++ b/packages/wrangler/src/pipelines/client.ts
@@ -1,6 +1,7 @@
 import assert from "node:assert";
 import { createHash } from "node:crypto";
 import http from "node:http";
+import { setTimeout as setTimeoutPromise } from "node:timers/promises";
 import { fetchResult } from "../cfetch";
 import { getCloudflareApiEnvironmentFromEnv } from "../environment-variables/misc-variables";
 import { UserError } from "../errors";
@@ -126,22 +127,15 @@ export async function generateR2ServiceToken(
 	pipelineName: string
 ): Promise<S3AccessKey> {
 	// TODO: Refactor into startHttpServerWithTimeout function and update `getOauthToken`
-	let server: http.Server;
-	let loginTimeoutHandle: ReturnType<typeof setTimeout>;
-	const timerPromise = new Promise<S3AccessKey>((_, reject) => {
-		loginTimeoutHandle = setTimeout(() => {
-			server.close();
-			clearTimeout(loginTimeoutHandle);
-			reject(
-				new UserError(
-					"Timed out waiting for authorization code, please try again."
-				)
-			);
-		}, 120000); // wait for 120 seconds for the user to authorize
-	});
+	const controller = new AbortController();
+	const signal = controller.signal;
+
+	// Create timeout promise to prevent hanging forever
+	const timeoutPromise = setTimeoutPromise(120000, "timeout", { signal });
 
-	const loginPromise = new Promise<S3AccessKey>((resolve, reject) => {
-		server = http.createServer(async (request, response) => {
+	// Create server promise to handle the callback and register the cleanup handler on the controller
+	const serverPromise = new Promise<S3AccessKey>((resolve, reject) => {
+		const server = http.createServer(async (request, response) => {
 			assert(request.url, "This request doesn't have a URL"); // This should never happen
 
 			if (request.method !== "GET") {
@@ -179,6 +173,10 @@ export async function generateR2ServiceToken(
 			response.end();
 		});
 
+		// Register cleanup handler
+		signal.addEventListener("abort", () => {
+			server.close();
+		});
 		server.listen(8976, "localhost");
 	});
 
@@ -192,7 +190,15 @@ export async function generateR2ServiceToken(
 	logger.log(`Opening a link in your default browser: ${urlToOpen}`);
 	await openInBrowser(urlToOpen);
 
-	return Promise.race([timerPromise, loginPromise]);
+	const result = await Promise.race([timeoutPromise, serverPromise]);
+	controller.abort();
+	if (result === "timeout") {
+		throw new UserError(
+			"Timed out waiting for authorization code, please try again."
+		);
+	}
+
+	return result as S3AccessKey;
 }
 
 // Get R2 bucket information from v4 API