Write auth details per-environment (#5933)

.changeset/new-frogs-decide.md

@@ -0,0 +1,5 @@
+---
+"wrangler": minor
+---
+
+feature: allow for writing authentication details per API environment. This allows someone targetting staging to have their staging auth details saved separately from production, this saves them logging in and out when switching environments.

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

@@ -1,8 +1,6 @@
 import fs from "node:fs";
-import path from "node:path";
 import { http, HttpResponse } from "msw";
-import { getGlobalWranglerConfigPath } from "../global-wrangler-config-path";
-import { USER_AUTH_CONFIG_FILE, writeAuthConfigFile } from "../user";
+import { getAuthConfigFilePath, writeAuthConfigFile } from "../user";
 import { mockConsoleMethods } from "./helpers/mock-console";
 import { msw } from "./helpers/msw";
 import { runInTempDir } from "./helpers/run-in-tmp";
@@ -23,10 +21,7 @@ describe("logout", () => {
 			refresh_token: "some-refresh-tok",
 		});
 		// Make sure that logout removed the config file containing the auth tokens.
-		const config = path.join(
-			getGlobalWranglerConfigPath(),
-			USER_AUTH_CONFIG_FILE
-		);
+		const config = getAuthConfigFilePath();
 		let counter = 0;
 		msw.use(
 			http.post(

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

@@ -1,13 +1,15 @@
 import { http, HttpResponse } from "msw";
 import { vi } from "vitest";
+import { getGlobalWranglerConfigPath } from "../global-wrangler-config-path";
 import { CI } from "../is-ci";
 import {
+	getAuthConfigFilePath,
 	loginOrRefreshIfRequired,
 	readAuthConfigFile,
 	requireAuth,
 	writeAuthConfigFile,
 } from "../user";
-import { mockConsoleMethods } from "./helpers/mock-console";
+import { mockConsoleMethods, normalizeSlashes } from "./helpers/mock-console";
 import { useMockIsTTY } from "./helpers/mock-istty";
 import {
 	mockExchangeRefreshTokenForAccessToken,
@@ -37,6 +39,10 @@ describe("User", () => {
 		isCISpy = vi.spyOn(CI, "isCI").mockReturnValue(false);
 	});
 
+	afterEach(() => {
+		delete process.env.WRANGLER_API_ENVIRONMENT;
+	});
+
 	describe("login", () => {
 		it("should login a user when `wrangler login` is run", async () => {
 			mockOAuthServerCallback("success");
@@ -75,6 +81,49 @@ describe("User", () => {
 				scopes: ["account:read"],
 			});
 		});
+
+		it("login works in a different environment", async () => {
+			process.env.WRANGLER_API_ENVIRONMENT = "staging";
+			mockOAuthServerCallback("success");
+
+			let counter = 0;
+			msw.use(
+				http.post(
+					"*/oauth2/token",
+					async () => {
+						counter += 1;
+
+						return HttpResponse.json({
+							access_token: "test-access-token",
+							expires_in: 100000,
+							refresh_token: "test-refresh-token",
+							scope: "account:read",
+						});
+					},
+					{ once: true }
+				)
+			);
+
+			await runWrangler("login");
+
+			expect(counter).toBe(1);
+			expect(std.out).toMatchInlineSnapshot(`
+				"Attempting to login via OAuth...
+				Opening a link in your default browser: https://dash.staging.cloudflare.com/oauth2/auth?response_type=code&client_id=4b2ea6cc-9421-4761-874b-ce550e0e3def&redirect_uri=http%3A%2F%2Flocalhost%3A8976%2Foauth%2Fcallback&scope=account%3Aread%20user%3Aread%20workers%3Awrite%20workers_kv%3Awrite%20workers_routes%3Awrite%20workers_scripts%3Awrite%20workers_tail%3Aread%20d1%3Awrite%20pages%3Awrite%20zone%3Aread%20ssl_certs%3Awrite%20constellation%3Awrite%20ai%3Awrite%20queues%3Awrite%20offline_access&state=MOCK_STATE_PARAM&code_challenge=MOCK_CODE_CHALLENGE&code_challenge_method=S256
+				Successfully logged in."
+			`);
+
+			expect(normalizeSlashes(getAuthConfigFilePath())).toBe(
+				normalizeSlashes(`${getGlobalWranglerConfigPath()}/config/staging.toml`)
+			);
+			expect(readAuthConfigFile()).toEqual<UserAuthConfig>({
+				api_token: undefined,
+				oauth_token: "test-access-token",
+				refresh_token: "test-refresh-token",
+				expiration_time: expect.any(String),
+				scopes: ["account:read"],
+			});
+		});
 	});
 
 	it("should handle errors for failed token refresh in a non-interactive environment", async () => {
@@ -119,4 +168,18 @@ describe("User", () => {
 		});
 		await expect(loginOrRefreshIfRequired()).resolves.toEqual(false);
 	});
+
+	it("should have auth per environment", async () => {
+		setIsTTY(false);
+		process.env.WRANGLER_API_ENVIRONMENT = "staging";
+
+		writeAuthConfigFile({
+			oauth_token: "hunter2",
+			refresh_token: "Order 66",
+		});
+
+		expect(normalizeSlashes(getAuthConfigFilePath())).toBe(
+			normalizeSlashes(`${getGlobalWranglerConfigPath()}/config/staging.toml`)
+		);
+	});
 });

packages/wrangler/src/user/user.ts

@@ -220,6 +220,7 @@ import {
 	saveToConfigCache,
 } from "../config-cache";
 import { NoDefaultValueProvided, select } from "../dialogs";
+import { getCloudflareApiEnvironmentFromEnv } from "../environment-variables/misc-variables";
 import { UserError } from "../errors";
 import { getGlobalWranglerConfigPath } from "../global-wrangler-config-path";
 import { CI } from "../is-ci";
@@ -306,7 +307,7 @@ interface AuthTokens {
  * The path to the config file that holds user authentication data,
  * relative to the user's home directory.
  */
-export const USER_AUTH_CONFIG_FILE = "config/default.toml";
+export const USER_AUTH_CONFIG_PATH = "config";
 
 /**
  * The data that may be read from the `USER_CONFIG_FILE`.
@@ -875,33 +876,32 @@ async function generatePKCECodes(): Promise<PKCECodes> {
 	return { codeChallenge, codeVerifier };
 }
 
+export function getAuthConfigFilePath() {
+	const environment = getCloudflareApiEnvironmentFromEnv();
+	const filePath = `${USER_AUTH_CONFIG_PATH}/${environment === "production" ? "default.toml" : `${environment}.toml`}`;
+
+	return path.join(getGlobalWranglerConfigPath(), filePath);
+}
+
 /**
  * Writes a a wrangler config file (auth credentials) to disk,
  * and updates the user auth state with the new credentials.
  */
 export function writeAuthConfigFile(config: UserAuthConfig) {
-	const authConfigFilePath = path.join(
-		getGlobalWranglerConfigPath(),
-		USER_AUTH_CONFIG_FILE
-	);
-	mkdirSync(path.dirname(authConfigFilePath), {
+	const configPath = getAuthConfigFilePath();
+
+	mkdirSync(path.dirname(configPath), {
 		recursive: true,
 	});
-	writeFileSync(
-		path.join(authConfigFilePath),
-		TOML.stringify(config as TOML.JsonMap),
-		{ encoding: "utf-8" }
-	);
+	writeFileSync(path.join(configPath), TOML.stringify(config as TOML.JsonMap), {
+		encoding: "utf-8",
+	});
 
 	reinitialiseAuthTokens();
 }
 
 export function readAuthConfigFile(): UserAuthConfig {
-	const authConfigFilePath = path.join(
-		getGlobalWranglerConfigPath(),
-		USER_AUTH_CONFIG_FILE
-	);
-	const toml = parseTOML(readFileSync(authConfigFilePath));
+	const toml = parseTOML(readFileSync(getAuthConfigFilePath()));
 	return toml;
 }
 
@@ -1106,7 +1106,7 @@ export async function logout(): Promise<void> {
 		},
 	});
 	await response.text(); // blank text? would be nice if it was something meaningful
-	rmSync(path.join(getGlobalWranglerConfigPath(), USER_AUTH_CONFIG_FILE));
+	rmSync(getAuthConfigFilePath());
 	logger.log(`Successfully logged out.`);
 }
 

packages/wrangler/turbo.json

@@ -51,7 +51,8 @@
 				"SENTRY_DSN",
 				"SYSTEMROOT",
 				"WRANGLER_DISABLE_REQUEST_BODY_DRAINING",
-				"WRANGLER_WORKER_REGISTRY_PORT"
+				"WRANGLER_WORKER_REGISTRY_PORT",
+				"WRANGLER_API_ENVIRONMENT"
 			]
 		},
 		"test:ci": {