feat: Use OAuth flow to generate R2 tokens for Pipelines

.changeset/odd-ducks-attack.md

@@ -0,0 +1,5 @@
+---
+"wrangler": patch
+---
+
+feat: Use OAuth flow to generate R2 tokens for Pipelines

packages/wrangler/package.json

@@ -69,6 +69,7 @@
 		"type:tests": "tsc -p ./src/__tests__/tsconfig.json && tsc -p ./e2e/tsconfig.json"
 	},
 	"dependencies": {
+		"@aws-sdk/client-s3": "^3.721.0",
 		"@cloudflare/kv-asset-handler": "workspace:*",
 		"@esbuild-plugins/node-globals-polyfill": "^0.2.3",
 		"@esbuild-plugins/node-modules-polyfill": "^0.2.2",

packages/wrangler/src/__tests__/pipelines.test.ts

@@ -46,76 +46,7 @@ describe("pipelines", () => {
 		endpoint: "https://0001.pipelines.cloudflarestorage.com",
 	} satisfies Pipeline;
 
-	function mockCreateR2Token(bucket: string) {
-		const requests = { count: 0 };
-		msw.use(
-			http.get(
-				"*/accounts/:accountId/r2/buckets/:bucket",
-				async ({ params }) => {
-					expect(params.accountId).toEqual("some-account-id");
-					expect(params.bucket).toEqual(bucket);
-					requests.count++;
-					return HttpResponse.json(
-						{
-							success: true,
-							errors: [],
-							messages: [],
-							result: null,
-						},
-						{ status: 200 }
-					);
-				},
-				{ once: true }
-			),
-			http.get(
-				"*/user/tokens/permission_groups",
-				async () => {
-					requests.count++;
-					return HttpResponse.json(
-						{
-							success: true,
-							errors: [],
-							messages: [],
-							result: [
-								{
-									id: "2efd5506f9c8494dacb1fa10a3e7d5b6",
-									name: "Workers R2 Storage Bucket Item Write",
-									description:
-										"Grants write access to Cloudflare R2 Bucket Scoped Storage",
-									scopes: ["com.cloudflare.edge.r2.bucket"],
-								},
-							],
-						},
-						{ status: 200 }
-					);
-				},
-				{ once: true }
-			),
-			http.post(
-				"*/user/tokens",
-				async () => {
-					requests.count++;
-					return HttpResponse.json(
-						{
-							success: true,
-							errors: [],
-							messages: [],
-							result: {
-								id: "service-token-id",
-								name: "my-service-token",
-								value: "my-secret-value",
-							},
-						},
-						{ status: 200 }
-					);
-				},
-				{ once: true }
-			)
-		);
-		return requests;
-	}
-
-	function mockCreeatR2TokenFailure(bucket: string) {
+	function mockCreateR2TokenFailure(bucket: string) {
 		const requests = { count: 0 };
 		msw.use(
 			http.get(
@@ -310,6 +241,7 @@ describe("pipelines", () => {
 		);
 		return requests;
 	}
+
 	beforeAll(() => {
 		__testSkipDelays();
 	});
@@ -383,15 +315,6 @@ describe("pipelines", () => {
 			`);
 		});
 
-		it("should create a pipeline", async () => {
-			const tokenReq = mockCreateR2Token("test-bucket");
-			const requests = mockCreateRequest("my-pipeline");
-			await runWrangler("pipelines create my-pipeline --r2 test-bucket");
-
-			expect(tokenReq.count).toEqual(3);
-			expect(requests.count).toEqual(1);
-		});
-
 		it("should create a pipeline with explicit credentials", async () => {
 			const requests = mockCreateRequest("my-pipeline");
 			await runWrangler(
@@ -401,7 +324,7 @@ describe("pipelines", () => {
 		});
 
 		it("should fail a missing bucket", async () => {
-			const requests = mockCreeatR2TokenFailure("bad-bucket");
+			const requests = mockCreateR2TokenFailure("bad-bucket");
 			await expect(
 				runWrangler("pipelines create bad-pipeline --r2 bad-bucket")
 			).rejects.toThrowError();
@@ -543,22 +466,21 @@ describe("pipelines", () => {
 
 		it("should update a pipeline with new bucket", async () => {
 			const pipeline: Pipeline = samplePipeline;
-			const tokenReq = mockCreateR2Token("new-bucket");
 			mockShowRequest(pipeline.name, pipeline);
 
 			const update = JSON.parse(JSON.stringify(pipeline));
 			update.destination.path.bucket = "new_bucket";
 			update.destination.credentials = {
 				endpoint: "https://some-account-id.r2.cloudflarestorage.com",
 				access_key_id: "service-token-id",
-				secret_access_key:
-					"be22cbae9c1585c7b61a92fdb75afd10babd535fb9b317f90ac9a9ca896d02d7",
+				secret_access_key: "my-secret-access-key",
 			};
 			const updateReq = mockUpdateRequest(update.name, update);
 
-			await runWrangler("pipelines update my-pipeline --r2 new-bucket");
+			await runWrangler(
+				"pipelines update my-pipeline --r2 new-bucket --access-key-id service-token-id --secret-access-key my-secret-access-key"
+			);
 
-			expect(tokenReq.count).toEqual(3);
 			expect(updateReq.count).toEqual(1);
 		});
 

packages/wrangler/src/pipelines/client.ts

@@ -1,5 +1,12 @@
+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";
+import { logger } from "../logger";
+import openInBrowser from "../open-in-browser";
 import type { R2BucketInfo } from "../r2/helpers";
 
 // ensure this is in sync with:
@@ -96,44 +103,102 @@ export type PermissionGroup = {
 	scopes: string[];
 };
 
-// Generate a Service Token to write to R2 for a pipeline
+interface S3AccessKey {
+	accessKeyId: string;
+	secretAccessKey: string;
+}
+
+/**
+ * Generate an R2 service token for the given account ID, bucket name, and pipeline name.
+ *
+ * This function kicks off its own OAuth process using the Workers Pipelines OAuth client requesting the scope
+ * `pipelines:setup`. Once the user confirms, our OAuth callback endpoint will validate the request, exchange the
+ * authorization code and return a bucket-scoped R2 token.
+ *
+ * This OAuth flow is distinct from the one used in `wrangler login` to ensure these tokens are generated server-side
+ * and that only the tokens of concern are returned to the user.
+ * @param accountId
+ * @param bucketName
+ * @param pipelineName
+ */
 export async function generateR2ServiceToken(
-	label: string,
 	accountId: string,
-	bucket: string
-): Promise<ServiceToken> {
-	const res = await fetchResult<PermissionGroup[]>(
-		`/user/tokens/permission_groups`,
-		{
-			method: "GET",
-		}
-	);
-	const perm = res.find(
-		(g) => g.name == "Workers R2 Storage Bucket Item Write"
-	);
-	if (!perm) {
-		throw new Error("Missing R2 Permissions");
-	}
+	bucketName: string,
+	pipelineName: string
+): Promise<S3AccessKey> {
+	// TODO: Refactor into startHttpServerWithTimeout function and update `getOauthToken`
+	const controller = new AbortController();
+	const signal = controller.signal;
 
-	// generate specific bucket write token for pipeline
-	const body = JSON.stringify({
-		policies: [
-			{
-				effect: "allow",
-				permission_groups: [{ id: perm.id }],
-				resources: {
-					[`com.cloudflare.edge.r2.bucket.${accountId}_default_${bucket}`]: "*",
-				},
-			},
-		],
-		name: label,
-	});
+	// Create timeout promise to prevent hanging forever
+	const timeoutPromise = setTimeoutPromise(120000, "timeout", { signal });
 
-	return await fetchResult<ServiceToken>(`/user/tokens`, {
-		method: "POST",
-		headers: API_HEADERS,
-		body,
+	// 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") {
+				response.writeHead(405);
+				response.end("Method not allowed.");
+				return;
+			}
+
+			const { pathname, searchParams } = new URL(
+				request.url,
+				`http://${request.headers.host}`
+			);
+
+			if (pathname !== "/") {
+				response.writeHead(404);
+				response.end("Not found.");
+				return;
+			}
+
+			// Retrieve values from the URL parameters
+			const accessKeyId = searchParams.get("access-key-id");
+			const secretAccessKey = searchParams.get("secret-access-key");
+
+			if (!accessKeyId || !secretAccessKey) {
+				reject(new UserError("Missing required URL parameters"));
+				return;
+			}
+
+			resolve({ accessKeyId, secretAccessKey } as S3AccessKey);
+			// Do a final redirect to "clear" the URL of the sensitive URL parameters that were returned.
+			response.writeHead(307, {
+				Location:
+					"https://welcome.developers.workers.dev/wrangler-oauth-consent-granted",
+			});
+			response.end();
+		});
+
+		// Register cleanup handler
+		signal.addEventListener("abort", () => {
+			server.close();
+		});
+		server.listen(8976, "localhost");
 	});
+
+	const env = getCloudflareApiEnvironmentFromEnv();
+	const oauthDomain =
+		env === "staging"
+			? "oauth.pipelines-staging.cloudflare.com"
+			: "oauth.pipelines.cloudflare.com";
+
+	const urlToOpen = `https://${oauthDomain}/oauth/login?accountId=${accountId}&bucketName=${bucketName}&pipelineName=${pipelineName}`;
+	logger.log(`Opening a link in your default browser: ${urlToOpen}`);
+	await openInBrowser(urlToOpen);
+
+	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