Skip to content

Commit cb586a5

Browse files
EhabYEhabY
committed
Review comments
1 parent acf26aa commit cb586a5

File tree

7 files changed

+115
-98
lines changed

7 files changed

+115
-98
lines changed

‎src/commands.ts‎

Lines changed: 10 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -79,7 +79,6 @@ export class Commands {
7979
*/
8080
public async login(args?: {
8181
url?: string;
82-
safeHostname?: string;
8382
autoLogin?: boolean;
8483
}): Promise<void> {
8584
if (this.deploymentManager.isAuthenticated()) {
@@ -94,25 +93,27 @@ export class Commands {
9493
currentDeployment?.url,
9594
);
9695
if (!url) {
97-
return;
96+
return; // The user aborted.
9897
}
9998

100-
// It is possible that we are trying to log into an old-style host, in which
101-
// case we want to write with the provided blank hostname instead of
102-
// generating one.
103-
const safeHostname = args?.safeHostname ?? toSafeHost(url);
99+
const safeHostname = toSafeHost(url);
104100
this.logger.info("Using hostname", safeHostname);
105101

106-
const result = await this.loginCoordinator.promptForLogin({
102+
const result = await this.loginCoordinator.ensureLoggedIn({
107103
safeHostname,
108104
url,
109105
autoLogin: args?.autoLogin,
110106
});
111107

112-
if (!result.success || !result.user || result.token === undefined) {
108+
if (!result.success) {
113109
return;
114110
}
115111

112+
if (!result.user) {
113+
// Login might have happened in another process/window so we do not have the user yet.
114+
result.user = await this.extensionClient.getAuthenticatedUser();
115+
}
116+
116117
await this.deploymentManager.changeDeployment({
117118
url,
118119
safeHostname,
@@ -171,7 +172,7 @@ export class Commands {
171172

172173
this.logger.info("Logging out");
173174

174-
await this.deploymentManager.logout();
175+
await this.deploymentManager.clearDeployment();
175176

176177
vscode.window
177178
.showInformationMessage("You've been logged out of Coder!", "Login")

‎src/core/cliManager.ts‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -757,7 +757,9 @@ export class CliManager {
757757
await fs.writeFile(tempPath, content);
758758
await fs.rename(tempPath, filePath);
759759
} catch (err) {
760-
await fs.rm(tempPath, { force: true }).catch(() => {});
760+
await fs.rm(tempPath, { force: true }).catch((rmErr) => {
761+
this.output.warn("Failed to delete temp file", tempPath, rmErr);
762+
});
761763
throw err;
762764
}
763765
}

‎src/core/secretsManager.ts‎

Lines changed: 12 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -98,9 +98,9 @@ export class SecretsManager {
9898
safeHostname: string,
9999
listener: (auth: SessionAuth | undefined) => void | Promise<void>,
100100
): Disposable {
101-
const key = `${SESSION_KEY_PREFIX}${safeHostname}`;
101+
const sessionKey = this.getSessionKey(safeHostname);
102102
return this.secrets.onDidChange(async (e) => {
103-
if (e.key !== key) {
103+
if (e.key !== sessionKey) {
104104
return;
105105
}
106106
const auth = await this.getSessionAuth(safeHostname);
@@ -115,14 +115,9 @@ export class SecretsManager {
115115
public async getSessionAuth(
116116
safeHostname: string,
117117
): Promise<SessionAuth | undefined> {
118-
if (!safeHostname) {
119-
return undefined;
120-
}
121-
118+
const sessionKey = this.getSessionKey(safeHostname);
122119
try {
123-
const data = await this.secrets.get(
124-
`${SESSION_KEY_PREFIX}${safeHostname}`,
125-
);
120+
const data = await this.secrets.get(sessionKey);
126121
if (!data) {
127122
return undefined;
128123
}
@@ -136,22 +131,18 @@ export class SecretsManager {
136131
safeHostname: string,
137132
auth: SessionAuth,
138133
): Promise<void> {
139-
if (!safeHostname) {
140-
return;
141-
}
142-
143-
await this.secrets.store(
144-
`${SESSION_KEY_PREFIX}${safeHostname}`,
145-
JSON.stringify(auth),
146-
);
134+
const sessionKey = this.getSessionKey(safeHostname);
135+
await this.secrets.store(sessionKey, JSON.stringify(auth));
147136
await this.recordDeploymentAccess(safeHostname);
148137
}
149138

150139
private async clearSessionAuth(safeHostname: string): Promise<void> {
151-
if (!safeHostname) {
152-
return;
153-
}
154-
await this.secrets.delete(`${SESSION_KEY_PREFIX}${safeHostname}`);
140+
const sessionKey = this.getSessionKey(safeHostname);
141+
await this.secrets.delete(sessionKey);
142+
}
143+
144+
private getSessionKey(safeHostname: string): string {
145+
return `${SESSION_KEY_PREFIX}${safeHostname || "<legacy>"}`;
155146
}
156147

157148
/**

‎src/deployment/deploymentManager.ts‎

Lines changed: 25 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -78,7 +78,7 @@ export class DeploymentManager implements vscode.Disposable {
7878
public async changeDeployment(
7979
deployment: DeploymentWithAuth & { user: User },
8080
): Promise<void> {
81-
this.setDeploymentCore(deployment);
81+
this.setDeployment(deployment);
8282

8383
this.refreshWorkspaces();
8484
await this.persistDeployment(deployment);
@@ -89,15 +89,15 @@ export class DeploymentManager implements vscode.Disposable {
8989
* Immediately tries to fetch user and upgrade to authenticated state.
9090
* Use this for startup or when you don't have the user yet.
9191
*/
92-
public async setDeploymentWithoutAuth(
92+
public async setDeploymentAndValidate(
9393
deployment: Deployment & { token?: string },
9494
): Promise<void> {
95-
this.setDeploymentCore({ ...deployment });
95+
this.setDeployment({ ...deployment });
9696

9797
await this.tryFetchAndUpgradeUser();
9898
}
9999

100-
private setDeploymentCore(deployment: DeploymentWithAuth): void {
100+
private setDeployment(deployment: DeploymentWithAuth): void {
101101
if (deployment.token === undefined) {
102102
this.client.setHost(deployment.url);
103103
} else {
@@ -109,9 +109,9 @@ export class DeploymentManager implements vscode.Disposable {
109109
}
110110

111111
/**
112-
* Log out from the current deployment.
112+
* Clears the current deployment.
113113
*/
114-
public async logout(): Promise<void> {
114+
public async clearDeployment(): Promise<void> {
115115
this.client.setCredentials(undefined, undefined);
116116

117117
this.authListenerDisposable?.dispose();
@@ -142,7 +142,7 @@ export class DeploymentManager implements vscode.Disposable {
142142
const auth = await this.secretsManager.getSessionAuth(
143143
deployment.safeHostname,
144144
);
145-
await this.setDeploymentWithoutAuth({
145+
await this.setDeploymentAndValidate({
146146
...deployment,
147147
token: auth?.token,
148148
});
@@ -162,17 +162,18 @@ export class DeploymentManager implements vscode.Disposable {
162162
this.authListenerDisposable = this.secretsManager.onDidChangeSessionAuth(
163163
safeHostname,
164164
async (auth) => {
165+
if (this.currentDeployment?.safeHostname !== safeHostname) {
166+
return;
167+
}
165168
if (auth) {
166-
if (this.currentDeployment?.safeHostname !== safeHostname) {
167-
return;
168-
}
169-
170-
this.client.setCredentials(this.currentDeployment.url, auth.token);
169+
this.client.setCredentials(auth.url, auth.token);
171170

172171
// If we don't have a user yet, try to fetch one
173172
if (!this.currentDeployment?.user) {
174173
await this.tryFetchAndUpgradeUser();
175174
}
175+
} else {
176+
await this.clearDeployment();
176177
}
177178
},
178179
);
@@ -186,9 +187,20 @@ export class DeploymentManager implements vscode.Disposable {
186187
return;
187188
}
188189

190+
const safeHostname = this.currentDeployment.safeHostname;
191+
189192
try {
190193
const user = await this.client.getAuthenticatedUser();
191-
this.currentDeployment = { ...this.currentDeployment, user };
194+
195+
// Re-validate deployment hasn't changed during await
196+
if (this.currentDeployment?.safeHostname !== safeHostname) {
197+
this.logger.debug(
198+
"Deployment changed during user fetch, discarding result",
199+
);
200+
return;
201+
}
202+
203+
this.currentDeployment.user = user;
192204
this.updateAuthContexts(user);
193205
this.refreshWorkspaces();
194206

‎src/extension.ts‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -313,7 +313,7 @@ export async function activate(ctx: vscode.ExtensionContext): Promise<void> {
313313
ctx.subscriptions.push(details);
314314

315315
// Will automatically fetch the user and upgrade the deployment
316-
await deploymentManager.setDeploymentWithoutAuth({
316+
await deploymentManager.setDeploymentAndValidate({
317317
safeHostname: details.safeHostname,
318318
url: details.url,
319319
token: details.token,
@@ -365,7 +365,7 @@ export async function activate(ctx: vscode.ExtensionContext): Promise<void> {
365365
output.info(`Initializing deployment: ${deployment.url}`);
366366
const auth = await secretsManager.getSessionAuth(deployment.safeHostname);
367367
deploymentManager
368-
.setDeploymentWithoutAuth({ ...deployment, token: auth?.token })
368+
.setDeploymentAndValidate({ ...deployment, token: auth?.token })
369369
.then(() => {
370370
if (deploymentManager.isAuthenticated()) {
371371
output.info("Credentials are valid");
@@ -478,7 +478,7 @@ async function setupDeploymentFromUri(
478478
}
479479

480480
// Will automatically fetch the user and upgrade the deployment
481-
await deploymentManager.setDeploymentWithoutAuth({
481+
await deploymentManager.setDeploymentAndValidate({
482482
safeHostname: safeHost,
483483
url,
484484
token,

‎src/login/loginCoordinator.ts‎

Lines changed: 41 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import { isAxiosError } from "axios";
12
import { getErrorMessage } from "coder/site/src/api/errors";
23
import * as vscode from "vscode";
34

@@ -13,11 +14,9 @@ import type { SecretsManager } from "../core/secretsManager";
1314
import type { Deployment } from "../deployment/types";
1415
import type { Logger } from "../logging/logger";
1516

16-
interface LoginResult {
17-
success: boolean;
18-
user?: User;
19-
token?: string;
20-
}
17+
type LoginResult =
18+
| { success: false }
19+
| { success: true; user?: User; token: string };
2120

2221
interface LoginOptions {
2322
safeHostname: string;
@@ -42,7 +41,7 @@ export class LoginCoordinator {
4241
* Direct login - for user-initiated login via commands.
4342
* Stores session auth and URL history on success.
4443
*/
45-
public async promptForLogin(
44+
public async ensureLoggedIn(
4645
options: LoginOptions & { url: string },
4746
): Promise<LoginResult> {
4847
const { safeHostname, url } = options;
@@ -61,11 +60,11 @@ export class LoginCoordinator {
6160
/**
6261
* Shows dialog then login - for system-initiated auth (remote).
6362
*/
64-
public async promptForLoginWithDialog(
63+
public async ensureLoggedInWithDialog(
6564
options: LoginOptions & { message?: string; detailPrefix?: string },
6665
): Promise<LoginResult> {
6766
const { safeHostname, url, detailPrefix, message } = options;
68-
return this.executeWithGuard(safeHostname, () => {
67+
return this.executeWithGuard(safeHostname, async () => {
6968
// Show dialog promise
7069
const dialogPromise = this.vscodeProposed.window
7170
.showErrorMessage(
@@ -103,15 +102,20 @@ export class LoginCoordinator {
103102
return result;
104103
} else {
105104
// User cancelled
106-
return { success: false };
105+
return { success: false } as const;
107106
}
108107
});
109108

110109
// Race between user clicking login and cross-window detection
111-
return Promise.race([
112-
dialogPromise,
113-
this.waitForCrossWindowLogin(safeHostname),
114-
]);
110+
const {
111+
promise: crossWindowPromise,
112+
dispose: disposeCrossWindowListener,
113+
} = this.waitForCrossWindowLogin(safeHostname);
114+
try {
115+
return await Promise.race([dialogPromise, crossWindowPromise]);
116+
} finally {
117+
disposeCrossWindowListener();
118+
}
115119
});
116120
}
117121

@@ -121,7 +125,7 @@ export class LoginCoordinator {
121125
url: string,
122126
): Promise<void> {
123127
// Empty token is valid for mTLS
124-
if (result.success && result.token !== undefined) {
128+
if (result.success) {
125129
await this.secretsManager.setSessionAuth(safeHostname, {
126130
url,
127131
token: result.token,
@@ -154,21 +158,28 @@ export class LoginCoordinator {
154158

155159
/**
156160
* Waits for login detected from another window.
161+
* Returns a promise and a dispose function to clean up the listener.
157162
*/
158-
private async waitForCrossWindowLogin(
159-
safeHostname: string,
160-
): Promise<LoginResult> {
161-
return new Promise((resolve) => {
162-
const disposable = this.secretsManager.onDidChangeSessionAuth(
163+
private waitForCrossWindowLogin(safeHostname: string): {
164+
promise: Promise<LoginResult>;
165+
dispose: () => void;
166+
} {
167+
let disposable: vscode.Disposable | undefined;
168+
const promise = new Promise<LoginResult>((resolve) => {
169+
disposable = this.secretsManager.onDidChangeSessionAuth(
163170
safeHostname,
164171
(auth) => {
165172
if (auth?.token) {
166-
disposable.dispose();
173+
disposable?.dispose();
167174
resolve({ success: true, token: auth.token });
168175
}
169176
},
170177
);
171178
});
179+
return {
180+
promise,
181+
dispose: () => disposable?.dispose(),
182+
};
172183
}
173184

174185
/**
@@ -197,15 +208,18 @@ export class LoginCoordinator {
197208

198209
// Attempt authentication with current credentials (token or mTLS)
199210
try {
200-
const user = await client.getAuthenticatedUser();
201-
// Return the token that was used (empty string for mTLS since
202-
// the `vscodessh` command currently always requires a token file)
203-
return { success: true, token: storedToken ?? "", user };
211+
if (!needsToken || storedToken) {
212+
const user = await client.getAuthenticatedUser();
213+
// Return the token that was used (empty string for mTLS since
214+
// the `vscodessh` command currently always requires a token file)
215+
return { success: true, token: storedToken ?? "", user };
216+
}
204217
} catch (err) {
205-
if (needsToken) {
206-
// For token auth: silently continue to prompt for new credentials
218+
const is401 = isAxiosError(err) && err.response?.status === 401;
219+
if (needsToken && is401) {
220+
// For token auth with 401: silently continue to prompt for new credentials
207221
} else {
208-
// For mTLS: show error and abort (no credentials to prompt for)
222+
// For mTLS or non-401 errors: show error and abort
209223
const message = getErrorMessage(err, "no response from the server");
210224
if (isAutoLogin) {
211225
this.logger.warn("Failed to log in to Coder server:", message);

0 commit comments

Comments
 (0)