Handle secretsManager comments part 1

Author: EhabY

src/core/cliManager.ts

@@ -732,13 +732,10 @@ export class CliManager {
 	 *
 	 * If the label is empty, read the old deployment-unaware config instead.
 	 */
-	private async updateTokenForCli(
-		label: string,
-		token: string | undefined | null,
-	) {
+	private async updateTokenForCli(label: string, token: string | null) {
 		if (token !== null) {
 			const tokenPath = this.pathResolver.getSessionTokenPath(label);
-			await this.atomicWriteFile(tokenPath, token ?? "");
+			await this.atomicWriteFile(tokenPath, token);
 		}
 	}
 

src/core/secretsManager.ts

@@ -3,7 +3,10 @@ import { toSafeHost } from "../util";
 import type { Memento, SecretStorage, Disposable } from "vscode";
 
 import type { Deployment } from "../deployment";
+import { Logger } from "../logging/logger";
 
+// Each deployment has its own key to ensure atomic operations (multiple windows
+// writing to a shared key could drop data) and to receive proper VS Code events.
 const SESSION_KEY_PREFIX = "coder.session.";
 
 const CURRENT_DEPLOYMENT_KEY = "coder.currentDeployment";
@@ -31,6 +34,7 @@ export class SecretsManager {
 	constructor(
 		private readonly secrets: SecretStorage,
 		private readonly memento: Memento,
+		private readonly logger: Logger,
 	) {}
 
 	/**
@@ -50,16 +54,16 @@ export class SecretsManager {
 	/**
 	 * Gets the current deployment from storage.
 	 */
-	public async getCurrentDeployment(): Promise<Deployment | undefined> {
+	public async getCurrentDeployment(): Promise<Deployment | null> {
 		try {
 			const data = await this.secrets.get(CURRENT_DEPLOYMENT_KEY);
 			if (!data) {
-				return undefined;
+				return null;
 			}
-			const parsed = JSON.parse(data) as { deployment: Deployment | null };
-			return parsed.deployment ?? undefined;
+			const parsed = JSON.parse(data) as CurrentDeploymentState;
+			return parsed.deployment;
 		} catch {
-			return undefined;
+			return null;
 		}
 	}
 
@@ -75,16 +79,14 @@ export class SecretsManager {
 				return;
 			}
 
+			const deployment = await this.getCurrentDeployment();
 			try {
-				const data = await this.secrets.get(CURRENT_DEPLOYMENT_KEY);
-				if (data) {
-					const parsed = JSON.parse(data) as {
-						deployment: Deployment | null;
-					};
-					await listener({ deployment: parsed.deployment });
-				}
-			} catch {
-				// Ignore parse errors
+				await listener({ deployment });
+			} catch (err) {
+				this.logger.error(
+					"Error in onDidChangeCurrentDeployment listener",
+					err,
+				);
 			}
 		});
 	}
@@ -102,7 +104,11 @@ export class SecretsManager {
 				return;
 			}
 			const auth = await this.getSessionAuth(label);
-			await listener(auth);
+			try {
+				await listener(auth);
+			} catch (err) {
+				this.logger.error("Error in onDidChangeSessionAuth listener", err);
+			}
 		});
 	}
 
@@ -122,16 +128,6 @@ export class SecretsManager {
 		}
 	}
 
-	public async getSessionToken(label: string): Promise<string | undefined> {
-		const auth = await this.getSessionAuth(label);
-		return auth?.token;
-	}
-
-	public async getUrl(label: string): Promise<string | undefined> {
-		const auth = await this.getSessionAuth(label);
-		return auth?.url;
-	}
-
 	public async setSessionAuth(label: string, auth: SessionAuth): Promise<void> {
 		await this.secrets.store(
 			`${SESSION_KEY_PREFIX}${label}`,
@@ -210,6 +206,7 @@ export class SecretsManager {
 
 		await this.setSessionAuth(label, { url: legacyUrl, token: oldToken });
 		await this.secrets.delete(LEGACY_SESSION_TOKEN_KEY);
+		await this.memento.update("url", undefined);
 
 		// Also set as current deployment if none exists
 		const currentDeployment = await this.getCurrentDeployment();

src/deployment/deploymentManager.ts

@@ -142,10 +142,13 @@ export class DeploymentManager implements vscode.Disposable {
 
 					this.logger.info("Deployment changed from another window");
 					if (deployment) {
-						const token = await this.secretsManager.getSessionToken(
+						const auth = await this.secretsManager.getSessionAuth(
 							deployment.label,
 						);
-						await this.setDeploymentWithoutAuth({ ...deployment, token });
+						await this.setDeploymentWithoutAuth({
+							...deployment,
+							token: auth?.token,
+						});
 					}
 				},
 			);

src/extension.ts

@@ -74,7 +74,7 @@ export async function activate(ctx: vscode.ExtensionContext): Promise<void> {
 	// in commands that operate on the current login.
 	const client = CoderApi.create(
 		deployment?.url || "",
-		await secretsManager.getSessionToken(deployment?.label ?? ""),
+		(await secretsManager.getSessionAuth(deployment?.label ?? ""))?.token,
 		output,
 	);
 	ctx.subscriptions.push(client);
@@ -362,9 +362,9 @@ export async function activate(ctx: vscode.ExtensionContext): Promise<void> {
 		contextManager.set("coder.loaded", true);
 	} else if (deployment) {
 		output.info(`Initializing deployment: ${deployment.url}`);
-		const storedToken = await secretsManager.getSessionToken(deployment.label);
+		const auth = await secretsManager.getSessionAuth(deployment.label);
 		deploymentManager
-			.setDeploymentWithoutAuth({ ...deployment, token: storedToken })
+			.setDeploymentWithoutAuth({ ...deployment, token: auth?.token })
 			.then(() => {
 				if (deploymentManager.isAuthenticated()) {
 					output.info("Credentials are valid");
@@ -468,7 +468,7 @@ async function setupDeploymentFromUri(
 	let token: string | undefined = params.get("token") ?? undefined;
 	if (token === undefined) {
 		if (needToken(vscode.workspace.getConfiguration())) {
-			token = await secretsManager.getSessionToken(label);
+			token = (await secretsManager.getSessionAuth(label))?.token;
 		} else {
 			token = "";
 		}

src/login/loginCoordinator.ts

@@ -81,11 +81,11 @@ export class LoginCoordinator {
 				.then(async (action) => {
 					if (action === "Login") {
 						// Proceed with the login flow, handling logging in from another window
-						const storedUrl = await this.secretsManager.getUrl(label);
+						const storedAuth = await this.secretsManager.getSessionAuth(label);
 						const newUrl = await maybeAskUrl(
 							this.mementoManager,
 							url,
-							storedUrl,
+							storedAuth?.url,
 						);
 						if (!newUrl) {
 							throw new Error("URL must be provided");
@@ -178,7 +178,8 @@ export class LoginCoordinator {
 
 		let storedToken: string | undefined;
 		if (needsToken) {
-			storedToken = await this.secretsManager.getSessionToken(deployment.label);
+			const auth = await this.secretsManager.getSessionAuth(deployment.label);
+			storedToken = auth?.token;
 			if (storedToken) {
 				client.setSessionToken(storedToken);
 			}

src/remote/remote.ts

@@ -97,10 +97,10 @@ export class Remote {
 		await this.migrateSessionToken(parts.label);
 
 		// Get the URL and token belonging to this host.
-		const baseUrlRaw = (await this.secretsManager.getUrl(parts.label)) ?? "";
-		const token =
-			(await this.secretsManager.getSessionToken(parts.label)) ?? "";
-		if (baseUrlRaw && token) {
+		const auth = await this.secretsManager.getSessionAuth(parts.label);
+		const baseUrlRaw = auth?.url ?? "";
+		const token = auth?.token;
+		if (baseUrlRaw && token !== undefined) {
 			await this.cliManager.configure(parts.label, baseUrlRaw, token);
 		}
 
@@ -533,7 +533,7 @@ export class Remote {
 		return {
 			label: parts.label,
 			url: baseUrlRaw,
-			token,
+			token: token ?? "",
 			dispose: () => {
 				disposables.forEach((d) => d.dispose());
 			},