Improve consistency in client-side login/logout experience

Author: EhabY

src/commands.ts

@@ -179,31 +179,32 @@ export class Commands {
 	}
 
 	/**
-	 * Log into the provided deployment.  If the deployment URL is not specified,
+	 * Log into the provided deployment. If the deployment URL is not specified,
 	 * ask for it first with a menu showing recent URLs along with the default URL
 	 * and CODER_URL, if those are set.
 	 */
-	public async login(...args: string[]): Promise<void> {
-		// Destructure would be nice but VS Code can pass undefined which errors.
-		const inputUrl = args[0];
-		const inputToken = args[1];
-		const inputLabel = args[2];
-		const isAutologin =
-			typeof args[3] === "undefined" ? false : Boolean(args[3]);
-
-		const url = await this.maybeAskUrl(inputUrl);
+	public async login(args?: {
+		url?: string;
+		token?: string;
+		label?: string;
+		autoLogin?: boolean;
+	}): Promise<void> {
+		const url = await this.maybeAskUrl(args?.url);
 		if (!url) {
 			return; // The user aborted.
 		}
 
 		// It is possible that we are trying to log into an old-style host, in which
 		// case we want to write with the provided blank label instead of generating
 		// a host label.
-		const label =
-			typeof inputLabel === "undefined" ? toSafeHost(url) : inputLabel;
+		const label = args?.label === undefined ? toSafeHost(url) : args?.label;
 
 		// Try to get a token from the user, if we need one, and their user.
-		const res = await this.maybeAskToken(url, inputToken, isAutologin);
+		const res = await this.maybeAskToken(
+			url,
+			args?.token,
+			args?.autoLogin === true,
+		);
 		if (!res) {
 			return; // The user aborted, or unable to auth.
 		}
@@ -257,21 +258,23 @@ export class Commands {
 	 */
 	private async maybeAskToken(
 		url: string,
-		token: string,
-		isAutologin: boolean,
+		token: string | undefined,
+		isAutoLogin: boolean,
 	): Promise<{ user: User; token: string } | null> {
 		const client = CoderApi.create(url, token, this.logger, () =>
 			vscode.workspace.getConfiguration(),
 		);
-		if (!needToken(vscode.workspace.getConfiguration())) {
-			try {
-				const user = await client.getAuthenticatedUser();
-				// For non-token auth, we write a blank token since the `vscodessh`
-				// command currently always requires a token file.
-				return { token: "", user };
-			} catch (err) {
+		const needsToken = needToken(vscode.workspace.getConfiguration());
+		try {
+			const user = await client.getAuthenticatedUser();
+			// For non-token auth, we write a blank token since the `vscodessh`
+			// command currently always requires a token file.
+			// For token auth, we have valid access so we can just return the user here
+			return { token: needsToken && token ? token : "", user };
+		} catch (err) {
+			if (!needToken(vscode.workspace.getConfiguration())) {
 				const message = getErrorMessage(err, "no response from the server");
-				if (isAutologin) {
+				if (isAutoLogin) {
 					this.logger.warn("Failed to log in to Coder server:", message);
 				} else {
 					this.vscodeProposed.window.showErrorMessage(
@@ -303,6 +306,9 @@ export class Commands {
 			value: token || (await this.secretsManager.getSessionToken()),
 			ignoreFocusOut: true,
 			validateInput: async (value) => {
+				if (!value) {
+					return null;
+				}
 				client.setSessionToken(value);
 				try {
 					user = await client.getAuthenticatedUser();
@@ -371,7 +377,10 @@ export class Commands {
 			// Sanity check; command should not be available if no url.
 			throw new Error("You are not logged in");
 		}
+		await this.forceLogout();
+	}
 
+	public async forceLogout(): Promise<void> {
 		// Clear from the REST client.  An empty url will indicate to other parts of
 		// the code that we are logged out.
 		this.restClient.setHost("");
@@ -390,7 +399,7 @@ export class Commands {
 			.showInformationMessage("You've been logged out of Coder!", "Login")
 			.then((action) => {
 				if (action === "Login") {
-					vscode.commands.executeCommand("coder.login");
+					this.login();
 				}
 			});
 

src/core/secretsManager.ts

@@ -1,4 +1,6 @@
-import type { SecretStorage } from "vscode";
+import type { SecretStorage, Disposable } from "vscode";
+
+const SESSION_TOKEN_KEY = "sessionToken";
 
 export class SecretsManager {
 	constructor(private readonly secrets: SecretStorage) {}
@@ -8,9 +10,9 @@ export class SecretsManager {
 	 */
 	public async setSessionToken(sessionToken?: string): Promise<void> {
 		if (!sessionToken) {
-			await this.secrets.delete("sessionToken");
+			await this.secrets.delete(SESSION_TOKEN_KEY);
 		} else {
-			await this.secrets.store("sessionToken", sessionToken);
+			await this.secrets.store(SESSION_TOKEN_KEY, sessionToken);
 		}
 	}
 
@@ -19,11 +21,22 @@ export class SecretsManager {
 	 */
 	public async getSessionToken(): Promise<string | undefined> {
 		try {
-			return await this.secrets.get("sessionToken");
+			return await this.secrets.get(SESSION_TOKEN_KEY);
 		} catch {
 			// The VS Code session store has become corrupt before, and
 			// will fail to get the session token...
 			return undefined;
 		}
 	}
+
+	/**
+	 * Subscribe to changes to the session token which can be used to indicate user login status.
+	 */
+	public onDidChangeSessionToken(listener: () => Promise<void>): Disposable {
+		return this.secrets.onDidChange((e) => {
+			if (e.key === SESSION_TOKEN_KEY) {
+				listener();
+			}
+		});
+	}
 }

src/extension.ts

@@ -299,6 +299,21 @@ export async function activate(ctx: vscode.ExtensionContext): Promise<void> {
 		commands.viewLogs.bind(commands),
 	);
 
+	ctx.subscriptions.push(
+		secretsManager.onDidChangeSessionToken(async () => {
+			const token = await secretsManager.getSessionToken();
+			const url = mementoManager.getUrl();
+			if (!token) {
+				output.info("Logging out");
+				await commands.forceLogout();
+			} else if (url) {
+				output.info("Logging in");
+				// Should login the user directly if the URL+Token are valid
+				await commands.login({ url, token });
+			}
+		}),
+	);
+
 	// Since the "onResolveRemoteAuthority:ssh-remote" activation event exists
 	// in package.json we're able to perform actions before the authority is
 	// resolved by the remote SSH extension.
@@ -410,13 +425,7 @@ export async function activate(ctx: vscode.ExtensionContext): Promise<void> {
 				cfg.get<string>("coder.defaultUrl")?.trim() ||
 				process.env.CODER_URL?.trim();
 			if (defaultUrl) {
-				vscode.commands.executeCommand(
-					"coder.login",
-					defaultUrl,
-					undefined,
-					undefined,
-					"true",
-				);
+				commands.login({ url: defaultUrl, autoLogin: true });
 			}
 		}
 	}

src/remote/remote.ts

@@ -257,12 +257,7 @@ export class Remote {
 				await this.closeRemote();
 			} else {
 				// Log in then try again.
-				await vscode.commands.executeCommand(
-					"coder.login",
-					baseUrlRaw,
-					undefined,
-					parts.label,
-				);
+				await this.commands.login({ url: baseUrlRaw, label: parts.label });
 				await this.setup(remoteAuthority, firstConnect);
 			}
 			return;
@@ -382,12 +377,7 @@ export class Remote {
 					if (!result) {
 						await this.closeRemote();
 					} else {
-						await vscode.commands.executeCommand(
-							"coder.login",
-							baseUrlRaw,
-							undefined,
-							parts.label,
-						);
+						await this.commands.login({ url: baseUrlRaw, label: parts.label });
 						await this.setup(remoteAuthority, firstConnect);
 					}
 					return;