Skip to content

Commit

Permalink
feat: Use OAuth flow to generate R2 tokens for Pipelines (#7534)
Browse files Browse the repository at this point in the history 
* feat: Use OAuth flow to generate R2 tokens for Pipelines

This commit changes the generateR2Tokens flow which will direct the user
to the web browser to perform a OAuth flow to grant the Workers
Pipelines client the ability to generate R2 tokens on behalf of the
user. This will only run if the user does not provide the credentials as
CLI parameters.

Due to requiring user interactivity, and reliance on the callbacks,
there is no easy way to support a "headless" mode for `wrangler pipelines
create` (or `update`) unless the user provides the tokens as arguments.
The same applies for testing this flow, which can only be done manually
at this time.

* fix: add forced delayed to allow r2 tokens time to sync

Create odd-ducks-attack.md

* Add docs around bespoke OAuth solution

* fix: Wait for R2 token to sync

After creating an R2 token, there is a slight delay before if can be
used. Previously we would sleep for some amount of time, but this method
is really sensitive to latency.

Instead, use the S3 SDK and try using the token until we exhaust all
attempts, or we finally succeed in using it. Each failure results in a
constant backoff of 1 second.

This commit does add the dependency `@aws-sdk/client-s3`.

* fix pnpm-lock.yaml

* 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.

---------

Co-authored-by: emily-shen <[email protected]>
  • Loading branch information
@cmackenzie1 @emily-shen
cmackenzie1 and emily-shen authored Jan 9, 2025
1 parent b8e5f63 commit 7c8ae1c
Show file tree
Hide file tree
Showing 6 changed files with 1,378 additions and 140 deletions.
5 changes: 5 additions & 0 deletions .changeset/odd-ducks-attack.md
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
---
"wrangler": patch
---

feat: Use OAuth flow to generate R2 tokens for Pipelines
1 change: 1 addition & 0 deletions packages/wrangler/package.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
92 changes: 7 additions & 85 deletions packages/wrangler/src/__tests__/pipelines.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -310,6 +241,7 @@ describe("pipelines", () => {
);
return requests;
}

beforeAll(() => {
__testSkipDelays();
});
Expand Down Expand Up @@ -380,15 +312,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(
Expand All @@ -398,7 +321,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();
Expand Down Expand Up @@ -540,22 +463,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);
});

Expand Down
131 changes: 98 additions & 33 deletions packages/wrangler/src/pipelines/client.ts
View file
Original file line number Diff line number Diff line change
@@ -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:
Expand Down Expand Up @@ -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
Expand Down
Loading

0 comments on commit 7c8ae1c

Please sign in to comment.