Skip to content

Commit 16da8fd

Browse files
EhabYEhabY
committed
Self-review of OAuth logic
1 parent 4a4a6ad commit 16da8fd

File tree

5 files changed

+142
-136
lines changed

5 files changed

+142
-136
lines changed

‎src/core/secretsManager.ts‎

Lines changed: 105 additions & 104 deletions
Original file line numberDiff line numberDiff line change
@@ -15,9 +15,10 @@ const SESSION_KEY_PREFIX = "coder.session.";
1515
const OAUTH_TOKENS_PREFIX = "coder.oauth.tokens.";
1616
const OAUTH_CLIENT_PREFIX = "coder.oauth.client.";
1717

18-
const CURRENT_DEPLOYMENT_KEY = "coder.currentDeployment";
1918
const OAUTH_CALLBACK_KEY = "coder.oauthCallback";
2019

20+
const CURRENT_DEPLOYMENT_KEY = "coder.currentDeployment";
21+
2122
const DEPLOYMENT_USAGE_KEY = "coder.deploymentUsage";
2223
const DEFAULT_MAX_DEPLOYMENTS = 10;
2324

@@ -115,38 +116,6 @@ export class SecretsManager {
115116
});
116117
}
117118

118-
/**
119-
* Write an OAuth callback result to secrets storage.
120-
* Used for cross-window communication when OAuth callback arrives in a different window.
121-
*/
122-
public async setOAuthCallback(data: OAuthCallbackData): Promise<void> {
123-
await this.secrets.store(OAUTH_CALLBACK_KEY, JSON.stringify(data));
124-
}
125-
126-
/**
127-
* Listen for OAuth callback results from any VS Code window.
128-
* The listener receives the state parameter, code (if success), and error (if failed).
129-
*/
130-
public onDidChangeOAuthCallback(
131-
listener: (data: OAuthCallbackData) => void,
132-
): Disposable {
133-
return this.secrets.onDidChange(async (e) => {
134-
if (e.key !== OAUTH_CALLBACK_KEY) {
135-
return;
136-
}
137-
138-
try {
139-
const data = await this.secrets.get(OAUTH_CALLBACK_KEY);
140-
if (data) {
141-
const parsed = JSON.parse(data) as OAuthCallbackData;
142-
listener(parsed);
143-
}
144-
} catch {
145-
// Ignore parse errors
146-
}
147-
});
148-
}
149-
150119
/**
151120
* Listen for changes to a specific deployment's session auth.
152121
*/
@@ -203,77 +172,6 @@ export class SecretsManager {
203172
return `${SESSION_KEY_PREFIX}${safeHostname || "<legacy>"}`;
204173
}
205174

206-
public async getOAuthTokens(
207-
safeHostname: string,
208-
): Promise<StoredOAuthTokens | undefined> {
209-
try {
210-
const data = await this.secrets.get(
211-
`${OAUTH_TOKENS_PREFIX}${safeHostname}`,
212-
);
213-
if (!data) {
214-
return undefined;
215-
}
216-
return JSON.parse(data) as StoredOAuthTokens;
217-
} catch {
218-
return undefined;
219-
}
220-
}
221-
222-
public async setOAuthTokens(
223-
safeHostname: string,
224-
tokens: StoredOAuthTokens,
225-
): Promise<void> {
226-
await this.secrets.store(
227-
`${OAUTH_TOKENS_PREFIX}${safeHostname}`,
228-
JSON.stringify(tokens),
229-
);
230-
await this.recordDeploymentAccess(safeHostname);
231-
}
232-
233-
public async clearOAuthTokens(safeHostname: string): Promise<void> {
234-
await this.secrets.delete(`${OAUTH_TOKENS_PREFIX}${safeHostname}`);
235-
}
236-
237-
public async getOAuthClientRegistration(
238-
safeHostname: string,
239-
): Promise<ClientRegistrationResponse | undefined> {
240-
try {
241-
const data = await this.secrets.get(
242-
`${OAUTH_CLIENT_PREFIX}${safeHostname}`,
243-
);
244-
if (!data) {
245-
return undefined;
246-
}
247-
return JSON.parse(data) as ClientRegistrationResponse;
248-
} catch {
249-
return undefined;
250-
}
251-
}
252-
253-
public async setOAuthClientRegistration(
254-
safeHostname: string,
255-
registration: ClientRegistrationResponse,
256-
): Promise<void> {
257-
await this.secrets.store(
258-
`${OAUTH_CLIENT_PREFIX}${safeHostname}`,
259-
JSON.stringify(registration),
260-
);
261-
await this.recordDeploymentAccess(safeHostname);
262-
}
263-
264-
public async clearOAuthClientRegistration(
265-
safeHostname: string,
266-
): Promise<void> {
267-
await this.secrets.delete(`${OAUTH_CLIENT_PREFIX}${safeHostname}`);
268-
}
269-
270-
public async clearOAuthData(safeHostname: string): Promise<void> {
271-
await Promise.all([
272-
this.clearOAuthTokens(safeHostname),
273-
this.clearOAuthClientRegistration(safeHostname),
274-
]);
275-
}
276-
277175
/**
278176
* Record that a deployment was accessed, moving it to the front of the LRU list.
279177
* Prunes deployments beyond maxCount, clearing their auth data.
@@ -359,4 +257,107 @@ export class SecretsManager {
359257

360258
return safeHostname;
361259
}
260+
261+
/**
262+
* Write an OAuth callback result to secrets storage.
263+
* Used for cross-window communication when OAuth callback arrives in a different window.
264+
*/
265+
public async setOAuthCallback(data: OAuthCallbackData): Promise<void> {
266+
await this.secrets.store(OAUTH_CALLBACK_KEY, JSON.stringify(data));
267+
}
268+
269+
/**
270+
* Listen for OAuth callback results from any VS Code window.
271+
* The listener receives the state parameter, code (if success), and error (if failed).
272+
*/
273+
public onDidChangeOAuthCallback(
274+
listener: (data: OAuthCallbackData) => void,
275+
): Disposable {
276+
return this.secrets.onDidChange(async (e) => {
277+
if (e.key !== OAUTH_CALLBACK_KEY) {
278+
return;
279+
}
280+
281+
try {
282+
const data = await this.secrets.get(OAUTH_CALLBACK_KEY);
283+
if (data) {
284+
const parsed = JSON.parse(data) as OAuthCallbackData;
285+
listener(parsed);
286+
}
287+
} catch {
288+
// Ignore parse errors
289+
}
290+
});
291+
}
292+
293+
public async getOAuthTokens(
294+
safeHostname: string,
295+
): Promise<StoredOAuthTokens | undefined> {
296+
try {
297+
const data = await this.secrets.get(
298+
`${OAUTH_TOKENS_PREFIX}${safeHostname}`,
299+
);
300+
if (!data) {
301+
return undefined;
302+
}
303+
return JSON.parse(data) as StoredOAuthTokens;
304+
} catch {
305+
return undefined;
306+
}
307+
}
308+
309+
public async setOAuthTokens(
310+
safeHostname: string,
311+
tokens: StoredOAuthTokens,
312+
): Promise<void> {
313+
await this.secrets.store(
314+
`${OAUTH_TOKENS_PREFIX}${safeHostname}`,
315+
JSON.stringify(tokens),
316+
);
317+
await this.recordDeploymentAccess(safeHostname);
318+
}
319+
320+
public async clearOAuthTokens(safeHostname: string): Promise<void> {
321+
await this.secrets.delete(`${OAUTH_TOKENS_PREFIX}${safeHostname}`);
322+
}
323+
324+
public async getOAuthClientRegistration(
325+
safeHostname: string,
326+
): Promise<ClientRegistrationResponse | undefined> {
327+
try {
328+
const data = await this.secrets.get(
329+
`${OAUTH_CLIENT_PREFIX}${safeHostname}`,
330+
);
331+
if (!data) {
332+
return undefined;
333+
}
334+
return JSON.parse(data) as ClientRegistrationResponse;
335+
} catch {
336+
return undefined;
337+
}
338+
}
339+
340+
public async setOAuthClientRegistration(
341+
safeHostname: string,
342+
registration: ClientRegistrationResponse,
343+
): Promise<void> {
344+
await this.secrets.store(
345+
`${OAUTH_CLIENT_PREFIX}${safeHostname}`,
346+
JSON.stringify(registration),
347+
);
348+
await this.recordDeploymentAccess(safeHostname);
349+
}
350+
351+
public async clearOAuthClientRegistration(
352+
safeHostname: string,
353+
): Promise<void> {
354+
await this.secrets.delete(`${OAUTH_CLIENT_PREFIX}${safeHostname}`);
355+
}
356+
357+
public async clearOAuthData(safeHostname: string): Promise<void> {
358+
await Promise.all([
359+
this.clearOAuthTokens(safeHostname),
360+
this.clearOAuthClientRegistration(safeHostname),
361+
]);
362+
}
362363
}

‎src/deployment/deploymentManager.ts‎

Lines changed: 4 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -89,12 +89,6 @@ export class DeploymentManager implements vscode.Disposable {
8989
public async setDeploymentIfValid(
9090
deployment: Deployment & { token?: string },
9191
): Promise<boolean> {
92-
// TODO used to trigger
93-
/**
94-
* this.oauthSessionManager.refreshIfAlmostExpired().catch((error) => {
95-
this.logger.warn("Setup token refresh failed:", error);
96-
});
97-
*/
9892
const auth = await this.secretsManager.getSessionAuth(
9993
deployment.safeHostname,
10094
);
@@ -134,11 +128,14 @@ export class DeploymentManager implements vscode.Disposable {
134128
} else {
135129
this.client.setCredentials(deployment.url, deployment.token);
136130
}
137-
await this.oauthSessionManager.setDeployment(deployment);
138131

132+
// Register auth listener before setDeployment so background token refresh
133+
// can update client credentials via the listener
139134
this.registerAuthListener();
140135
this.updateAuthContexts();
141136
this.refreshWorkspaces();
137+
138+
await this.oauthSessionManager.setDeployment(deployment);
142139
await this.persistDeployment(deployment);
143140
}
144141

@@ -158,13 +155,6 @@ export class DeploymentManager implements vscode.Disposable {
158155
await this.secretsManager.setCurrentDeployment(undefined);
159156
}
160157

161-
/**
162-
* Clear OAuth state for a deployment when switching to token auth.
163-
*/
164-
public async clearOAuthState(label: string): Promise<void> {
165-
await this.oauthSessionManager.clearOAuthState(label);
166-
}
167-
168158
public dispose(): void {
169159
this.#authListenerDisposable?.dispose();
170160
this.#crossWindowSyncDisposable?.dispose();

‎src/extension.ts‎

Lines changed: 4 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -172,11 +172,7 @@ export async function activate(ctx: vscode.ExtensionContext): Promise<void> {
172172
throw new Error("workspace must be specified as a query parameter");
173173
}
174174

175-
await setupDeploymentFromUri(
176-
params,
177-
serviceContainer,
178-
deploymentManager,
179-
);
175+
await setupDeploymentFromUri(params, serviceContainer);
180176

181177
await commands.open(
182178
owner,
@@ -230,11 +226,7 @@ export async function activate(ctx: vscode.ExtensionContext): Promise<void> {
230226
);
231227
}
232228

233-
await setupDeploymentFromUri(
234-
params,
235-
serviceContainer,
236-
deploymentManager,
237-
);
229+
await setupDeploymentFromUri(params, serviceContainer);
238230

239231
await commands.openDevContainer(
240232
workspaceOwner,
@@ -460,7 +452,6 @@ async function showTreeViewSearch(id: string): Promise<void> {
460452
async function setupDeploymentFromUri(
461453
params: URLSearchParams,
462454
serviceContainer: ServiceContainer,
463-
deploymentManager: DeploymentManager,
464455
): Promise<void> {
465456
const secretsManager = serviceContainer.getSecretsManager();
466457
const mementoManager = serviceContainer.getMementoManager();
@@ -495,8 +486,8 @@ async function setupDeploymentFromUri(
495486
}
496487
} else {
497488
await secretsManager.setSessionAuth(safeHost, { url, token });
498-
// Clear OAuth state since we're using a non-OAuth token
499-
await deploymentManager.clearOAuthState(safeHost);
489+
// Clear OAuth tokens since we're using a non-OAuth token
490+
await secretsManager.clearOAuthTokens(safeHost);
500491
}
501492
}
502493

‎src/login/loginCoordinator.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -251,7 +251,7 @@ export class LoginCoordinator {
251251
const result = await this.loginWithToken(client);
252252
if (result.success) {
253253
// Clear OAuth state since user explicitly chose token auth
254-
await oauthSessionManager.clearOAuthState(deployment.safeHostname);
254+
await oauthSessionManager.clearOAuthState();
255255
}
256256
return result;
257257
}

‎src/oauth/sessionManager.ts‎

Lines changed: 28 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -152,8 +152,7 @@ export class OAuthSessionManager implements vscode.Disposable {
152152
required_scopes: DEFAULT_OAUTH_SCOPES,
153153
},
154154
);
155-
this.clearInMemoryTokens();
156-
await this.secretsManager.clearOAuthTokens(this.deployment.safeHostname);
155+
await this.clearOAuthState();
157156
return;
158157
}
159158

@@ -322,6 +321,19 @@ export class OAuthSessionManager implements vscode.Disposable {
322321
this.deployment = deployment;
323322
this.clearInMemoryTokens();
324323
await this.loadTokens();
324+
325+
// Refresh tokens if needed to prevent 401s
326+
if (this.isTokenExpired()) {
327+
try {
328+
await this.refreshToken();
329+
} catch (error) {
330+
this.logger.warn("Token refresh failed (expired):", error);
331+
}
332+
} else {
333+
this.refreshIfAlmostExpired().catch((error) => {
334+
this.logger.warn("Background token refresh failed:", error);
335+
});
336+
}
325337
}
326338

327339
public clearDeployment(): void {
@@ -695,6 +707,16 @@ export class OAuthSessionManager implements vscode.Disposable {
695707
return timeUntilExpiry < TOKEN_REFRESH_THRESHOLD_MS;
696708
}
697709

710+
/**
711+
* Check if token is expired.
712+
*/
713+
private isTokenExpired(): boolean {
714+
if (!this.storedTokens) {
715+
return false;
716+
}
717+
return Date.now() >= this.storedTokens.expiry_timestamp;
718+
}
719+
698720
public async revokeRefreshToken(): Promise<void> {
699721
if (!this.storedTokens?.refresh_token) {
700722
this.logger.info("No refresh token to revoke");
@@ -754,9 +776,11 @@ export class OAuthSessionManager implements vscode.Disposable {
754776
* Clears in-memory state and OAuth tokens from storage.
755777
* Preserves client registration for potential future OAuth use.
756778
*/
757-
public async clearOAuthState(label: string): Promise<void> {
779+
public async clearOAuthState(): Promise<void> {
758780
this.clearInMemoryTokens();
759-
await this.secretsManager.clearOAuthTokens(label);
781+
if (this.deployment) {
782+
await this.secretsManager.clearOAuthTokens(this.deployment.safeHostname);
783+
}
760784
}
761785

762786
/**

0 commit comments

Comments
 (0)