From ee74d713dbfe605c3c19aa96bead6f011f34c5e4 Mon Sep 17 00:00:00 2001 From: Ehab Younes Date: Tue, 21 Oct 2025 16:43:21 +0300 Subject: [PATCH 01/20] Add simple OAuth command --- package.json | 5 + src/core/secretsManager.ts | 26 ++ src/extension.ts | 14 +- src/oauth.ts | 492 +++++++++++++++++++++++++++++++++++++ 4 files changed, 536 insertions(+), 1 deletion(-) create mode 100644 src/oauth.ts diff --git a/package.json b/package.json index bd60a54c..d7f6f6e0 100644 --- a/package.json +++ b/package.json @@ -255,6 +255,11 @@ "title": "Search", "category": "Coder", "icon": "$(search)" + }, + { + "command": "coder.oauth.testAuth", + "title": "Test OAuth Auth", + "category": "Coder" } ], "menus": { diff --git a/src/core/secretsManager.ts b/src/core/secretsManager.ts index 94827b15..94982040 100644 --- a/src/core/secretsManager.ts +++ b/src/core/secretsManager.ts @@ -4,6 +4,8 @@ const SESSION_TOKEN_KEY = "sessionToken"; const LOGIN_STATE_KEY = "loginState"; +const OAUTH_CLIENT_REGISTRATION_KEY = "oauthClientRegistration"; + export enum AuthAction { LOGIN, LOGOUT, @@ -70,4 +72,28 @@ export class SecretsManager { } }); } + + /** + * Store OAuth client registration data. + */ + public async setOAuthClientRegistration( + registration: string | undefined, + ): Promise { + if (registration) { + await this.secrets.store(OAUTH_CLIENT_REGISTRATION_KEY, registration); + } else { + await this.secrets.delete(OAUTH_CLIENT_REGISTRATION_KEY); + } + } + + /** + * Get OAuth client registration data. + */ + public async getOAuthClientRegistration(): Promise { + try { + return await this.secrets.get(OAUTH_CLIENT_REGISTRATION_KEY); + } catch { + return undefined; + } + } } diff --git a/src/extension.ts b/src/extension.ts index 974cbe7d..411f6eae 100644 --- a/src/extension.ts +++ b/src/extension.ts @@ -13,6 +13,7 @@ import { Commands } from "./commands"; import { ServiceContainer } from "./core/container"; import { AuthAction } from "./core/secretsManager"; import { CertificateError, getErrorDetail } from "./error"; +import { activateCoderOAuth, CALLBACK_PATH } from "./oauth"; import { maybeAskUrl } from "./promptUtils"; import { Remote } from "./remote/remote"; import { getRemoteSshExtension } from "./remote/sshExtension"; @@ -116,11 +117,22 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { ctx.subscriptions, ); + const oauthHelper = activateCoderOAuth(client, secretsManager, output, ctx); + // Handle vscode:// URIs. const uriHandler = vscode.window.registerUriHandler({ handleUri: async (uri) => { - const cliManager = serviceContainer.getCliManager(); const params = new URLSearchParams(uri.query); + + if (uri.path === CALLBACK_PATH) { + const code = params.get("code"); + const state = params.get("state"); + const error = params.get("error"); + oauthHelper.handleCallback(code, state, error); + return; + } + + const cliManager = serviceContainer.getCliManager(); if (uri.path === "/open") { const owner = params.get("owner"); const workspace = params.get("workspace"); diff --git a/src/oauth.ts b/src/oauth.ts new file mode 100644 index 00000000..4575451d --- /dev/null +++ b/src/oauth.ts @@ -0,0 +1,492 @@ +import { createHash, randomBytes } from "node:crypto"; +import * as vscode from "vscode"; + +import { type CoderApi } from "./api/coderApi"; +import { type SecretsManager } from "./core/secretsManager"; + +import type { Logger } from "./logging/logger"; + +export const CALLBACK_PATH = "/oauth/callback"; + +interface ClientRegistrationRequest { + redirect_uris: string[]; + application_type: "native" | "web"; + grant_types: string[]; + response_types: string[]; + client_name: string; + token_endpoint_auth_method: + | "none" + | "client_secret_post" + | "client_secret_basic"; +} + +interface ClientRegistrationResponse { + client_id: string; + client_secret?: string; + client_id_issued_at?: number; + client_secret_expires_at?: number; + redirect_uris: string[]; + grant_types: string[]; +} + +interface OAuthServerMetadata { + issuer: string; + authorization_endpoint: string; + token_endpoint: string; + registration_endpoint?: string; + response_types_supported?: string[]; + grant_types_supported?: string[]; + code_challenge_methods_supported?: string[]; +} + +interface TokenResponse { + access_token: string; + token_type: string; + expires_in?: number; + refresh_token?: string; + scope?: string; +} + +/** + * Generate PKCE verifier and challenge (RFC 7636) + */ +function generatePKCE(): { verifier: string; challenge: string } { + const verifier = randomBytes(32).toString("base64url"); + const challenge = createHash("sha256").update(verifier).digest("base64url"); + return { verifier, challenge }; +} + +/** + * OAuth helper for Coder authentication + */ +class CoderOAuthHelper { + private _clientId: string | undefined; + private _clientRegistration: ClientRegistrationResponse | undefined; + private _cachedMetadata: OAuthServerMetadata | undefined; + private _pendingAuthResolve: + | ((value: { code: string; verifier: string }) => void) + | undefined; + private _pendingAuthReject: ((reason: Error) => void) | undefined; + private _expectedState: string | undefined; + private _pendingVerifier: string | undefined; + + private readonly extensionId: string; + + constructor( + private readonly client: CoderApi, + private readonly secretsManager: SecretsManager, + private readonly logger: Logger, + context: vscode.ExtensionContext, + ) { + this.loadClientRegistration(); + this.extensionId = context.extension.id; + } + + /** + * Discover OAuth server endpoints using RFC 8414 + * Caches result in memory for the session + * Throws error if server returns 404 (OAuth not supported) + */ + private async discoverOAuthEndpoints(): Promise { + if (this._cachedMetadata) { + return this._cachedMetadata; + } + + this.logger.info("Discovering OAuth endpoints..."); + + const response = await this.client + .getAxiosInstance() + .request({ + url: "/.well-known/oauth-authorization-server", + method: "GET", + headers: { + Accept: "application/json", + }, + }); + + const metadata = response.data; + + // Validate required fields + if ( + !metadata.authorization_endpoint || + !metadata.token_endpoint || + !metadata.issuer + ) { + throw new Error( + "OAuth server metadata missing required endpoints: " + + JSON.stringify(metadata), + ); + } + + this._cachedMetadata = metadata; + this.logger.info("OAuth endpoints discovered:", { + authorization: metadata.authorization_endpoint, + token: metadata.token_endpoint, + registration: metadata.registration_endpoint, + }); + + return metadata; + } + + /** + * Get redirect URI. + */ + private getRedirectUri(): string { + return `${vscode.env.uriScheme}://${this.extensionId}${CALLBACK_PATH}`; + } + + /** + * Load stored client registration from SecretsManager + */ + private async loadClientRegistration(): Promise { + try { + const stored = await this.secretsManager.getOAuthClientRegistration(); + if (stored) { + const registration = JSON.parse(stored) as ClientRegistrationResponse; + this._clientRegistration = registration; + this._clientId = registration.client_id; + this.logger.info("Loaded existing OAuth client:", this._clientId); + } + } catch (error) { + this.logger.error("Failed to load client registration:", error); + } + } + + /** + * Save client registration to SecretsManager + */ + private async saveClientRegistration( + registration: ClientRegistrationResponse, + ): Promise { + try { + await this.secretsManager.setOAuthClientRegistration( + JSON.stringify(registration), + ); + this._clientRegistration = registration; + this._clientId = registration.client_id; + this.logger.info("Saved OAuth client registration:", this._clientId); + } catch (error) { + this.logger.error("Failed to save client registration:", error); + } + } + + /** + * Clear stored client registration from SecretsManager + */ + async clearClientRegistration(): Promise { + await this.secretsManager.setOAuthClientRegistration(undefined); + this._clientRegistration = undefined; + this._clientId = undefined; + this.logger.info("Cleared OAuth client registration"); + } + + /** + * Register OAuth client dynamically (RFC 7591) + * Uses discovered registration endpoint from OAuth server metadata + */ + async registerClient(): Promise { + const redirectUri = this.getRedirectUri(); + + // Check if we need a new registration + if (this._clientId && this._clientRegistration) { + if (this._clientRegistration.redirect_uris.includes(redirectUri)) { + this.logger.info("Using existing client registration:", this._clientId); + return this._clientId; + } + this.logger.info("Redirect URI changed, re-registering client"); + } + + // Discover endpoints - will throw if 404 + const metadata = await this.discoverOAuthEndpoints(); + + if (!metadata.registration_endpoint) { + throw new Error( + "Server does not support dynamic client registration (no registration_endpoint in metadata)", + ); + } + + const registrationRequest: ClientRegistrationRequest = { + redirect_uris: [redirectUri], + application_type: "native", + grant_types: ["authorization_code"], + response_types: ["code"], + client_name: "VS Code Coder Extension", + token_endpoint_auth_method: "client_secret_post", + }; + + this.logger.info( + "Registering OAuth client at:", + metadata.registration_endpoint, + ); + + const response = await this.client + .getAxiosInstance() + .request({ + url: metadata.registration_endpoint, + method: "POST", + headers: { + "Content-Type": "application/json", + Accept: "application/json", + }, + data: registrationRequest, + }); + + await this.saveClientRegistration(response.data); + return response.data.client_id; + } + + /** + * Generate OAuth authorization URL with PKCE + * Uses discovered authorization endpoint + */ + private generateAuthUrl( + metadata: OAuthServerMetadata, + clientId: string, + state: string, + challenge: string, + scope = "all", + ): string { + const params = new URLSearchParams({ + client_id: clientId, + response_type: "code", + redirect_uri: this.getRedirectUri(), + scope: scope, + state: state, + code_challenge: challenge, + code_challenge_method: "S256", + }); + + const url = `${metadata.authorization_endpoint}?${params.toString()}`; + + this.logger.info("OAuth Authorization URL:", url); + this.logger.info("Client ID:", clientId); + this.logger.info("Redirect URI:", this.getRedirectUri()); + this.logger.info("Scope:", scope); + + return url; + } + + /** + * Start OAuth authorization flow + * Returns a promise that resolves when the callback is received with the authorization code + */ + async startAuthFlow( + scope = "all", + ): Promise<{ code: string; verifier: string }> { + // Discover endpoints first - will throw if 404 + const metadata = await this.discoverOAuthEndpoints(); + + // Register client + const clientId = await this.registerClient(); + + // Generate PKCE and state (kept in closure) + const state = randomBytes(16).toString("base64url"); + const { verifier, challenge } = generatePKCE(); + + // Build auth URL with discovered endpoints + const authUrl = this.generateAuthUrl( + metadata, + clientId, + state, + challenge, + scope, + ); + + // Create promise that waits for callback + return new Promise<{ code: string; verifier: string }>( + (resolve, reject) => { + const timeout = setTimeout( + () => { + this._pendingAuthResolve = undefined; + this._pendingAuthReject = undefined; + this._expectedState = undefined; + this._pendingVerifier = undefined; + reject(new Error("OAuth flow timed out after 5 minutes")); + }, + 5 * 60 * 1000, + ); + + // Store resolvers, state, and verifier for callback handler + this._pendingAuthResolve = (result) => { + clearTimeout(timeout); + this._pendingAuthResolve = undefined; + this._pendingAuthReject = undefined; + this._expectedState = undefined; + this._pendingVerifier = undefined; + resolve(result); + }; + + this._pendingAuthReject = (error) => { + clearTimeout(timeout); + this._pendingAuthResolve = undefined; + this._pendingAuthReject = undefined; + this._expectedState = undefined; + this._pendingVerifier = undefined; + reject(error); + }; + + this._expectedState = state; + this._pendingVerifier = verifier; + + vscode.env.openExternal(vscode.Uri.parse(authUrl)).then( + () => {}, + (error) => { + if (error instanceof Error) { + this._pendingAuthReject?.(error); + } else { + this._pendingAuthReject?.(new Error("Failed to open browser")); + } + }, + ); + }, + ); + } + + /** + * Handle OAuth callback from URI handler + * Called by extension.ts when vscode:// callback is received + */ + handleCallback( + code: string | null, + state: string | null, + error: string | null, + ): void { + if (!this._pendingAuthResolve || !this._pendingAuthReject) { + this.logger.warn("Received OAuth callback but no pending auth flow"); + return; + } + + if (error) { + this._pendingAuthReject(new Error(`OAuth error: ${error}`)); + return; + } + + if (!code) { + this._pendingAuthReject(new Error("No authorization code received")); + return; + } + + if (!state) { + this._pendingAuthReject(new Error("No state received")); + return; + } + + // Get verifier from pending flow + const verifier = this._pendingVerifier; + if (!verifier) { + this._pendingAuthReject(new Error("No PKCE verifier found")); + return; + } + + this._pendingAuthResolve({ code, verifier }); + } + + /** + * Exchange authorization code for access token + * Uses discovered token endpoint and PKCE verifier + */ + async exchangeCodeForToken( + code: string, + verifier: string, + ): Promise { + // Discover endpoints - will throw if 404 + const metadata = await this.discoverOAuthEndpoints(); + + if (!this._clientRegistration) { + throw new Error("No client registration found"); + } + + this.logger.info("Exchanging authorization code for token"); + + const tokenRequest = new URLSearchParams({ + grant_type: "authorization_code", + code: code, + redirect_uri: this.getRedirectUri(), + client_id: this._clientRegistration.client_id, + code_verifier: verifier, + }); + + // Add client secret if present + if (this._clientRegistration.client_secret) { + tokenRequest.append( + "client_secret", + this._clientRegistration.client_secret, + ); + } + + const response = await this.client + .getAxiosInstance() + .request({ + url: metadata.token_endpoint, + method: "POST", + headers: { + "Content-Type": "application/x-www-form-urlencoded", + Accept: "application/json", + }, + data: tokenRequest.toString(), + }); + + this.logger.info("Token exchange successful"); + return response.data; + } + + getClientId(): string | undefined { + return this._clientId; + } +} + +/** + * Activate OAuth functionality + * Returns the OAuth helper instance for use by URI handler + */ +export function activateCoderOAuth( + client: CoderApi, + secretsManager: SecretsManager, + logger: Logger, + context: vscode.ExtensionContext, +): CoderOAuthHelper { + const oauthHelper = new CoderOAuthHelper( + client, + secretsManager, + logger, + context, + ); + + // Register command to test OAuth flow + context.subscriptions.push( + vscode.commands.registerCommand("coder.oauth.testAuth", async () => { + try { + // Start OAuth flow and wait for callback + const { code, verifier } = await oauthHelper.startAuthFlow(); + logger.info( + "Authorization code received:", + code.substring(0, 8) + "...", + ); + + // Exchange code for token + const tokenResponse = await oauthHelper.exchangeCodeForToken( + code, + verifier, + ); + + vscode.window.showInformationMessage( + `OAuth flow completed! Access token received (expires in ${tokenResponse.expires_in}s)`, + ); + logger.info("OAuth flow completed:", { + token_type: tokenResponse.token_type, + expires_in: tokenResponse.expires_in, + scope: tokenResponse.scope, + }); + + client.setSessionToken(tokenResponse.access_token); + const response = await client.getWorkspaces({ q: "owner:me" }); + logger.info(response.workspaces.map((w) => w.name).toString()); + } catch (error) { + vscode.window.showErrorMessage(`OAuth flow failed: ${error}`); + logger.error("OAuth flow failed:", error); + } + }), + ); + + return oauthHelper; +} From 563929949c449c74a9dffa78b5fe346eab921d4e Mon Sep 17 00:00:00 2001 From: Ehab Younes Date: Tue, 21 Oct 2025 23:37:53 +0300 Subject: [PATCH 02/20] Refactoring into seperate files --- src/core/secretsManager.ts | 29 ++- src/extension.ts | 10 +- src/oauth.ts | 492 ------------------------------------- src/oauth/oauthHelper.ts | 465 +++++++++++++++++++++++++++++++++++ src/oauth/types.ts | 58 +++++ src/oauth/utils.ts | 28 +++ 6 files changed, 580 insertions(+), 502 deletions(-) delete mode 100644 src/oauth.ts create mode 100644 src/oauth/oauthHelper.ts create mode 100644 src/oauth/types.ts create mode 100644 src/oauth/utils.ts diff --git a/src/core/secretsManager.ts b/src/core/secretsManager.ts index 94982040..4821b962 100644 --- a/src/core/secretsManager.ts +++ b/src/core/secretsManager.ts @@ -1,3 +1,5 @@ +import { type ClientRegistrationResponse } from "../oauth/types"; + import type { SecretStorage, Disposable } from "vscode"; const SESSION_TOKEN_KEY = "sessionToken"; @@ -19,10 +21,10 @@ export class SecretsManager { * Set or unset the last used token. */ public async setSessionToken(sessionToken?: string): Promise { - if (!sessionToken) { - await this.secrets.delete(SESSION_TOKEN_KEY); - } else { + if (sessionToken) { await this.secrets.store(SESSION_TOKEN_KEY, sessionToken); + } else { + await this.secrets.delete(SESSION_TOKEN_KEY); } } @@ -77,10 +79,13 @@ export class SecretsManager { * Store OAuth client registration data. */ public async setOAuthClientRegistration( - registration: string | undefined, + registration: ClientRegistrationResponse | undefined, ): Promise { if (registration) { - await this.secrets.store(OAUTH_CLIENT_REGISTRATION_KEY, registration); + await this.secrets.store( + OAUTH_CLIENT_REGISTRATION_KEY, + JSON.stringify(registration), + ); } else { await this.secrets.delete(OAUTH_CLIENT_REGISTRATION_KEY); } @@ -89,11 +94,19 @@ export class SecretsManager { /** * Get OAuth client registration data. */ - public async getOAuthClientRegistration(): Promise { + public async getOAuthClientRegistration(): Promise< + ClientRegistrationResponse | undefined + > { try { - return await this.secrets.get(OAUTH_CLIENT_REGISTRATION_KEY); + const stringifiedResponse = await this.secrets.get( + OAUTH_CLIENT_REGISTRATION_KEY, + ); + if (stringifiedResponse) { + return JSON.parse(stringifiedResponse) as ClientRegistrationResponse; + } } catch { - return undefined; + // Do nothing } + return undefined; } } diff --git a/src/extension.ts b/src/extension.ts index 411f6eae..f350a6e5 100644 --- a/src/extension.ts +++ b/src/extension.ts @@ -13,7 +13,8 @@ import { Commands } from "./commands"; import { ServiceContainer } from "./core/container"; import { AuthAction } from "./core/secretsManager"; import { CertificateError, getErrorDetail } from "./error"; -import { activateCoderOAuth, CALLBACK_PATH } from "./oauth"; +import { activateCoderOAuth } from "./oauth/oauthHelper"; +import { CALLBACK_PATH } from "./oauth/utils"; import { maybeAskUrl } from "./promptUtils"; import { Remote } from "./remote/remote"; import { getRemoteSshExtension } from "./remote/sshExtension"; @@ -117,7 +118,12 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { ctx.subscriptions, ); - const oauthHelper = activateCoderOAuth(client, secretsManager, output, ctx); + const oauthHelper = await activateCoderOAuth( + client, + secretsManager, + output, + ctx, + ); // Handle vscode:// URIs. const uriHandler = vscode.window.registerUriHandler({ diff --git a/src/oauth.ts b/src/oauth.ts deleted file mode 100644 index 4575451d..00000000 --- a/src/oauth.ts +++ /dev/null @@ -1,492 +0,0 @@ -import { createHash, randomBytes } from "node:crypto"; -import * as vscode from "vscode"; - -import { type CoderApi } from "./api/coderApi"; -import { type SecretsManager } from "./core/secretsManager"; - -import type { Logger } from "./logging/logger"; - -export const CALLBACK_PATH = "/oauth/callback"; - -interface ClientRegistrationRequest { - redirect_uris: string[]; - application_type: "native" | "web"; - grant_types: string[]; - response_types: string[]; - client_name: string; - token_endpoint_auth_method: - | "none" - | "client_secret_post" - | "client_secret_basic"; -} - -interface ClientRegistrationResponse { - client_id: string; - client_secret?: string; - client_id_issued_at?: number; - client_secret_expires_at?: number; - redirect_uris: string[]; - grant_types: string[]; -} - -interface OAuthServerMetadata { - issuer: string; - authorization_endpoint: string; - token_endpoint: string; - registration_endpoint?: string; - response_types_supported?: string[]; - grant_types_supported?: string[]; - code_challenge_methods_supported?: string[]; -} - -interface TokenResponse { - access_token: string; - token_type: string; - expires_in?: number; - refresh_token?: string; - scope?: string; -} - -/** - * Generate PKCE verifier and challenge (RFC 7636) - */ -function generatePKCE(): { verifier: string; challenge: string } { - const verifier = randomBytes(32).toString("base64url"); - const challenge = createHash("sha256").update(verifier).digest("base64url"); - return { verifier, challenge }; -} - -/** - * OAuth helper for Coder authentication - */ -class CoderOAuthHelper { - private _clientId: string | undefined; - private _clientRegistration: ClientRegistrationResponse | undefined; - private _cachedMetadata: OAuthServerMetadata | undefined; - private _pendingAuthResolve: - | ((value: { code: string; verifier: string }) => void) - | undefined; - private _pendingAuthReject: ((reason: Error) => void) | undefined; - private _expectedState: string | undefined; - private _pendingVerifier: string | undefined; - - private readonly extensionId: string; - - constructor( - private readonly client: CoderApi, - private readonly secretsManager: SecretsManager, - private readonly logger: Logger, - context: vscode.ExtensionContext, - ) { - this.loadClientRegistration(); - this.extensionId = context.extension.id; - } - - /** - * Discover OAuth server endpoints using RFC 8414 - * Caches result in memory for the session - * Throws error if server returns 404 (OAuth not supported) - */ - private async discoverOAuthEndpoints(): Promise { - if (this._cachedMetadata) { - return this._cachedMetadata; - } - - this.logger.info("Discovering OAuth endpoints..."); - - const response = await this.client - .getAxiosInstance() - .request({ - url: "/.well-known/oauth-authorization-server", - method: "GET", - headers: { - Accept: "application/json", - }, - }); - - const metadata = response.data; - - // Validate required fields - if ( - !metadata.authorization_endpoint || - !metadata.token_endpoint || - !metadata.issuer - ) { - throw new Error( - "OAuth server metadata missing required endpoints: " + - JSON.stringify(metadata), - ); - } - - this._cachedMetadata = metadata; - this.logger.info("OAuth endpoints discovered:", { - authorization: metadata.authorization_endpoint, - token: metadata.token_endpoint, - registration: metadata.registration_endpoint, - }); - - return metadata; - } - - /** - * Get redirect URI. - */ - private getRedirectUri(): string { - return `${vscode.env.uriScheme}://${this.extensionId}${CALLBACK_PATH}`; - } - - /** - * Load stored client registration from SecretsManager - */ - private async loadClientRegistration(): Promise { - try { - const stored = await this.secretsManager.getOAuthClientRegistration(); - if (stored) { - const registration = JSON.parse(stored) as ClientRegistrationResponse; - this._clientRegistration = registration; - this._clientId = registration.client_id; - this.logger.info("Loaded existing OAuth client:", this._clientId); - } - } catch (error) { - this.logger.error("Failed to load client registration:", error); - } - } - - /** - * Save client registration to SecretsManager - */ - private async saveClientRegistration( - registration: ClientRegistrationResponse, - ): Promise { - try { - await this.secretsManager.setOAuthClientRegistration( - JSON.stringify(registration), - ); - this._clientRegistration = registration; - this._clientId = registration.client_id; - this.logger.info("Saved OAuth client registration:", this._clientId); - } catch (error) { - this.logger.error("Failed to save client registration:", error); - } - } - - /** - * Clear stored client registration from SecretsManager - */ - async clearClientRegistration(): Promise { - await this.secretsManager.setOAuthClientRegistration(undefined); - this._clientRegistration = undefined; - this._clientId = undefined; - this.logger.info("Cleared OAuth client registration"); - } - - /** - * Register OAuth client dynamically (RFC 7591) - * Uses discovered registration endpoint from OAuth server metadata - */ - async registerClient(): Promise { - const redirectUri = this.getRedirectUri(); - - // Check if we need a new registration - if (this._clientId && this._clientRegistration) { - if (this._clientRegistration.redirect_uris.includes(redirectUri)) { - this.logger.info("Using existing client registration:", this._clientId); - return this._clientId; - } - this.logger.info("Redirect URI changed, re-registering client"); - } - - // Discover endpoints - will throw if 404 - const metadata = await this.discoverOAuthEndpoints(); - - if (!metadata.registration_endpoint) { - throw new Error( - "Server does not support dynamic client registration (no registration_endpoint in metadata)", - ); - } - - const registrationRequest: ClientRegistrationRequest = { - redirect_uris: [redirectUri], - application_type: "native", - grant_types: ["authorization_code"], - response_types: ["code"], - client_name: "VS Code Coder Extension", - token_endpoint_auth_method: "client_secret_post", - }; - - this.logger.info( - "Registering OAuth client at:", - metadata.registration_endpoint, - ); - - const response = await this.client - .getAxiosInstance() - .request({ - url: metadata.registration_endpoint, - method: "POST", - headers: { - "Content-Type": "application/json", - Accept: "application/json", - }, - data: registrationRequest, - }); - - await this.saveClientRegistration(response.data); - return response.data.client_id; - } - - /** - * Generate OAuth authorization URL with PKCE - * Uses discovered authorization endpoint - */ - private generateAuthUrl( - metadata: OAuthServerMetadata, - clientId: string, - state: string, - challenge: string, - scope = "all", - ): string { - const params = new URLSearchParams({ - client_id: clientId, - response_type: "code", - redirect_uri: this.getRedirectUri(), - scope: scope, - state: state, - code_challenge: challenge, - code_challenge_method: "S256", - }); - - const url = `${metadata.authorization_endpoint}?${params.toString()}`; - - this.logger.info("OAuth Authorization URL:", url); - this.logger.info("Client ID:", clientId); - this.logger.info("Redirect URI:", this.getRedirectUri()); - this.logger.info("Scope:", scope); - - return url; - } - - /** - * Start OAuth authorization flow - * Returns a promise that resolves when the callback is received with the authorization code - */ - async startAuthFlow( - scope = "all", - ): Promise<{ code: string; verifier: string }> { - // Discover endpoints first - will throw if 404 - const metadata = await this.discoverOAuthEndpoints(); - - // Register client - const clientId = await this.registerClient(); - - // Generate PKCE and state (kept in closure) - const state = randomBytes(16).toString("base64url"); - const { verifier, challenge } = generatePKCE(); - - // Build auth URL with discovered endpoints - const authUrl = this.generateAuthUrl( - metadata, - clientId, - state, - challenge, - scope, - ); - - // Create promise that waits for callback - return new Promise<{ code: string; verifier: string }>( - (resolve, reject) => { - const timeout = setTimeout( - () => { - this._pendingAuthResolve = undefined; - this._pendingAuthReject = undefined; - this._expectedState = undefined; - this._pendingVerifier = undefined; - reject(new Error("OAuth flow timed out after 5 minutes")); - }, - 5 * 60 * 1000, - ); - - // Store resolvers, state, and verifier for callback handler - this._pendingAuthResolve = (result) => { - clearTimeout(timeout); - this._pendingAuthResolve = undefined; - this._pendingAuthReject = undefined; - this._expectedState = undefined; - this._pendingVerifier = undefined; - resolve(result); - }; - - this._pendingAuthReject = (error) => { - clearTimeout(timeout); - this._pendingAuthResolve = undefined; - this._pendingAuthReject = undefined; - this._expectedState = undefined; - this._pendingVerifier = undefined; - reject(error); - }; - - this._expectedState = state; - this._pendingVerifier = verifier; - - vscode.env.openExternal(vscode.Uri.parse(authUrl)).then( - () => {}, - (error) => { - if (error instanceof Error) { - this._pendingAuthReject?.(error); - } else { - this._pendingAuthReject?.(new Error("Failed to open browser")); - } - }, - ); - }, - ); - } - - /** - * Handle OAuth callback from URI handler - * Called by extension.ts when vscode:// callback is received - */ - handleCallback( - code: string | null, - state: string | null, - error: string | null, - ): void { - if (!this._pendingAuthResolve || !this._pendingAuthReject) { - this.logger.warn("Received OAuth callback but no pending auth flow"); - return; - } - - if (error) { - this._pendingAuthReject(new Error(`OAuth error: ${error}`)); - return; - } - - if (!code) { - this._pendingAuthReject(new Error("No authorization code received")); - return; - } - - if (!state) { - this._pendingAuthReject(new Error("No state received")); - return; - } - - // Get verifier from pending flow - const verifier = this._pendingVerifier; - if (!verifier) { - this._pendingAuthReject(new Error("No PKCE verifier found")); - return; - } - - this._pendingAuthResolve({ code, verifier }); - } - - /** - * Exchange authorization code for access token - * Uses discovered token endpoint and PKCE verifier - */ - async exchangeCodeForToken( - code: string, - verifier: string, - ): Promise { - // Discover endpoints - will throw if 404 - const metadata = await this.discoverOAuthEndpoints(); - - if (!this._clientRegistration) { - throw new Error("No client registration found"); - } - - this.logger.info("Exchanging authorization code for token"); - - const tokenRequest = new URLSearchParams({ - grant_type: "authorization_code", - code: code, - redirect_uri: this.getRedirectUri(), - client_id: this._clientRegistration.client_id, - code_verifier: verifier, - }); - - // Add client secret if present - if (this._clientRegistration.client_secret) { - tokenRequest.append( - "client_secret", - this._clientRegistration.client_secret, - ); - } - - const response = await this.client - .getAxiosInstance() - .request({ - url: metadata.token_endpoint, - method: "POST", - headers: { - "Content-Type": "application/x-www-form-urlencoded", - Accept: "application/json", - }, - data: tokenRequest.toString(), - }); - - this.logger.info("Token exchange successful"); - return response.data; - } - - getClientId(): string | undefined { - return this._clientId; - } -} - -/** - * Activate OAuth functionality - * Returns the OAuth helper instance for use by URI handler - */ -export function activateCoderOAuth( - client: CoderApi, - secretsManager: SecretsManager, - logger: Logger, - context: vscode.ExtensionContext, -): CoderOAuthHelper { - const oauthHelper = new CoderOAuthHelper( - client, - secretsManager, - logger, - context, - ); - - // Register command to test OAuth flow - context.subscriptions.push( - vscode.commands.registerCommand("coder.oauth.testAuth", async () => { - try { - // Start OAuth flow and wait for callback - const { code, verifier } = await oauthHelper.startAuthFlow(); - logger.info( - "Authorization code received:", - code.substring(0, 8) + "...", - ); - - // Exchange code for token - const tokenResponse = await oauthHelper.exchangeCodeForToken( - code, - verifier, - ); - - vscode.window.showInformationMessage( - `OAuth flow completed! Access token received (expires in ${tokenResponse.expires_in}s)`, - ); - logger.info("OAuth flow completed:", { - token_type: tokenResponse.token_type, - expires_in: tokenResponse.expires_in, - scope: tokenResponse.scope, - }); - - client.setSessionToken(tokenResponse.access_token); - const response = await client.getWorkspaces({ q: "owner:me" }); - logger.info(response.workspaces.map((w) => w.name).toString()); - } catch (error) { - vscode.window.showErrorMessage(`OAuth flow failed: ${error}`); - logger.error("OAuth flow failed:", error); - } - }), - ); - - return oauthHelper; -} diff --git a/src/oauth/oauthHelper.ts b/src/oauth/oauthHelper.ts new file mode 100644 index 00000000..54c87c15 --- /dev/null +++ b/src/oauth/oauthHelper.ts @@ -0,0 +1,465 @@ +import * as vscode from "vscode"; + +import { type CoderApi } from "../api/coderApi"; +import { type SecretsManager } from "../core/secretsManager"; + +import { CALLBACK_PATH, generatePKCE, generateState } from "./utils"; + +import type { Logger } from "../logging/logger"; + +import type { + AuthorizationRequestParams, + ClientRegistrationRequest, + ClientRegistrationResponse, + OAuthServerMetadata, + TokenRequestParams, + TokenResponse, +} from "./types"; + +const OAUTH_GRANT_TYPE = "authorization_code" as const; +const OAUTH_RESPONSE_TYPE = "code" as const; +const OAUTH_AUTH_METHOD = "client_secret_post" as const; +const PKCE_CHALLENGE_METHOD = "S256" as const; +const CLIENT_NAME = "VS Code Coder Extension"; + +export class CoderOAuthHelper { + private clientRegistration: ClientRegistrationResponse | undefined; + private cachedMetadata: OAuthServerMetadata | undefined; + private pendingAuthResolve: + | ((value: { code: string; verifier: string }) => void) + | undefined; + private pendingAuthReject: ((reason: Error) => void) | undefined; + private expectedState: string | undefined; + private pendingVerifier: string | undefined; + + private readonly extensionId: string; + + static async create( + client: CoderApi, + secretsManager: SecretsManager, + logger: Logger, + context: vscode.ExtensionContext, + ): Promise { + const helper = new CoderOAuthHelper( + client, + secretsManager, + logger, + context, + ); + await helper.loadRegistration(); + return helper; + } + private constructor( + private readonly client: CoderApi, + private readonly secretsManager: SecretsManager, + private readonly logger: Logger, + context: vscode.ExtensionContext, + ) { + this.extensionId = context.extension.id; + } + + private async getMetadata(): Promise { + if (this.cachedMetadata) { + return this.cachedMetadata; + } + + this.logger.info("Discovering OAuth endpoints..."); + + const response = await this.client + .getAxiosInstance() + .request({ + url: "/.well-known/oauth-authorization-server", + method: "GET", + headers: { + Accept: "application/json", + }, + }); + + const metadata = response.data; + + if ( + !metadata.authorization_endpoint || + !metadata.token_endpoint || + !metadata.issuer + ) { + throw new Error( + "OAuth server metadata missing required endpoints: " + + JSON.stringify(metadata), + ); + } + + if ( + !includesAllTypes(metadata.grant_types_supported, [ + OAUTH_GRANT_TYPE, + "refresh_token", + ]) + ) { + throw new Error( + `Server does not support required grant types: authorization_code, refresh_token. Supported: ${metadata.grant_types_supported?.join(", ") || "none"}`, + ); + } + + if ( + !includesAllTypes(metadata.response_types_supported, [ + OAUTH_RESPONSE_TYPE, + ]) + ) { + throw new Error( + `Server does not support required response type: code. Supported: ${metadata.response_types_supported?.join(", ") || "none"}`, + ); + } + + if ( + !includesAllTypes(metadata.token_endpoint_auth_methods_supported, [ + OAUTH_AUTH_METHOD, + ]) + ) { + throw new Error( + `Server does not support required auth method: client_secret_post. Supported: ${metadata.token_endpoint_auth_methods_supported?.join(", ") || "none"}`, + ); + } + + if ( + !includesAllTypes(metadata.code_challenge_methods_supported, [ + PKCE_CHALLENGE_METHOD, + ]) + ) { + throw new Error( + `Server does not support required PKCE method: S256. Supported: ${metadata.code_challenge_methods_supported?.join(", ") || "none"}`, + ); + } + + this.cachedMetadata = metadata; + this.logger.info("OAuth endpoints discovered:", { + authorization: metadata.authorization_endpoint, + token: metadata.token_endpoint, + registration: metadata.registration_endpoint, + }); + + return metadata; + } + + private getRedirectUri(): string { + return `${vscode.env.uriScheme}://${this.extensionId}${CALLBACK_PATH}`; + } + + private async loadRegistration(): Promise { + const registration = await this.secretsManager.getOAuthClientRegistration(); + if (registration) { + this.clientRegistration = registration; + this.logger.info("Loaded existing OAuth client:", registration.client_id); + } + } + + private async saveRegistration( + registration: ClientRegistrationResponse, + ): Promise { + await this.secretsManager.setOAuthClientRegistration(registration); + this.clientRegistration = registration; + this.logger.info( + "Saved OAuth client registration:", + registration.client_id, + ); + } + + async clearClientRegistration(): Promise { + await this.secretsManager.setOAuthClientRegistration(undefined); + this.clientRegistration = undefined; + this.logger.info("Cleared OAuth client registration"); + } + + async registerClient(): Promise { + const redirectUri = this.getRedirectUri(); + + if (this.clientRegistration?.client_id) { + const clientId = this.clientRegistration.client_id; + if (this.clientRegistration.redirect_uris.includes(redirectUri)) { + this.logger.info("Using existing client registration:", clientId); + return clientId; + } + this.logger.info("Redirect URI changed, re-registering client"); + } + + const metadata = await this.getMetadata(); + + if (!metadata.registration_endpoint) { + throw new Error( + "Server does not support dynamic client registration (no registration_endpoint in metadata)", + ); + } + + const registrationRequest: ClientRegistrationRequest = { + redirect_uris: [redirectUri], + application_type: "native", + grant_types: [OAUTH_GRANT_TYPE], + response_types: [OAUTH_RESPONSE_TYPE], + client_name: CLIENT_NAME, + token_endpoint_auth_method: OAUTH_AUTH_METHOD, + }; + + const response = await this.client + .getAxiosInstance() + .request({ + url: metadata.registration_endpoint, + method: "POST", + headers: { + "Content-Type": "application/json", + Accept: "application/json", + }, + data: registrationRequest, + }); + + await this.saveRegistration(response.data); + return response.data.client_id; + } + + private buildAuthorizationUrl( + metadata: OAuthServerMetadata, + clientId: string, + state: string, + challenge: string, + scope = "all", + ): string { + if ( + metadata.scopes_supported && + !metadata.scopes_supported.includes(scope) + ) { + this.logger.warn( + `Requested scope "${scope}" not in server's supported scopes. Server may still accept it.`, + { supported_scopes: metadata.scopes_supported }, + ); + } + + const params: AuthorizationRequestParams = { + client_id: clientId, + response_type: OAUTH_RESPONSE_TYPE, + redirect_uri: this.getRedirectUri(), + scope, + state, + code_challenge: challenge, + code_challenge_method: PKCE_CHALLENGE_METHOD, + }; + + const url = `${metadata.authorization_endpoint}?${new URLSearchParams(params as unknown as Record).toString()}`; + + this.logger.info("OAuth Authorization URL:", url); + this.logger.info("Client ID:", clientId); + this.logger.info("Redirect URI:", this.getRedirectUri()); + this.logger.info("Scope:", scope); + + return url; + } + + async startAuthorization( + scope = "all", + ): Promise<{ code: string; verifier: string }> { + const metadata = await this.getMetadata(); + const clientId = await this.registerClient(); + const state = generateState(); + const { verifier, challenge } = generatePKCE(); + + const authUrl = this.buildAuthorizationUrl( + metadata, + clientId, + state, + challenge, + scope, + ); + + return new Promise<{ code: string; verifier: string }>( + (resolve, reject) => { + const timeout = setTimeout( + () => { + this.clearPendingAuth(); + reject(new Error("OAuth flow timed out after 5 minutes")); + }, + 5 * 60 * 1000, + ); + + const clearPromise = () => { + clearTimeout(timeout); + this.clearPendingAuth(); + }; + + this.pendingAuthResolve = (result) => { + clearPromise(); + resolve(result); + }; + + this.pendingAuthReject = (error) => { + clearPromise(); + reject(error); + }; + + this.expectedState = state; + this.pendingVerifier = verifier; + + vscode.env.openExternal(vscode.Uri.parse(authUrl)).then( + () => {}, + (error) => { + if (error instanceof Error) { + this.pendingAuthReject?.(error); + } else { + this.pendingAuthReject?.(new Error("Failed to open browser")); + } + }, + ); + }, + ); + } + + private clearPendingAuth(): void { + this.pendingAuthResolve = undefined; + this.pendingAuthReject = undefined; + this.expectedState = undefined; + this.pendingVerifier = undefined; + } + + handleCallback( + code: string | null, + state: string | null, + error: string | null, + ): void { + if (!this.pendingAuthResolve || !this.pendingAuthReject) { + this.logger.warn("Received OAuth callback but no pending auth flow"); + return; + } + + if (error) { + this.pendingAuthReject(new Error(`OAuth error: ${error}`)); + return; + } + + if (!code) { + this.pendingAuthReject(new Error("No authorization code received")); + return; + } + + if (!state) { + this.pendingAuthReject(new Error("No state received")); + return; + } + + if (state !== this.expectedState) { + this.pendingAuthReject( + new Error("State mismatch - possible CSRF attack"), + ); + return; + } + + const verifier = this.pendingVerifier; + if (!verifier) { + this.pendingAuthReject(new Error("No PKCE verifier found")); + return; + } + + this.pendingAuthResolve({ code, verifier }); + } + + async exchangeToken(code: string, verifier: string): Promise { + const metadata = await this.getMetadata(); + + if (!this.clientRegistration) { + throw new Error("No client registration found"); + } + + this.logger.info("Exchanging authorization code for token"); + + const params: TokenRequestParams = { + grant_type: OAUTH_GRANT_TYPE, + code, + redirect_uri: this.getRedirectUri(), + client_id: this.clientRegistration.client_id, + code_verifier: verifier, + }; + + if (this.clientRegistration.client_secret) { + params.client_secret = this.clientRegistration.client_secret; + } + + const tokenRequest = new URLSearchParams( + params as unknown as Record, + ); + + const response = await this.client + .getAxiosInstance() + .request({ + url: metadata.token_endpoint, + method: "POST", + headers: { + "Content-Type": "application/x-www-form-urlencoded", + Accept: "application/json", + }, + data: tokenRequest.toString(), + }); + + this.logger.info("Token exchange successful"); + return response.data; + } + + getClientId(): string | undefined { + return this.clientRegistration?.client_id; + } +} + +function includesAllTypes( + arr: string[] | undefined, + requiredTypes: readonly string[], +): boolean { + if (arr === undefined) { + // Supported types are not sent by the server so just assume everything is allowed + return true; + } + + return requiredTypes.every((type) => arr.includes(type)); +} + +/** + * Activates OAuth support for the Coder extension. + * Initializes the OAuth helper and registers the test auth command. + */ +export async function activateCoderOAuth( + client: CoderApi, + secretsManager: SecretsManager, + logger: Logger, + context: vscode.ExtensionContext, +): Promise { + const oauthHelper = await CoderOAuthHelper.create( + client, + secretsManager, + logger, + context, + ); + + context.subscriptions.push( + vscode.commands.registerCommand("coder.oauth.testAuth", async () => { + try { + const { code, verifier } = await oauthHelper.startAuthorization(); + logger.info( + "Authorization code received:", + code.substring(0, 8) + "...", + ); + + const tokenResponse = await oauthHelper.exchangeToken(code, verifier); + + vscode.window.showInformationMessage( + `OAuth flow completed! Access token received (expires in ${tokenResponse.expires_in}s)`, + ); + logger.info("OAuth flow completed:", { + token_type: tokenResponse.token_type, + expires_in: tokenResponse.expires_in, + scope: tokenResponse.scope, + }); + + client.setSessionToken(tokenResponse.access_token); + const response = await client.getWorkspaces({ q: "owner:me" }); + logger.info(response.workspaces.map((w) => w.name).toString()); + } catch (error) { + vscode.window.showErrorMessage(`OAuth flow failed: ${error}`); + logger.error("OAuth flow failed:", error); + } + }), + ); + + return oauthHelper; +} diff --git a/src/oauth/types.ts b/src/oauth/types.ts new file mode 100644 index 00000000..493e3707 --- /dev/null +++ b/src/oauth/types.ts @@ -0,0 +1,58 @@ +export interface ClientRegistrationRequest { + redirect_uris: string[]; + token_endpoint_auth_method: "client_secret_post"; + application_type: "native" | "web"; + grant_types: string[]; + response_types: string[]; + client_name?: string; + client_uri?: string; + scope?: string[]; +} + +export interface ClientRegistrationResponse { + client_id: string; + client_secret?: string; + client_id_issued_at?: number; + client_secret_expires_at?: number; + redirect_uris: string[]; + grant_types: string[]; +} + +export interface OAuthServerMetadata { + issuer: string; + authorization_endpoint: string; + token_endpoint: string; + registration_endpoint?: string; + response_types_supported?: string[]; + grant_types_supported?: string[]; + code_challenge_methods_supported?: string[]; + scopes_supported?: string[]; + token_endpoint_auth_methods_supported?: string[]; +} + +export interface TokenResponse { + access_token: string; + token_type: string; + expires_in?: number; + refresh_token?: string; + scope?: string; +} + +export interface AuthorizationRequestParams { + client_id: string; + response_type: "code"; + redirect_uri: string; + scope: string; + state: string; + code_challenge: string; + code_challenge_method: "S256"; +} + +export interface TokenRequestParams { + grant_type: "authorization_code"; + code: string; + redirect_uri: string; + client_id: string; + code_verifier: string; + client_secret?: string; +} diff --git a/src/oauth/utils.ts b/src/oauth/utils.ts new file mode 100644 index 00000000..7d66a139 --- /dev/null +++ b/src/oauth/utils.ts @@ -0,0 +1,28 @@ +import { createHash, randomBytes } from "node:crypto"; + +/** + * OAuth callback path for handling authorization responses (RFC 6749). + */ +export const CALLBACK_PATH = "/oauth/callback"; + +export interface PKCEChallenge { + verifier: string; + challenge: string; +} + +/** + * Generates a PKCE challenge pair (RFC 7636). + * Creates a code verifier and its SHA256 challenge for secure OAuth flows. + */ +export function generatePKCE(): PKCEChallenge { + const verifier = randomBytes(32).toString("base64url"); + const challenge = createHash("sha256").update(verifier).digest("base64url"); + return { verifier, challenge }; +} + +/** + * Generates a cryptographically secure state parameter to prevent CSRF attacks (RFC 6749). + */ +export function generateState(): string { + return randomBytes(16).toString("base64url"); +} From 9bceb8952e6be92b01b42accf8776a804c326a0b Mon Sep 17 00:00:00 2001 From: Ehab Younes Date: Wed, 22 Oct 2025 10:33:01 +0300 Subject: [PATCH 03/20] Update types --- src/oauth/types.ts | 133 ++++++++++++++++++++++++++++++++++++++++----- 1 file changed, 119 insertions(+), 14 deletions(-) diff --git a/src/oauth/types.ts b/src/oauth/types.ts index 493e3707..6ecaa0ff 100644 --- a/src/oauth/types.ts +++ b/src/oauth/types.ts @@ -1,53 +1,112 @@ +// OAuth 2.1 Grant Types +export type GrantType = + | "authorization_code" + | "refresh_token" + | "client_credentials"; + +// OAuth 2.1 Response Types +export type ResponseType = "code"; + +// Token Endpoint Authentication Methods +export type TokenEndpointAuthMethod = + | "client_secret_post" + | "client_secret_basic" + | "none"; + +// Application Types +export type ApplicationType = "native" | "web"; + +// PKCE Code Challenge Methods (OAuth 2.1 requires S256) +export type CodeChallengeMethod = "S256"; + +// Token Types +export type TokenType = "Bearer" | "DPoP"; + +// Client Registration Request (RFC 7591 + OAuth 2.1) export interface ClientRegistrationRequest { redirect_uris: string[]; - token_endpoint_auth_method: "client_secret_post"; - application_type: "native" | "web"; - grant_types: string[]; - response_types: string[]; + token_endpoint_auth_method: TokenEndpointAuthMethod; + application_type: ApplicationType; + grant_types: GrantType[]; + response_types: ResponseType[]; client_name?: string; client_uri?: string; - scope?: string[]; + logo_uri?: string; + scope?: string; + contacts?: string[]; + tos_uri?: string; + policy_uri?: string; + jwks_uri?: string; + software_id?: string; + software_version?: string; } +// Client Registration Response (RFC 7591) export interface ClientRegistrationResponse { client_id: string; client_secret?: string; client_id_issued_at?: number; client_secret_expires_at?: number; redirect_uris: string[]; - grant_types: string[]; + token_endpoint_auth_method: TokenEndpointAuthMethod; + application_type?: ApplicationType; + grant_types: GrantType[]; + response_types: ResponseType[]; + client_name?: string; + client_uri?: string; + logo_uri?: string; + scope?: string; + contacts?: string[]; + tos_uri?: string; + policy_uri?: string; + jwks_uri?: string; + software_id?: string; + software_version?: string; + registration_client_uri?: string; + registration_access_token?: string; } +// OAuth 2.1 Authorization Server Metadata (RFC 8414) export interface OAuthServerMetadata { issuer: string; authorization_endpoint: string; token_endpoint: string; registration_endpoint?: string; - response_types_supported?: string[]; - grant_types_supported?: string[]; - code_challenge_methods_supported?: string[]; + jwks_uri?: string; + response_types_supported: ResponseType[]; + grant_types_supported?: GrantType[]; + code_challenge_methods_supported: CodeChallengeMethod[]; scopes_supported?: string[]; - token_endpoint_auth_methods_supported?: string[]; + token_endpoint_auth_methods_supported?: TokenEndpointAuthMethod[]; + revocation_endpoint?: string; + revocation_endpoint_auth_methods_supported?: TokenEndpointAuthMethod[]; + introspection_endpoint?: string; + introspection_endpoint_auth_methods_supported?: TokenEndpointAuthMethod[]; + service_documentation?: string; + ui_locales_supported?: string[]; } +// Token Response (RFC 6749 Section 5.1) export interface TokenResponse { access_token: string; - token_type: string; + token_type: TokenType; expires_in?: number; refresh_token?: string; scope?: string; } +// Authorization Request Parameters (OAuth 2.1) export interface AuthorizationRequestParams { client_id: string; - response_type: "code"; + response_type: ResponseType; redirect_uri: string; - scope: string; + scope?: string; state: string; code_challenge: string; - code_challenge_method: "S256"; + code_challenge_method: CodeChallengeMethod; } +// Token Request Parameters - Authorization Code Grant (OAuth 2.1) export interface TokenRequestParams { grant_type: "authorization_code"; code: string; @@ -56,3 +115,49 @@ export interface TokenRequestParams { code_verifier: string; client_secret?: string; } + +// Token Request Parameters - Refresh Token Grant +export interface RefreshTokenRequestParams { + grant_type: "refresh_token"; + refresh_token: string; + client_id: string; + client_secret?: string; + scope?: string; +} + +// Token Request Parameters - Client Credentials Grant +export interface ClientCredentialsRequestParams { + grant_type: "client_credentials"; + client_id: string; + client_secret: string; + scope?: string; +} + +// Union type for all token request types +export type TokenRequestParamsUnion = + | TokenRequestParams + | RefreshTokenRequestParams + | ClientCredentialsRequestParams; + +// Token Revocation Request (RFC 7009) +export interface TokenRevocationRequest { + token: string; + token_type_hint?: "access_token" | "refresh_token"; + client_id: string; + client_secret?: string; +} + +// Error Response (RFC 6749 Section 5.2) +export interface OAuthErrorResponse { + error: + | "invalid_request" + | "invalid_client" + | "invalid_grant" + | "unauthorized_client" + | "unsupported_grant_type" + | "invalid_scope" + | "server_error" + | "temporarily_unavailable"; + error_description?: string; + error_uri?: string; +} From 01bd44e2c47cdfd16b9f12546931a9b150170c94 Mon Sep 17 00:00:00 2001 From: Ehab Younes Date: Wed, 22 Oct 2025 16:18:31 +0300 Subject: [PATCH 04/20] Smol refactorings --- src/oauth/oauthHelper.ts | 83 ++++++++++++++++------------------------ 1 file changed, 34 insertions(+), 49 deletions(-) diff --git a/src/oauth/oauthHelper.ts b/src/oauth/oauthHelper.ts index 54c87c15..bbd0a82c 100644 --- a/src/oauth/oauthHelper.ts +++ b/src/oauth/oauthHelper.ts @@ -16,12 +16,15 @@ import type { TokenResponse, } from "./types"; -const OAUTH_GRANT_TYPE = "authorization_code" as const; -const OAUTH_RESPONSE_TYPE = "code" as const; -const OAUTH_AUTH_METHOD = "client_secret_post" as const; +const AUTH_GRANT_TYPE = "authorization_code" as const; +const REFRESH_GRANT_TYPE = "refresh_token" as const; +const RESPONSE_TYPE = "code" as const; +const OAUTH_METHOD = "client_secret_post" as const; const PKCE_CHALLENGE_METHOD = "S256" as const; const CLIENT_NAME = "VS Code Coder Extension"; +const REQUIRED_GRANT_TYPES = [AUTH_GRANT_TYPE, REFRESH_GRANT_TYPE] as const; + export class CoderOAuthHelper { private clientRegistration: ClientRegistrationResponse | undefined; private cachedMetadata: OAuthServerMetadata | undefined; @@ -46,7 +49,7 @@ export class CoderOAuthHelper { logger, context, ); - await helper.loadRegistration(); + await helper.loadClientRegistration(); return helper; } private constructor( @@ -67,13 +70,7 @@ export class CoderOAuthHelper { const response = await this.client .getAxiosInstance() - .request({ - url: "/.well-known/oauth-authorization-server", - method: "GET", - headers: { - Accept: "application/json", - }, - }); + .get("/.well-known/oauth-authorization-server"); const metadata = response.data; @@ -89,33 +86,26 @@ export class CoderOAuthHelper { } if ( - !includesAllTypes(metadata.grant_types_supported, [ - OAUTH_GRANT_TYPE, - "refresh_token", - ]) + !includesAllTypes(metadata.grant_types_supported, REQUIRED_GRANT_TYPES) ) { throw new Error( - `Server does not support required grant types: authorization_code, refresh_token. Supported: ${metadata.grant_types_supported?.join(", ") || "none"}`, + `Server does not support required grant types: ${REQUIRED_GRANT_TYPES.join(", ")}. Supported: ${metadata.grant_types_supported?.join(", ") || "none"}`, ); } - if ( - !includesAllTypes(metadata.response_types_supported, [ - OAUTH_RESPONSE_TYPE, - ]) - ) { + if (!includesAllTypes(metadata.response_types_supported, [RESPONSE_TYPE])) { throw new Error( - `Server does not support required response type: code. Supported: ${metadata.response_types_supported?.join(", ") || "none"}`, + `Server does not support required response type: ${RESPONSE_TYPE}. Supported: ${metadata.response_types_supported?.join(", ") || "none"}`, ); } if ( !includesAllTypes(metadata.token_endpoint_auth_methods_supported, [ - OAUTH_AUTH_METHOD, + OAUTH_METHOD, ]) ) { throw new Error( - `Server does not support required auth method: client_secret_post. Supported: ${metadata.token_endpoint_auth_methods_supported?.join(", ") || "none"}`, + `Server does not support required auth method: ${OAUTH_METHOD}. Supported: ${metadata.token_endpoint_auth_methods_supported?.join(", ") || "none"}`, ); } @@ -125,7 +115,7 @@ export class CoderOAuthHelper { ]) ) { throw new Error( - `Server does not support required PKCE method: S256. Supported: ${metadata.code_challenge_methods_supported?.join(", ") || "none"}`, + `Server does not support required PKCE method: ${PKCE_CHALLENGE_METHOD}. Supported: ${metadata.code_challenge_methods_supported?.join(", ") || "none"}`, ); } @@ -143,7 +133,7 @@ export class CoderOAuthHelper { return `${vscode.env.uriScheme}://${this.extensionId}${CALLBACK_PATH}`; } - private async loadRegistration(): Promise { + private async loadClientRegistration(): Promise { const registration = await this.secretsManager.getOAuthClientRegistration(); if (registration) { this.clientRegistration = registration; @@ -151,7 +141,7 @@ export class CoderOAuthHelper { } } - private async saveRegistration( + private async saveClientRegistration( registration: ClientRegistrationResponse, ): Promise { await this.secretsManager.setOAuthClientRegistration(registration); @@ -191,25 +181,21 @@ export class CoderOAuthHelper { const registrationRequest: ClientRegistrationRequest = { redirect_uris: [redirectUri], application_type: "native", - grant_types: [OAUTH_GRANT_TYPE], - response_types: [OAUTH_RESPONSE_TYPE], + grant_types: [AUTH_GRANT_TYPE], + response_types: [RESPONSE_TYPE], client_name: CLIENT_NAME, - token_endpoint_auth_method: OAUTH_AUTH_METHOD, + token_endpoint_auth_method: OAUTH_METHOD, }; const response = await this.client .getAxiosInstance() - .request({ - url: metadata.registration_endpoint, - method: "POST", - headers: { - "Content-Type": "application/json", - Accept: "application/json", - }, - data: registrationRequest, - }); + .post( + metadata.registration_endpoint, + registrationRequest, + ); + + await this.saveClientRegistration(response.data); - await this.saveRegistration(response.data); return response.data.client_id; } @@ -232,7 +218,7 @@ export class CoderOAuthHelper { const params: AuthorizationRequestParams = { client_id: clientId, - response_type: OAUTH_RESPONSE_TYPE, + response_type: RESPONSE_TYPE, redirect_uri: this.getRedirectUri(), scope, state, @@ -268,12 +254,15 @@ export class CoderOAuthHelper { return new Promise<{ code: string; verifier: string }>( (resolve, reject) => { + const timeoutMins = 5; const timeout = setTimeout( () => { this.clearPendingAuth(); - reject(new Error("OAuth flow timed out after 5 minutes")); + reject( + new Error(`OAuth flow timed out after ${timeoutMins} minutes`), + ); }, - 5 * 60 * 1000, + timeoutMins * 60 * 1000, ); const clearPromise = () => { @@ -366,7 +355,7 @@ export class CoderOAuthHelper { this.logger.info("Exchanging authorization code for token"); const params: TokenRequestParams = { - grant_type: OAUTH_GRANT_TYPE, + grant_type: AUTH_GRANT_TYPE, code, redirect_uri: this.getRedirectUri(), client_id: this.clientRegistration.client_id, @@ -383,14 +372,10 @@ export class CoderOAuthHelper { const response = await this.client .getAxiosInstance() - .request({ - url: metadata.token_endpoint, - method: "POST", + .post(metadata.token_endpoint, tokenRequest.toString(), { headers: { "Content-Type": "application/x-www-form-urlencoded", - Accept: "application/json", }, - data: tokenRequest.toString(), }); this.logger.info("Token exchange successful"); From b7f232530e7a2d75c4352d23e47e9ffd7f8c12b2 Mon Sep 17 00:00:00 2001 From: Ehab Younes Date: Fri, 24 Oct 2025 15:16:31 +0300 Subject: [PATCH 05/20] Add token refresh and revocation --- package.json | 9 +- src/core/secretsManager.ts | 46 +++++- src/oauth/oauthHelper.ts | 300 ++++++++++++++++++++++++++++++++++--- 3 files changed, 331 insertions(+), 24 deletions(-) diff --git a/package.json b/package.json index d7f6f6e0..add68824 100644 --- a/package.json +++ b/package.json @@ -257,8 +257,13 @@ "icon": "$(search)" }, { - "command": "coder.oauth.testAuth", - "title": "Test OAuth Auth", + "command": "coder.oauth.login", + "title": "OAuth Login", + "category": "Coder" + }, + { + "command": "coder.oauth.logout", + "title": "OAuth Logout", "category": "Coder" } ], diff --git a/src/core/secretsManager.ts b/src/core/secretsManager.ts index 4821b962..b108af0a 100644 --- a/src/core/secretsManager.ts +++ b/src/core/secretsManager.ts @@ -1,4 +1,7 @@ -import { type ClientRegistrationResponse } from "../oauth/types"; +import { + type TokenResponse, + type ClientRegistrationResponse, +} from "../oauth/types"; import type { SecretStorage, Disposable } from "vscode"; @@ -8,6 +11,12 @@ const LOGIN_STATE_KEY = "loginState"; const OAUTH_CLIENT_REGISTRATION_KEY = "oauthClientRegistration"; +const OAUTH_TOKENS_KEY = "oauthTokens"; + +export type StoredOAuthTokens = Omit & { + expiry_timestamp: number; +}; + export enum AuthAction { LOGIN, LOGOUT, @@ -109,4 +118,39 @@ export class SecretsManager { } return undefined; } + + /** + * Store OAuth token data including expiry timestamp. + */ + public async setOAuthTokens( + tokens: StoredOAuthTokens | undefined, + ): Promise { + if (tokens) { + await this.secrets.store(OAUTH_TOKENS_KEY, JSON.stringify(tokens)); + } else { + await this.secrets.delete(OAUTH_TOKENS_KEY); + } + } + + /** + * Get stored OAuth token data. + */ + public async getOAuthTokens(): Promise { + try { + const stringifiedTokens = await this.secrets.get(OAUTH_TOKENS_KEY); + if (stringifiedTokens) { + return JSON.parse(stringifiedTokens) as StoredOAuthTokens; + } + } catch { + // Do nothing + } + return undefined; + } + + /** + * Clear OAuth token data. + */ + public async clearOAuthTokens(): Promise { + await this.secrets.delete(OAUTH_TOKENS_KEY); + } } diff --git a/src/oauth/oauthHelper.ts b/src/oauth/oauthHelper.ts index bbd0a82c..03154fc7 100644 --- a/src/oauth/oauthHelper.ts +++ b/src/oauth/oauthHelper.ts @@ -1,7 +1,10 @@ import * as vscode from "vscode"; import { type CoderApi } from "../api/coderApi"; -import { type SecretsManager } from "../core/secretsManager"; +import { + type StoredOAuthTokens, + type SecretsManager, +} from "../core/secretsManager"; import { CALLBACK_PATH, generatePKCE, generateState } from "./utils"; @@ -12,8 +15,10 @@ import type { ClientRegistrationRequest, ClientRegistrationResponse, OAuthServerMetadata, + RefreshTokenRequestParams, TokenRequestParams, TokenResponse, + TokenRevocationRequest, } from "./types"; const AUTH_GRANT_TYPE = "authorization_code" as const; @@ -25,6 +30,9 @@ const CLIENT_NAME = "VS Code Coder Extension"; const REQUIRED_GRANT_TYPES = [AUTH_GRANT_TYPE, REFRESH_GRANT_TYPE] as const; +// Token refresh timing constants +const TOKEN_REFRESH_THRESHOLD_MS = 5 * 60 * 1000; // 5 minutes before expiry + export class CoderOAuthHelper { private clientRegistration: ClientRegistrationResponse | undefined; private cachedMetadata: OAuthServerMetadata | undefined; @@ -34,6 +42,8 @@ export class CoderOAuthHelper { private pendingAuthReject: ((reason: Error) => void) | undefined; private expectedState: string | undefined; private pendingVerifier: string | undefined; + private storedTokens: StoredOAuthTokens | undefined; + private refreshTimer: NodeJS.Timeout | undefined; private readonly extensionId: string; @@ -50,6 +60,7 @@ export class CoderOAuthHelper { context, ); await helper.loadClientRegistration(); + await helper.loadTokens(); return helper; } private constructor( @@ -120,10 +131,11 @@ export class CoderOAuthHelper { } this.cachedMetadata = metadata; - this.logger.info("OAuth endpoints discovered:", { + this.logger.debug("OAuth endpoints discovered:", { authorization: metadata.authorization_endpoint, token: metadata.token_endpoint, registration: metadata.registration_endpoint, + revocation: metadata.revocation_endpoint, }); return metadata; @@ -141,6 +153,20 @@ export class CoderOAuthHelper { } } + private async loadTokens(): Promise { + const tokens = await this.secretsManager.getOAuthTokens(); + if (tokens) { + this.storedTokens = tokens; + this.logger.info("Loaded stored OAuth tokens", { + expires_at: new Date(tokens.expiry_timestamp).toISOString(), + }); + + if (tokens.refresh_token) { + this.startRefreshTimer(); + } + } + } + private async saveClientRegistration( registration: ClientRegistrationResponse, ): Promise { @@ -178,9 +204,10 @@ export class CoderOAuthHelper { ); } + // "web" type since VS Code Secrets API allows secure client_secret storage (confidential client). const registrationRequest: ClientRegistrationRequest = { redirect_uris: [redirectUri], - application_type: "native", + application_type: "web", grant_types: [AUTH_GRANT_TYPE], response_types: [RESPONSE_TYPE], client_name: CLIENT_NAME, @@ -228,10 +255,11 @@ export class CoderOAuthHelper { const url = `${metadata.authorization_endpoint}?${new URLSearchParams(params as unknown as Record).toString()}`; - this.logger.info("OAuth Authorization URL:", url); - this.logger.info("Client ID:", clientId); - this.logger.info("Redirect URI:", this.getRedirectUri()); - this.logger.info("Scope:", scope); + this.logger.debug("Building OAuth authorization URL:", { + client_id: clientId, + redirect_uri: this.getRedirectUri(), + scope, + }); return url; } @@ -366,25 +394,234 @@ export class CoderOAuthHelper { params.client_secret = this.clientRegistration.client_secret; } - const tokenRequest = new URLSearchParams( - params as unknown as Record, - ); + const tokenRequest = toUrlSearchParams(params); const response = await this.client .getAxiosInstance() - .post(metadata.token_endpoint, tokenRequest.toString(), { + .post(metadata.token_endpoint, tokenRequest, { headers: { "Content-Type": "application/x-www-form-urlencoded", }, }); this.logger.info("Token exchange successful"); + + await this.saveTokens(response.data); + return response.data; } getClientId(): string | undefined { return this.clientRegistration?.client_id; } + + /** + * Refresh the access token using the stored refresh token. + */ + private async refreshToken(): Promise { + if (!this.storedTokens?.refresh_token) { + throw new Error("No refresh token available"); + } + + if (!this.clientRegistration) { + throw new Error("No client registration found"); + } + + const metadata = await this.getMetadata(); + + this.logger.debug("Refreshing access token"); + + const params: RefreshTokenRequestParams = { + grant_type: REFRESH_GRANT_TYPE, + refresh_token: this.storedTokens.refresh_token, + client_id: this.clientRegistration.client_id, + }; + + if (this.clientRegistration.client_secret) { + params.client_secret = this.clientRegistration.client_secret; + } + + const tokenRequest = toUrlSearchParams(params); + + const response = await this.client + .getAxiosInstance() + .post(metadata.token_endpoint, tokenRequest, { + headers: { + "Content-Type": "application/x-www-form-urlencoded", + }, + }); + + this.logger.debug("Token refresh successful"); + + await this.saveTokens(response.data); + + return response.data; + } + + /** + * Save token response to secrets storage and restart the refresh timer. + */ + private async saveTokens(tokenResponse: TokenResponse): Promise { + const expiryTimestamp = tokenResponse.expires_in + ? Date.now() + tokenResponse.expires_in * 1000 + : Date.now() + 3600 * 1000; // Default to 1 hour if not specified + + const tokens: StoredOAuthTokens = { + ...tokenResponse, + expiry_timestamp: expiryTimestamp, + }; + + this.storedTokens = tokens; + await this.secretsManager.setOAuthTokens(tokens); + + this.logger.info("Tokens saved", { + expires_at: new Date(expiryTimestamp).toISOString(), + }); + + // Restart timer with new expiry (creates self-perpetuating refresh cycle) + this.startRefreshTimer(); + } + + /** + * Start the background token refresh timer. + * Sets a timeout to fire exactly when the token is 5 minutes from expiry. + */ + private startRefreshTimer(): void { + this.stopRefreshTimer(); + + if (!this.storedTokens?.refresh_token) { + this.logger.debug("No refresh token available, skipping timer setup"); + return; + } + + const now = Date.now(); + const timeUntilRefresh = + this.storedTokens.expiry_timestamp - TOKEN_REFRESH_THRESHOLD_MS - now; + + // If token is already expired or expires very soon, refresh immediately + if (timeUntilRefresh <= 0) { + this.logger.info("Token needs immediate refresh"); + this.refreshToken().catch((error) => { + this.logger.error("Immediate token refresh failed:", error); + }); + return; + } + + // Set timeout to fire exactly when token is 5 minutes from expiry + this.refreshTimer = setTimeout(() => { + this.logger.debug("Token refresh timer fired, refreshing token..."); + this.refreshToken().catch((error) => { + this.logger.error("Scheduled token refresh failed:", error); + }); + }, timeUntilRefresh); + + this.logger.debug("Token refresh timer scheduled", { + fires_at: new Date(now + timeUntilRefresh).toISOString(), + fires_in: timeUntilRefresh / 1000, + }); + } + + /** + * Stop the background token refresh timer. + */ + private stopRefreshTimer(): void { + if (this.refreshTimer) { + clearTimeout(this.refreshTimer); + this.refreshTimer = undefined; + this.logger.debug("Background token refresh timer stopped"); + } + } + + /** + * Revoke a token using the OAuth server's revocation endpoint. + */ + private async revokeToken( + token: string, + tokenTypeHint?: "access_token" | "refresh_token", + ): Promise { + if (!this.clientRegistration) { + throw new Error("No client registration found"); + } + + const metadata = await this.getMetadata(); + + if (!metadata.revocation_endpoint) { + this.logger.warn( + "Server does not support token revocation (no revocation_endpoint)", + ); + return; + } + + this.logger.info("Revoking token", { tokenTypeHint }); + + const params: TokenRevocationRequest = { + token, + client_id: this.clientRegistration.client_id, + }; + + if (tokenTypeHint) { + params.token_type_hint = tokenTypeHint; + } + + if (this.clientRegistration.client_secret) { + params.client_secret = this.clientRegistration.client_secret; + } + + const revocationRequest = toUrlSearchParams(params); + + try { + await this.client + .getAxiosInstance() + .post(metadata.revocation_endpoint, revocationRequest, { + headers: { + "Content-Type": "application/x-www-form-urlencoded", + }, + }); + + this.logger.info("Token revocation successful"); + } catch (error) { + this.logger.error("Token revocation failed:", error); + throw error; + } + } + + /** + * Logout by revoking tokens and clearing all OAuth data. + */ + async logout(): Promise { + this.stopRefreshTimer(); + + // Revoke refresh token (which also invalidates access token per RFC 7009) + if (this.storedTokens?.refresh_token) { + try { + await this.revokeToken( + this.storedTokens.refresh_token, + "refresh_token", + ); + } catch (error) { + this.logger.warn("Token revocation failed during logout:", error); + } + } + + // Clear stored tokens + await this.secretsManager.clearOAuthTokens(); + this.storedTokens = undefined; + + // Clear client registration + await this.clearClientRegistration(); + + // Trigger logout state change for other windows + // await this.secretsManager.triggerLoginStateChange("logout"); + + this.logger.info("Logout complete"); + } + + /** + * Cleanup method to be called when disposing the helper. + */ + dispose(): void { + this.stopRefreshTimer(); + } } function includesAllTypes( @@ -399,6 +636,20 @@ function includesAllTypes( return requiredTypes.every((type) => arr.includes(type)); } +/** + * Converts an object with string properties to Record, + * filtering out undefined values for use with URLSearchParams. + */ +function toUrlSearchParams(obj: object): URLSearchParams { + const params = Object.fromEntries( + Object.entries(obj).filter( + ([, value]) => value !== undefined && typeof value === "string", + ), + ) as Record; + + return new URLSearchParams(params); +} + /** * Activates OAuth support for the Coder extension. * Initializes the OAuth helper and registers the test auth command. @@ -417,33 +668,40 @@ export async function activateCoderOAuth( ); context.subscriptions.push( - vscode.commands.registerCommand("coder.oauth.testAuth", async () => { + vscode.commands.registerCommand("coder.oauth.login", async () => { try { const { code, verifier } = await oauthHelper.startAuthorization(); - logger.info( - "Authorization code received:", - code.substring(0, 8) + "...", - ); const tokenResponse = await oauthHelper.exchangeToken(code, verifier); - vscode.window.showInformationMessage( - `OAuth flow completed! Access token received (expires in ${tokenResponse.expires_in}s)`, - ); logger.info("OAuth flow completed:", { token_type: tokenResponse.token_type, expires_in: tokenResponse.expires_in, scope: tokenResponse.scope, }); + vscode.window.showInformationMessage( + `OAuth flow completed! Access token received (expires in ${tokenResponse.expires_in}s)`, + ); + + // Test API call to verify token works client.setSessionToken(tokenResponse.access_token); - const response = await client.getWorkspaces({ q: "owner:me" }); - logger.info(response.workspaces.map((w) => w.name).toString()); + await client.getWorkspaces({ q: "owner:me" }); } catch (error) { vscode.window.showErrorMessage(`OAuth flow failed: ${error}`); logger.error("OAuth flow failed:", error); } }), + vscode.commands.registerCommand("coder.oauth.logout", async () => { + try { + await oauthHelper.logout(); + vscode.window.showInformationMessage("Successfully logged out"); + logger.info("User logged out via OAuth"); + } catch (error) { + vscode.window.showErrorMessage(`Logout failed: ${error}`); + logger.error("OAuth logout failed:", error); + } + }), ); return oauthHelper; From 56d77ab747fda04fabc76d70d3a7bb2cbe9d17ed Mon Sep 17 00:00:00 2001 From: Ehab Younes Date: Fri, 24 Oct 2025 16:15:05 +0300 Subject: [PATCH 06/20] Add proper scopes --- src/oauth/oauthHelper.ts | 96 ++++++++++++++++++++++++++++++++++------ 1 file changed, 82 insertions(+), 14 deletions(-) diff --git a/src/oauth/oauthHelper.ts b/src/oauth/oauthHelper.ts index 03154fc7..de369ba8 100644 --- a/src/oauth/oauthHelper.ts +++ b/src/oauth/oauthHelper.ts @@ -30,8 +30,28 @@ const CLIENT_NAME = "VS Code Coder Extension"; const REQUIRED_GRANT_TYPES = [AUTH_GRANT_TYPE, REFRESH_GRANT_TYPE] as const; -// Token refresh timing constants -const TOKEN_REFRESH_THRESHOLD_MS = 5 * 60 * 1000; // 5 minutes before expiry +// Token refresh timing constants (5 minutes before expiry) +const TOKEN_REFRESH_THRESHOLD_MS = 5 * 60 * 1000; + +/** + * Minimal scopes required by the VS Code extension: + * - workspace:read: List and read workspace details + * - workspace:update: Update workspace version + * - workspace:start: Start stopped workspaces + * - workspace:ssh: SSH configuration for remote connections + * - workspace:application_connect: Connect to workspace agents/apps + * - template:read: Read templates and versions + * - user:read_personal: Read authenticated user info + */ +const DEFAULT_OAUTH_SCOPES = [ + "workspace:read", + "workspace:update", + "workspace:start", + "workspace:ssh", + "workspace:application_connect", + "template:read", + "user:read_personal", +].join(" "); export class CoderOAuthHelper { private clientRegistration: ClientRegistrationResponse | undefined; @@ -156,9 +176,22 @@ export class CoderOAuthHelper { private async loadTokens(): Promise { const tokens = await this.secretsManager.getOAuthTokens(); if (tokens) { + if (!this.hasRequiredScopes(tokens.scope)) { + this.logger.warn( + "Stored token missing required scopes, clearing tokens", + { + stored_scope: tokens.scope, + required_scopes: DEFAULT_OAUTH_SCOPES, + }, + ); + await this.secretsManager.clearOAuthTokens(); + return; + } + this.storedTokens = tokens; this.logger.info("Loaded stored OAuth tokens", { expires_at: new Date(tokens.expiry_timestamp).toISOString(), + scope: tokens.scope, }); if (tokens.refresh_token) { @@ -167,6 +200,40 @@ export class CoderOAuthHelper { } } + /** + * Check if granted scopes cover all required scopes. + * Supports wildcard scopes like "workspace:*" which grant all "workspace:" prefixed scopes. + */ + private hasRequiredScopes(grantedScope: string | undefined): boolean { + if (!grantedScope) { + return false; + } + + const grantedScopes = new Set(grantedScope.split(" ")); + const requiredScopes = DEFAULT_OAUTH_SCOPES.split(" "); + + for (const required of requiredScopes) { + // Check exact match + if (grantedScopes.has(required)) { + continue; + } + + // Check wildcard match (e.g., "workspace:*" grants "workspace:read") + const colonIndex = required.indexOf(":"); + if (colonIndex !== -1) { + const prefix = required.substring(0, colonIndex); + const wildcard = `${prefix}:*`; + if (grantedScopes.has(wildcard)) { + continue; + } + } + + return false; + } + + return true; + } + private async saveClientRegistration( registration: ClientRegistrationResponse, ): Promise { @@ -231,16 +298,19 @@ export class CoderOAuthHelper { clientId: string, state: string, challenge: string, - scope = "all", + scope: string, ): string { - if ( - metadata.scopes_supported && - !metadata.scopes_supported.includes(scope) - ) { - this.logger.warn( - `Requested scope "${scope}" not in server's supported scopes. Server may still accept it.`, - { supported_scopes: metadata.scopes_supported }, + if (metadata.scopes_supported) { + const requestedScopes = scope.split(" "); + const unsupportedScopes = requestedScopes.filter( + (s) => !metadata.scopes_supported?.includes(s), ); + if (unsupportedScopes.length > 0) { + this.logger.warn( + `Requested scopes not in server's supported scopes: ${unsupportedScopes.join(", ")}. Server may still accept them.`, + { supported_scopes: metadata.scopes_supported }, + ); + } } const params: AuthorizationRequestParams = { @@ -264,9 +334,7 @@ export class CoderOAuthHelper { return url; } - async startAuthorization( - scope = "all", - ): Promise<{ code: string; verifier: string }> { + async startAuthorization(): Promise<{ code: string; verifier: string }> { const metadata = await this.getMetadata(); const clientId = await this.registerClient(); const state = generateState(); @@ -277,7 +345,7 @@ export class CoderOAuthHelper { clientId, state, challenge, - scope, + DEFAULT_OAUTH_SCOPES, ); return new Promise<{ code: string; verifier: string }>( From 2e1454c17b49108e3fb67fe5914e6d3161882b4f Mon Sep 17 00:00:00 2001 From: Ehab Younes Date: Sat, 25 Oct 2025 00:21:54 +0300 Subject: [PATCH 07/20] Hook up into the login/logout logic of the extension --- package.json | 10 -- src/commands.ts | 215 +++++++++++++++++++++++++++++++------ src/core/secretsManager.ts | 14 +++ src/extension.ts | 27 ++++- src/oauth/oauthHelper.ts | 95 ++++++---------- 5 files changed, 252 insertions(+), 109 deletions(-) diff --git a/package.json b/package.json index add68824..bd60a54c 100644 --- a/package.json +++ b/package.json @@ -255,16 +255,6 @@ "title": "Search", "category": "Coder", "icon": "$(search)" - }, - { - "command": "coder.oauth.login", - "title": "OAuth Login", - "category": "Coder" - }, - { - "command": "coder.oauth.logout", - "title": "OAuth Logout", - "category": "Coder" } ], "menus": { diff --git a/src/commands.ts b/src/commands.ts index 384b4d79..c9f33909 100644 --- a/src/commands.ts +++ b/src/commands.ts @@ -19,6 +19,7 @@ import { type SecretsManager } from "./core/secretsManager"; import { CertificateError } from "./error"; import { getGlobalFlags } from "./globalFlags"; import { type Logger } from "./logging/logger"; +import { type CoderOAuthHelper } from "./oauth/oauthHelper"; import { maybeAskAgent, maybeAskUrl } from "./promptUtils"; import { escapeCommandArg, toRemoteAuthority, toSafeHost } from "./util"; import { @@ -49,6 +50,7 @@ export class Commands { public constructor( serviceContainer: ServiceContainer, private readonly restClient: Api, + private readonly oauthHelper: CoderOAuthHelper, ) { this.vscodeProposed = serviceContainer.getVsCodeProposed(); this.logger = serviceContainer.getLogger(); @@ -60,59 +62,119 @@ export class Commands { } /** - * 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. + * Check if server supports OAuth by attempting to fetch the well-known endpoint. */ - public async login(args?: { - url?: string; - token?: string; - label?: string; - autoLogin?: boolean; - }): Promise { - if (this.contextManager.get("coder.authenticated")) { - return; + private async checkOAuthSupport(client: CoderApi): Promise { + try { + await client + .getAxiosInstance() + .get("/.well-known/oauth-authorization-server"); + this.logger.debug("Server supports OAuth"); + return true; + } catch (error) { + this.logger.debug("Server does not support OAuth:", error); + return false; } - this.logger.info("Logging in"); + } - const url = await maybeAskUrl(this.mementoManager, args?.url); - if (!url) { - return; // The user aborted. - } + /** + * Ask user to choose between OAuth and legacy API token authentication. + */ + private async askAuthMethod(): Promise<"oauth" | "legacy" | undefined> { + const choice = await vscode.window.showQuickPick( + [ + { + label: "$(key) OAuth (Recommended)", + detail: "Secure authentication with automatic token refresh", + value: "oauth", + }, + { + label: "$(lock) API Token", + detail: "Use a manually created API key", + value: "legacy", + }, + ], + { + title: "Choose Authentication Method", + placeHolder: "How would you like to authenticate?", + ignoreFocusOut: true, + }, + ); - // 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 = args?.label === undefined ? toSafeHost(url) : args.label; + return choice?.value as "oauth" | "legacy" | undefined; + } - // Try to get a token from the user, if we need one, and their user. - const autoLogin = args?.autoLogin === true; - const res = await this.maybeAskToken(url, args?.token, autoLogin); - if (!res) { - return; // The user aborted, or unable to auth. + /** + * Authenticate using OAuth flow. + * Returns the access token and authenticated user, or null if failed/cancelled. + */ + private async loginWithOAuth( + url: string, + ): Promise<{ user: User; token: string } | null> { + try { + this.logger.info("Starting OAuth authentication"); + + // Start OAuth authorization flow + const { code, verifier } = await this.oauthHelper.startAuthorization(url); + + // Exchange authorization code for tokens + const tokenResponse = await this.oauthHelper.exchangeToken( + code, + verifier, + ); + + // Validate token by fetching user + const client = CoderApi.create( + url, + tokenResponse.access_token, + this.logger, + ); + const user = await client.getAuthenticatedUser(); + + this.logger.info("OAuth authentication successful"); + + return { + token: tokenResponse.access_token, + user, + }; + } catch (error) { + this.logger.error("OAuth authentication failed:", error); + vscode.window.showErrorMessage( + `OAuth authentication failed: ${getErrorMessage(error, "Unknown error")}`, + ); + return null; } + } - // The URL is good and the token is either good or not required; authorize - // the global client. + /** + * Complete the login process by storing credentials and updating context. + */ + private async completeLogin( + url: string, + label: string, + token: string, + user: User, + ): Promise { + // Authorize the global client this.restClient.setHost(url); - this.restClient.setSessionToken(res.token); + this.restClient.setSessionToken(token); - // Store these to be used in later sessions. + // Store for later sessions await this.mementoManager.setUrl(url); - await this.secretsManager.setSessionToken(res.token); + await this.secretsManager.setSessionToken(token); - // Store on disk to be used by the cli. - await this.cliManager.configure(label, url, res.token); + // Store on disk for CLI + await this.cliManager.configure(label, url, token); - // These contexts control various menu items and the sidebar. + // Update contexts this.contextManager.set("coder.authenticated", true); - if (res.user.roles.find((role) => role.name === "owner")) { + if (user.roles.find((role) => role.name === "owner")) { this.contextManager.set("coder.isOwner", true); } vscode.window .showInformationMessage( - `Welcome to Coder, ${res.user.username}!`, + `Welcome to Coder, ${user.username}!`, { detail: "You can now use the Coder extension to manage your Coder instance.", @@ -130,6 +192,73 @@ export class Commands { vscode.commands.executeCommand("coder.refreshWorkspaces"); } + /** + * 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?: { + url?: string; + token?: string; + label?: string; + autoLogin?: boolean; + }): Promise { + if (this.contextManager.get("coder.authenticated")) { + return; + } + this.logger.info("Logging in"); + + const url = await maybeAskUrl(this.mementoManager, args?.url); + if (!url) { + return; // The user aborted. + } + + const label = args?.label ?? toSafeHost(url); + const autoLogin = args?.autoLogin === true; + + // Check if we have an existing valid legacy token + const existingToken = await this.secretsManager.getSessionToken(); + const client = CoderApi.create(url, existingToken, this.logger); + if (existingToken && !args?.token) { + try { + const user = await client.getAuthenticatedUser(); + this.logger.info("Using existing valid session token"); + await this.completeLogin(url, label, existingToken, user); + return; + } catch { + this.logger.debug("Existing token invalid, clearing it"); + await this.secretsManager.setSessionToken(); + } + } + + // Check if server supports OAuth + const supportsOAuth = await this.checkOAuthSupport(client); + + if (supportsOAuth && !autoLogin) { + const choice = await this.askAuthMethod(); + if (!choice) { + return; + } + + if (choice === "oauth") { + const res = await this.loginWithOAuth(url); + if (!res) { + return; + } + await this.completeLogin(url, label, res.token, res.user); + return; + } + } + + // Use legacy token flow (existing behavior) + const res = await this.maybeAskToken(url, args?.token, autoLogin); + if (!res) { + return; + } + + await this.completeLogin(url, label, res.token, res.user); + } + /** * If necessary, ask for a token, and keep asking until the token has been * validated. Return the token and user that was fetched to validate the @@ -255,6 +384,22 @@ export class Commands { // Sanity check; command should not be available if no url. throw new Error("You are not logged in"); } + + // Check if using OAuth + const hasOAuthTokens = await this.secretsManager.getOAuthTokens(); + if (hasOAuthTokens) { + this.logger.info("Logging out via OAuth"); + try { + await this.oauthHelper.logout(); + } catch (error) { + this.logger.warn( + "OAuth logout failed, continuing with cleanup:", + error, + ); + } + } + + // Continue with standard logout (clears sessionToken, contexts, etc) await this.forceLogout(); } diff --git a/src/core/secretsManager.ts b/src/core/secretsManager.ts index b108af0a..02681132 100644 --- a/src/core/secretsManager.ts +++ b/src/core/secretsManager.ts @@ -84,6 +84,20 @@ export class SecretsManager { }); } + /** + * Listens for session token changes. + */ + public onDidChangeSessionToken( + listener: (token: string | undefined) => Promise, + ): Disposable { + return this.secrets.onDidChange(async (e) => { + if (e.key === SESSION_TOKEN_KEY) { + const token = await this.getSessionToken(); + await listener(token); + } + }); + } + /** * Store OAuth client registration data. */ diff --git a/src/extension.ts b/src/extension.ts index f350a6e5..cf5314fd 100644 --- a/src/extension.ts +++ b/src/extension.ts @@ -119,11 +119,34 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { ); const oauthHelper = await activateCoderOAuth( - client, + url || "", secretsManager, output, ctx, ); + ctx.subscriptions.push(oauthHelper); + + // Listen for session token changes and sync state across all components + ctx.subscriptions.push( + secretsManager.onDidChangeSessionToken(async (token) => { + if (!token) { + output.debug("Session token cleared"); + client.setSessionToken(""); + return; + } + + output.debug("Session token changed, syncing state"); + + client.setSessionToken(token); + const url = mementoManager.getUrl(); + if (url) { + const cliManager = serviceContainer.getCliManager(); + // TODO label might not match? + await cliManager.configure(toSafeHost(url), url, token); + output.debug("Updated CLI config with new token"); + } + }), + ); // Handle vscode:// URIs. const uriHandler = vscode.window.registerUriHandler({ @@ -290,7 +313,7 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { // Register globally available commands. Many of these have visibility // controlled by contexts, see `when` in the package.json. - const commands = new Commands(serviceContainer, client); + const commands = new Commands(serviceContainer, client, oauthHelper); ctx.subscriptions.push( vscode.commands.registerCommand( "coder.login", diff --git a/src/oauth/oauthHelper.ts b/src/oauth/oauthHelper.ts index de369ba8..efdc9b8b 100644 --- a/src/oauth/oauthHelper.ts +++ b/src/oauth/oauthHelper.ts @@ -1,6 +1,6 @@ import * as vscode from "vscode"; -import { type CoderApi } from "../api/coderApi"; +import { CoderApi } from "../api/coderApi"; import { type StoredOAuthTokens, type SecretsManager, @@ -53,7 +53,9 @@ const DEFAULT_OAUTH_SCOPES = [ "user:read_personal", ].join(" "); -export class CoderOAuthHelper { +export class CoderOAuthHelper implements vscode.Disposable { + private readonly client: CoderApi; + private clientRegistration: ClientRegistrationResponse | undefined; private cachedMetadata: OAuthServerMetadata | undefined; private pendingAuthResolve: @@ -68,13 +70,13 @@ export class CoderOAuthHelper { private readonly extensionId: string; static async create( - client: CoderApi, + baseUrl: string, secretsManager: SecretsManager, logger: Logger, context: vscode.ExtensionContext, ): Promise { const helper = new CoderOAuthHelper( - client, + baseUrl, secretsManager, logger, context, @@ -84,11 +86,12 @@ export class CoderOAuthHelper { return helper; } private constructor( - private readonly client: CoderApi, + baseUrl: string, private readonly secretsManager: SecretsManager, private readonly logger: Logger, context: vscode.ExtensionContext, ) { + this.client = CoderApi.create(baseUrl, undefined, logger); this.extensionId = context.extension.id; } @@ -193,6 +196,7 @@ export class CoderOAuthHelper { expires_at: new Date(tokens.expiry_timestamp).toISOString(), scope: tokens.scope, }); + this.client.setSessionToken(tokens.access_token); if (tokens.refresh_token) { this.startRefreshTimer(); @@ -334,7 +338,11 @@ export class CoderOAuthHelper { return url; } - async startAuthorization(): Promise<{ code: string; verifier: string }> { + async startAuthorization( + url: string, + ): Promise<{ code: string; verifier: string }> { + this.client.setHost(url); + const metadata = await this.getMetadata(); const clientId = await this.registerClient(); const state = generateState(); @@ -541,6 +549,8 @@ export class CoderOAuthHelper { this.storedTokens = tokens; await this.secretsManager.setOAuthTokens(tokens); + await this.secretsManager.setSessionToken(tokenResponse.access_token); + this.client.setSessionToken(tokens.access_token); this.logger.info("Tokens saved", { expires_at: new Date(expiryTimestamp).toISOString(), @@ -552,7 +562,7 @@ export class CoderOAuthHelper { /** * Start the background token refresh timer. - * Sets a timeout to fire exactly when the token is 5 minutes from expiry. + * Sets a timeout to fire when the token is close to expiry. */ private startRefreshTimer(): void { this.stopRefreshTimer(); @@ -671,24 +681,29 @@ export class CoderOAuthHelper { } } - // Clear stored tokens await this.secretsManager.clearOAuthTokens(); this.storedTokens = undefined; - - // Clear client registration await this.clearClientRegistration(); - // Trigger logout state change for other windows - // await this.secretsManager.triggerLoginStateChange("logout"); - - this.logger.info("Logout complete"); + this.logger.info("OAuth logout complete"); } /** - * Cleanup method to be called when disposing the helper. + * Clears all in-memory state and rejects any pending operations. */ dispose(): void { this.stopRefreshTimer(); + + if (this.pendingAuthReject) { + this.pendingAuthReject(new Error("OAuth helper disposed")); + } + this.clearPendingAuth(); + + this.storedTokens = undefined; + this.clientRegistration = undefined; + this.cachedMetadata = undefined; + + this.logger.debug("OAuth helper disposed, all state cleared"); } } @@ -720,57 +735,13 @@ function toUrlSearchParams(obj: object): URLSearchParams { /** * Activates OAuth support for the Coder extension. - * Initializes the OAuth helper and registers the test auth command. + * Initializes and returns the OAuth helper. */ export async function activateCoderOAuth( - client: CoderApi, + baseUrl: string, secretsManager: SecretsManager, logger: Logger, context: vscode.ExtensionContext, ): Promise { - const oauthHelper = await CoderOAuthHelper.create( - client, - secretsManager, - logger, - context, - ); - - context.subscriptions.push( - vscode.commands.registerCommand("coder.oauth.login", async () => { - try { - const { code, verifier } = await oauthHelper.startAuthorization(); - - const tokenResponse = await oauthHelper.exchangeToken(code, verifier); - - logger.info("OAuth flow completed:", { - token_type: tokenResponse.token_type, - expires_in: tokenResponse.expires_in, - scope: tokenResponse.scope, - }); - - vscode.window.showInformationMessage( - `OAuth flow completed! Access token received (expires in ${tokenResponse.expires_in}s)`, - ); - - // Test API call to verify token works - client.setSessionToken(tokenResponse.access_token); - await client.getWorkspaces({ q: "owner:me" }); - } catch (error) { - vscode.window.showErrorMessage(`OAuth flow failed: ${error}`); - logger.error("OAuth flow failed:", error); - } - }), - vscode.commands.registerCommand("coder.oauth.logout", async () => { - try { - await oauthHelper.logout(); - vscode.window.showInformationMessage("Successfully logged out"); - logger.info("User logged out via OAuth"); - } catch (error) { - vscode.window.showErrorMessage(`Logout failed: ${error}`); - logger.error("OAuth logout failed:", error); - } - }), - ); - - return oauthHelper; + return CoderOAuthHelper.create(baseUrl, secretsManager, logger, context); } From a0f90e61adffbff427a7e00d02d9e2f75114ea25 Mon Sep 17 00:00:00 2001 From: Ehab Younes Date: Mon, 27 Oct 2025 12:12:35 +0300 Subject: [PATCH 08/20] WIP refactoring --- src/commands.ts | 319 ++++++------ src/core/mementoManager.ts | 2 +- src/core/secretsManager.ts | 12 +- src/extension.ts | 10 +- src/oauth/clientRegistry.ts | 111 +++++ src/oauth/metadataClient.ts | 137 ++++++ src/oauth/oauthHelper.ts | 747 ----------------------------- src/oauth/sessionManager.ts | 636 ++++++++++++++++++++++++ src/oauth/tokenRefreshScheduler.ts | 65 +++ src/oauth/utils.ts | 14 + src/remote/remote.ts | 10 +- 11 files changed, 1124 insertions(+), 939 deletions(-) create mode 100644 src/oauth/clientRegistry.ts create mode 100644 src/oauth/metadataClient.ts delete mode 100644 src/oauth/oauthHelper.ts create mode 100644 src/oauth/sessionManager.ts create mode 100644 src/oauth/tokenRefreshScheduler.ts diff --git a/src/commands.ts b/src/commands.ts index c9f33909..4e16e617 100644 --- a/src/commands.ts +++ b/src/commands.ts @@ -19,7 +19,8 @@ import { type SecretsManager } from "./core/secretsManager"; import { CertificateError } from "./error"; import { getGlobalFlags } from "./globalFlags"; import { type Logger } from "./logging/logger"; -import { type CoderOAuthHelper } from "./oauth/oauthHelper"; +import { OAuthMetadataClient } from "./oauth/metadataClient"; +import { type OAuthSessionManager } from "./oauth/sessionManager"; import { maybeAskAgent, maybeAskUrl } from "./promptUtils"; import { escapeCommandArg, toRemoteAuthority, toSafeHost } from "./util"; import { @@ -28,6 +29,8 @@ import { WorkspaceTreeItem, } from "./workspace/workspacesProvider"; +type AuthMethod = "oauth" | "legacy"; + export class Commands { private readonly vscodeProposed: typeof vscode; private readonly logger: Logger; @@ -50,7 +53,7 @@ export class Commands { public constructor( serviceContainer: ServiceContainer, private readonly restClient: Api, - private readonly oauthHelper: CoderOAuthHelper, + private readonly oauthSessionManager: OAuthSessionManager, ) { this.vscodeProposed = serviceContainer.getVsCodeProposed(); this.logger = serviceContainer.getLogger(); @@ -62,119 +65,59 @@ export class Commands { } /** - * Check if server supports OAuth by attempting to fetch the well-known endpoint. + * 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. */ - private async checkOAuthSupport(client: CoderApi): Promise { - try { - await client - .getAxiosInstance() - .get("/.well-known/oauth-authorization-server"); - this.logger.debug("Server supports OAuth"); - return true; - } catch (error) { - this.logger.debug("Server does not support OAuth:", error); - return false; + public async login(args?: { + url?: string; + token?: string; + label?: string; + autoLogin?: boolean; + }): Promise { + if (this.contextManager.get("coder.authenticated")) { + return; } - } - - /** - * Ask user to choose between OAuth and legacy API token authentication. - */ - private async askAuthMethod(): Promise<"oauth" | "legacy" | undefined> { - const choice = await vscode.window.showQuickPick( - [ - { - label: "$(key) OAuth (Recommended)", - detail: "Secure authentication with automatic token refresh", - value: "oauth", - }, - { - label: "$(lock) API Token", - detail: "Use a manually created API key", - value: "legacy", - }, - ], - { - title: "Choose Authentication Method", - placeHolder: "How would you like to authenticate?", - ignoreFocusOut: true, - }, - ); - - return choice?.value as "oauth" | "legacy" | undefined; - } - - /** - * Authenticate using OAuth flow. - * Returns the access token and authenticated user, or null if failed/cancelled. - */ - private async loginWithOAuth( - url: string, - ): Promise<{ user: User; token: string } | null> { - try { - this.logger.info("Starting OAuth authentication"); - - // Start OAuth authorization flow - const { code, verifier } = await this.oauthHelper.startAuthorization(url); - - // Exchange authorization code for tokens - const tokenResponse = await this.oauthHelper.exchangeToken( - code, - verifier, - ); + this.logger.info("Logging in"); - // Validate token by fetching user - const client = CoderApi.create( - url, - tokenResponse.access_token, - this.logger, - ); - const user = await client.getAuthenticatedUser(); + const url = await maybeAskUrl(this.mementoManager, args?.url); + if (!url) { + return; // The user aborted. + } - this.logger.info("OAuth authentication successful"); + // 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 = args?.label ?? toSafeHost(url); + // Try to get a token from the user, if we need one, and their user. + const autoLogin = args?.autoLogin === true; - return { - token: tokenResponse.access_token, - user, - }; - } catch (error) { - this.logger.error("OAuth authentication failed:", error); - vscode.window.showErrorMessage( - `OAuth authentication failed: ${getErrorMessage(error, "Unknown error")}`, - ); - return null; + const res = await this.maybeAskToken(url, args?.token, autoLogin); + if (!res) { + return; // The user aborted, or unable to auth. } - } - /** - * Complete the login process by storing credentials and updating context. - */ - private async completeLogin( - url: string, - label: string, - token: string, - user: User, - ): Promise { - // Authorize the global client + // The URL is good and the token is either good or not required; authorize + // the global client. this.restClient.setHost(url); - this.restClient.setSessionToken(token); + this.restClient.setSessionToken(res.token); - // Store for later sessions + // Store these to be used in later sessions. await this.mementoManager.setUrl(url); - await this.secretsManager.setSessionToken(token); + await this.secretsManager.setSessionToken(res.token); - // Store on disk for CLI - await this.cliManager.configure(label, url, token); + // Store on disk to be used by the cli. + await this.cliManager.configure(label, url, res.token); - // Update contexts + // These contexts control various menu items and the sidebar. this.contextManager.set("coder.authenticated", true); - if (user.roles.find((role) => role.name === "owner")) { + if (res.user.roles.some((role) => role.name === "owner")) { this.contextManager.set("coder.isOwner", true); } vscode.window .showInformationMessage( - `Welcome to Coder, ${user.username}!`, + `Welcome to Coder, ${res.user.username}!`, { detail: "You can now use the Coder extension to manage your Coder instance.", @@ -192,73 +135,6 @@ export class Commands { vscode.commands.executeCommand("coder.refreshWorkspaces"); } - /** - * 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?: { - url?: string; - token?: string; - label?: string; - autoLogin?: boolean; - }): Promise { - if (this.contextManager.get("coder.authenticated")) { - return; - } - this.logger.info("Logging in"); - - const url = await maybeAskUrl(this.mementoManager, args?.url); - if (!url) { - return; // The user aborted. - } - - const label = args?.label ?? toSafeHost(url); - const autoLogin = args?.autoLogin === true; - - // Check if we have an existing valid legacy token - const existingToken = await this.secretsManager.getSessionToken(); - const client = CoderApi.create(url, existingToken, this.logger); - if (existingToken && !args?.token) { - try { - const user = await client.getAuthenticatedUser(); - this.logger.info("Using existing valid session token"); - await this.completeLogin(url, label, existingToken, user); - return; - } catch { - this.logger.debug("Existing token invalid, clearing it"); - await this.secretsManager.setSessionToken(); - } - } - - // Check if server supports OAuth - const supportsOAuth = await this.checkOAuthSupport(client); - - if (supportsOAuth && !autoLogin) { - const choice = await this.askAuthMethod(); - if (!choice) { - return; - } - - if (choice === "oauth") { - const res = await this.loginWithOAuth(url); - if (!res) { - return; - } - await this.completeLogin(url, label, res.token, res.user); - return; - } - } - - // Use legacy token flow (existing behavior) - const res = await this.maybeAskToken(url, args?.token, autoLogin); - if (!res) { - return; - } - - await this.completeLogin(url, label, res.token, res.user); - } - /** * If necessary, ask for a token, and keep asking until the token has been * validated. Return the token and user that was fetched to validate the @@ -298,6 +174,64 @@ export class Commands { } } + // Check if server supports OAuth + const supportsOAuth = await this.checkOAuthSupport(client); + + let choice: AuthMethod | undefined = "legacy"; + if (supportsOAuth) { + choice = await this.askAuthMethod(); + } + + if (choice === "oauth") { + return this.loginWithOAuth(url, client); + } else if (choice === "legacy") { + return this.loginWithToken(url, token, client); + } + + // User aborted. + return null; + } + + private async checkOAuthSupport(client: CoderApi): Promise { + const metadataClient = new OAuthMetadataClient( + client.getAxiosInstance(), + this.logger, + ); + return metadataClient.checkOAuthSupport(); + } + + /** + * Ask user to choose between OAuth and legacy API token authentication. + */ + private async askAuthMethod(): Promise { + const choice = await vscode.window.showQuickPick( + [ + { + label: "$(key) OAuth (Recommended)", + detail: "Secure authentication with automatic token refresh", + value: "oauth" as const, + }, + { + label: "$(lock) API Token", + detail: "Use a manually created API key", + value: "legacy" as const, + }, + ], + { + title: "Choose Authentication Method", + placeHolder: "How would you like to authenticate?", + ignoreFocusOut: true, + }, + ); + + return choice?.value; + } + + private async loginWithToken( + url: string, + token: string | undefined, + client: CoderApi, + ): Promise<{ user: User; token: string } | null> { // This prompt is for convenience; do not error if they close it since // they may already have a token or already have the page opened. await vscode.env.openExternal(vscode.Uri.parse(`${url}/cli-auth`)); @@ -342,12 +276,52 @@ export class Commands { }, }); - if (validatedToken && user) { - return { token: validatedToken, user }; + if (user === undefined || validatedToken === undefined) { + return null; } - // User aborted. - return null; + return { user, token: validatedToken }; + } + + /** + * Authenticate using OAuth flow. + * Returns the access token and authenticated user, or null if failed/cancelled. + */ + private async loginWithOAuth( + url: string, + client: CoderApi, + ): Promise<{ user: User; token: string } | null> { + try { + this.logger.info("Starting OAuth authentication"); + + // Start OAuth authorization flow + // TODO just pass the client here and do all the neccessary steps (If we are already logged in we'd have the right token and the OAuth client registration saved). + const { code, verifier } = + await this.oauthSessionManager.startAuthorization(url); + + // Exchange authorization code for tokens + const tokenResponse = await this.oauthSessionManager.exchangeToken( + code, + verifier, + ); + + // Validate token by fetching user + client.setSessionToken(tokenResponse.access_token); + const user = await client.getAuthenticatedUser(); + + this.logger.info("OAuth authentication successful"); + + return { + token: tokenResponse.access_token, + user, + }; + } catch (error) { + this.logger.error("OAuth authentication failed:", error); + vscode.window.showErrorMessage( + `OAuth authentication failed: ${getErrorMessage(error, "Unknown error")}`, + ); + return null; + } } /** @@ -390,7 +364,7 @@ export class Commands { if (hasOAuthTokens) { this.logger.info("Logging out via OAuth"); try { - await this.oauthHelper.logout(); + await this.oauthSessionManager.logout(); } catch (error) { this.logger.warn( "OAuth logout failed, continuing with cleanup:", @@ -524,7 +498,7 @@ export class Commands { true, ); } else { - throw new Error("Unable to open unknown sidebar item"); + throw new TypeError("Unable to open unknown sidebar item"); } } else { // If there is no tree item, then the user manually ran this command. @@ -570,7 +544,7 @@ export class Commands { configDir, ); terminal.sendText( - `${escapeCommandArg(binary)}${` ${globalFlags.join(" ")}`} ssh ${app.workspace_name}`, + `${escapeCommandArg(binary)} ${globalFlags.join(" ")} ssh ${app.workspace_name}`, ); await new Promise((resolve) => setTimeout(resolve, 5000)); terminal.sendText(app.command ?? ""); @@ -669,7 +643,7 @@ export class Commands { workspaceAgent, ); - const hostPath = localWorkspaceFolder ? localWorkspaceFolder : undefined; + const hostPath = localWorkspaceFolder || undefined; const configFile = hostPath && localConfigFile ? { @@ -770,7 +744,6 @@ export class Commands { if (ex instanceof CertificateError) { ex.showNotification(); } - return; }); }); quickPick.show(); diff --git a/src/core/mementoManager.ts b/src/core/mementoManager.ts index f79be46c..a317ffe5 100644 --- a/src/core/mementoManager.ts +++ b/src/core/mementoManager.ts @@ -13,7 +13,7 @@ export class MementoManager { * If the URL is falsey, then remove it as the last used URL and do not touch * the history. */ - public async setUrl(url?: string): Promise { + public async setUrl(url: string | undefined): Promise { await this.memento.update("url", url); if (url) { const history = this.withUrlHistory(url); diff --git a/src/core/secretsManager.ts b/src/core/secretsManager.ts index 02681132..d16292f1 100644 --- a/src/core/secretsManager.ts +++ b/src/core/secretsManager.ts @@ -15,6 +15,7 @@ const OAUTH_TOKENS_KEY = "oauthTokens"; export type StoredOAuthTokens = Omit & { expiry_timestamp: number; + deployment_url: string; }; export enum AuthAction { @@ -29,7 +30,9 @@ export class SecretsManager { /** * Set or unset the last used token. */ - public async setSessionToken(sessionToken?: string): Promise { + public async setSessionToken( + sessionToken: string | undefined, + ): Promise { if (sessionToken) { await this.secrets.store(SESSION_TOKEN_KEY, sessionToken); } else { @@ -160,11 +163,4 @@ export class SecretsManager { } return undefined; } - - /** - * Clear OAuth token data. - */ - public async clearOAuthTokens(): Promise { - await this.secrets.delete(OAUTH_TOKENS_KEY); - } } diff --git a/src/extension.ts b/src/extension.ts index cf5314fd..38823378 100644 --- a/src/extension.ts +++ b/src/extension.ts @@ -13,7 +13,7 @@ import { Commands } from "./commands"; import { ServiceContainer } from "./core/container"; import { AuthAction } from "./core/secretsManager"; import { CertificateError, getErrorDetail } from "./error"; -import { activateCoderOAuth } from "./oauth/oauthHelper"; +import { OAuthSessionManager } from "./oauth/sessionManager"; import { CALLBACK_PATH } from "./oauth/utils"; import { maybeAskUrl } from "./promptUtils"; import { Remote } from "./remote/remote"; @@ -118,13 +118,13 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { ctx.subscriptions, ); - const oauthHelper = await activateCoderOAuth( + const oauthSessionManager = await OAuthSessionManager.create( url || "", secretsManager, output, ctx, ); - ctx.subscriptions.push(oauthHelper); + ctx.subscriptions.push(oauthSessionManager); // Listen for session token changes and sync state across all components ctx.subscriptions.push( @@ -157,7 +157,7 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { const code = params.get("code"); const state = params.get("state"); const error = params.get("error"); - oauthHelper.handleCallback(code, state, error); + oauthSessionManager.handleCallback(code, state, error); return; } @@ -313,7 +313,7 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { // Register globally available commands. Many of these have visibility // controlled by contexts, see `when` in the package.json. - const commands = new Commands(serviceContainer, client, oauthHelper); + const commands = new Commands(serviceContainer, client, oauthSessionManager); ctx.subscriptions.push( vscode.commands.registerCommand( "coder.login", diff --git a/src/oauth/clientRegistry.ts b/src/oauth/clientRegistry.ts new file mode 100644 index 00000000..91b4949f --- /dev/null +++ b/src/oauth/clientRegistry.ts @@ -0,0 +1,111 @@ +import type { AxiosInstance } from "axios"; + +import type { SecretsManager } from "../core/secretsManager"; +import type { Logger } from "../logging/logger"; + +import type { + ClientRegistrationRequest, + ClientRegistrationResponse, + OAuthServerMetadata, +} from "./types"; + +const AUTH_GRANT_TYPE = "authorization_code" as const; +const RESPONSE_TYPE = "code" as const; +const OAUTH_METHOD = "client_secret_post" as const; +const CLIENT_NAME = "VS Code Coder Extension"; + +/** + * Manages OAuth client registration and persistence. + */ +export class OAuthClientRegistry { + private registration: ClientRegistrationResponse | undefined; + + constructor( + private readonly axiosInstance: AxiosInstance, + private readonly secretsManager: SecretsManager, + private readonly logger: Logger, + ) {} + + /** + * Load existing client registration from secure storage. + * Should be called during initialization. + */ + async load(): Promise { + const registration = await this.secretsManager.getOAuthClientRegistration(); + if (registration) { + this.registration = registration; + this.logger.info("Loaded existing OAuth client:", registration.client_id); + } + } + + /** + * Get the current client registration if one exists. + */ + get(): ClientRegistrationResponse | undefined { + return this.registration; + } + + /** + * Register a new OAuth client or return existing if still valid. + * Re-registers if redirect URI has changed. + */ + async register( + metadata: OAuthServerMetadata, + redirectUri: string, + ): Promise { + if (this.registration?.client_id) { + if (this.registration.redirect_uris.includes(redirectUri)) { + this.logger.info( + "Using existing client registration:", + this.registration.client_id, + ); + return this.registration; + } + this.logger.info("Redirect URI changed, re-registering client"); + } + + if (!metadata.registration_endpoint) { + throw new Error("Server does not support dynamic client registration"); + } + + // "web" type since VS Code Secrets API allows secure client_secret storage (confidential client) + const registrationRequest: ClientRegistrationRequest = { + redirect_uris: [redirectUri], + application_type: "web", + grant_types: [AUTH_GRANT_TYPE], + response_types: [RESPONSE_TYPE], + client_name: CLIENT_NAME, + token_endpoint_auth_method: OAUTH_METHOD, + }; + + const response = await this.axiosInstance.post( + metadata.registration_endpoint, + registrationRequest, + ); + + await this.save(response.data); + + return response.data; + } + + /** + * Save client registration to secure storage. + */ + private async save(registration: ClientRegistrationResponse): Promise { + await this.secretsManager.setOAuthClientRegistration(registration); + this.registration = registration; + this.logger.info( + "Saved OAuth client registration:", + registration.client_id, + ); + } + + /** + * Clear the current client registration from memory and storage. + */ + async clear(): Promise { + await this.secretsManager.setOAuthClientRegistration(undefined); + this.registration = undefined; + this.logger.info("Cleared OAuth client registration"); + } +} diff --git a/src/oauth/metadataClient.ts b/src/oauth/metadataClient.ts new file mode 100644 index 00000000..7f3227dc --- /dev/null +++ b/src/oauth/metadataClient.ts @@ -0,0 +1,137 @@ +import type { AxiosInstance } from "axios"; + +import type { Logger } from "../logging/logger"; + +import type { OAuthServerMetadata } from "./types"; + +const OAUTH_DISCOVERY_ENDPOINT = "/.well-known/oauth-authorization-server"; + +const AUTH_GRANT_TYPE = "authorization_code" as const; +const REFRESH_GRANT_TYPE = "refresh_token" as const; +const RESPONSE_TYPE = "code" as const; +const OAUTH_METHOD = "client_secret_post" as const; +const PKCE_CHALLENGE_METHOD = "S256" as const; + +const REQUIRED_GRANT_TYPES = [AUTH_GRANT_TYPE, REFRESH_GRANT_TYPE] as const; + +/** + * Client for discovering and validating OAuth server metadata. + */ +export class OAuthMetadataClient { + constructor( + private readonly axiosInstance: AxiosInstance, + private readonly logger: Logger, + ) {} + + /** + * Check if a server supports OAuth by attempting to fetch the well-known endpoint. + */ + async checkOAuthSupport(): Promise { + try { + await this.axiosInstance.get(OAUTH_DISCOVERY_ENDPOINT); + this.logger.debug("Server supports OAuth"); + return true; + } catch (error) { + this.logger.debug("Server does not support OAuth:", error); + return false; + } + } + + /** + * Fetch and validate OAuth server metadata. + * Throws detailed errors if server doesn't meet OAuth 2.1 requirements. + */ + async getMetadata(): Promise { + this.logger.info("Discovering OAuth endpoints..."); + + const response = await this.axiosInstance.get( + OAUTH_DISCOVERY_ENDPOINT, + ); + + const metadata = response.data; + + this.validateRequiredEndpoints(metadata); + this.validateGrantTypes(metadata); + this.validateResponseTypes(metadata); + this.validateAuthMethods(metadata); + this.validatePKCEMethods(metadata); + + this.logger.debug("OAuth endpoints discovered:", { + authorization: metadata.authorization_endpoint, + token: metadata.token_endpoint, + registration: metadata.registration_endpoint, + revocation: metadata.revocation_endpoint, + }); + + return metadata; + } + + private validateRequiredEndpoints(metadata: OAuthServerMetadata): void { + if ( + !metadata.authorization_endpoint || + !metadata.token_endpoint || + !metadata.issuer + ) { + throw new Error( + "OAuth server metadata missing required endpoints: " + + JSON.stringify(metadata), + ); + } + } + + private validateGrantTypes(metadata: OAuthServerMetadata): void { + if ( + !includesAllTypes(metadata.grant_types_supported, REQUIRED_GRANT_TYPES) + ) { + throw new Error( + `Server does not support required grant types: ${REQUIRED_GRANT_TYPES.join(", ")}. Supported: ${metadata.grant_types_supported?.join(", ") || "none"}`, + ); + } + } + + private validateResponseTypes(metadata: OAuthServerMetadata): void { + if (!includesAllTypes(metadata.response_types_supported, [RESPONSE_TYPE])) { + throw new Error( + `Server does not support required response type: ${RESPONSE_TYPE}. Supported: ${metadata.response_types_supported?.join(", ") || "none"}`, + ); + } + } + + private validateAuthMethods(metadata: OAuthServerMetadata): void { + if ( + !includesAllTypes(metadata.token_endpoint_auth_methods_supported, [ + OAUTH_METHOD, + ]) + ) { + throw new Error( + `Server does not support required auth method: ${OAUTH_METHOD}. Supported: ${metadata.token_endpoint_auth_methods_supported?.join(", ") || "none"}`, + ); + } + } + + private validatePKCEMethods(metadata: OAuthServerMetadata): void { + if ( + !includesAllTypes(metadata.code_challenge_methods_supported, [ + PKCE_CHALLENGE_METHOD, + ]) + ) { + throw new Error( + `Server does not support required PKCE method: ${PKCE_CHALLENGE_METHOD}. Supported: ${metadata.code_challenge_methods_supported?.join(", ") || "none"}`, + ); + } + } +} + +/** + * Check if an array includes all required types. + * If the array is undefined, returns true (server didn't specify, assume all allowed). + */ +function includesAllTypes( + arr: string[] | undefined, + requiredTypes: readonly string[], +): boolean { + if (arr === undefined) { + return true; + } + return requiredTypes.every((type) => arr.includes(type)); +} diff --git a/src/oauth/oauthHelper.ts b/src/oauth/oauthHelper.ts deleted file mode 100644 index efdc9b8b..00000000 --- a/src/oauth/oauthHelper.ts +++ /dev/null @@ -1,747 +0,0 @@ -import * as vscode from "vscode"; - -import { CoderApi } from "../api/coderApi"; -import { - type StoredOAuthTokens, - type SecretsManager, -} from "../core/secretsManager"; - -import { CALLBACK_PATH, generatePKCE, generateState } from "./utils"; - -import type { Logger } from "../logging/logger"; - -import type { - AuthorizationRequestParams, - ClientRegistrationRequest, - ClientRegistrationResponse, - OAuthServerMetadata, - RefreshTokenRequestParams, - TokenRequestParams, - TokenResponse, - TokenRevocationRequest, -} from "./types"; - -const AUTH_GRANT_TYPE = "authorization_code" as const; -const REFRESH_GRANT_TYPE = "refresh_token" as const; -const RESPONSE_TYPE = "code" as const; -const OAUTH_METHOD = "client_secret_post" as const; -const PKCE_CHALLENGE_METHOD = "S256" as const; -const CLIENT_NAME = "VS Code Coder Extension"; - -const REQUIRED_GRANT_TYPES = [AUTH_GRANT_TYPE, REFRESH_GRANT_TYPE] as const; - -// Token refresh timing constants (5 minutes before expiry) -const TOKEN_REFRESH_THRESHOLD_MS = 5 * 60 * 1000; - -/** - * Minimal scopes required by the VS Code extension: - * - workspace:read: List and read workspace details - * - workspace:update: Update workspace version - * - workspace:start: Start stopped workspaces - * - workspace:ssh: SSH configuration for remote connections - * - workspace:application_connect: Connect to workspace agents/apps - * - template:read: Read templates and versions - * - user:read_personal: Read authenticated user info - */ -const DEFAULT_OAUTH_SCOPES = [ - "workspace:read", - "workspace:update", - "workspace:start", - "workspace:ssh", - "workspace:application_connect", - "template:read", - "user:read_personal", -].join(" "); - -export class CoderOAuthHelper implements vscode.Disposable { - private readonly client: CoderApi; - - private clientRegistration: ClientRegistrationResponse | undefined; - private cachedMetadata: OAuthServerMetadata | undefined; - private pendingAuthResolve: - | ((value: { code: string; verifier: string }) => void) - | undefined; - private pendingAuthReject: ((reason: Error) => void) | undefined; - private expectedState: string | undefined; - private pendingVerifier: string | undefined; - private storedTokens: StoredOAuthTokens | undefined; - private refreshTimer: NodeJS.Timeout | undefined; - - private readonly extensionId: string; - - static async create( - baseUrl: string, - secretsManager: SecretsManager, - logger: Logger, - context: vscode.ExtensionContext, - ): Promise { - const helper = new CoderOAuthHelper( - baseUrl, - secretsManager, - logger, - context, - ); - await helper.loadClientRegistration(); - await helper.loadTokens(); - return helper; - } - private constructor( - baseUrl: string, - private readonly secretsManager: SecretsManager, - private readonly logger: Logger, - context: vscode.ExtensionContext, - ) { - this.client = CoderApi.create(baseUrl, undefined, logger); - this.extensionId = context.extension.id; - } - - private async getMetadata(): Promise { - if (this.cachedMetadata) { - return this.cachedMetadata; - } - - this.logger.info("Discovering OAuth endpoints..."); - - const response = await this.client - .getAxiosInstance() - .get("/.well-known/oauth-authorization-server"); - - const metadata = response.data; - - if ( - !metadata.authorization_endpoint || - !metadata.token_endpoint || - !metadata.issuer - ) { - throw new Error( - "OAuth server metadata missing required endpoints: " + - JSON.stringify(metadata), - ); - } - - if ( - !includesAllTypes(metadata.grant_types_supported, REQUIRED_GRANT_TYPES) - ) { - throw new Error( - `Server does not support required grant types: ${REQUIRED_GRANT_TYPES.join(", ")}. Supported: ${metadata.grant_types_supported?.join(", ") || "none"}`, - ); - } - - if (!includesAllTypes(metadata.response_types_supported, [RESPONSE_TYPE])) { - throw new Error( - `Server does not support required response type: ${RESPONSE_TYPE}. Supported: ${metadata.response_types_supported?.join(", ") || "none"}`, - ); - } - - if ( - !includesAllTypes(metadata.token_endpoint_auth_methods_supported, [ - OAUTH_METHOD, - ]) - ) { - throw new Error( - `Server does not support required auth method: ${OAUTH_METHOD}. Supported: ${metadata.token_endpoint_auth_methods_supported?.join(", ") || "none"}`, - ); - } - - if ( - !includesAllTypes(metadata.code_challenge_methods_supported, [ - PKCE_CHALLENGE_METHOD, - ]) - ) { - throw new Error( - `Server does not support required PKCE method: ${PKCE_CHALLENGE_METHOD}. Supported: ${metadata.code_challenge_methods_supported?.join(", ") || "none"}`, - ); - } - - this.cachedMetadata = metadata; - this.logger.debug("OAuth endpoints discovered:", { - authorization: metadata.authorization_endpoint, - token: metadata.token_endpoint, - registration: metadata.registration_endpoint, - revocation: metadata.revocation_endpoint, - }); - - return metadata; - } - - private getRedirectUri(): string { - return `${vscode.env.uriScheme}://${this.extensionId}${CALLBACK_PATH}`; - } - - private async loadClientRegistration(): Promise { - const registration = await this.secretsManager.getOAuthClientRegistration(); - if (registration) { - this.clientRegistration = registration; - this.logger.info("Loaded existing OAuth client:", registration.client_id); - } - } - - private async loadTokens(): Promise { - const tokens = await this.secretsManager.getOAuthTokens(); - if (tokens) { - if (!this.hasRequiredScopes(tokens.scope)) { - this.logger.warn( - "Stored token missing required scopes, clearing tokens", - { - stored_scope: tokens.scope, - required_scopes: DEFAULT_OAUTH_SCOPES, - }, - ); - await this.secretsManager.clearOAuthTokens(); - return; - } - - this.storedTokens = tokens; - this.logger.info("Loaded stored OAuth tokens", { - expires_at: new Date(tokens.expiry_timestamp).toISOString(), - scope: tokens.scope, - }); - this.client.setSessionToken(tokens.access_token); - - if (tokens.refresh_token) { - this.startRefreshTimer(); - } - } - } - - /** - * Check if granted scopes cover all required scopes. - * Supports wildcard scopes like "workspace:*" which grant all "workspace:" prefixed scopes. - */ - private hasRequiredScopes(grantedScope: string | undefined): boolean { - if (!grantedScope) { - return false; - } - - const grantedScopes = new Set(grantedScope.split(" ")); - const requiredScopes = DEFAULT_OAUTH_SCOPES.split(" "); - - for (const required of requiredScopes) { - // Check exact match - if (grantedScopes.has(required)) { - continue; - } - - // Check wildcard match (e.g., "workspace:*" grants "workspace:read") - const colonIndex = required.indexOf(":"); - if (colonIndex !== -1) { - const prefix = required.substring(0, colonIndex); - const wildcard = `${prefix}:*`; - if (grantedScopes.has(wildcard)) { - continue; - } - } - - return false; - } - - return true; - } - - private async saveClientRegistration( - registration: ClientRegistrationResponse, - ): Promise { - await this.secretsManager.setOAuthClientRegistration(registration); - this.clientRegistration = registration; - this.logger.info( - "Saved OAuth client registration:", - registration.client_id, - ); - } - - async clearClientRegistration(): Promise { - await this.secretsManager.setOAuthClientRegistration(undefined); - this.clientRegistration = undefined; - this.logger.info("Cleared OAuth client registration"); - } - - async registerClient(): Promise { - const redirectUri = this.getRedirectUri(); - - if (this.clientRegistration?.client_id) { - const clientId = this.clientRegistration.client_id; - if (this.clientRegistration.redirect_uris.includes(redirectUri)) { - this.logger.info("Using existing client registration:", clientId); - return clientId; - } - this.logger.info("Redirect URI changed, re-registering client"); - } - - const metadata = await this.getMetadata(); - - if (!metadata.registration_endpoint) { - throw new Error( - "Server does not support dynamic client registration (no registration_endpoint in metadata)", - ); - } - - // "web" type since VS Code Secrets API allows secure client_secret storage (confidential client). - const registrationRequest: ClientRegistrationRequest = { - redirect_uris: [redirectUri], - application_type: "web", - grant_types: [AUTH_GRANT_TYPE], - response_types: [RESPONSE_TYPE], - client_name: CLIENT_NAME, - token_endpoint_auth_method: OAUTH_METHOD, - }; - - const response = await this.client - .getAxiosInstance() - .post( - metadata.registration_endpoint, - registrationRequest, - ); - - await this.saveClientRegistration(response.data); - - return response.data.client_id; - } - - private buildAuthorizationUrl( - metadata: OAuthServerMetadata, - clientId: string, - state: string, - challenge: string, - scope: string, - ): string { - if (metadata.scopes_supported) { - const requestedScopes = scope.split(" "); - const unsupportedScopes = requestedScopes.filter( - (s) => !metadata.scopes_supported?.includes(s), - ); - if (unsupportedScopes.length > 0) { - this.logger.warn( - `Requested scopes not in server's supported scopes: ${unsupportedScopes.join(", ")}. Server may still accept them.`, - { supported_scopes: metadata.scopes_supported }, - ); - } - } - - const params: AuthorizationRequestParams = { - client_id: clientId, - response_type: RESPONSE_TYPE, - redirect_uri: this.getRedirectUri(), - scope, - state, - code_challenge: challenge, - code_challenge_method: PKCE_CHALLENGE_METHOD, - }; - - const url = `${metadata.authorization_endpoint}?${new URLSearchParams(params as unknown as Record).toString()}`; - - this.logger.debug("Building OAuth authorization URL:", { - client_id: clientId, - redirect_uri: this.getRedirectUri(), - scope, - }); - - return url; - } - - async startAuthorization( - url: string, - ): Promise<{ code: string; verifier: string }> { - this.client.setHost(url); - - const metadata = await this.getMetadata(); - const clientId = await this.registerClient(); - const state = generateState(); - const { verifier, challenge } = generatePKCE(); - - const authUrl = this.buildAuthorizationUrl( - metadata, - clientId, - state, - challenge, - DEFAULT_OAUTH_SCOPES, - ); - - return new Promise<{ code: string; verifier: string }>( - (resolve, reject) => { - const timeoutMins = 5; - const timeout = setTimeout( - () => { - this.clearPendingAuth(); - reject( - new Error(`OAuth flow timed out after ${timeoutMins} minutes`), - ); - }, - timeoutMins * 60 * 1000, - ); - - const clearPromise = () => { - clearTimeout(timeout); - this.clearPendingAuth(); - }; - - this.pendingAuthResolve = (result) => { - clearPromise(); - resolve(result); - }; - - this.pendingAuthReject = (error) => { - clearPromise(); - reject(error); - }; - - this.expectedState = state; - this.pendingVerifier = verifier; - - vscode.env.openExternal(vscode.Uri.parse(authUrl)).then( - () => {}, - (error) => { - if (error instanceof Error) { - this.pendingAuthReject?.(error); - } else { - this.pendingAuthReject?.(new Error("Failed to open browser")); - } - }, - ); - }, - ); - } - - private clearPendingAuth(): void { - this.pendingAuthResolve = undefined; - this.pendingAuthReject = undefined; - this.expectedState = undefined; - this.pendingVerifier = undefined; - } - - handleCallback( - code: string | null, - state: string | null, - error: string | null, - ): void { - if (!this.pendingAuthResolve || !this.pendingAuthReject) { - this.logger.warn("Received OAuth callback but no pending auth flow"); - return; - } - - if (error) { - this.pendingAuthReject(new Error(`OAuth error: ${error}`)); - return; - } - - if (!code) { - this.pendingAuthReject(new Error("No authorization code received")); - return; - } - - if (!state) { - this.pendingAuthReject(new Error("No state received")); - return; - } - - if (state !== this.expectedState) { - this.pendingAuthReject( - new Error("State mismatch - possible CSRF attack"), - ); - return; - } - - const verifier = this.pendingVerifier; - if (!verifier) { - this.pendingAuthReject(new Error("No PKCE verifier found")); - return; - } - - this.pendingAuthResolve({ code, verifier }); - } - - async exchangeToken(code: string, verifier: string): Promise { - const metadata = await this.getMetadata(); - - if (!this.clientRegistration) { - throw new Error("No client registration found"); - } - - this.logger.info("Exchanging authorization code for token"); - - const params: TokenRequestParams = { - grant_type: AUTH_GRANT_TYPE, - code, - redirect_uri: this.getRedirectUri(), - client_id: this.clientRegistration.client_id, - code_verifier: verifier, - }; - - if (this.clientRegistration.client_secret) { - params.client_secret = this.clientRegistration.client_secret; - } - - const tokenRequest = toUrlSearchParams(params); - - const response = await this.client - .getAxiosInstance() - .post(metadata.token_endpoint, tokenRequest, { - headers: { - "Content-Type": "application/x-www-form-urlencoded", - }, - }); - - this.logger.info("Token exchange successful"); - - await this.saveTokens(response.data); - - return response.data; - } - - getClientId(): string | undefined { - return this.clientRegistration?.client_id; - } - - /** - * Refresh the access token using the stored refresh token. - */ - private async refreshToken(): Promise { - if (!this.storedTokens?.refresh_token) { - throw new Error("No refresh token available"); - } - - if (!this.clientRegistration) { - throw new Error("No client registration found"); - } - - const metadata = await this.getMetadata(); - - this.logger.debug("Refreshing access token"); - - const params: RefreshTokenRequestParams = { - grant_type: REFRESH_GRANT_TYPE, - refresh_token: this.storedTokens.refresh_token, - client_id: this.clientRegistration.client_id, - }; - - if (this.clientRegistration.client_secret) { - params.client_secret = this.clientRegistration.client_secret; - } - - const tokenRequest = toUrlSearchParams(params); - - const response = await this.client - .getAxiosInstance() - .post(metadata.token_endpoint, tokenRequest, { - headers: { - "Content-Type": "application/x-www-form-urlencoded", - }, - }); - - this.logger.debug("Token refresh successful"); - - await this.saveTokens(response.data); - - return response.data; - } - - /** - * Save token response to secrets storage and restart the refresh timer. - */ - private async saveTokens(tokenResponse: TokenResponse): Promise { - const expiryTimestamp = tokenResponse.expires_in - ? Date.now() + tokenResponse.expires_in * 1000 - : Date.now() + 3600 * 1000; // Default to 1 hour if not specified - - const tokens: StoredOAuthTokens = { - ...tokenResponse, - expiry_timestamp: expiryTimestamp, - }; - - this.storedTokens = tokens; - await this.secretsManager.setOAuthTokens(tokens); - await this.secretsManager.setSessionToken(tokenResponse.access_token); - this.client.setSessionToken(tokens.access_token); - - this.logger.info("Tokens saved", { - expires_at: new Date(expiryTimestamp).toISOString(), - }); - - // Restart timer with new expiry (creates self-perpetuating refresh cycle) - this.startRefreshTimer(); - } - - /** - * Start the background token refresh timer. - * Sets a timeout to fire when the token is close to expiry. - */ - private startRefreshTimer(): void { - this.stopRefreshTimer(); - - if (!this.storedTokens?.refresh_token) { - this.logger.debug("No refresh token available, skipping timer setup"); - return; - } - - const now = Date.now(); - const timeUntilRefresh = - this.storedTokens.expiry_timestamp - TOKEN_REFRESH_THRESHOLD_MS - now; - - // If token is already expired or expires very soon, refresh immediately - if (timeUntilRefresh <= 0) { - this.logger.info("Token needs immediate refresh"); - this.refreshToken().catch((error) => { - this.logger.error("Immediate token refresh failed:", error); - }); - return; - } - - // Set timeout to fire exactly when token is 5 minutes from expiry - this.refreshTimer = setTimeout(() => { - this.logger.debug("Token refresh timer fired, refreshing token..."); - this.refreshToken().catch((error) => { - this.logger.error("Scheduled token refresh failed:", error); - }); - }, timeUntilRefresh); - - this.logger.debug("Token refresh timer scheduled", { - fires_at: new Date(now + timeUntilRefresh).toISOString(), - fires_in: timeUntilRefresh / 1000, - }); - } - - /** - * Stop the background token refresh timer. - */ - private stopRefreshTimer(): void { - if (this.refreshTimer) { - clearTimeout(this.refreshTimer); - this.refreshTimer = undefined; - this.logger.debug("Background token refresh timer stopped"); - } - } - - /** - * Revoke a token using the OAuth server's revocation endpoint. - */ - private async revokeToken( - token: string, - tokenTypeHint?: "access_token" | "refresh_token", - ): Promise { - if (!this.clientRegistration) { - throw new Error("No client registration found"); - } - - const metadata = await this.getMetadata(); - - if (!metadata.revocation_endpoint) { - this.logger.warn( - "Server does not support token revocation (no revocation_endpoint)", - ); - return; - } - - this.logger.info("Revoking token", { tokenTypeHint }); - - const params: TokenRevocationRequest = { - token, - client_id: this.clientRegistration.client_id, - }; - - if (tokenTypeHint) { - params.token_type_hint = tokenTypeHint; - } - - if (this.clientRegistration.client_secret) { - params.client_secret = this.clientRegistration.client_secret; - } - - const revocationRequest = toUrlSearchParams(params); - - try { - await this.client - .getAxiosInstance() - .post(metadata.revocation_endpoint, revocationRequest, { - headers: { - "Content-Type": "application/x-www-form-urlencoded", - }, - }); - - this.logger.info("Token revocation successful"); - } catch (error) { - this.logger.error("Token revocation failed:", error); - throw error; - } - } - - /** - * Logout by revoking tokens and clearing all OAuth data. - */ - async logout(): Promise { - this.stopRefreshTimer(); - - // Revoke refresh token (which also invalidates access token per RFC 7009) - if (this.storedTokens?.refresh_token) { - try { - await this.revokeToken( - this.storedTokens.refresh_token, - "refresh_token", - ); - } catch (error) { - this.logger.warn("Token revocation failed during logout:", error); - } - } - - await this.secretsManager.clearOAuthTokens(); - this.storedTokens = undefined; - await this.clearClientRegistration(); - - this.logger.info("OAuth logout complete"); - } - - /** - * Clears all in-memory state and rejects any pending operations. - */ - dispose(): void { - this.stopRefreshTimer(); - - if (this.pendingAuthReject) { - this.pendingAuthReject(new Error("OAuth helper disposed")); - } - this.clearPendingAuth(); - - this.storedTokens = undefined; - this.clientRegistration = undefined; - this.cachedMetadata = undefined; - - this.logger.debug("OAuth helper disposed, all state cleared"); - } -} - -function includesAllTypes( - arr: string[] | undefined, - requiredTypes: readonly string[], -): boolean { - if (arr === undefined) { - // Supported types are not sent by the server so just assume everything is allowed - return true; - } - - return requiredTypes.every((type) => arr.includes(type)); -} - -/** - * Converts an object with string properties to Record, - * filtering out undefined values for use with URLSearchParams. - */ -function toUrlSearchParams(obj: object): URLSearchParams { - const params = Object.fromEntries( - Object.entries(obj).filter( - ([, value]) => value !== undefined && typeof value === "string", - ), - ) as Record; - - return new URLSearchParams(params); -} - -/** - * Activates OAuth support for the Coder extension. - * Initializes and returns the OAuth helper. - */ -export async function activateCoderOAuth( - baseUrl: string, - secretsManager: SecretsManager, - logger: Logger, - context: vscode.ExtensionContext, -): Promise { - return CoderOAuthHelper.create(baseUrl, secretsManager, logger, context); -} diff --git a/src/oauth/sessionManager.ts b/src/oauth/sessionManager.ts new file mode 100644 index 00000000..f478184f --- /dev/null +++ b/src/oauth/sessionManager.ts @@ -0,0 +1,636 @@ +import axios, { type AxiosInstance } from "axios"; +import * as vscode from "vscode"; + +import { OAuthClientRegistry } from "./clientRegistry"; +import { OAuthMetadataClient } from "./metadataClient"; +import { OAuthTokenRefreshScheduler } from "./tokenRefreshScheduler"; +import { + CALLBACK_PATH, + generatePKCE, + generateState, + toUrlSearchParams, +} from "./utils"; + +import type { SecretsManager, StoredOAuthTokens } from "../core/secretsManager"; +import type { Logger } from "../logging/logger"; + +import type { + OAuthServerMetadata, + RefreshTokenRequestParams, + TokenRequestParams, + TokenResponse, + TokenRevocationRequest, +} from "./types"; + +const AUTH_GRANT_TYPE = "authorization_code" as const; +const REFRESH_GRANT_TYPE = "refresh_token" as const; +const RESPONSE_TYPE = "code" as const; +const PKCE_CHALLENGE_METHOD = "S256" as const; + +/** + * Minimal scopes required by the VS Code extension. + */ +const DEFAULT_OAUTH_SCOPES = [ + "workspace:read", + "workspace:update", + "workspace:start", + "workspace:ssh", + "workspace:application_connect", + "template:read", + "user:read_personal", +].join(" "); + +/** + * Manages OAuth session lifecycle for a Coder deployment. + * Coordinates authorization flow, token management, and automatic refresh. + */ +export class OAuthSessionManager implements vscode.Disposable { + private readonly extensionId: string; + private readonly refreshScheduler: OAuthTokenRefreshScheduler; + + private metadataClient: OAuthMetadataClient; + private clientRegistry: OAuthClientRegistry; + + private metadata: OAuthServerMetadata | undefined; + private storedTokens: StoredOAuthTokens | undefined; + + // Pending authorization flow state + private pendingAuthResolve: + | ((value: { code: string; verifier: string }) => void) + | undefined; + private pendingAuthReject: ((reason: Error) => void) | undefined; + private expectedState: string | undefined; + private pendingVerifier: string | undefined; + + /** + * Create and initialize a new OAuth session manager. + */ + static async create( + deploymentUrl: string, + secretsManager: SecretsManager, + logger: Logger, + context: vscode.ExtensionContext, + ): Promise { + const manager = new OAuthSessionManager( + deploymentUrl, + secretsManager, + logger, + context, + ); + await manager.initialize(); + return manager; + } + + private constructor( + private deploymentUrl: string, + private readonly secretsManager: SecretsManager, + private readonly logger: Logger, + context: vscode.ExtensionContext, + ) { + this.extensionId = context.extension.id; + + const axiosInstance = this.createAxiosInstance(); + + this.metadataClient = new OAuthMetadataClient(axiosInstance, logger); + this.clientRegistry = new OAuthClientRegistry( + axiosInstance, + secretsManager, + logger, + ); + this.refreshScheduler = new OAuthTokenRefreshScheduler(async () => { + await this.refreshToken(); + }, logger); + } + + /** + * Create axios instance for the current deployment URL. + */ + private createAxiosInstance(): AxiosInstance { + return axios.create({ + baseURL: this.deploymentUrl, + }); + } + + /** + * Initialize the session manager by loading persisted state. + */ + private async initialize(): Promise { + await this.clientRegistry.load(); + await this.loadTokens(); + } + + /** + * Load stored tokens and start refresh timer if applicable. + * Validates that tokens belong to the current deployment URL. + */ + private async loadTokens(): Promise { + const tokens = await this.secretsManager.getOAuthTokens(); + if (!tokens) { + return; + } + + // Validate URL match (only if we have a deploymentUrl set) + if ( + this.deploymentUrl && + tokens.deployment_url && + tokens.deployment_url !== this.deploymentUrl + ) { + this.logger.warn("Stored tokens for different deployment, clearing", { + stored: tokens.deployment_url, + current: this.deploymentUrl, + }); + await this.clearStaleData(); + return; + } + + if (!this.hasRequiredScopes(tokens.scope)) { + this.logger.warn( + "Stored token missing required scopes, clearing tokens", + { + stored_scope: tokens.scope, + required_scopes: DEFAULT_OAUTH_SCOPES, + }, + ); + await this.secretsManager.setOAuthTokens(undefined); + return; + } + + this.storedTokens = tokens; + this.logger.info("Loaded stored OAuth tokens", { + expires_at: new Date(tokens.expiry_timestamp).toISOString(), + scope: tokens.scope, + deployment: tokens.deployment_url, + }); + + if (tokens.refresh_token) { + this.refreshScheduler.schedule(tokens); + } + } + + /** + * Clear stale data when tokens don't match current deployment. + */ + private async clearStaleData(): Promise { + this.refreshScheduler.stop(); + await this.secretsManager.setOAuthTokens(undefined); + await this.clientRegistry.clear(); + } + + /** + * Clear all state when switching to a new deployment URL. + */ + private async clearForNewUrl(): Promise { + this.refreshScheduler.stop(); + this.metadata = undefined; + this.storedTokens = undefined; + await this.secretsManager.setOAuthTokens(undefined); + await this.clientRegistry.clear(); + } + + /** + * Check if granted scopes cover all required scopes. + * Supports wildcard scopes like "workspace:*". + */ + private hasRequiredScopes(grantedScope: string | undefined): boolean { + if (!grantedScope) { + return false; + } + + const grantedScopes = new Set(grantedScope.split(" ")); + const requiredScopes = DEFAULT_OAUTH_SCOPES.split(" "); + + for (const required of requiredScopes) { + if (grantedScopes.has(required)) { + continue; + } + + // Check wildcard match (e.g., "workspace:*" grants "workspace:read") + const colonIndex = required.indexOf(":"); + if (colonIndex !== -1) { + const prefix = required.substring(0, colonIndex); + const wildcard = `${prefix}:*`; + if (grantedScopes.has(wildcard)) { + continue; + } + } + + return false; + } + + return true; + } + + /** + * Get the redirect URI for OAuth callbacks. + */ + private getRedirectUri(): string { + return `${vscode.env.uriScheme}://${this.extensionId}${CALLBACK_PATH}`; + } + + /** + * Get OAuth server metadata, fetching if not already cached. + */ + private async getMetadata(): Promise { + this.metadata ??= await this.metadataClient.getMetadata(); + return this.metadata; + } + + /** + * Build authorization URL with all required OAuth 2.1 parameters. + */ + private buildAuthorizationUrl( + metadata: OAuthServerMetadata, + clientId: string, + state: string, + challenge: string, + ): string { + if (metadata.scopes_supported) { + const requestedScopes = DEFAULT_OAUTH_SCOPES.split(" "); + const unsupportedScopes = requestedScopes.filter( + (s) => !metadata.scopes_supported?.includes(s), + ); + if (unsupportedScopes.length > 0) { + this.logger.warn( + `Requested scopes not in server's supported scopes: ${unsupportedScopes.join(", ")}. Server may still accept them.`, + { supported_scopes: metadata.scopes_supported }, + ); + } + } + + const params = new URLSearchParams({ + client_id: clientId, + response_type: RESPONSE_TYPE, + redirect_uri: this.getRedirectUri(), + scope: DEFAULT_OAUTH_SCOPES, + state, + code_challenge: challenge, + code_challenge_method: PKCE_CHALLENGE_METHOD, + }); + + const url = `${metadata.authorization_endpoint}?${params.toString()}`; + + this.logger.debug("Built OAuth authorization URL:", { + client_id: clientId, + redirect_uri: this.getRedirectUri(), + scope: DEFAULT_OAUTH_SCOPES, + }); + + return url; + } + + /** + * Start OAuth authorization flow. + * Opens browser for user authentication and waits for callback. + * Returns authorization code and PKCE verifier on success. + * + * @param url Coder deployment URL to authenticate against + */ + async startAuthorization( + url: string, + ): Promise<{ code: string; verifier: string }> { + if (this.deploymentUrl !== url) { + this.logger.info("Deployment URL changed, clearing cached state", { + old: this.deploymentUrl, + new: url, + }); + await this.clearForNewUrl(); + this.deploymentUrl = url; + + // Recreate components with new axios instance for new URL + const axiosInstance = this.createAxiosInstance(); + this.metadataClient = new OAuthMetadataClient(axiosInstance, this.logger); + this.clientRegistry = new OAuthClientRegistry( + axiosInstance, + this.secretsManager, + this.logger, + ); + } + + // Clear cached metadata (may be stale) + this.metadata = undefined; + + const metadata = await this.getMetadata(); + const registration = await this.clientRegistry.register( + metadata, + this.getRedirectUri(), + ); + const state = generateState(); + const { verifier, challenge } = generatePKCE(); + + const authUrl = this.buildAuthorizationUrl( + metadata, + registration.client_id, + state, + challenge, + ); + + return new Promise<{ code: string; verifier: string }>( + (resolve, reject) => { + const timeoutMins = 5; + const timeout = setTimeout( + () => { + this.clearPendingAuth(); + reject( + new Error(`OAuth flow timed out after ${timeoutMins} minutes`), + ); + }, + timeoutMins * 60 * 1000, + ); + + const clearPromise = () => { + clearTimeout(timeout); + this.clearPendingAuth(); + }; + + this.pendingAuthResolve = (result) => { + clearPromise(); + resolve(result); + }; + + this.pendingAuthReject = (error) => { + clearPromise(); + reject(error); + }; + + this.expectedState = state; + this.pendingVerifier = verifier; + + vscode.env.openExternal(vscode.Uri.parse(authUrl)).then( + () => {}, + (error) => { + if (error instanceof Error) { + this.pendingAuthReject?.(error); + } else { + this.pendingAuthReject?.(new Error("Failed to open browser")); + } + }, + ); + }, + ); + } + + /** + * Clear pending authorization flow state. + */ + private clearPendingAuth(): void { + this.pendingAuthResolve = undefined; + this.pendingAuthReject = undefined; + this.expectedState = undefined; + this.pendingVerifier = undefined; + } + + /** + * Handle OAuth callback from browser redirect. + * Validates state and resolves pending authorization promise. + */ + handleCallback( + code: string | null, + state: string | null, + error: string | null, + ): void { + if (!this.pendingAuthResolve || !this.pendingAuthReject) { + this.logger.warn("Received OAuth callback but no pending auth flow"); + return; + } + + if (error) { + this.pendingAuthReject(new Error(`OAuth error: ${error}`)); + return; + } + + if (!code) { + this.pendingAuthReject(new Error("No authorization code received")); + return; + } + + if (!state) { + this.pendingAuthReject(new Error("No state received")); + return; + } + + if (state !== this.expectedState) { + this.pendingAuthReject( + new Error("State mismatch - possible CSRF attack"), + ); + return; + } + + const verifier = this.pendingVerifier; + if (!verifier) { + this.pendingAuthReject(new Error("No PKCE verifier found")); + return; + } + + this.pendingAuthResolve({ code, verifier }); + } + + /** + * Exchange authorization code for access token. + */ + async exchangeToken(code: string, verifier: string): Promise { + const metadata = await this.getMetadata(); + const registration = this.clientRegistry.get(); + + if (!registration) { + throw new Error("No client registration found"); + } + + this.logger.info("Exchanging authorization code for token"); + + const params: TokenRequestParams = { + grant_type: AUTH_GRANT_TYPE, + code, + redirect_uri: this.getRedirectUri(), + client_id: registration.client_id, + client_secret: registration.client_secret, + code_verifier: verifier, + }; + + const tokenRequest = toUrlSearchParams(params); + + const axiosInstance = this.createAxiosInstance(); + const response = await axiosInstance.post( + metadata.token_endpoint, + tokenRequest, + { + headers: { + "Content-Type": "application/x-www-form-urlencoded", + }, + }, + ); + + this.logger.info("Token exchange successful"); + + await this.saveTokens(response.data); + + return response.data; + } + + /** + * Refresh the access token using the stored refresh token. + */ + private async refreshToken(): Promise { + if (!this.storedTokens?.refresh_token) { + throw new Error("No refresh token available"); + } + + const registration = this.clientRegistry.get(); + if (!registration) { + throw new Error("No client registration found"); + } + + const metadata = await this.getMetadata(); + + this.logger.debug("Refreshing access token"); + + const params: RefreshTokenRequestParams = { + grant_type: REFRESH_GRANT_TYPE, + refresh_token: this.storedTokens.refresh_token, + client_id: registration.client_id, + client_secret: registration.client_secret, + }; + + const tokenRequest = toUrlSearchParams(params); + + const axiosInstance = this.createAxiosInstance(); + const response = await axiosInstance.post( + metadata.token_endpoint, + tokenRequest, + { + headers: { + "Content-Type": "application/x-www-form-urlencoded", + }, + }, + ); + + this.logger.debug("Token refresh successful"); + + await this.saveTokens(response.data); + + return response.data; + } + + /** + * Save token response to storage and schedule automatic refresh. + * Also triggers event via secretsManager to update global client. + */ + private async saveTokens(tokenResponse: TokenResponse): Promise { + const expiryTimestamp = tokenResponse.expires_in + ? Date.now() + tokenResponse.expires_in * 1000 + : Date.now() + 3600 * 1000; // TODO Default to 1 hour + + const tokens: StoredOAuthTokens = { + ...tokenResponse, + deployment_url: this.deploymentUrl, + expiry_timestamp: expiryTimestamp, + }; + + this.storedTokens = tokens; + await this.secretsManager.setOAuthTokens(tokens); + + // Trigger event to update global client (works for login & background refresh) + // TODO Add a setting to check if we have OAuth or token setup so we can start the background refresh + await this.secretsManager.setSessionToken(tokenResponse.access_token); + + this.logger.info("Tokens saved", { + expires_at: new Date(expiryTimestamp).toISOString(), + deployment: this.deploymentUrl, + }); + + // Schedule automatic refresh + this.refreshScheduler.schedule(tokens); + } + + /** + * Revoke a token using the OAuth server's revocation endpoint. + */ + private async revokeToken(token: string): Promise { + const registration = this.clientRegistry.get(); + if (!registration) { + throw new Error("No client registration found"); + } + + const metadata = await this.getMetadata(); + + if (!metadata.revocation_endpoint) { + this.logger.warn( + "Server does not support token revocation (no revocation_endpoint)", + ); + return; + } + + this.logger.info("Revoking refresh token"); + + const params: TokenRevocationRequest = { + token, + client_id: registration.client_id, + client_secret: registration.client_secret, + token_type_hint: "refresh_token", + }; + + const revocationRequest = toUrlSearchParams(params); + + try { + const axiosInstance = this.createAxiosInstance(); + await axiosInstance.post( + metadata.revocation_endpoint, + revocationRequest, + { + headers: { + "Content-Type": "application/x-www-form-urlencoded", + }, + }, + ); + + this.logger.info("Token revocation successful"); + } catch (error) { + this.logger.error("Token revocation failed:", error); + throw error; + } + } + + /** + * Logout by revoking tokens and clearing all OAuth data. + */ + async logout(): Promise { + this.refreshScheduler.stop(); + + // Revoke refresh token (which also invalidates access token per RFC 7009) + if (this.storedTokens?.refresh_token) { + try { + await this.revokeToken(this.storedTokens.refresh_token); + } catch (error) { + this.logger.warn("Token revocation failed during logout:", error); + } + } + + await this.secretsManager.setOAuthTokens(undefined); + this.storedTokens = undefined; + await this.clientRegistry.clear(); + + this.logger.info("OAuth logout complete"); + } + + /** + * Get the client ID if registered. + */ + getClientId(): string | undefined { + return this.clientRegistry.get()?.client_id; + } + + /** + * Clears all in-memory state and rejects any pending operations. + */ + dispose(): void { + this.refreshScheduler.stop(); + + if (this.pendingAuthReject) { + this.pendingAuthReject(new Error("OAuth session manager disposed")); + } + this.clearPendingAuth(); + this.storedTokens = undefined; + this.metadata = undefined; + + this.logger.debug("OAuth session manager disposed"); + } +} diff --git a/src/oauth/tokenRefreshScheduler.ts b/src/oauth/tokenRefreshScheduler.ts new file mode 100644 index 00000000..3eeabb9e --- /dev/null +++ b/src/oauth/tokenRefreshScheduler.ts @@ -0,0 +1,65 @@ +import type { StoredOAuthTokens } from "../core/secretsManager"; +import type { Logger } from "../logging/logger"; + +// Token refresh timing constant +const TOKEN_REFRESH_THRESHOLD_MS = 20 * 60 * 1000; + +/** + * Manages automatic token refresh scheduling. + * Calculates optimal refresh timing and triggers refresh callbacks. + */ +export class OAuthTokenRefreshScheduler { + private refreshTimer: NodeJS.Timeout | undefined; + + constructor( + private readonly refreshCallback: () => Promise, + private readonly logger: Logger, + ) {} + + /** + * Schedule automatic token refresh based on token expiry. + */ + schedule(tokens: StoredOAuthTokens): void { + this.stop(); + + if (!tokens.refresh_token) { + this.logger.debug("No refresh token available, skipping timer setup"); + return; + } + + const now = Date.now(); + const timeUntilRefresh = + tokens.expiry_timestamp - TOKEN_REFRESH_THRESHOLD_MS - now; + + if (timeUntilRefresh <= 0) { + this.logger.info("Token needs immediate refresh"); + this.refreshCallback().catch((error) => { + this.logger.error("Immediate token refresh failed:", error); + }); + return; + } + + this.refreshTimer = setTimeout(() => { + this.logger.debug("Token refresh timer fired, refreshing token..."); + this.refreshCallback().catch((error) => { + this.logger.error("Scheduled token refresh failed:", error); + }); + }, timeUntilRefresh); + + this.logger.debug("Token refresh timer scheduled", { + fires_at: new Date(now + timeUntilRefresh).toISOString(), + fires_in_seconds: Math.round(timeUntilRefresh / 1000), + }); + } + + /** + * Stop the background token refresh timer. + */ + stop(): void { + if (this.refreshTimer) { + clearTimeout(this.refreshTimer); + this.refreshTimer = undefined; + this.logger.debug("Token refresh timer stopped"); + } + } +} diff --git a/src/oauth/utils.ts b/src/oauth/utils.ts index 7d66a139..61beeb50 100644 --- a/src/oauth/utils.ts +++ b/src/oauth/utils.ts @@ -26,3 +26,17 @@ export function generatePKCE(): PKCEChallenge { export function generateState(): string { return randomBytes(16).toString("base64url"); } + +/** + * Converts an object with string properties to URLSearchParams, + * filtering out undefined values for use with OAuth requests. + */ +export function toUrlSearchParams(obj: object): URLSearchParams { + const params = Object.fromEntries( + Object.entries(obj).filter( + ([, value]) => value !== undefined && typeof value === "string", + ), + ) as Record; + + return new URLSearchParams(params); +} diff --git a/src/remote/remote.ts b/src/remote/remote.ts index 4193e46d..6af22390 100644 --- a/src/remote/remote.ts +++ b/src/remote/remote.ts @@ -148,14 +148,14 @@ export class Remote { if (result.type === "login") { return this.setup(remoteAuthority, firstConnect, remoteSshExtensionId); - } else if (!result.userChoice) { - // User declined to log in. - await this.closeRemote(); - return; - } else { + } else if (result.userChoice === "Log In") { // Log in then try again. await this.commands.login({ url: baseUrlRaw, label: parts.label }); return this.setup(remoteAuthority, firstConnect, remoteSshExtensionId); + } else { + // User declined to log in. + await this.closeRemote(); + return; } }; From ac4784555fd3bab237e64f5b3bb6828cf72c84cb Mon Sep 17 00:00:00 2001 From: Ehab Younes Date: Mon, 27 Oct 2025 12:37:25 +0300 Subject: [PATCH 09/20] Simplify the OAuth flow --- src/commands.ts | 55 +++---- src/extension.ts | 5 +- src/oauth/clientRegistry.ts | 111 -------------- src/oauth/sessionManager.ts | 287 +++++++++++++++++++++--------------- 4 files changed, 191 insertions(+), 267 deletions(-) delete mode 100644 src/oauth/clientRegistry.ts diff --git a/src/commands.ts b/src/commands.ts index 4e16e617..2e43fbbc 100644 --- a/src/commands.ts +++ b/src/commands.ts @@ -183,13 +183,14 @@ export class Commands { } if (choice === "oauth") { - return this.loginWithOAuth(url, client); + return this.loginWithOAuth(client); } else if (choice === "legacy") { - return this.loginWithToken(url, token, client); + const initialToken = + token || (await this.secretsManager.getSessionToken()); + return this.loginWithToken(client, initialToken); } - // User aborted. - return null; + return null; // User aborted. } private async checkOAuthSupport(client: CoderApi): Promise { @@ -228,10 +229,13 @@ export class Commands { } private async loginWithToken( - url: string, - token: string | undefined, client: CoderApi, + initialToken: string | undefined, ): Promise<{ user: User; token: string } | null> { + const url = client.getAxiosInstance().defaults.baseURL; + if (!url) { + throw new Error("No base URL set on REST client"); + } // This prompt is for convenience; do not error if they close it since // they may already have a token or already have the page opened. await vscode.env.openExternal(vscode.Uri.parse(`${url}/cli-auth`)); @@ -244,7 +248,7 @@ export class Commands { title: "Coder API Key", password: true, placeHolder: "Paste your API key.", - value: token || (await this.secretsManager.getSessionToken()), + value: initialToken, ignoreFocusOut: true, validateInput: async (value) => { if (!value) { @@ -288,29 +292,17 @@ export class Commands { * Returns the access token and authenticated user, or null if failed/cancelled. */ private async loginWithOAuth( - url: string, client: CoderApi, ): Promise<{ user: User; token: string } | null> { try { this.logger.info("Starting OAuth authentication"); - // Start OAuth authorization flow - // TODO just pass the client here and do all the neccessary steps (If we are already logged in we'd have the right token and the OAuth client registration saved). - const { code, verifier } = - await this.oauthSessionManager.startAuthorization(url); - - // Exchange authorization code for tokens - const tokenResponse = await this.oauthSessionManager.exchangeToken( - code, - verifier, - ); + const tokenResponse = await this.oauthSessionManager.login(client); // Validate token by fetching user client.setSessionToken(tokenResponse.access_token); const user = await client.getAuthenticatedUser(); - this.logger.info("OAuth authentication successful"); - return { token: tokenResponse.access_token, user, @@ -359,9 +351,19 @@ export class Commands { throw new Error("You are not logged in"); } + await this.forceLogout(); + } + + public async forceLogout(): Promise { + if (!this.contextManager.get("coder.authenticated")) { + return; + } + this.logger.info("Logging out"); + // Check if using OAuth - const hasOAuthTokens = await this.secretsManager.getOAuthTokens(); - if (hasOAuthTokens) { + const isOAuthLoggedIn = + await this.oauthSessionManager.isLoggedInWithOAuth(); + if (isOAuthLoggedIn) { this.logger.info("Logging out via OAuth"); try { await this.oauthSessionManager.logout(); @@ -373,15 +375,6 @@ export class Commands { } } - // Continue with standard logout (clears sessionToken, contexts, etc) - await this.forceLogout(); - } - - public async forceLogout(): Promise { - if (!this.contextManager.get("coder.authenticated")) { - return; - } - this.logger.info("Logging out"); // Clear from the REST client. An empty url will indicate to other parts of // the code that we are logged out. this.restClient.setHost(""); diff --git a/src/extension.ts b/src/extension.ts index 38823378..61f9afd9 100644 --- a/src/extension.ts +++ b/src/extension.ts @@ -129,19 +129,18 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { // Listen for session token changes and sync state across all components ctx.subscriptions.push( secretsManager.onDidChangeSessionToken(async (token) => { + client.setSessionToken(token ?? ""); if (!token) { output.debug("Session token cleared"); - client.setSessionToken(""); return; } output.debug("Session token changed, syncing state"); - client.setSessionToken(token); const url = mementoManager.getUrl(); if (url) { const cliManager = serviceContainer.getCliManager(); - // TODO label might not match? + // TODO label might not match the one in remote? await cliManager.configure(toSafeHost(url), url, token); output.debug("Updated CLI config with new token"); } diff --git a/src/oauth/clientRegistry.ts b/src/oauth/clientRegistry.ts deleted file mode 100644 index 91b4949f..00000000 --- a/src/oauth/clientRegistry.ts +++ /dev/null @@ -1,111 +0,0 @@ -import type { AxiosInstance } from "axios"; - -import type { SecretsManager } from "../core/secretsManager"; -import type { Logger } from "../logging/logger"; - -import type { - ClientRegistrationRequest, - ClientRegistrationResponse, - OAuthServerMetadata, -} from "./types"; - -const AUTH_GRANT_TYPE = "authorization_code" as const; -const RESPONSE_TYPE = "code" as const; -const OAUTH_METHOD = "client_secret_post" as const; -const CLIENT_NAME = "VS Code Coder Extension"; - -/** - * Manages OAuth client registration and persistence. - */ -export class OAuthClientRegistry { - private registration: ClientRegistrationResponse | undefined; - - constructor( - private readonly axiosInstance: AxiosInstance, - private readonly secretsManager: SecretsManager, - private readonly logger: Logger, - ) {} - - /** - * Load existing client registration from secure storage. - * Should be called during initialization. - */ - async load(): Promise { - const registration = await this.secretsManager.getOAuthClientRegistration(); - if (registration) { - this.registration = registration; - this.logger.info("Loaded existing OAuth client:", registration.client_id); - } - } - - /** - * Get the current client registration if one exists. - */ - get(): ClientRegistrationResponse | undefined { - return this.registration; - } - - /** - * Register a new OAuth client or return existing if still valid. - * Re-registers if redirect URI has changed. - */ - async register( - metadata: OAuthServerMetadata, - redirectUri: string, - ): Promise { - if (this.registration?.client_id) { - if (this.registration.redirect_uris.includes(redirectUri)) { - this.logger.info( - "Using existing client registration:", - this.registration.client_id, - ); - return this.registration; - } - this.logger.info("Redirect URI changed, re-registering client"); - } - - if (!metadata.registration_endpoint) { - throw new Error("Server does not support dynamic client registration"); - } - - // "web" type since VS Code Secrets API allows secure client_secret storage (confidential client) - const registrationRequest: ClientRegistrationRequest = { - redirect_uris: [redirectUri], - application_type: "web", - grant_types: [AUTH_GRANT_TYPE], - response_types: [RESPONSE_TYPE], - client_name: CLIENT_NAME, - token_endpoint_auth_method: OAUTH_METHOD, - }; - - const response = await this.axiosInstance.post( - metadata.registration_endpoint, - registrationRequest, - ); - - await this.save(response.data); - - return response.data; - } - - /** - * Save client registration to secure storage. - */ - private async save(registration: ClientRegistrationResponse): Promise { - await this.secretsManager.setOAuthClientRegistration(registration); - this.registration = registration; - this.logger.info( - "Saved OAuth client registration:", - registration.client_id, - ); - } - - /** - * Clear the current client registration from memory and storage. - */ - async clear(): Promise { - await this.secretsManager.setOAuthClientRegistration(undefined); - this.registration = undefined; - this.logger.info("Cleared OAuth client registration"); - } -} diff --git a/src/oauth/sessionManager.ts b/src/oauth/sessionManager.ts index f478184f..e41d5112 100644 --- a/src/oauth/sessionManager.ts +++ b/src/oauth/sessionManager.ts @@ -1,7 +1,8 @@ -import axios, { type AxiosInstance } from "axios"; +import { type AxiosInstance } from "axios"; import * as vscode from "vscode"; -import { OAuthClientRegistry } from "./clientRegistry"; +import { CoderApi } from "../api/coderApi"; + import { OAuthMetadataClient } from "./metadataClient"; import { OAuthTokenRefreshScheduler } from "./tokenRefreshScheduler"; import { @@ -15,6 +16,8 @@ import type { SecretsManager, StoredOAuthTokens } from "../core/secretsManager"; import type { Logger } from "../logging/logger"; import type { + ClientRegistrationRequest, + ClientRegistrationResponse, OAuthServerMetadata, RefreshTokenRequestParams, TokenRequestParams, @@ -48,10 +51,6 @@ export class OAuthSessionManager implements vscode.Disposable { private readonly extensionId: string; private readonly refreshScheduler: OAuthTokenRefreshScheduler; - private metadataClient: OAuthMetadataClient; - private clientRegistry: OAuthClientRegistry; - - private metadata: OAuthServerMetadata | undefined; private storedTokens: StoredOAuthTokens | undefined; // Pending authorization flow state @@ -77,7 +76,7 @@ export class OAuthSessionManager implements vscode.Disposable { logger, context, ); - await manager.initialize(); + await manager.loadTokens(); return manager; } @@ -88,37 +87,11 @@ export class OAuthSessionManager implements vscode.Disposable { context: vscode.ExtensionContext, ) { this.extensionId = context.extension.id; - - const axiosInstance = this.createAxiosInstance(); - - this.metadataClient = new OAuthMetadataClient(axiosInstance, logger); - this.clientRegistry = new OAuthClientRegistry( - axiosInstance, - secretsManager, - logger, - ); this.refreshScheduler = new OAuthTokenRefreshScheduler(async () => { await this.refreshToken(); }, logger); } - /** - * Create axios instance for the current deployment URL. - */ - private createAxiosInstance(): AxiosInstance { - return axios.create({ - baseURL: this.deploymentUrl, - }); - } - - /** - * Initialize the session manager by loading persisted state. - */ - private async initialize(): Promise { - await this.clientRegistry.load(); - await this.loadTokens(); - } - /** * Load stored tokens and start refresh timer if applicable. * Validates that tokens belong to the current deployment URL. @@ -129,19 +102,15 @@ export class OAuthSessionManager implements vscode.Disposable { return; } - // Validate URL match (only if we have a deploymentUrl set) - if ( - this.deploymentUrl && - tokens.deployment_url && - tokens.deployment_url !== this.deploymentUrl - ) { + if (this.deploymentUrl && tokens.deployment_url !== this.deploymentUrl) { this.logger.warn("Stored tokens for different deployment, clearing", { stored: tokens.deployment_url, current: this.deploymentUrl, }); - await this.clearStaleData(); + await this.clearTokenState(); return; } + this.deploymentUrl = tokens.deployment_url; if (!this.hasRequiredScopes(tokens.scope)) { this.logger.warn( @@ -156,11 +125,7 @@ export class OAuthSessionManager implements vscode.Disposable { } this.storedTokens = tokens; - this.logger.info("Loaded stored OAuth tokens", { - expires_at: new Date(tokens.expiry_timestamp).toISOString(), - scope: tokens.scope, - deployment: tokens.deployment_url, - }); + this.logger.info(`Loaded stored OAuth tokens for ${tokens.deployment_url}`); if (tokens.refresh_token) { this.refreshScheduler.schedule(tokens); @@ -170,21 +135,11 @@ export class OAuthSessionManager implements vscode.Disposable { /** * Clear stale data when tokens don't match current deployment. */ - private async clearStaleData(): Promise { + private async clearTokenState(): Promise { this.refreshScheduler.stop(); - await this.secretsManager.setOAuthTokens(undefined); - await this.clientRegistry.clear(); - } - - /** - * Clear all state when switching to a new deployment URL. - */ - private async clearForNewUrl(): Promise { - this.refreshScheduler.stop(); - this.metadata = undefined; this.storedTokens = undefined; await this.secretsManager.setOAuthTokens(undefined); - await this.clientRegistry.clear(); + await this.secretsManager.setOAuthClientRegistration(undefined); } /** @@ -193,7 +148,8 @@ export class OAuthSessionManager implements vscode.Disposable { */ private hasRequiredScopes(grantedScope: string | undefined): boolean { if (!grantedScope) { - return false; + // TODO server always returns empty scopes + return true; } const grantedScopes = new Set(grantedScope.split(" ")); @@ -228,11 +184,126 @@ export class OAuthSessionManager implements vscode.Disposable { } /** - * Get OAuth server metadata, fetching if not already cached. + * Prepare common OAuth operation setup: CoderApi, metadata, and registration. + * Used by refresh and revoke operations to reduce duplication. + */ + private async prepareOAuthOperation( + deploymentUrl: string, + token?: string, + ): Promise<{ + axiosInstance: AxiosInstance; + metadata: OAuthServerMetadata; + registration: ClientRegistrationResponse; + }> { + const client = CoderApi.create(deploymentUrl, token, this.logger); + const axiosInstance = client.getAxiosInstance(); + + const metadataClient = new OAuthMetadataClient(axiosInstance, this.logger); + const metadata = await metadataClient.getMetadata(); + + const registration = await this.secretsManager.getOAuthClientRegistration(); + if (!registration) { + throw new Error("No client registration found"); + } + + return { axiosInstance, metadata, registration }; + } + + /** + * Register OAuth client or return existing if still valid. + * Re-registers if redirect URI has changed. + */ + private async registerClient( + axiosInstance: AxiosInstance, + metadata: OAuthServerMetadata, + ): Promise { + const redirectUri = this.getRedirectUri(); + + const existing = await this.secretsManager.getOAuthClientRegistration(); + if (existing?.client_id) { + if (existing.redirect_uris.includes(redirectUri)) { + this.logger.info( + "Using existing client registration:", + existing.client_id, + ); + return existing; + } + this.logger.info("Redirect URI changed, re-registering client"); + } + + if (!metadata.registration_endpoint) { + throw new Error("Server does not support dynamic client registration"); + } + + const registrationRequest: ClientRegistrationRequest = { + redirect_uris: [redirectUri], + application_type: "web", + grant_types: ["authorization_code"], + response_types: ["code"], + client_name: "VS Code Coder Extension", + token_endpoint_auth_method: "client_secret_post", + }; + + const response = await axiosInstance.post( + metadata.registration_endpoint, + registrationRequest, + ); + + await this.secretsManager.setOAuthClientRegistration(response.data); + this.logger.info( + "Saved OAuth client registration:", + response.data.client_id, + ); + + return response.data; + } + + /** + * Simplified OAuth login flow that handles the entire process. + * Fetches metadata, registers client, starts authorization, and exchanges tokens. + * + * @param client CoderApi instance for the deployment to authenticate against + * @returns TokenResponse containing access token and optional refresh token */ - private async getMetadata(): Promise { - this.metadata ??= await this.metadataClient.getMetadata(); - return this.metadata; + async login(client: CoderApi): Promise { + const baseUrl = client.getAxiosInstance().defaults.baseURL; + if (!baseUrl) { + throw new Error("CoderApi instance has no base URL set"); + } + if (this.deploymentUrl !== baseUrl) { + this.logger.info("Deployment URL changed, clearing cached state", { + old: this.deploymentUrl, + new: baseUrl, + }); + await this.clearTokenState(); + this.deploymentUrl = baseUrl; + } + + this.logger.info("Starting OAuth login flow"); + + const axiosInstance = client.getAxiosInstance(); + const metadataClient = new OAuthMetadataClient(axiosInstance, this.logger); + const metadata = await metadataClient.getMetadata(); + + // Only register the client on login + const registration = await this.registerClient(axiosInstance, metadata); + + const { code, verifier } = await this.startAuthorization( + metadata, + registration, + ); + + const tokenResponse = await this.exchangeToken( + code, + verifier, + axiosInstance, + metadata, + registration, + ); + + this.logger.info("OAuth login flow completed successfully"); + + return tokenResponse; } /** @@ -282,38 +353,11 @@ export class OAuthSessionManager implements vscode.Disposable { * Start OAuth authorization flow. * Opens browser for user authentication and waits for callback. * Returns authorization code and PKCE verifier on success. - * - * @param url Coder deployment URL to authenticate against */ - async startAuthorization( - url: string, + private async startAuthorization( + metadata: OAuthServerMetadata, + registration: ClientRegistrationResponse, ): Promise<{ code: string; verifier: string }> { - if (this.deploymentUrl !== url) { - this.logger.info("Deployment URL changed, clearing cached state", { - old: this.deploymentUrl, - new: url, - }); - await this.clearForNewUrl(); - this.deploymentUrl = url; - - // Recreate components with new axios instance for new URL - const axiosInstance = this.createAxiosInstance(); - this.metadataClient = new OAuthMetadataClient(axiosInstance, this.logger); - this.clientRegistry = new OAuthClientRegistry( - axiosInstance, - this.secretsManager, - this.logger, - ); - } - - // Clear cached metadata (may be stale) - this.metadata = undefined; - - const metadata = await this.getMetadata(); - const registration = await this.clientRegistry.register( - metadata, - this.getRedirectUri(), - ); const state = generateState(); const { verifier, challenge } = generatePKCE(); @@ -382,6 +426,8 @@ export class OAuthSessionManager implements vscode.Disposable { /** * Handle OAuth callback from browser redirect. * Validates state and resolves pending authorization promise. + * + * // TODO this has to work across windows! */ handleCallback( code: string | null, @@ -427,14 +473,13 @@ export class OAuthSessionManager implements vscode.Disposable { /** * Exchange authorization code for access token. */ - async exchangeToken(code: string, verifier: string): Promise { - const metadata = await this.getMetadata(); - const registration = this.clientRegistry.get(); - - if (!registration) { - throw new Error("No client registration found"); - } - + private async exchangeToken( + code: string, + verifier: string, + axiosInstance: AxiosInstance, + metadata: OAuthServerMetadata, + registration: ClientRegistrationResponse, + ): Promise { this.logger.info("Exchanging authorization code for token"); const params: TokenRequestParams = { @@ -448,7 +493,6 @@ export class OAuthSessionManager implements vscode.Disposable { const tokenRequest = toUrlSearchParams(params); - const axiosInstance = this.createAxiosInstance(); const response = await axiosInstance.post( metadata.token_endpoint, tokenRequest, @@ -474,12 +518,11 @@ export class OAuthSessionManager implements vscode.Disposable { throw new Error("No refresh token available"); } - const registration = this.clientRegistry.get(); - if (!registration) { - throw new Error("No client registration found"); - } - - const metadata = await this.getMetadata(); + const { axiosInstance, metadata, registration } = + await this.prepareOAuthOperation( + this.deploymentUrl, + this.storedTokens.access_token, + ); this.logger.debug("Refreshing access token"); @@ -492,7 +535,6 @@ export class OAuthSessionManager implements vscode.Disposable { const tokenRequest = toUrlSearchParams(params); - const axiosInstance = this.createAxiosInstance(); const response = await axiosInstance.post( metadata.token_endpoint, tokenRequest, @@ -545,12 +587,11 @@ export class OAuthSessionManager implements vscode.Disposable { * Revoke a token using the OAuth server's revocation endpoint. */ private async revokeToken(token: string): Promise { - const registration = this.clientRegistry.get(); - if (!registration) { - throw new Error("No client registration found"); - } - - const metadata = await this.getMetadata(); + const { axiosInstance, metadata, registration } = + await this.prepareOAuthOperation( + this.deploymentUrl, + this.storedTokens?.access_token, + ); if (!metadata.revocation_endpoint) { this.logger.warn( @@ -571,7 +612,6 @@ export class OAuthSessionManager implements vscode.Disposable { const revocationRequest = toUrlSearchParams(params); try { - const axiosInstance = this.createAxiosInstance(); await axiosInstance.post( metadata.revocation_endpoint, revocationRequest, @@ -604,18 +644,22 @@ export class OAuthSessionManager implements vscode.Disposable { } } - await this.secretsManager.setOAuthTokens(undefined); - this.storedTokens = undefined; - await this.clientRegistry.clear(); + await this.clearTokenState(); this.logger.info("OAuth logout complete"); } /** - * Get the client ID if registered. + * Check if currently logged in with OAuth. + * Returns true only if valid OAuth tokens exist for the current deployment. */ - getClientId(): string | undefined { - return this.clientRegistry.get()?.client_id; + async isLoggedInWithOAuth(): Promise { + const tokens = await this.secretsManager.getOAuthTokens(); + if (!tokens) { + return false; + } + + return this.deploymentUrl === tokens.deployment_url; } /** @@ -629,7 +673,6 @@ export class OAuthSessionManager implements vscode.Disposable { } this.clearPendingAuth(); this.storedTokens = undefined; - this.metadata = undefined; this.logger.debug("OAuth session manager disposed"); } From a5a409516731e4f26a2d9d8cabcde08f84497fb1 Mon Sep 17 00:00:00 2001 From: Ehab Younes Date: Tue, 28 Oct 2025 18:41:17 +0300 Subject: [PATCH 10/20] Add error handling --- src/api/coderApi.ts | 98 +++++++++++++- src/commands.ts | 31 +---- src/extension.ts | 21 +-- src/oauth/errors.ts | 173 +++++++++++++++++++++++ src/oauth/metadataClient.ts | 2 +- src/oauth/sessionManager.ts | 211 ++++++++++++++++++----------- src/oauth/tokenRefreshScheduler.ts | 65 --------- 7 files changed, 414 insertions(+), 187 deletions(-) create mode 100644 src/oauth/errors.ts delete mode 100644 src/oauth/tokenRefreshScheduler.ts diff --git a/src/api/coderApi.ts b/src/api/coderApi.ts index 04c696be..d9c30868 100644 --- a/src/api/coderApi.ts +++ b/src/api/coderApi.ts @@ -3,6 +3,7 @@ import { type AxiosInstance, type AxiosHeaders, type AxiosResponseTransformer, + isAxiosError, } from "axios"; import { Api } from "coder/site/src/api/api"; import { @@ -31,6 +32,12 @@ import { HttpClientLogLevel, } from "../logging/types"; import { sizeOf } from "../logging/utils"; +import { + parseOAuthError, + requiresReAuthentication, + isNetworkError, +} from "../oauth/errors"; +import { type OAuthSessionManager } from "../oauth/sessionManager"; import { HttpStatusCode } from "../websocket/codes"; import { type UnidirectionalStream, @@ -72,6 +79,7 @@ export class CoderApi extends Api { baseUrl: string, token: string | undefined, output: Logger, + oauthSessionManager?: OAuthSessionManager, ): CoderApi { const client = new CoderApi(output); client.setHost(baseUrl); @@ -79,7 +87,7 @@ export class CoderApi extends Api { client.setSessionToken(token); } - setupInterceptors(client, output); + setupInterceptors(client, output, oauthSessionManager); return client; } @@ -390,7 +398,11 @@ export class CoderApi extends Api { /** * Set up logging and request interceptors for the CoderApi instance. */ -function setupInterceptors(client: CoderApi, output: Logger): void { +function setupInterceptors( + client: CoderApi, + output: Logger, + oauthSessionManager?: OAuthSessionManager, +): void { addLoggingInterceptors(client.getAxiosInstance(), output); client.getAxiosInstance().interceptors.request.use(async (config) => { @@ -428,6 +440,11 @@ function setupInterceptors(client: CoderApi, output: Logger): void { } }, ); + + // OAuth token refresh interceptors + if (oauthSessionManager) { + addOAuthInterceptors(client, output, oauthSessionManager); + } } function addLoggingInterceptors(client: AxiosInstance, logger: Logger) { @@ -457,7 +474,7 @@ function addLoggingInterceptors(client: AxiosInstance, logger: Logger) { }, (error: unknown) => { logError(logger, error, getLogLevel()); - return Promise.reject(error); + throw error; }, ); @@ -468,7 +485,80 @@ function addLoggingInterceptors(client: AxiosInstance, logger: Logger) { }, (error: unknown) => { logError(logger, error, getLogLevel()); - return Promise.reject(error); + throw error; + }, + ); +} + +/** + * Add OAuth token refresh interceptors. + * Success interceptor: proactively refreshes token when approaching expiry. + * Error interceptor: reactively refreshes token on 401/403 responses. + */ +function addOAuthInterceptors( + client: CoderApi, + logger: Logger, + oauthSessionManager: OAuthSessionManager, +) { + client.getAxiosInstance().interceptors.response.use( + // Success response interceptor: proactive token refresh + (response) => { + if (oauthSessionManager.shouldRefreshToken()) { + logger.debug( + "Token approaching expiry, triggering proactive refresh in background", + ); + + // Fire-and-forget: don't await, don't block response + oauthSessionManager.refreshToken().catch((error) => { + logger.warn("Background token refresh failed:", error); + }); + } + + return response; + }, + // Error response interceptor: reactive token refresh on 401/403 + async (error: unknown) => { + if (!isAxiosError(error)) { + throw error; + } + + const status = error.response?.status; + if (status !== 401 && status !== 403) { + throw error; + } + + if (!oauthSessionManager.isLoggedInWithOAuth()) { + throw error; + } + + logger.info(`Received ${status} response, attempting token refresh`); + + try { + const newTokens = await oauthSessionManager.refreshToken(); + client.setSessionToken(newTokens.access_token); + + logger.info("Token refresh successful, updated session token"); + } catch (refreshError) { + logger.error("Token refresh failed:", refreshError); + + const oauthError = parseOAuthError(refreshError); + if (oauthError && requiresReAuthentication(oauthError)) { + logger.error( + `OAuth error requires re-authentication: ${oauthError.errorCode}`, + ); + + oauthSessionManager + .showReAuthenticationModal(oauthError) + .catch((err) => { + logger.error("Failed to show re-auth modal:", err); + }); + } else if (isNetworkError(refreshError)) { + logger.warn( + "Token refresh failed due to network error, will retry later", + ); + } + } + throw error; }, ); } diff --git a/src/commands.ts b/src/commands.ts index 2e43fbbc..bf01cf5b 100644 --- a/src/commands.ts +++ b/src/commands.ts @@ -360,20 +360,10 @@ export class Commands { } this.logger.info("Logging out"); - // Check if using OAuth - const isOAuthLoggedIn = - await this.oauthSessionManager.isLoggedInWithOAuth(); - if (isOAuthLoggedIn) { - this.logger.info("Logging out via OAuth"); - try { - await this.oauthSessionManager.logout(); - } catch (error) { - this.logger.warn( - "OAuth logout failed, continuing with cleanup:", - error, - ); - } - } + // Fire and forget + this.oauthSessionManager.logout().catch((error) => { + this.logger.warn("OAuth logout failed, continuing with cleanup:", error); + }); // Clear from the REST client. An empty url will indicate to other parts of // the code that we are logged out. @@ -545,19 +535,6 @@ export class Commands { }, ); } - // Check if app has a URL to open - if (app.url) { - return vscode.window.withProgress( - { - location: vscode.ProgressLocation.Notification, - title: `Opening ${app.name || "application"} in browser...`, - cancellable: false, - }, - async () => { - await vscode.env.openExternal(vscode.Uri.parse(app.url!)); - }, - ); - } // If no URL or command, show information about the app status vscode.window.showInformationMessage(`${app.name}`, { diff --git a/src/extension.ts b/src/extension.ts index 61f9afd9..e564bcc4 100644 --- a/src/extension.ts +++ b/src/extension.ts @@ -65,14 +65,24 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { // Try to clear this flag ASAP const isFirstConnect = await mementoManager.getAndClearFirstConnect(); + const url = mementoManager.getUrl(); + + // Create OAuth session manager before the main client + const oauthSessionManager = await OAuthSessionManager.create( + url || "", + serviceContainer, + ctx, + ); + ctx.subscriptions.push(oauthSessionManager); + // This client tracks the current login and will be used through the life of // the plugin to poll workspaces for the current login, as well as being used // in commands that operate on the current login. - const url = mementoManager.getUrl(); const client = CoderApi.create( url || "", await secretsManager.getSessionToken(), output, + oauthSessionManager, ); const myWorkspacesProvider = new WorkspaceProvider( @@ -118,14 +128,6 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { ctx.subscriptions, ); - const oauthSessionManager = await OAuthSessionManager.create( - url || "", - secretsManager, - output, - ctx, - ); - ctx.subscriptions.push(oauthSessionManager); - // Listen for session token changes and sync state across all components ctx.subscriptions.push( secretsManager.onDidChangeSessionToken(async (token) => { @@ -407,6 +409,7 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { remoteSshExtension.id, ); if (details) { + // TODO if the URL is different then we need to update the OAuth session!!! (Centralize this logic) ctx.subscriptions.push(details); // Authenticate the plugin client which is used in the sidebar to display // workspaces belonging to this deployment. diff --git a/src/oauth/errors.ts b/src/oauth/errors.ts new file mode 100644 index 00000000..67c7cd47 --- /dev/null +++ b/src/oauth/errors.ts @@ -0,0 +1,173 @@ +import { isAxiosError } from "axios"; + +import type { OAuthErrorResponse } from "./types"; + +/** + * Base class for OAuth errors + */ +export class OAuthError extends Error { + constructor( + message: string, + public readonly errorCode: string, + public readonly description?: string, + public readonly errorUri?: string, + ) { + super(message); + this.name = "OAuthError"; + } +} + +/** + * Refresh token is invalid, expired, or revoked. Requires re-authentication. + */ +export class InvalidGrantError extends OAuthError { + constructor(description?: string, errorUri?: string) { + super( + "OAuth refresh token is invalid, expired, or revoked", + "invalid_grant", + description, + errorUri, + ); + this.name = "InvalidGrantError"; + } +} + +/** + * Client credentials are invalid. Requires re-registration. + */ +export class InvalidClientError extends OAuthError { + constructor(description?: string, errorUri?: string) { + super( + "OAuth client credentials are invalid", + "invalid_client", + description, + errorUri, + ); + this.name = "InvalidClientError"; + } +} + +/** + * Invalid request error - malformed OAuth request + */ +export class InvalidRequestError extends OAuthError { + constructor(description?: string, errorUri?: string) { + super( + "OAuth request is malformed or invalid", + "invalid_request", + description, + errorUri, + ); + this.name = "InvalidRequestError"; + } +} + +/** + * Client is not authorized for this grant type. + */ +export class UnauthorizedClientError extends OAuthError { + constructor(description?: string, errorUri?: string) { + super( + "OAuth client is not authorized for this grant type", + "unauthorized_client", + description, + errorUri, + ); + this.name = "UnauthorizedClientError"; + } +} + +/** + * Unsupported grant type error. + */ +export class UnsupportedGrantTypeError extends OAuthError { + constructor(description?: string, errorUri?: string) { + super( + "OAuth grant type is not supported", + "unsupported_grant_type", + description, + errorUri, + ); + this.name = "UnsupportedGrantTypeError"; + } +} + +/** + * Invalid scope error. + */ +export class InvalidScopeError extends OAuthError { + constructor(description?: string, errorUri?: string) { + super( + "OAuth scope is invalid, unknown, malformed, or exceeds the scope granted by the resource owner", + "invalid_scope", + description, + errorUri, + ); + this.name = "InvalidScopeError"; + } +} + +/** + * Parses an axios error to extract OAuth error information + * Returns an OAuthError instance if the error is OAuth-related, otherwise returns null + */ +export function parseOAuthError(error: unknown): OAuthError | null { + if (!isAxiosError(error)) { + return null; + } + + const data = error.response?.data; + + if (!isOAuthErrorResponse(data)) { + return null; + } + + const { error: errorCode, error_description, error_uri } = data; + + switch (errorCode) { + case "invalid_grant": + return new InvalidGrantError(error_description, error_uri); + case "invalid_client": + return new InvalidClientError(error_description, error_uri); + case "invalid_request": + return new InvalidRequestError(error_description, error_uri); + case "unauthorized_client": + return new UnauthorizedClientError(error_description, error_uri); + case "unsupported_grant_type": + return new UnsupportedGrantTypeError(error_description, error_uri); + case "invalid_scope": + return new InvalidScopeError(error_description, error_uri); + default: + return new OAuthError( + `OAuth error: ${errorCode}`, + errorCode, + error_description, + error_uri, + ); + } +} + +function isOAuthErrorResponse(data: unknown): data is OAuthErrorResponse { + return ( + data !== null && + typeof data === "object" && + "error" in data && + typeof data.error === "string" + ); +} + +/** + * Checks if an error requires re-authentication + */ +export function requiresReAuthentication(error: OAuthError): boolean { + return ( + error instanceof InvalidGrantError || error instanceof InvalidClientError + ); +} + +/** + * Checks if an error is a network/connectivity error + */ +export function isNetworkError(error: unknown): boolean { + return isAxiosError(error) && !error.response && Boolean(error.request); +} diff --git a/src/oauth/metadataClient.ts b/src/oauth/metadataClient.ts index 7f3227dc..98568525 100644 --- a/src/oauth/metadataClient.ts +++ b/src/oauth/metadataClient.ts @@ -42,7 +42,7 @@ export class OAuthMetadataClient { * Throws detailed errors if server doesn't meet OAuth 2.1 requirements. */ async getMetadata(): Promise { - this.logger.info("Discovering OAuth endpoints..."); + this.logger.debug("Discovering OAuth endpoints..."); const response = await this.axiosInstance.get( OAUTH_DISCOVERY_ENDPOINT, diff --git a/src/oauth/sessionManager.ts b/src/oauth/sessionManager.ts index e41d5112..77fcec63 100644 --- a/src/oauth/sessionManager.ts +++ b/src/oauth/sessionManager.ts @@ -1,10 +1,11 @@ import { type AxiosInstance } from "axios"; import * as vscode from "vscode"; +import { type ServiceContainer } from "src/core/container"; + import { CoderApi } from "../api/coderApi"; import { OAuthMetadataClient } from "./metadataClient"; -import { OAuthTokenRefreshScheduler } from "./tokenRefreshScheduler"; import { CALLBACK_PATH, generatePKCE, @@ -15,6 +16,7 @@ import { import type { SecretsManager, StoredOAuthTokens } from "../core/secretsManager"; import type { Logger } from "../logging/logger"; +import type { OAuthError } from "./errors"; import type { ClientRegistrationRequest, ClientRegistrationResponse, @@ -30,6 +32,16 @@ const REFRESH_GRANT_TYPE = "refresh_token" as const; const RESPONSE_TYPE = "code" as const; const PKCE_CHALLENGE_METHOD = "S256" as const; +/** + * Token refresh threshold: refresh when token expires in less than this time + */ +const TOKEN_REFRESH_THRESHOLD_MS = 10 * 60 * 1000; + +/** + * Minimum time between refresh attempts to prevent thrashing + */ +const REFRESH_THROTTLE_MS = 30 * 1000; + /** * Minimal scopes required by the VS Code extension. */ @@ -48,10 +60,9 @@ const DEFAULT_OAUTH_SCOPES = [ * Coordinates authorization flow, token management, and automatic refresh. */ export class OAuthSessionManager implements vscode.Disposable { - private readonly extensionId: string; - private readonly refreshScheduler: OAuthTokenRefreshScheduler; - private storedTokens: StoredOAuthTokens | undefined; + private refreshInProgress = false; + private lastRefreshAttempt = 0; // Pending authorization flow state private pendingAuthResolve: @@ -66,15 +77,15 @@ export class OAuthSessionManager implements vscode.Disposable { */ static async create( deploymentUrl: string, - secretsManager: SecretsManager, - logger: Logger, + container: ServiceContainer, context: vscode.ExtensionContext, ): Promise { const manager = new OAuthSessionManager( deploymentUrl, - secretsManager, - logger, - context, + container.getSecretsManager(), + container.getLogger(), + container.getVsCodeProposed(), + context.extension.id, ); await manager.loadTokens(); return manager; @@ -84,16 +95,12 @@ export class OAuthSessionManager implements vscode.Disposable { private deploymentUrl: string, private readonly secretsManager: SecretsManager, private readonly logger: Logger, - context: vscode.ExtensionContext, - ) { - this.extensionId = context.extension.id; - this.refreshScheduler = new OAuthTokenRefreshScheduler(async () => { - await this.refreshToken(); - }, logger); - } + private readonly vscodeProposed: typeof vscode, + private readonly extensionId: string, + ) {} /** - * Load stored tokens and start refresh timer if applicable. + * Load stored tokens from storage. * Validates that tokens belong to the current deployment URL. */ private async loadTokens(): Promise { @@ -126,18 +133,15 @@ export class OAuthSessionManager implements vscode.Disposable { this.storedTokens = tokens; this.logger.info(`Loaded stored OAuth tokens for ${tokens.deployment_url}`); - - if (tokens.refresh_token) { - this.refreshScheduler.schedule(tokens); - } } /** * Clear stale data when tokens don't match current deployment. */ private async clearTokenState(): Promise { - this.refreshScheduler.stop(); this.storedTokens = undefined; + this.refreshInProgress = false; + this.lastRefreshAttempt = 0; await this.secretsManager.setOAuthTokens(undefined); await this.secretsManager.setOAuthClientRegistration(undefined); } @@ -270,7 +274,7 @@ export class OAuthSessionManager implements vscode.Disposable { if (!baseUrl) { throw new Error("CoderApi instance has no base URL set"); } - if (this.deploymentUrl !== baseUrl) { + if (this.deploymentUrl && this.deploymentUrl !== baseUrl) { this.logger.info("Deployment URL changed, clearing cached state", { old: this.deploymentUrl, new: baseUrl, @@ -279,8 +283,6 @@ export class OAuthSessionManager implements vscode.Disposable { this.deploymentUrl = baseUrl; } - this.logger.info("Starting OAuth login flow"); - const axiosInstance = client.getAxiosInstance(); const metadataClient = new OAuthMetadataClient(axiosInstance, this.logger); const metadata = await metadataClient.getMetadata(); @@ -512,48 +514,60 @@ export class OAuthSessionManager implements vscode.Disposable { /** * Refresh the access token using the stored refresh token. + * Uses a mutex to prevent concurrent refresh attempts. */ - private async refreshToken(): Promise { + async refreshToken(): Promise { + if (this.refreshInProgress) { + throw new Error("Token refresh already in progress"); + } + if (!this.storedTokens?.refresh_token) { throw new Error("No refresh token available"); } - const { axiosInstance, metadata, registration } = - await this.prepareOAuthOperation( - this.deploymentUrl, - this.storedTokens.access_token, - ); + this.refreshInProgress = true; + this.lastRefreshAttempt = Date.now(); - this.logger.debug("Refreshing access token"); + try { + const { axiosInstance, metadata, registration } = + await this.prepareOAuthOperation( + this.deploymentUrl, + this.storedTokens.access_token, + ); - const params: RefreshTokenRequestParams = { - grant_type: REFRESH_GRANT_TYPE, - refresh_token: this.storedTokens.refresh_token, - client_id: registration.client_id, - client_secret: registration.client_secret, - }; + this.logger.debug("Refreshing access token"); - const tokenRequest = toUrlSearchParams(params); + const params: RefreshTokenRequestParams = { + grant_type: REFRESH_GRANT_TYPE, + refresh_token: this.storedTokens.refresh_token, + client_id: registration.client_id, + client_secret: registration.client_secret, + }; - const response = await axiosInstance.post( - metadata.token_endpoint, - tokenRequest, - { - headers: { - "Content-Type": "application/x-www-form-urlencoded", + const tokenRequest = toUrlSearchParams(params); + + const response = await axiosInstance.post( + metadata.token_endpoint, + tokenRequest, + { + headers: { + "Content-Type": "application/x-www-form-urlencoded", + }, }, - }, - ); + ); - this.logger.debug("Token refresh successful"); + this.logger.debug("Token refresh successful"); - await this.saveTokens(response.data); + await this.saveTokens(response.data); - return response.data; + return response.data; + } finally { + this.refreshInProgress = false; + } } /** - * Save token response to storage and schedule automatic refresh. + * Save token response to storage. * Also triggers event via secretsManager to update global client. */ private async saveTokens(tokenResponse: TokenResponse): Promise { @@ -571,34 +585,54 @@ export class OAuthSessionManager implements vscode.Disposable { await this.secretsManager.setOAuthTokens(tokens); // Trigger event to update global client (works for login & background refresh) - // TODO Add a setting to check if we have OAuth or token setup so we can start the background refresh await this.secretsManager.setSessionToken(tokenResponse.access_token); this.logger.info("Tokens saved", { expires_at: new Date(expiryTimestamp).toISOString(), deployment: this.deploymentUrl, }); + } + + /** + * Check if token should be refreshed. + * Returns true if: + * 1. Token expires in less than TOKEN_REFRESH_THRESHOLD_MS + * 2. Last refresh attempt was more than REFRESH_THROTTLE_MS ago + * 3. No refresh is currently in progress + */ + shouldRefreshToken(): boolean { + if ( + !this.isLoggedInWithOAuth() || + !this.storedTokens?.refresh_token || + this.refreshInProgress + ) { + return false; + } - // Schedule automatic refresh - this.refreshScheduler.schedule(tokens); + const now = Date.now(); + if (now - this.lastRefreshAttempt < REFRESH_THROTTLE_MS) { + return false; + } + + const timeUntilExpiry = this.storedTokens.expiry_timestamp - now; + return timeUntilExpiry < TOKEN_REFRESH_THRESHOLD_MS; } /** * Revoke a token using the OAuth server's revocation endpoint. */ - private async revokeToken(token: string): Promise { + private async revokeToken( + token: string, + tokenTypeHint: "access_token" | "refresh_token" = "refresh_token", + ): Promise { const { axiosInstance, metadata, registration } = await this.prepareOAuthOperation( this.deploymentUrl, this.storedTokens?.access_token, ); - if (!metadata.revocation_endpoint) { - this.logger.warn( - "Server does not support token revocation (no revocation_endpoint)", - ); - return; - } + const revocationEndpoint = + metadata.revocation_endpoint || `${metadata.issuer}/oauth2/revoke`; this.logger.info("Revoking refresh token"); @@ -606,21 +640,17 @@ export class OAuthSessionManager implements vscode.Disposable { token, client_id: registration.client_id, client_secret: registration.client_secret, - token_type_hint: "refresh_token", + token_type_hint: tokenTypeHint, }; const revocationRequest = toUrlSearchParams(params); try { - await axiosInstance.post( - metadata.revocation_endpoint, - revocationRequest, - { - headers: { - "Content-Type": "application/x-www-form-urlencoded", - }, + await axiosInstance.post(revocationEndpoint, revocationRequest, { + headers: { + "Content-Type": "application/x-www-form-urlencoded", }, - ); + }); this.logger.info("Token revocation successful"); } catch (error) { @@ -633,7 +663,9 @@ export class OAuthSessionManager implements vscode.Disposable { * Logout by revoking tokens and clearing all OAuth data. */ async logout(): Promise { - this.refreshScheduler.stop(); + if (!this.isLoggedInWithOAuth()) { + return; + } // Revoke refresh token (which also invalidates access token per RFC 7009) if (this.storedTokens?.refresh_token) { @@ -650,29 +682,46 @@ export class OAuthSessionManager implements vscode.Disposable { } /** - * Check if currently logged in with OAuth. - * Returns true only if valid OAuth tokens exist for the current deployment. + * Returns true if (valid or invalid) OAuth tokens exist for the current deployment. */ - async isLoggedInWithOAuth(): Promise { - const tokens = await this.secretsManager.getOAuthTokens(); - if (!tokens) { - return false; - } + isLoggedInWithOAuth(): boolean { + return this.storedTokens !== undefined; + } - return this.deploymentUrl === tokens.deployment_url; + /** + * Show a modal dialog to the user when OAuth re-authentication is required. + * This is called when the refresh token is invalid or the client credentials are invalid. + */ + async showReAuthenticationModal(error: OAuthError): Promise { + const errorMessage = + error.description || + "Your session is no longer valid. This could be due to token expiration or revocation."; + + // Log out first to clear invalid tokens + await vscode.commands.executeCommand("coder.logout"); + + const action = await this.vscodeProposed.window.showErrorMessage( + `Authentication Error`, + { modal: true, useCustom: true, detail: errorMessage }, + "Log in again", + ); + + if (action === "Log in again") { + await vscode.commands.executeCommand("coder.login"); + } } /** * Clears all in-memory state and rejects any pending operations. */ dispose(): void { - this.refreshScheduler.stop(); - if (this.pendingAuthReject) { this.pendingAuthReject(new Error("OAuth session manager disposed")); } this.clearPendingAuth(); this.storedTokens = undefined; + this.refreshInProgress = false; + this.lastRefreshAttempt = 0; this.logger.debug("OAuth session manager disposed"); } diff --git a/src/oauth/tokenRefreshScheduler.ts b/src/oauth/tokenRefreshScheduler.ts deleted file mode 100644 index 3eeabb9e..00000000 --- a/src/oauth/tokenRefreshScheduler.ts +++ /dev/null @@ -1,65 +0,0 @@ -import type { StoredOAuthTokens } from "../core/secretsManager"; -import type { Logger } from "../logging/logger"; - -// Token refresh timing constant -const TOKEN_REFRESH_THRESHOLD_MS = 20 * 60 * 1000; - -/** - * Manages automatic token refresh scheduling. - * Calculates optimal refresh timing and triggers refresh callbacks. - */ -export class OAuthTokenRefreshScheduler { - private refreshTimer: NodeJS.Timeout | undefined; - - constructor( - private readonly refreshCallback: () => Promise, - private readonly logger: Logger, - ) {} - - /** - * Schedule automatic token refresh based on token expiry. - */ - schedule(tokens: StoredOAuthTokens): void { - this.stop(); - - if (!tokens.refresh_token) { - this.logger.debug("No refresh token available, skipping timer setup"); - return; - } - - const now = Date.now(); - const timeUntilRefresh = - tokens.expiry_timestamp - TOKEN_REFRESH_THRESHOLD_MS - now; - - if (timeUntilRefresh <= 0) { - this.logger.info("Token needs immediate refresh"); - this.refreshCallback().catch((error) => { - this.logger.error("Immediate token refresh failed:", error); - }); - return; - } - - this.refreshTimer = setTimeout(() => { - this.logger.debug("Token refresh timer fired, refreshing token..."); - this.refreshCallback().catch((error) => { - this.logger.error("Scheduled token refresh failed:", error); - }); - }, timeUntilRefresh); - - this.logger.debug("Token refresh timer scheduled", { - fires_at: new Date(now + timeUntilRefresh).toISOString(), - fires_in_seconds: Math.round(timeUntilRefresh / 1000), - }); - } - - /** - * Stop the background token refresh timer. - */ - stop(): void { - if (this.refreshTimer) { - clearTimeout(this.refreshTimer); - this.refreshTimer = undefined; - this.logger.debug("Token refresh timer stopped"); - } - } -} From 27fdb3ec8c5a97481c72b4eab789b1d11370b3ed Mon Sep 17 00:00:00 2001 From: Ehab Younes Date: Thu, 30 Oct 2025 18:50:47 +0300 Subject: [PATCH 11/20] Allow callback handling in multiple VS Code windows --- src/core/secretsManager.ts | 40 ++++++++++++ src/extension.ts | 2 +- src/oauth/sessionManager.ts | 120 +++++++++++++----------------------- 3 files changed, 85 insertions(+), 77 deletions(-) diff --git a/src/core/secretsManager.ts b/src/core/secretsManager.ts index d16292f1..cdda9c72 100644 --- a/src/core/secretsManager.ts +++ b/src/core/secretsManager.ts @@ -13,11 +13,19 @@ const OAUTH_CLIENT_REGISTRATION_KEY = "oauthClientRegistration"; const OAUTH_TOKENS_KEY = "oauthTokens"; +const OAUTH_CALLBACK_KEY = "coder.oauthCallback"; + export type StoredOAuthTokens = Omit & { expiry_timestamp: number; deployment_url: string; }; +interface OAuthCallbackData { + state: string; + code: string | null; + error: string | null; +} + export enum AuthAction { LOGIN, LOGOUT, @@ -163,4 +171,36 @@ export class SecretsManager { } return undefined; } + + /** + * Write an OAuth callback result to secrets storage. + * Used for cross-window communication when OAuth callback arrives in a different window. + */ + public async setOAuthCallback(data: OAuthCallbackData): Promise { + await this.secrets.store(OAUTH_CALLBACK_KEY, JSON.stringify(data)); + } + + /** + * Listen for OAuth callback results from any VS Code window. + * The listener receives the state parameter, code (if success), and error (if failed). + */ + public onDidChangeOAuthCallback( + listener: (data: OAuthCallbackData) => void, + ): Disposable { + return this.secrets.onDidChange(async (e) => { + if (e.key !== OAUTH_CALLBACK_KEY) { + return; + } + + try { + const data = await this.secrets.get(OAUTH_CALLBACK_KEY); + if (data) { + const parsed = JSON.parse(data) as OAuthCallbackData; + listener(parsed); + } + } catch { + // Ignore parse errors + } + }); + } } diff --git a/src/extension.ts b/src/extension.ts index e564bcc4..837b8d64 100644 --- a/src/extension.ts +++ b/src/extension.ts @@ -158,7 +158,7 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { const code = params.get("code"); const state = params.get("state"); const error = params.get("error"); - oauthSessionManager.handleCallback(code, state, error); + await oauthSessionManager.handleCallback(code, state, error); return; } diff --git a/src/oauth/sessionManager.ts b/src/oauth/sessionManager.ts index 77fcec63..04733b5f 100644 --- a/src/oauth/sessionManager.ts +++ b/src/oauth/sessionManager.ts @@ -64,13 +64,7 @@ export class OAuthSessionManager implements vscode.Disposable { private refreshInProgress = false; private lastRefreshAttempt = 0; - // Pending authorization flow state - private pendingAuthResolve: - | ((value: { code: string; verifier: string }) => void) - | undefined; private pendingAuthReject: ((reason: Error) => void) | undefined; - private expectedState: string | undefined; - private pendingVerifier: string | undefined; /** * Create and initialize a new OAuth session manager. @@ -370,12 +364,12 @@ export class OAuthSessionManager implements vscode.Disposable { challenge, ); - return new Promise<{ code: string; verifier: string }>( + const callbackPromise = new Promise<{ code: string; verifier: string }>( (resolve, reject) => { const timeoutMins = 5; - const timeout = setTimeout( + const timeoutHandle = setTimeout( () => { - this.clearPendingAuth(); + cleanup(); reject( new Error(`OAuth flow timed out after ${timeoutMins} minutes`), ); @@ -383,93 +377,67 @@ export class OAuthSessionManager implements vscode.Disposable { timeoutMins * 60 * 1000, ); - const clearPromise = () => { - clearTimeout(timeout); - this.clearPendingAuth(); - }; - - this.pendingAuthResolve = (result) => { - clearPromise(); - resolve(result); - }; - - this.pendingAuthReject = (error) => { - clearPromise(); - reject(error); - }; + const listener = this.secretsManager.onDidChangeOAuthCallback( + ({ state: callbackState, code, error }) => { + if (callbackState !== state) { + return; + } - this.expectedState = state; - this.pendingVerifier = verifier; + cleanup(); - vscode.env.openExternal(vscode.Uri.parse(authUrl)).then( - () => {}, - (error) => { - if (error instanceof Error) { - this.pendingAuthReject?.(error); + if (error) { + reject(new Error(`OAuth error: ${error}`)); + } else if (code) { + resolve({ code, verifier }); } else { - this.pendingAuthReject?.(new Error("Failed to open browser")); + reject(new Error("No authorization code received")); } }, ); + + const cleanup = () => { + clearTimeout(timeoutHandle); + listener.dispose(); + }; + + this.pendingAuthReject = (error) => { + cleanup(); + reject(error); + }; }, ); - } - /** - * Clear pending authorization flow state. - */ - private clearPendingAuth(): void { - this.pendingAuthResolve = undefined; - this.pendingAuthReject = undefined; - this.expectedState = undefined; - this.pendingVerifier = undefined; + try { + await vscode.env.openExternal(vscode.Uri.parse(authUrl)); + } catch (error) { + throw error instanceof Error + ? error + : new Error("Failed to open browser"); + } + + return callbackPromise; } /** * Handle OAuth callback from browser redirect. - * Validates state and resolves pending authorization promise. - * - * // TODO this has to work across windows! + * Writes the callback result to secrets storage, triggering the waiting window to proceed. */ - handleCallback( + async handleCallback( code: string | null, state: string | null, error: string | null, - ): void { - if (!this.pendingAuthResolve || !this.pendingAuthReject) { - this.logger.warn("Received OAuth callback but no pending auth flow"); - return; - } - - if (error) { - this.pendingAuthReject(new Error(`OAuth error: ${error}`)); - return; - } - - if (!code) { - this.pendingAuthReject(new Error("No authorization code received")); - return; - } - + ): Promise { if (!state) { - this.pendingAuthReject(new Error("No state received")); - return; - } - - if (state !== this.expectedState) { - this.pendingAuthReject( - new Error("State mismatch - possible CSRF attack"), - ); + this.logger.warn("Received OAuth callback with no state parameter"); return; } - const verifier = this.pendingVerifier; - if (!verifier) { - this.pendingAuthReject(new Error("No PKCE verifier found")); - return; + try { + await this.secretsManager.setOAuthCallback({ state, code, error }); + this.logger.debug("OAuth callback processed successfully"); + } catch (err) { + this.logger.error("Failed to process OAuth callback:", err); } - - this.pendingAuthResolve({ code, verifier }); } /** @@ -712,13 +680,13 @@ export class OAuthSessionManager implements vscode.Disposable { } /** - * Clears all in-memory state and rejects any pending operations. + * Clears all in-memory state. */ dispose(): void { if (this.pendingAuthReject) { this.pendingAuthReject(new Error("OAuth session manager disposed")); } - this.clearPendingAuth(); + this.pendingAuthReject = undefined; this.storedTokens = undefined; this.refreshInProgress = false; this.lastRefreshAttempt = 0; From c17cd28fd6cafb4c3eb5f0057120836731b0a683 Mon Sep 17 00:00:00 2001 From: Ehab Younes Date: Tue, 11 Nov 2025 21:50:30 +0300 Subject: [PATCH 12/20] Improve UX --- src/commands.ts | 89 +++++++++++-------------------------- src/oauth/metadataClient.ts | 10 ++--- src/oauth/sessionManager.ts | 9 +++- src/promptUtils.ts | 55 +++++++++++++++++++++++ 4 files changed, 95 insertions(+), 68 deletions(-) diff --git a/src/commands.ts b/src/commands.ts index bf01cf5b..4cf0289a 100644 --- a/src/commands.ts +++ b/src/commands.ts @@ -19,9 +19,8 @@ import { type SecretsManager } from "./core/secretsManager"; import { CertificateError } from "./error"; import { getGlobalFlags } from "./globalFlags"; import { type Logger } from "./logging/logger"; -import { OAuthMetadataClient } from "./oauth/metadataClient"; import { type OAuthSessionManager } from "./oauth/sessionManager"; -import { maybeAskAgent, maybeAskUrl } from "./promptUtils"; +import { maybeAskAgent, maybeAskUrl, maybeAskAuthMethod } from "./promptUtils"; import { escapeCommandArg, toRemoteAuthority, toSafeHost } from "./util"; import { AgentTreeItem, @@ -29,8 +28,6 @@ import { WorkspaceTreeItem, } from "./workspace/workspacesProvider"; -type AuthMethod = "oauth" | "legacy"; - export class Commands { private readonly vscodeProposed: typeof vscode; private readonly logger: Logger; @@ -92,7 +89,7 @@ export class Commands { // Try to get a token from the user, if we need one, and their user. const autoLogin = args?.autoLogin === true; - const res = await this.maybeAskToken(url, args?.token, autoLogin); + const res = await this.attemptLogin(url, args?.token, autoLogin); if (!res) { return; // The user aborted, or unable to auth. } @@ -136,12 +133,12 @@ export class Commands { } /** - * If necessary, ask for a token, and keep asking until the token has been - * validated. Return the token and user that was fetched to validate the - * token. Null means the user aborted or we were unable to authenticate with - * mTLS (in the latter case, an error notification will have been displayed). + * Attempt to authenticate using OAuth, token, or mTLS. If necessary, prompts + * for authentication method and credentials. Returns the token and user upon + * successful authentication. Null means the user aborted or authentication + * failed (in which case an error notification will have been displayed). */ - private async maybeAskToken( + private async attemptLogin( url: string, token: string | undefined, isAutoLogin: boolean, @@ -174,58 +171,18 @@ export class Commands { } } - // Check if server supports OAuth - const supportsOAuth = await this.checkOAuthSupport(client); - - let choice: AuthMethod | undefined = "legacy"; - if (supportsOAuth) { - choice = await this.askAuthMethod(); - } - - if (choice === "oauth") { - return this.loginWithOAuth(client); - } else if (choice === "legacy") { - const initialToken = - token || (await this.secretsManager.getSessionToken()); - return this.loginWithToken(client, initialToken); + const authMethod = await maybeAskAuthMethod(client); + switch (authMethod) { + case "oauth": + return this.loginWithOAuth(client); + case "legacy": { + const initialToken = + token || (await this.secretsManager.getSessionToken()); + return this.loginWithToken(client, initialToken); + } + case undefined: + return null; // User aborted } - - return null; // User aborted. - } - - private async checkOAuthSupport(client: CoderApi): Promise { - const metadataClient = new OAuthMetadataClient( - client.getAxiosInstance(), - this.logger, - ); - return metadataClient.checkOAuthSupport(); - } - - /** - * Ask user to choose between OAuth and legacy API token authentication. - */ - private async askAuthMethod(): Promise { - const choice = await vscode.window.showQuickPick( - [ - { - label: "$(key) OAuth (Recommended)", - detail: "Secure authentication with automatic token refresh", - value: "oauth" as const, - }, - { - label: "$(lock) API Token", - detail: "Use a manually created API key", - value: "legacy" as const, - }, - ], - { - title: "Choose Authentication Method", - placeHolder: "How would you like to authenticate?", - ignoreFocusOut: true, - }, - ); - - return choice?.value; } private async loginWithToken( @@ -297,7 +254,15 @@ export class Commands { try { this.logger.info("Starting OAuth authentication"); - const tokenResponse = await this.oauthSessionManager.login(client); + const tokenResponse = await vscode.window.withProgress( + { + location: vscode.ProgressLocation.Notification, + title: "Authenticating", + cancellable: false, + }, + async (progress) => + await this.oauthSessionManager.login(client, progress), + ); // Validate token by fetching user client.setSessionToken(tokenResponse.access_token); diff --git a/src/oauth/metadataClient.ts b/src/oauth/metadataClient.ts index 98568525..149d64fa 100644 --- a/src/oauth/metadataClient.ts +++ b/src/oauth/metadataClient.ts @@ -26,13 +26,13 @@ export class OAuthMetadataClient { /** * Check if a server supports OAuth by attempting to fetch the well-known endpoint. */ - async checkOAuthSupport(): Promise { + public static async checkOAuthSupport( + axiosInstance: AxiosInstance, + ): Promise { try { - await this.axiosInstance.get(OAUTH_DISCOVERY_ENDPOINT); - this.logger.debug("Server supports OAuth"); + await axiosInstance.get(OAUTH_DISCOVERY_ENDPOINT); return true; - } catch (error) { - this.logger.debug("Server does not support OAuth:", error); + } catch { return false; } } diff --git a/src/oauth/sessionManager.ts b/src/oauth/sessionManager.ts index 04733b5f..c44ba331 100644 --- a/src/oauth/sessionManager.ts +++ b/src/oauth/sessionManager.ts @@ -263,7 +263,10 @@ export class OAuthSessionManager implements vscode.Disposable { * @param client CoderApi instance for the deployment to authenticate against * @returns TokenResponse containing access token and optional refresh token */ - async login(client: CoderApi): Promise { + async login( + client: CoderApi, + progress: vscode.Progress<{ message?: string; increment?: number }>, + ): Promise { const baseUrl = client.getAxiosInstance().defaults.baseURL; if (!baseUrl) { throw new Error("CoderApi instance has no base URL set"); @@ -282,13 +285,16 @@ export class OAuthSessionManager implements vscode.Disposable { const metadata = await metadataClient.getMetadata(); // Only register the client on login + progress.report({ message: "registering client...", increment: 10 }); const registration = await this.registerClient(axiosInstance, metadata); + progress.report({ message: "waiting for authorization...", increment: 30 }); const { code, verifier } = await this.startAuthorization( metadata, registration, ); + progress.report({ message: "exchanging token...", increment: 30 }); const tokenResponse = await this.exchangeToken( code, verifier, @@ -297,6 +303,7 @@ export class OAuthSessionManager implements vscode.Disposable { registration, ); + progress.report({ increment: 30 }); this.logger.info("OAuth login flow completed successfully"); return tokenResponse; diff --git a/src/promptUtils.ts b/src/promptUtils.ts index 4d058f12..ce488753 100644 --- a/src/promptUtils.ts +++ b/src/promptUtils.ts @@ -1,7 +1,11 @@ import { type WorkspaceAgent } from "coder/site/src/api/typesGenerated"; import * as vscode from "vscode"; +import { type CoderApi } from "./api/coderApi"; import { type MementoManager } from "./core/mementoManager"; +import { OAuthMetadataClient } from "./oauth/metadataClient"; + +type AuthMethod = "oauth" | "legacy"; /** * Find the requested agent if specified, otherwise return the agent if there @@ -129,3 +133,54 @@ export async function maybeAskUrl( } return url; } + +export async function maybeAskAuthMethod( + client: CoderApi, +): Promise { + // Check if server supports OAuth with progress indication + const supportsOAuth = await vscode.window.withProgress( + { + location: vscode.ProgressLocation.Notification, + title: "Checking authentication methods", + cancellable: false, + }, + async () => { + return await OAuthMetadataClient.checkOAuthSupport( + client.getAxiosInstance(), + ); + }, + ); + + if (supportsOAuth) { + return await askAuthMethod(); + } else { + return "legacy"; + } +} + +/** + * Ask user to choose between OAuth and legacy API token authentication. + */ +async function askAuthMethod(): Promise { + const choice = await vscode.window.showQuickPick( + [ + { + label: "OAuth (Recommended)", + description: "Secure authentication with automatic token refresh", + value: "oauth" as const, + }, + { + label: "Session Token (Legacy)", + description: "Generate and paste a session token manually", + value: "legacy" as const, + }, + ], + { + title: "Select authentication method", + placeHolder: "How would you like to authenticate?", + ignoreFocusOut: true, + }, + ); + + return choice?.value; +} From 8aa1e768b150c4fee2dd59c6c730a378a40ad7c6 Mon Sep 17 00:00:00 2001 From: Ehab Younes Date: Wed, 19 Nov 2025 18:55:09 +0300 Subject: [PATCH 13/20] Better error handling on 400/401/403 due to auth --- src/api/coderApi.ts | 110 +++++++++++++++++++++++------------- src/oauth/errors.ts | 7 --- src/oauth/sessionManager.ts | 80 ++++++++++++++------------ 3 files changed, 116 insertions(+), 81 deletions(-) diff --git a/src/api/coderApi.ts b/src/api/coderApi.ts index d9c30868..f543667a 100644 --- a/src/api/coderApi.ts +++ b/src/api/coderApi.ts @@ -4,6 +4,7 @@ import { type AxiosHeaders, type AxiosResponseTransformer, isAxiosError, + type AxiosError, } from "axios"; import { Api } from "coder/site/src/api/api"; import { @@ -32,11 +33,7 @@ import { HttpClientLogLevel, } from "../logging/types"; import { sizeOf } from "../logging/utils"; -import { - parseOAuthError, - requiresReAuthentication, - isNetworkError, -} from "../oauth/errors"; +import { parseOAuthError, requiresReAuthentication } from "../oauth/errors"; import { type OAuthSessionManager } from "../oauth/sessionManager"; import { HttpStatusCode } from "../websocket/codes"; import { @@ -493,7 +490,7 @@ function addLoggingInterceptors(client: AxiosInstance, logger: Logger) { /** * Add OAuth token refresh interceptors. * Success interceptor: proactively refreshes token when approaching expiry. - * Error interceptor: reactively refreshes token on 401/403 responses. + * Error interceptor: reactively refreshes token on 401 responses. */ function addOAuthInterceptors( client: CoderApi, @@ -516,53 +513,90 @@ function addOAuthInterceptors( return response; }, - // Error response interceptor: reactive token refresh on 401/403 + // Error response interceptor: reactive token refresh on 401 async (error: unknown) => { if (!isAxiosError(error)) { throw error; } - const status = error.response?.status; - if (status !== 401 && status !== 403) { - throw error; + if (error.config) { + const config = error.config as { + _oauthRetryAttempted?: boolean; + }; + if (config._oauthRetryAttempted) { + throw error; + } } - if (!oauthSessionManager.isLoggedInWithOAuth()) { + const status = error.response?.status; + + // These could indicate permanent auth failures that won't be fixed by token refresh + if (status === 400 || status === 403) { + handlePossibleOAuthError(error, logger, oauthSessionManager); throw error; + } else if (status === 401) { + return handle401Error(error, client, logger, oauthSessionManager); } - logger.info(`Received ${status} response, attempting token refresh`); - - try { - const newTokens = await oauthSessionManager.refreshToken(); - client.setSessionToken(newTokens.access_token); - - logger.info("Token refresh successful, updated session token"); - } catch (refreshError) { - logger.error("Token refresh failed:", refreshError); - - const oauthError = parseOAuthError(refreshError); - if (oauthError && requiresReAuthentication(oauthError)) { - logger.error( - `OAuth error requires re-authentication: ${oauthError.errorCode}`, - ); - - oauthSessionManager - .showReAuthenticationModal(oauthError) - .catch((err) => { - logger.error("Failed to show re-auth modal:", err); - }); - } else if (isNetworkError(refreshError)) { - logger.warn( - "Token refresh failed due to network error, will retry later", - ); - } - } throw error; }, ); } +function handlePossibleOAuthError( + error: unknown, + logger: Logger, + oauthSessionManager: OAuthSessionManager, +): void { + const oauthError = parseOAuthError(error); + if (oauthError && requiresReAuthentication(oauthError)) { + logger.error( + `OAuth error requires re-authentication: ${oauthError.errorCode}`, + ); + + oauthSessionManager.showReAuthenticationModal(oauthError).catch((err) => { + logger.error("Failed to show re-auth modal:", err); + }); + } +} + +async function handle401Error( + error: AxiosError, + client: CoderApi, + logger: Logger, + oauthSessionManager: OAuthSessionManager, +): Promise { + if (!oauthSessionManager.isLoggedInWithOAuth()) { + throw error; + } + + logger.info("Received 401 response, attempting token refresh"); + + try { + const newTokens = await oauthSessionManager.refreshToken(); + client.setSessionToken(newTokens.access_token); + + logger.info("Token refresh successful, retrying request"); + + // Retry the original request with the new token + if (error.config) { + const config = error.config as RequestConfigWithMeta & { + _oauthRetryAttempted?: boolean; + }; + config._oauthRetryAttempted = true; + config.headers[coderSessionTokenHeader] = newTokens.access_token; + return client.getAxiosInstance().request(config); + } + + throw error; + } catch (refreshError) { + logger.error("Token refresh failed:", refreshError); + + handlePossibleOAuthError(refreshError, logger, oauthSessionManager); + throw error; + } +} + function wrapRequestTransform( transformer: AxiosResponseTransformer | AxiosResponseTransformer[], config: RequestConfigWithMeta, diff --git a/src/oauth/errors.ts b/src/oauth/errors.ts index 67c7cd47..9b7ee3ac 100644 --- a/src/oauth/errors.ts +++ b/src/oauth/errors.ts @@ -164,10 +164,3 @@ export function requiresReAuthentication(error: OAuthError): boolean { error instanceof InvalidGrantError || error instanceof InvalidClientError ); } - -/** - * Checks if an error is a network/connectivity error - */ -export function isNetworkError(error: unknown): boolean { - return isAxiosError(error) && !error.response && Boolean(error.request); -} diff --git a/src/oauth/sessionManager.ts b/src/oauth/sessionManager.ts index c44ba331..ffa5f899 100644 --- a/src/oauth/sessionManager.ts +++ b/src/oauth/sessionManager.ts @@ -61,7 +61,7 @@ const DEFAULT_OAUTH_SCOPES = [ */ export class OAuthSessionManager implements vscode.Disposable { private storedTokens: StoredOAuthTokens | undefined; - private refreshInProgress = false; + private refreshPromise: Promise | null = null; private lastRefreshAttempt = 0; private pendingAuthReject: ((reason: Error) => void) | undefined; @@ -134,7 +134,7 @@ export class OAuthSessionManager implements vscode.Disposable { */ private async clearTokenState(): Promise { this.storedTokens = undefined; - this.refreshInProgress = false; + this.refreshPromise = null; this.lastRefreshAttempt = 0; await this.secretsManager.setOAuthTokens(undefined); await this.secretsManager.setOAuthClientRegistration(undefined); @@ -489,56 +489,64 @@ export class OAuthSessionManager implements vscode.Disposable { /** * Refresh the access token using the stored refresh token. - * Uses a mutex to prevent concurrent refresh attempts. + * Uses a shared promise to handle concurrent refresh attempts. */ async refreshToken(): Promise { - if (this.refreshInProgress) { - throw new Error("Token refresh already in progress"); + // If a refresh is already in progress, return the existing promise + if (this.refreshPromise) { + this.logger.debug( + "Token refresh already in progress, waiting for result", + ); + return this.refreshPromise; } if (!this.storedTokens?.refresh_token) { throw new Error("No refresh token available"); } - this.refreshInProgress = true; + const refreshToken = this.storedTokens.refresh_token; + const accessToken = this.storedTokens.access_token; + this.lastRefreshAttempt = Date.now(); - try { - const { axiosInstance, metadata, registration } = - await this.prepareOAuthOperation( - this.deploymentUrl, - this.storedTokens.access_token, - ); + // Create and store the refresh promise + this.refreshPromise = (async () => { + try { + const { axiosInstance, metadata, registration } = + await this.prepareOAuthOperation(this.deploymentUrl, accessToken); - this.logger.debug("Refreshing access token"); + this.logger.debug("Refreshing access token"); - const params: RefreshTokenRequestParams = { - grant_type: REFRESH_GRANT_TYPE, - refresh_token: this.storedTokens.refresh_token, - client_id: registration.client_id, - client_secret: registration.client_secret, - }; + const params: RefreshTokenRequestParams = { + grant_type: REFRESH_GRANT_TYPE, + refresh_token: refreshToken, + client_id: registration.client_id, + client_secret: registration.client_secret, + }; - const tokenRequest = toUrlSearchParams(params); + const tokenRequest = toUrlSearchParams(params); - const response = await axiosInstance.post( - metadata.token_endpoint, - tokenRequest, - { - headers: { - "Content-Type": "application/x-www-form-urlencoded", + const response = await axiosInstance.post( + metadata.token_endpoint, + tokenRequest, + { + headers: { + "Content-Type": "application/x-www-form-urlencoded", + }, }, - }, - ); + ); - this.logger.debug("Token refresh successful"); + this.logger.debug("Token refresh successful"); - await this.saveTokens(response.data); + await this.saveTokens(response.data); - return response.data; - } finally { - this.refreshInProgress = false; - } + return response.data; + } finally { + this.refreshPromise = null; + } + })(); + + return this.refreshPromise; } /** @@ -579,7 +587,7 @@ export class OAuthSessionManager implements vscode.Disposable { if ( !this.isLoggedInWithOAuth() || !this.storedTokens?.refresh_token || - this.refreshInProgress + this.refreshPromise !== null ) { return false; } @@ -695,7 +703,7 @@ export class OAuthSessionManager implements vscode.Disposable { } this.pendingAuthReject = undefined; this.storedTokens = undefined; - this.refreshInProgress = false; + this.refreshPromise = null; this.lastRefreshAttempt = 0; this.logger.debug("OAuth session manager disposed"); From fe5ad6216e8712f279fd0519382919ee828edc4d Mon Sep 17 00:00:00 2001 From: Ehab Younes Date: Sun, 23 Nov 2025 18:28:53 +0300 Subject: [PATCH 14/20] Handle per deployment login and centralized login logic --- src/commands.ts | 213 +++---------- src/core/container.ts | 11 + src/core/secretsManager.ts | 422 +++++++++++++++++++------- src/extension.ts | 138 ++++++--- src/login/loginCoordinator.ts | 305 +++++++++++++++++++ src/oauth/sessionManager.ts | 118 ++++--- src/promptUtils.ts | 9 +- src/remote/remote.ts | 342 +++++++++++---------- test/unit/core/cliManager.test.ts | 46 --- test/unit/core/secretsManager.test.ts | 88 ++++-- 10 files changed, 1095 insertions(+), 597 deletions(-) create mode 100644 src/login/loginCoordinator.ts diff --git a/src/commands.ts b/src/commands.ts index 4cf0289a..04450600 100644 --- a/src/commands.ts +++ b/src/commands.ts @@ -1,15 +1,11 @@ import { type Api } from "coder/site/src/api/api"; -import { getErrorMessage } from "coder/site/src/api/errors"; import { - type User, type Workspace, type WorkspaceAgent, } from "coder/site/src/api/typesGenerated"; import * as vscode from "vscode"; import { createWorkspaceIdentifier, extractAgents } from "./api/api-helper"; -import { CoderApi } from "./api/coderApi"; -import { needToken } from "./api/utils"; import { type CliManager } from "./core/cliManager"; import { type ServiceContainer } from "./core/container"; import { type ContextManager } from "./core/contextManager"; @@ -19,8 +15,9 @@ import { type SecretsManager } from "./core/secretsManager"; import { CertificateError } from "./error"; import { getGlobalFlags } from "./globalFlags"; import { type Logger } from "./logging/logger"; +import { type LoginCoordinator } from "./login/loginCoordinator"; import { type OAuthSessionManager } from "./oauth/sessionManager"; -import { maybeAskAgent, maybeAskUrl, maybeAskAuthMethod } from "./promptUtils"; +import { maybeAskAgent, maybeAskUrl } from "./promptUtils"; import { escapeCommandArg, toRemoteAuthority, toSafeHost } from "./util"; import { AgentTreeItem, @@ -36,6 +33,8 @@ export class Commands { private readonly secretsManager: SecretsManager; private readonly cliManager: CliManager; private readonly contextManager: ContextManager; + private readonly loginCoordinator: LoginCoordinator; + // These will only be populated when actively connected to a workspace and are // used in commands. Because commands can be executed by the user, it is not // possible to pass in arguments, so we have to store the current workspace @@ -59,6 +58,7 @@ export class Commands { this.secretsManager = serviceContainer.getSecretsManager(); this.cliManager = serviceContainer.getCliManager(); this.contextManager = serviceContainer.getContextManager(); + this.loginCoordinator = serviceContainer.getLoginCoordinator(); } /** @@ -79,42 +79,49 @@ export class Commands { const url = await maybeAskUrl(this.mementoManager, args?.url); if (!url) { - return; // The user aborted. + return; } // 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 = args?.label ?? toSafeHost(url); - // Try to get a token from the user, if we need one, and their user. - const autoLogin = args?.autoLogin === true; + this.logger.info("Using deployment label", label); + + const result = await this.loginCoordinator.promptForLogin({ + url, + label, + autoLogin: args?.autoLogin, + oauthSessionManager: this.oauthSessionManager, + }); - const res = await this.attemptLogin(url, args?.token, autoLogin); - if (!res) { - return; // The user aborted, or unable to auth. + if (!result.success || !result.user || !result.token) { + return; } - // The URL is good and the token is either good or not required; authorize - // the global client. + // Authorize the global client this.restClient.setHost(url); - this.restClient.setSessionToken(res.token); + this.restClient.setSessionToken(result.token); - // Store these to be used in later sessions. + // Store for later sessions await this.mementoManager.setUrl(url); - await this.secretsManager.setSessionToken(res.token); + await this.secretsManager.setSessionToken(label, { + url, + sessionToken: result.token, + }); - // Store on disk to be used by the cli. - await this.cliManager.configure(label, url, res.token); + // Store on disk for CLI + await this.cliManager.configure(label, url, result.token); - // These contexts control various menu items and the sidebar. + // Update contexts this.contextManager.set("coder.authenticated", true); - if (res.user.roles.some((role) => role.name === "owner")) { + if (result.user.roles.some((role) => role.name === "owner")) { this.contextManager.set("coder.isOwner", true); } vscode.window .showInformationMessage( - `Welcome to Coder, ${res.user.username}!`, + `Welcome to Coder, ${result.user.username}!`, { detail: "You can now use the Coder extension to manage your Coder instance.", @@ -127,160 +134,10 @@ export class Commands { } }); - await this.secretsManager.triggerLoginStateChange("login"); - // Fetch workspaces for the new deployment. + await this.secretsManager.triggerLoginStateChange(label, "login"); vscode.commands.executeCommand("coder.refreshWorkspaces"); } - /** - * Attempt to authenticate using OAuth, token, or mTLS. If necessary, prompts - * for authentication method and credentials. Returns the token and user upon - * successful authentication. Null means the user aborted or authentication - * failed (in which case an error notification will have been displayed). - */ - private async attemptLogin( - url: string, - token: string | undefined, - isAutoLogin: boolean, - ): Promise<{ user: User; token: string } | null> { - const client = CoderApi.create(url, token, this.logger); - const needsToken = needToken(vscode.workspace.getConfiguration()); - if (!needsToken || token) { - 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) { - const message = getErrorMessage(err, "no response from the server"); - if (isAutoLogin) { - this.logger.warn("Failed to log in to Coder server:", message); - } else { - this.vscodeProposed.window.showErrorMessage( - "Failed to log in to Coder server", - { - detail: message, - modal: true, - useCustom: true, - }, - ); - } - // Invalid certificate, most likely. - return null; - } - } - - const authMethod = await maybeAskAuthMethod(client); - switch (authMethod) { - case "oauth": - return this.loginWithOAuth(client); - case "legacy": { - const initialToken = - token || (await this.secretsManager.getSessionToken()); - return this.loginWithToken(client, initialToken); - } - case undefined: - return null; // User aborted - } - } - - private async loginWithToken( - client: CoderApi, - initialToken: string | undefined, - ): Promise<{ user: User; token: string } | null> { - const url = client.getAxiosInstance().defaults.baseURL; - if (!url) { - throw new Error("No base URL set on REST client"); - } - // This prompt is for convenience; do not error if they close it since - // they may already have a token or already have the page opened. - await vscode.env.openExternal(vscode.Uri.parse(`${url}/cli-auth`)); - - // For token auth, start with the existing token in the prompt or the last - // used token. Once submitted, if there is a failure we will keep asking - // the user for a new token until they quit. - let user: User | undefined; - const validatedToken = await vscode.window.showInputBox({ - title: "Coder API Key", - password: true, - placeHolder: "Paste your API key.", - value: initialToken, - ignoreFocusOut: true, - validateInput: async (value) => { - if (!value) { - return null; - } - client.setSessionToken(value); - try { - user = await client.getAuthenticatedUser(); - } catch (err) { - // For certificate errors show both a notification and add to the - // text under the input box, since users sometimes miss the - // notification. - if (err instanceof CertificateError) { - err.showNotification(); - - return { - message: err.x509Err || err.message, - severity: vscode.InputBoxValidationSeverity.Error, - }; - } - // This could be something like the header command erroring or an - // invalid session token. - const message = getErrorMessage(err, "no response from the server"); - return { - message: "Failed to authenticate: " + message, - severity: vscode.InputBoxValidationSeverity.Error, - }; - } - }, - }); - - if (user === undefined || validatedToken === undefined) { - return null; - } - - return { user, token: validatedToken }; - } - - /** - * Authenticate using OAuth flow. - * Returns the access token and authenticated user, or null if failed/cancelled. - */ - private async loginWithOAuth( - client: CoderApi, - ): Promise<{ user: User; token: string } | null> { - try { - this.logger.info("Starting OAuth authentication"); - - const tokenResponse = await vscode.window.withProgress( - { - location: vscode.ProgressLocation.Notification, - title: "Authenticating", - cancellable: false, - }, - async (progress) => - await this.oauthSessionManager.login(client, progress), - ); - - // Validate token by fetching user - client.setSessionToken(tokenResponse.access_token); - const user = await client.getAuthenticatedUser(); - - return { - token: tokenResponse.access_token, - user, - }; - } catch (error) { - this.logger.error("OAuth authentication failed:", error); - vscode.window.showErrorMessage( - `OAuth authentication failed: ${getErrorMessage(error, "Unknown error")}`, - ); - return null; - } - } - /** * View the logs for the currently connected workspace. */ @@ -316,15 +173,16 @@ export class Commands { throw new Error("You are not logged in"); } - await this.forceLogout(); + await this.forceLogout(toSafeHost(url)); } - public async forceLogout(): Promise { + public async forceLogout(label: string): Promise { if (!this.contextManager.get("coder.authenticated")) { return; } - this.logger.info("Logging out"); + this.logger.info(`Logging out of deployment: ${label}`); + // Only clear REST client and UI context if logging out of current deployment // Fire and forget this.oauthSessionManager.logout().catch((error) => { this.logger.warn("OAuth logout failed, continuing with cleanup:", error); @@ -337,7 +195,7 @@ export class Commands { // Clear from memory. await this.mementoManager.setUrl(undefined); - await this.secretsManager.setSessionToken(undefined); + await this.secretsManager.setSessionToken(label, undefined); this.contextManager.set("coder.authenticated", false); vscode.window @@ -348,9 +206,10 @@ export class Commands { } }); - await this.secretsManager.triggerLoginStateChange("logout"); // This will result in clearing the workspace list. vscode.commands.executeCommand("coder.refreshWorkspaces"); + + await this.secretsManager.triggerLoginStateChange(label, "logout"); } /** diff --git a/src/core/container.ts b/src/core/container.ts index a8f938ea..d83ae679 100644 --- a/src/core/container.ts +++ b/src/core/container.ts @@ -1,6 +1,7 @@ import * as vscode from "vscode"; import { type Logger } from "../logging/logger"; +import { LoginCoordinator } from "../login/loginCoordinator"; import { CliManager } from "./cliManager"; import { ContextManager } from "./contextManager"; @@ -19,6 +20,7 @@ export class ServiceContainer implements vscode.Disposable { private readonly secretsManager: SecretsManager; private readonly cliManager: CliManager; private readonly contextManager: ContextManager; + private readonly loginCoordinator: LoginCoordinator; constructor( context: vscode.ExtensionContext, @@ -37,6 +39,11 @@ export class ServiceContainer implements vscode.Disposable { this.pathResolver, ); this.contextManager = new ContextManager(); + this.loginCoordinator = new LoginCoordinator( + this.secretsManager, + this.vscodeProposed, + this.logger, + ); } getVsCodeProposed(): typeof vscode { @@ -67,6 +74,10 @@ export class ServiceContainer implements vscode.Disposable { return this.contextManager; } + getLoginCoordinator(): LoginCoordinator { + return this.loginCoordinator; + } + /** * Dispose of all services and clean up resources. */ diff --git a/src/core/secretsManager.ts b/src/core/secretsManager.ts index cdda9c72..f4abdb45 100644 --- a/src/core/secretsManager.ts +++ b/src/core/secretsManager.ts @@ -2,6 +2,7 @@ import { type TokenResponse, type ClientRegistrationResponse, } from "../oauth/types"; +import { toSafeHost } from "../util"; import type { SecretStorage, Disposable } from "vscode"; @@ -9,17 +10,29 @@ const SESSION_TOKEN_KEY = "sessionToken"; const LOGIN_STATE_KEY = "loginState"; -const OAUTH_CLIENT_REGISTRATION_KEY = "oauthClientRegistration"; - -const OAUTH_TOKENS_KEY = "oauthTokens"; - const OAUTH_CALLBACK_KEY = "coder.oauthCallback"; +const SESSION_AUTH_MAP_KEY = "coder.sessionAuthMap"; +const OAUTH_DATA_MAP_KEY = "coder.oauthDataMap"; + export type StoredOAuthTokens = Omit & { expiry_timestamp: number; deployment_url: string; }; +export interface SessionAuth { + url: string; + sessionToken: string; +} + +export interface OAuthData { + oauthClientRegistration?: ClientRegistrationResponse; + oauthTokens?: StoredOAuthTokens; +} + +export type SessionAuthMap = Record; +export type OAuthDataMap = Record; + interface OAuthCallbackData { state: string; code: string | null; @@ -33,174 +46,379 @@ export enum AuthAction { } export class SecretsManager { - constructor(private readonly secrets: SecretStorage) {} - /** - * Set or unset the last used token. + * Track previous session tokens to detect actual changes. + * Maps label -> previous sessionToken value. */ - public async setSessionToken( - sessionToken: string | undefined, - ): Promise { - if (sessionToken) { - await this.secrets.store(SESSION_TOKEN_KEY, sessionToken); - } else { - await this.secrets.delete(SESSION_TOKEN_KEY); - } - } + private readonly previousSessionTokens = new Map< + string, + string | undefined + >(); - /** - * Get the last used token. - */ - public async getSessionToken(): Promise { - try { - 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; - } + constructor(private readonly secrets: SecretStorage) { + // Initialize previous session tokens + this.getSessionAuthMap().then((map) => { + for (const [label, auth] of Object.entries(map)) { + this.previousSessionTokens.set(label, auth.sessionToken); + } + }); } /** * Triggers a login/logout event that propagates across all VS Code windows. * Uses the secrets storage onDidChange event as a cross-window communication mechanism. - * Appends a timestamp to ensure the value always changes, guaranteeing the event fires. + * Stores JSON with action, label, and timestamp to ensure the value always changes. */ public async triggerLoginStateChange( + label: string, action: "login" | "logout", ): Promise { - const date = new Date().toISOString(); - await this.secrets.store(LOGIN_STATE_KEY, `${action}-${date}`); + const loginState = { + action, + label, + timestamp: new Date().toISOString(), + }; + await this.secrets.store(LOGIN_STATE_KEY, JSON.stringify(loginState)); } /** * Listens for login/logout events from any VS Code window. * The secrets storage onDidChange event fires across all windows, enabling cross-window sync. + * Parses JSON to extract action and label. */ public onDidChangeLoginState( - listener: (state: AuthAction) => Promise, + listener: (state: AuthAction, label: string) => Promise, ): Disposable { return this.secrets.onDidChange(async (e) => { if (e.key === LOGIN_STATE_KEY) { - const state = await this.secrets.get(LOGIN_STATE_KEY); - if (state?.startsWith("login")) { - listener(AuthAction.LOGIN); - } else if (state?.startsWith("logout")) { - listener(AuthAction.LOGOUT); - } else { - // Secret was deleted or is invalid - listener(AuthAction.INVALID); + const stateStr = await this.secrets.get(LOGIN_STATE_KEY); + if (!stateStr) { + await listener(AuthAction.INVALID, ""); + return; + } + + try { + const parsed = JSON.parse(stateStr) as { + action: string; + label: string; + timestamp: string; + }; + + if (parsed.action === "login") { + await listener(AuthAction.LOGIN, parsed.label); + } else if (parsed.action === "logout") { + await listener(AuthAction.LOGOUT, parsed.label); + } else { + await listener(AuthAction.INVALID, parsed.label); + } + } catch { + // Invalid JSON, treat as invalid state + await listener(AuthAction.INVALID, ""); + } + } + }); + } + + /** + * Write an OAuth callback result to secrets storage. + * Used for cross-window communication when OAuth callback arrives in a different window. + */ + public async setOAuthCallback(data: OAuthCallbackData): Promise { + await this.secrets.store(OAUTH_CALLBACK_KEY, JSON.stringify(data)); + } + + /** + * Listen for OAuth callback results from any VS Code window. + * The listener receives the state parameter, code (if success), and error (if failed). + */ + public onDidChangeOAuthCallback( + listener: (data: OAuthCallbackData) => void, + ): Disposable { + return this.secrets.onDidChange(async (e) => { + if (e.key !== OAUTH_CALLBACK_KEY) { + return; + } + + try { + const data = await this.secrets.get(OAUTH_CALLBACK_KEY); + if (data) { + const parsed = JSON.parse(data) as OAuthCallbackData; + listener(parsed); } + } catch { + // Ignore parse errors } }); } /** - * Listens for session token changes. + * Listen for changes to a specific deployment's session auth. + * Only fires when the session token actually changes for this deployment. + * OAuth token/registration changes will NOT trigger this listener. */ - public onDidChangeSessionToken( - listener: (token: string | undefined) => Promise, + public onDidChangeDeploymentAuth( + label: string, + listener: (auth: SessionAuth | undefined) => void | Promise, + ): Disposable { + return this.onDidChangeSessionAuthMap(async (map) => { + const auth = map[label]; + const newToken = auth?.sessionToken ?? ""; + const previousToken = this.previousSessionTokens.get(label) ?? ""; + + // Only fire listener if session token actually changed + if (newToken !== previousToken) { + this.previousSessionTokens.set(label, newToken); + await listener(auth); + } + }); + } + + /** + * Listen for changes to the session auth map across all deployments. + * Fires whenever any deployment's session auth is updated. + */ + private onDidChangeSessionAuthMap( + listener: (map: SessionAuthMap) => void | Promise, ): Disposable { return this.secrets.onDidChange(async (e) => { - if (e.key === SESSION_TOKEN_KEY) { - const token = await this.getSessionToken(); - await listener(token); + if (e.key !== SESSION_AUTH_MAP_KEY) { + return; + } + + try { + const map = await this.getSessionAuthMap(); + await listener(map); + } catch { + // Ignore errors in listener } }); } /** - * Store OAuth client registration data. + * Get session token for a specific deployment. + */ + public async getSessionToken(label: string): Promise { + const map = await this.getSessionAuthMap(); + return map[label]?.sessionToken; + } + + /** + * Set session token for a specific deployment. + */ + public async setSessionToken( + label: string, + auth: { url: string; sessionToken: string } | undefined, + ): Promise { + await this.updateSessionAuthMap((map) => { + if (auth === undefined) { + const newMap = { ...map }; + delete newMap[label]; + return newMap; + } + + return { + ...map, + [label]: auth, + }; + }); + } + + /** + * Get OAuth tokens for a specific deployment. + */ + public async getOAuthTokens( + label: string, + ): Promise { + const map = await this.getOAuthDataMap(); + return map[label]?.oauthTokens; + } + + /** + * Set OAuth tokens for a specific deployment. + */ + public async setOAuthTokens( + label: string, + tokens: StoredOAuthTokens | undefined, + ): Promise { + await this.updateOAuthDataMap((map) => { + const existing = map[label] || {}; + return { + ...map, + [label]: { + ...existing, + oauthTokens: tokens, + }, + }; + }); + } + + /** + * Get OAuth client registration for a specific deployment. + */ + public async getOAuthClientRegistration( + label: string, + ): Promise { + const map = await this.getOAuthDataMap(); + return map[label]?.oauthClientRegistration; + } + + /** + * Set OAuth client registration for a specific deployment. + * Creates the OAuth data entry if it doesn't exist. */ public async setOAuthClientRegistration( + label: string, registration: ClientRegistrationResponse | undefined, ): Promise { - if (registration) { - await this.secrets.store( - OAUTH_CLIENT_REGISTRATION_KEY, - JSON.stringify(registration), - ); - } else { - await this.secrets.delete(OAUTH_CLIENT_REGISTRATION_KEY); - } + await this.updateOAuthDataMap((map) => { + const existing = map[label] || {}; + return { + ...map, + [label]: { + ...existing, + oauthClientRegistration: registration, + }, + }; + }); + } + + public async clearOAuthData(label: string): Promise { + await this.updateOAuthDataMap((map) => { + const newMap = { ...map }; + delete newMap[label]; + return newMap; + }); } /** - * Get OAuth client registration data. + * Get the session auth map for all deployments. */ - public async getOAuthClientRegistration(): Promise< - ClientRegistrationResponse | undefined - > { + private async getSessionAuthMap(): Promise { try { - const stringifiedResponse = await this.secrets.get( - OAUTH_CLIENT_REGISTRATION_KEY, - ); - if (stringifiedResponse) { - return JSON.parse(stringifiedResponse) as ClientRegistrationResponse; + const data = await this.secrets.get(SESSION_AUTH_MAP_KEY); + if (data) { + return JSON.parse(data) as SessionAuthMap; } } catch { - // Do nothing + // Ignore parse errors } - return undefined; + return {}; } /** - * Store OAuth token data including expiry timestamp. + * Set the session auth map for all deployments. */ - public async setOAuthTokens( - tokens: StoredOAuthTokens | undefined, - ): Promise { - if (tokens) { - await this.secrets.store(OAUTH_TOKENS_KEY, JSON.stringify(tokens)); - } else { - await this.secrets.delete(OAUTH_TOKENS_KEY); - } + private async setSessionAuthMap(map: SessionAuthMap): Promise { + await this.secrets.store(SESSION_AUTH_MAP_KEY, JSON.stringify(map)); } /** - * Get stored OAuth token data. + * Get the OAuth data map for all deployments. */ - public async getOAuthTokens(): Promise { + private async getOAuthDataMap(): Promise { try { - const stringifiedTokens = await this.secrets.get(OAUTH_TOKENS_KEY); - if (stringifiedTokens) { - return JSON.parse(stringifiedTokens) as StoredOAuthTokens; + const data = await this.secrets.get(OAUTH_DATA_MAP_KEY); + if (data) { + return JSON.parse(data) as OAuthDataMap; } } catch { - // Do nothing + // Ignore parse errors } - return undefined; + return {}; } /** - * Write an OAuth callback result to secrets storage. - * Used for cross-window communication when OAuth callback arrives in a different window. + * Set the OAuth data map for all deployments. */ - public async setOAuthCallback(data: OAuthCallbackData): Promise { - await this.secrets.store(OAUTH_CALLBACK_KEY, JSON.stringify(data)); + private async setOAuthDataMap(map: OAuthDataMap): Promise { + await this.secrets.store(OAUTH_DATA_MAP_KEY, JSON.stringify(map)); } /** - * Listen for OAuth callback results from any VS Code window. - * The listener receives the state parameter, code (if success), and error (if failed). + * Promise used for synchronizing session auth map updates. */ - public onDidChangeOAuthCallback( - listener: (data: OAuthCallbackData) => void, - ): Disposable { - return this.secrets.onDidChange(async (e) => { - if (e.key !== OAUTH_CALLBACK_KEY) { - return; + private sessionAuthUpdatePromise: Promise = Promise.resolve(); + + /** + * Promise used for synchronizing OAuth data map updates. + */ + private oauthDataUpdatePromise: Promise = Promise.resolve(); + + /** + * Atomically update the session auth map using a synchronized updater function. + * All write operations should go through this method to prevent race conditions. + */ + private async updateSessionAuthMap( + updater: (map: SessionAuthMap) => SessionAuthMap, + ): Promise { + this.sessionAuthUpdatePromise = this.sessionAuthUpdatePromise.then( + async () => { + const currentMap = await this.getSessionAuthMap(); + const newMap = updater(currentMap); + await this.setSessionAuthMap(newMap); + }, + ); + + return this.sessionAuthUpdatePromise; + } + + /** + * Atomically update the OAuth data map using a synchronized updater function. + * All write operations should go through this method to prevent race conditions. + */ + private async updateOAuthDataMap( + updater: (map: OAuthDataMap) => OAuthDataMap, + ): Promise { + this.oauthDataUpdatePromise = this.oauthDataUpdatePromise.then(async () => { + const currentMap = await this.getOAuthDataMap(); + const newMap = updater(currentMap); + await this.setOAuthDataMap(newMap); + }); + + return this.oauthDataUpdatePromise; + } + + /** + * Migrate from old flat storage format to new label-based format. + * This is a one-time operation that runs on extension activation. + * + * @param url The deployment URL to use for generating the label + * @returns true if migration was performed, false if already migrated or nothing to migrate + */ + public async migrateFromLegacyStorage(url: string): Promise { + try { + // Check if already migrated (new map exists and has data) + const existingMap = await this.getSessionAuthMap(); + if (Object.keys(existingMap).length > 0) { + return false; // Already migrated } - try { - const data = await this.secrets.get(OAUTH_CALLBACK_KEY); - if (data) { - const parsed = JSON.parse(data) as OAuthCallbackData; - listener(parsed); - } - } catch { - // Ignore parse errors + // Directly access old session token from flat storage + const oldToken = await this.secrets.get(SESSION_TOKEN_KEY); + if (!oldToken) { + return false; // Nothing to migrate } - }); + + // Generate label from URL + const label = toSafeHost(url); + + // Create new session auth map with migrated token + const sessionAuthMap: SessionAuthMap = { + [label]: { + url: url, + sessionToken: oldToken, + }, + }; + + // Write new map to secrets + await this.setSessionAuthMap(sessionAuthMap); + + // Delete old session token key + await this.secrets.delete(SESSION_TOKEN_KEY); + + return true; // Migration successful + } catch (error) { + throw new Error(`Auth storage migration failed: ${error}`); + } } } diff --git a/src/extension.ts b/src/extension.ts index 837b8d64..38475a35 100644 --- a/src/extension.ts +++ b/src/extension.ts @@ -62,16 +62,21 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { const secretsManager = serviceContainer.getSecretsManager(); const contextManager = serviceContainer.getContextManager(); + // Migrate auth storage from old flat format to new label-based format + await migrateAuthStorage(serviceContainer); + // Try to clear this flag ASAP const isFirstConnect = await mementoManager.getAndClearFirstConnect(); const url = mementoManager.getUrl(); + const label = url ? toSafeHost(url) : ""; - // Create OAuth session manager before the main client + // Create OAuth session manager with login coordinator const oauthSessionManager = await OAuthSessionManager.create( url || "", + label, serviceContainer, - ctx, + ctx.extension.id, ); ctx.subscriptions.push(oauthSessionManager); @@ -80,7 +85,7 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { // in commands that operate on the current login. const client = CoderApi.create( url || "", - await secretsManager.getSessionToken(), + (await secretsManager.getSessionToken(label)) || "", output, oauthSessionManager, ); @@ -128,26 +133,39 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { ctx.subscriptions, ); - // Listen for session token changes and sync state across all components - ctx.subscriptions.push( - secretsManager.onDidChangeSessionToken(async (token) => { - client.setSessionToken(token ?? ""); - if (!token) { - output.debug("Session token cleared"); - return; - } + // Listen for deployment auth changes and sync state across all components + // This listener is re-registered when the user logs into a different deployment + let authChangeDisposable: vscode.Disposable | undefined; + const registerAuthListener = (deploymentLabel: string) => { + authChangeDisposable?.dispose(); - output.debug("Session token changed, syncing state"); + if (!deploymentLabel) { + return; + } - const url = mementoManager.getUrl(); - if (url) { - const cliManager = serviceContainer.getCliManager(); - // TODO label might not match the one in remote? - await cliManager.configure(toSafeHost(url), url, token); - output.debug("Updated CLI config with new token"); - } - }), - ); + output.debug("Registering auth listener for deployment", deploymentLabel); + authChangeDisposable = secretsManager.onDidChangeDeploymentAuth( + deploymentLabel, + async (auth) => { + const token = auth?.sessionToken ?? ""; + client.setSessionToken(token); + + // Update authentication context for current deployment + contextManager.set("coder.authenticated", auth !== undefined); + + output.debug("Session token changed, syncing state"); + + if (auth?.url) { + const cliManager = serviceContainer.getCliManager(); + await cliManager.configure(deploymentLabel, auth.url, token); + output.debug("Updated CLI config with new token"); + } + }, + ); + }; + + registerAuthListener(label); + ctx.subscriptions.push({ dispose: () => authChangeDisposable?.dispose() }); // Handle vscode:// URIs. const uriHandler = vscode.window.registerUriHandler({ @@ -184,11 +202,7 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { // queries will default to localhost) so ask for it if missing. // Pre-populate in case we do have the right URL so the user can just // hit enter and move on. - const url = await maybeAskUrl( - mementoManager, - params.get("url"), - mementoManager.getUrl(), - ); + const url = await maybeAskUrl(mementoManager, params.get("url")); if (url) { client.setHost(url); await mementoManager.setUrl(url); @@ -208,9 +222,13 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { ? params.get("token") : (params.get("token") ?? ""); + const label = toSafeHost(url); if (token) { client.setSessionToken(token); - await secretsManager.setSessionToken(token); + await secretsManager.setSessionToken(label, { + url, + sessionToken: token, + }); } // Store on disk to be used by the cli. @@ -268,11 +286,7 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { // queries will default to localhost) so ask for it if missing. // Pre-populate in case we do have the right URL so the user can just // hit enter and move on. - const url = await maybeAskUrl( - mementoManager, - params.get("url"), - mementoManager.getUrl(), - ); + const url = await maybeAskUrl(mementoManager, params.get("url")); if (url) { client.setHost(url); await mementoManager.setUrl(url); @@ -369,22 +383,32 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { ), ); - const remote = new Remote(serviceContainer, commands, ctx.extensionMode); + const remote = new Remote(serviceContainer, commands, ctx); ctx.subscriptions.push( - secretsManager.onDidChangeLoginState(async (state) => { + secretsManager.onDidChangeLoginState(async (state, label) => { switch (state) { case AuthAction.LOGIN: { - const token = await secretsManager.getSessionToken(); const url = mementoManager.getUrl(); + const token = await secretsManager.getSessionToken(label); // Should login the user directly if the URL+Token are valid await commands.login({ url, token }); - // Resolve any pending login detection promises - remote.resolveLoginDetected(); + + // Re-register auth listener for the new deployment + registerAuthListener(label); + + // Update OAuth session manager to match the new deployment + if (url) { + await oauthSessionManager.setDeployment(label, url); + } break; } case AuthAction.LOGOUT: - await commands.forceLogout(); + await commands.forceLogout(label); + + // Dispose auth listener when logged out + authChangeDisposable?.dispose(); + authChangeDisposable = undefined; break; case AuthAction.INVALID: break; @@ -408,13 +432,14 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { isFirstConnect, remoteSshExtension.id, ); + // TODO this is weird, why do we change the host here?? if (details) { - // TODO if the URL is different then we need to update the OAuth session!!! (Centralize this logic) ctx.subscriptions.push(details); // Authenticate the plugin client which is used in the sidebar to display // workspaces belonging to this deployment. client.setHost(details.url); client.setSessionToken(details.token); + await oauthSessionManager.setDeployment(details.label, details.url); } } catch (ex) { if (ex instanceof CertificateError) { @@ -503,6 +528,41 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { } } +/** + * Migrates old flat storage (sessionToken) to new label-based map storage. + * This is a one-time operation that runs on extension activation. + */ +async function migrateAuthStorage( + serviceContainer: ServiceContainer, +): Promise { + const secretsManager = serviceContainer.getSecretsManager(); + const mementoManager = serviceContainer.getMementoManager(); + const output = serviceContainer.getLogger(); + + try { + // Get deployment URL from memento + const url = mementoManager.getUrl(); + if (!url) { + output.info("No URL configured, skipping migration"); + return; + } + + // Perform migration using SecretsManager method + const migrated = await secretsManager.migrateFromLegacyStorage(url); + + if (migrated) { + output.info( + `Successfully migrated auth storage to label-based format (label: ${toSafeHost(url)})`, + ); + } + } catch (error) { + output.error( + `Auth storage migration failed: ${error}. You may need to log in again.`, + ); + // Don't throw - allow extension to continue + } +} + async function showTreeViewSearch(id: string): Promise { await vscode.commands.executeCommand(`${id}.focus`); await vscode.commands.executeCommand("list.find"); diff --git a/src/login/loginCoordinator.ts b/src/login/loginCoordinator.ts new file mode 100644 index 00000000..a923b12c --- /dev/null +++ b/src/login/loginCoordinator.ts @@ -0,0 +1,305 @@ +import { getErrorMessage } from "coder/site/src/api/errors"; +import * as vscode from "vscode"; + +import { CoderApi } from "../api/coderApi"; +import { needToken } from "../api/utils"; +import { type SecretsManager } from "../core/secretsManager"; +import { CertificateError } from "../error"; +import { type Logger } from "../logging/logger"; +import { maybeAskAuthMethod } from "../promptUtils"; + +import type { User } from "coder/site/src/api/typesGenerated"; + +import type { OAuthSessionManager } from "../oauth/sessionManager"; + +export interface LoginResult { + success: boolean; + user?: User; + token?: string; +} + +export interface LoginOptions { + url: string; + label: string; + oauthSessionManager: OAuthSessionManager; + message?: string; + detailPrefix?: string; + autoLogin?: boolean; +} + +/** + * Coordinates login prompts across windows and prevents duplicate dialogs. + */ +export class LoginCoordinator { + private readonly inProgressLogins = new Map>(); + + constructor( + private readonly secretsManager: SecretsManager, + private readonly vscodeProposed: typeof vscode, + private readonly logger: Logger, + ) {} + + /** + * Direct login - for user-initiated login via commands. + */ + public async promptForLogin( + options: Omit, + ): Promise { + return this.executeWithGuard(options.label, async () => { + const { url, label, oauthSessionManager } = options; + const existingToken = await this.secretsManager.getSessionToken(label); + return this.attemptLogin( + label, + url, + existingToken, + options.autoLogin ?? false, + oauthSessionManager, + ); + }); + } + + /** + * Shows dialog then login - for system-initiated auth (remote, OAuth refresh). + */ + public async promptForLoginWithDialog( + options: LoginOptions, + ): Promise { + return this.executeWithGuard(options.label, () => { + const { + url, + label, + detailPrefix: reason, + message, + oauthSessionManager, + } = options; + + // Show dialog promise + const dialogPromise = this.vscodeProposed.window + .showErrorMessage( + message || "Authentication Required", + { + modal: true, + useCustom: true, + detail: + (reason || `Authentication needed for ${label}`) + + " If you've already logged in, you may close this dialog.", + }, + "Login", + ) + .then(async (action) => { + if (action === "Login") { + // User clicked login - proceed with login flow + const existingToken = + await this.secretsManager.getSessionToken(label); + return this.attemptLogin( + label, + url, + existingToken, + false, + oauthSessionManager, + ); + } else { + // User cancelled + return { success: false }; + } + }); + + // Race between user clicking login and cross-window detection + return Promise.race([dialogPromise, this.waitForCrossWindowLogin(label)]); + }); + } + + /** + * Same-window guard wrapper. + */ + private async executeWithGuard( + label: string, + executeFn: () => Promise, + ): Promise { + const existingLogin = this.inProgressLogins.get(label); + if (existingLogin) { + return existingLogin; + } + + const loginPromise = executeFn(); + this.inProgressLogins.set(label, loginPromise); + + try { + return await loginPromise; + } finally { + this.inProgressLogins.delete(label); + } + } + + /** + * Waits for login detected from another window. + */ + private async waitForCrossWindowLogin(label: string): Promise { + return new Promise((resolve) => { + const disposable = this.secretsManager.onDidChangeDeploymentAuth( + label, + (auth) => { + if (auth?.sessionToken) { + disposable.dispose(); + resolve({ success: true }); + } + }, + ); + }); + } + + /** + * Attempt to authenticate using OAuth, token, or mTLS. If necessary, prompts + * for authentication method and credentials. Returns the token and user upon + * successful authentication. Null means the user aborted or authentication + * failed (in which case an error notification will have been displayed). + */ + private async attemptLogin( + label: string, + url: string, + token: string | undefined, + isAutoLogin: boolean, + oauthSessionManager: OAuthSessionManager, + ): Promise { + const client = CoderApi.create(url, token, this.logger); + const needsToken = needToken(vscode.workspace.getConfiguration()); + if (!needsToken || token) { + 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 { success: true, token: needsToken && token ? token : "", user }; + } catch (err) { + const message = getErrorMessage(err, "no response from the server"); + if (isAutoLogin) { + this.logger.warn("Failed to log in to Coder server:", message); + } else { + this.vscodeProposed.window.showErrorMessage( + "Failed to log in to Coder server", + { + detail: message, + modal: true, + useCustom: true, + }, + ); + } + // Invalid certificate, most likely. + return { success: false }; + } + } + + const authMethod = await maybeAskAuthMethod(client); + switch (authMethod) { + case "oauth": + return this.loginWithOAuth(client, oauthSessionManager, label); + case "legacy": { + const initialToken = + token || (await this.secretsManager.getSessionToken(label)); + return this.loginWithToken(client, initialToken); + } + case undefined: + return { success: false }; // User aborted + } + } + + /** + * Session token authentication flow. + */ + private async loginWithToken( + client: CoderApi, + initialToken: string | undefined, + ): Promise { + const url = client.getAxiosInstance().defaults.baseURL; + if (!url) { + throw new Error("No base URL set on REST client"); + } + // This prompt is for convenience; do not error if they close it since + // they may already have a token or already have the page opened. + await vscode.env.openExternal(vscode.Uri.parse(`${url}/cli-auth`)); + + // For token auth, start with the existing token in the prompt or the last + // used token. Once submitted, if there is a failure we will keep asking + // the user for a new token until they quit. + let user: User | undefined; + const validatedToken = await vscode.window.showInputBox({ + title: "Coder API Key", + password: true, + placeHolder: "Paste your API key.", + value: initialToken, + ignoreFocusOut: true, + validateInput: async (value) => { + if (!value) { + return null; + } + client.setSessionToken(value); + try { + user = await client.getAuthenticatedUser(); + } catch (err) { + // For certificate errors show both a notification and add to the + // text under the input box, since users sometimes miss the + // notification. + if (err instanceof CertificateError) { + err.showNotification(); + return { + message: err.x509Err || err.message, + severity: vscode.InputBoxValidationSeverity.Error, + }; + } + // This could be something like the header command erroring or an + // invalid session token. + const message = getErrorMessage(err, "no response from the server"); + return { + message: "Failed to authenticate: " + message, + severity: vscode.InputBoxValidationSeverity.Error, + }; + } + }, + }); + + if (user === undefined || validatedToken === undefined) { + return { success: false }; + } + + return { success: true, user, token: validatedToken }; + } + + /** + * OAuth authentication flow. + */ + private async loginWithOAuth( + client: CoderApi, + oauthSessionManager: OAuthSessionManager, + label: string, + ): Promise { + try { + this.logger.info("Starting OAuth authentication"); + + const tokenResponse = await vscode.window.withProgress( + { + location: vscode.ProgressLocation.Notification, + title: "Authenticating", + cancellable: false, + }, + async (progress) => + await oauthSessionManager.login(client, label, progress), + ); + + // Validate token by fetching user + client.setSessionToken(tokenResponse.access_token); + const user = await client.getAuthenticatedUser(); + + return { + success: true, + token: tokenResponse.access_token, + user, + }; + } catch (error) { + this.logger.error("OAuth authentication failed:", error); + vscode.window.showErrorMessage( + `OAuth authentication failed: ${getErrorMessage(error, "Unknown error")}`, + ); + return { success: false }; + } + } +} diff --git a/src/oauth/sessionManager.ts b/src/oauth/sessionManager.ts index ffa5f899..8bfb7f56 100644 --- a/src/oauth/sessionManager.ts +++ b/src/oauth/sessionManager.ts @@ -4,6 +4,7 @@ import * as vscode from "vscode"; import { type ServiceContainer } from "src/core/container"; import { CoderApi } from "../api/coderApi"; +import { type LoginCoordinator } from "../login/loginCoordinator"; import { OAuthMetadataClient } from "./metadataClient"; import { @@ -37,6 +38,11 @@ const PKCE_CHALLENGE_METHOD = "S256" as const; */ const TOKEN_REFRESH_THRESHOLD_MS = 10 * 60 * 1000; +/** + * Default expiry time for OAuth access tokens when the server doesn't provide one. + */ +const ACCESS_TOKEN_DEFAULT_EXPIRY_MS = 60 * 60 * 1000; + /** * Minimum time between refresh attempts to prevent thrashing */ @@ -69,17 +75,19 @@ export class OAuthSessionManager implements vscode.Disposable { /** * Create and initialize a new OAuth session manager. */ - static async create( + public static async create( deploymentUrl: string, + label: string, container: ServiceContainer, - context: vscode.ExtensionContext, + extensionId: string, ): Promise { const manager = new OAuthSessionManager( deploymentUrl, + label, container.getSecretsManager(), container.getLogger(), - container.getVsCodeProposed(), - context.extension.id, + container.getLoginCoordinator(), + extensionId, ); await manager.loadTokens(); return manager; @@ -87,9 +95,10 @@ export class OAuthSessionManager implements vscode.Disposable { private constructor( private deploymentUrl: string, + private label: string, private readonly secretsManager: SecretsManager, private readonly logger: Logger, - private readonly vscodeProposed: typeof vscode, + private readonly loginCoordinator: LoginCoordinator, private readonly extensionId: string, ) {} @@ -98,7 +107,7 @@ export class OAuthSessionManager implements vscode.Disposable { * Validates that tokens belong to the current deployment URL. */ private async loadTokens(): Promise { - const tokens = await this.secretsManager.getOAuthTokens(); + const tokens = await this.secretsManager.getOAuthTokens(this.label); if (!tokens) { return; } @@ -121,7 +130,7 @@ export class OAuthSessionManager implements vscode.Disposable { required_scopes: DEFAULT_OAUTH_SCOPES, }, ); - await this.secretsManager.setOAuthTokens(undefined); + await this.secretsManager.setOAuthTokens(this.label, undefined); return; } @@ -129,15 +138,16 @@ export class OAuthSessionManager implements vscode.Disposable { this.logger.info(`Loaded stored OAuth tokens for ${tokens.deployment_url}`); } - /** - * Clear stale data when tokens don't match current deployment. - */ private async clearTokenState(): Promise { + this.clearInMemoryTokens(); + await this.secretsManager.setOAuthTokens(this.label, undefined); + await this.secretsManager.setOAuthClientRegistration(this.label, undefined); + } + + private clearInMemoryTokens(): void { this.storedTokens = undefined; this.refreshPromise = null; this.lastRefreshAttempt = 0; - await this.secretsManager.setOAuthTokens(undefined); - await this.secretsManager.setOAuthClientRegistration(undefined); } /** @@ -199,7 +209,9 @@ export class OAuthSessionManager implements vscode.Disposable { const metadataClient = new OAuthMetadataClient(axiosInstance, this.logger); const metadata = await metadataClient.getMetadata(); - const registration = await this.secretsManager.getOAuthClientRegistration(); + const registration = await this.secretsManager.getOAuthClientRegistration( + this.label, + ); if (!registration) { throw new Error("No client registration found"); } @@ -217,7 +229,9 @@ export class OAuthSessionManager implements vscode.Disposable { ): Promise { const redirectUri = this.getRedirectUri(); - const existing = await this.secretsManager.getOAuthClientRegistration(); + const existing = await this.secretsManager.getOAuthClientRegistration( + this.label, + ); if (existing?.client_id) { if (existing.redirect_uris.includes(redirectUri)) { this.logger.info( @@ -247,7 +261,10 @@ export class OAuthSessionManager implements vscode.Disposable { registrationRequest, ); - await this.secretsManager.setOAuthClientRegistration(response.data); + await this.secretsManager.setOAuthClientRegistration( + this.label, + response.data, + ); this.logger.info( "Saved OAuth client registration:", response.data.client_id, @@ -256,15 +273,23 @@ export class OAuthSessionManager implements vscode.Disposable { return response.data; } + public async setDeployment(label: string, url: string): Promise { + this.logger.debug("Switching OAuth deployment", { label, url }); + this.label = label; + this.deploymentUrl = url; + this.clearInMemoryTokens(); + await this.loadTokens(); + } + /** * Simplified OAuth login flow that handles the entire process. * Fetches metadata, registers client, starts authorization, and exchanges tokens. * - * @param client CoderApi instance for the deployment to authenticate against * @returns TokenResponse containing access token and optional refresh token */ - async login( + public async login( client: CoderApi, + label: string, progress: vscode.Progress<{ message?: string; increment?: number }>, ): Promise { const baseUrl = client.getAxiosInstance().defaults.baseURL; @@ -276,9 +301,17 @@ export class OAuthSessionManager implements vscode.Disposable { old: this.deploymentUrl, new: baseUrl, }); - await this.clearTokenState(); + this.clearInMemoryTokens(); this.deploymentUrl = baseUrl; } + if (this.label && this.label !== label) { + this.logger.info("Deployment label changed, clearing cached state", { + old: this.label, + new: label, + }); + this.clearInMemoryTokens(); + this.label = label; + } const axiosInstance = client.getAxiosInstance(); const metadataClient = new OAuthMetadataClient(axiosInstance, this.logger); @@ -429,7 +462,7 @@ export class OAuthSessionManager implements vscode.Disposable { * Handle OAuth callback from browser redirect. * Writes the callback result to secrets storage, triggering the waiting window to proceed. */ - async handleCallback( + public async handleCallback( code: string | null, state: string | null, error: string | null, @@ -491,7 +524,7 @@ export class OAuthSessionManager implements vscode.Disposable { * Refresh the access token using the stored refresh token. * Uses a shared promise to handle concurrent refresh attempts. */ - async refreshToken(): Promise { + public async refreshToken(): Promise { // If a refresh is already in progress, return the existing promise if (this.refreshPromise) { this.logger.debug( @@ -556,7 +589,7 @@ export class OAuthSessionManager implements vscode.Disposable { private async saveTokens(tokenResponse: TokenResponse): Promise { const expiryTimestamp = tokenResponse.expires_in ? Date.now() + tokenResponse.expires_in * 1000 - : Date.now() + 3600 * 1000; // TODO Default to 1 hour + : Date.now() + ACCESS_TOKEN_DEFAULT_EXPIRY_MS; const tokens: StoredOAuthTokens = { ...tokenResponse, @@ -565,10 +598,11 @@ export class OAuthSessionManager implements vscode.Disposable { }; this.storedTokens = tokens; - await this.secretsManager.setOAuthTokens(tokens); - - // Trigger event to update global client (works for login & background refresh) - await this.secretsManager.setSessionToken(tokenResponse.access_token); + await this.secretsManager.setOAuthTokens(this.label, tokens); + await this.secretsManager.setSessionToken(this.label, { + url: this.deploymentUrl, + sessionToken: tokenResponse.access_token, + }); this.logger.info("Tokens saved", { expires_at: new Date(expiryTimestamp).toISOString(), @@ -583,7 +617,7 @@ export class OAuthSessionManager implements vscode.Disposable { * 2. Last refresh attempt was more than REFRESH_THROTTLE_MS ago * 3. No refresh is currently in progress */ - shouldRefreshToken(): boolean { + public shouldRefreshToken(): boolean { if ( !this.isLoggedInWithOAuth() || !this.storedTokens?.refresh_token || @@ -645,7 +679,7 @@ export class OAuthSessionManager implements vscode.Disposable { /** * Logout by revoking tokens and clearing all OAuth data. */ - async logout(): Promise { + public async logout(): Promise { if (!this.isLoggedInWithOAuth()) { return; } @@ -667,37 +701,43 @@ export class OAuthSessionManager implements vscode.Disposable { /** * Returns true if (valid or invalid) OAuth tokens exist for the current deployment. */ - isLoggedInWithOAuth(): boolean { + public isLoggedInWithOAuth(): boolean { return this.storedTokens !== undefined; } /** * Show a modal dialog to the user when OAuth re-authentication is required. * This is called when the refresh token is invalid or the client credentials are invalid. + * Clears tokens directly and lets listeners handle updates. */ - async showReAuthenticationModal(error: OAuthError): Promise { + public async showReAuthenticationModal(error: OAuthError): Promise { const errorMessage = error.description || "Your session is no longer valid. This could be due to token expiration or revocation."; - // Log out first to clear invalid tokens - await vscode.commands.executeCommand("coder.logout"); + // Clear invalid tokens - listeners will handle updates automatically + await this.clearTokenState(); + await this.secretsManager.setSessionToken(this.label, undefined); - const action = await this.vscodeProposed.window.showErrorMessage( - `Authentication Error`, - { modal: true, useCustom: true, detail: errorMessage }, - "Log in again", - ); + const result = await this.loginCoordinator.promptForLoginWithDialog({ + url: this.deploymentUrl, + label: this.label, + detailPrefix: errorMessage, + oauthSessionManager: this, + }); - if (action === "Log in again") { - await vscode.commands.executeCommand("coder.login"); + if (result.token) { + await this.secretsManager.setSessionToken(this.label, { + url: this.deploymentUrl, + sessionToken: result.token, + }); } } /** * Clears all in-memory state. */ - dispose(): void { + public dispose(): void { if (this.pendingAuthReject) { this.pendingAuthReject(new Error("OAuth session manager disposed")); } diff --git a/src/promptUtils.ts b/src/promptUtils.ts index ce488753..8f8369f5 100644 --- a/src/promptUtils.ts +++ b/src/promptUtils.ts @@ -65,7 +65,6 @@ export async function maybeAskAgent( */ async function askURL( mementoManager: MementoManager, - selection?: string, ): Promise { const defaultURL = vscode.workspace .getConfiguration() @@ -73,7 +72,10 @@ async function askURL( ?.trim(); const quickPick = vscode.window.createQuickPick(); quickPick.value = - selection || defaultURL || process.env.CODER_URL?.trim() || ""; + mementoManager.getUrl() || + defaultURL || + process.env.CODER_URL?.trim() || + ""; quickPick.placeholder = "https://example.coder.com"; quickPick.title = "Enter the URL of your Coder deployment."; @@ -115,9 +117,8 @@ async function askURL( export async function maybeAskUrl( mementoManager: MementoManager, providedUrl: string | undefined | null, - lastUsedUrl?: string, ): Promise { - let url = providedUrl || (await askURL(mementoManager, lastUsedUrl)); + let url = providedUrl || (await askURL(mementoManager)); if (!url) { // User aborted. return undefined; diff --git a/src/remote/remote.ts b/src/remote/remote.ts index 6af22390..2655ff8d 100644 --- a/src/remote/remote.ts +++ b/src/remote/remote.ts @@ -26,10 +26,14 @@ import * as cliUtils from "../core/cliUtils"; import { type ServiceContainer } from "../core/container"; import { type ContextManager } from "../core/contextManager"; import { type PathResolver } from "../core/pathResolver"; +import { type SecretsManager } from "../core/secretsManager"; import { featureSetForVersion, type FeatureSet } from "../featureSet"; import { getGlobalFlags } from "../globalFlags"; import { Inbox } from "../inbox"; import { type Logger } from "../logging/logger"; +import { type LoginCoordinator } from "../login/loginCoordinator"; +import { OAuthSessionManager } from "../oauth/sessionManager"; +import { maybeAskUrl } from "../promptUtils"; import { AuthorityPrefix, escapeCommandArg, @@ -44,6 +48,7 @@ import { computeSSHProperties, sshSupportsSetEnv } from "./sshSupport"; import { WorkspaceStateMachine } from "./workspaceStateMachine"; export interface RemoteDetails extends vscode.Disposable { + label: string; url: string; token: string; } @@ -55,48 +60,21 @@ export class Remote { private readonly pathResolver: PathResolver; private readonly cliManager: CliManager; private readonly contextManager: ContextManager; - - // Used to race between the login dialog and logging in from a different window - private loginDetectedResolver: (() => void) | undefined; - private loginDetectedRejector: ((reason?: Error) => void) | undefined; - private loginDetectedPromise: Promise = Promise.resolve(); + private readonly secretsManager: SecretsManager; + private readonly loginCoordinator: LoginCoordinator; public constructor( - serviceContainer: ServiceContainer, + private readonly serviceContainer: ServiceContainer, private readonly commands: Commands, - private readonly mode: vscode.ExtensionMode, + private readonly extensionContext: vscode.ExtensionContext, ) { this.vscodeProposed = serviceContainer.getVsCodeProposed(); this.logger = serviceContainer.getLogger(); this.pathResolver = serviceContainer.getPathResolver(); this.cliManager = serviceContainer.getCliManager(); this.contextManager = serviceContainer.getContextManager(); - } - - /** - * Creates a new promise that will be resolved when login is detected in another window. - */ - private createLoginDetectionPromise(): void { - if (this.loginDetectedRejector) { - this.loginDetectedRejector( - new Error("Login detection cancelled - new login attempt started"), - ); - } - this.loginDetectedPromise = new Promise((resolve, reject) => { - this.loginDetectedResolver = resolve; - this.loginDetectedRejector = reject; - }); - } - - /** - * Resolves the current login detection promise if one exists. - */ - public resolveLoginDetected(): void { - if (this.loginDetectedResolver) { - this.loginDetectedResolver(); - this.loginDetectedResolver = undefined; - this.loginDetectedRejector = undefined; - } + this.secretsManager = serviceContainer.getSecretsManager(); + this.loginCoordinator = serviceContainer.getLoginCoordinator(); } /** @@ -125,153 +103,196 @@ export class Remote { parts.label, ); - const showLoginDialog = async (message: string) => { - this.createLoginDetectionPromise(); - const dialogPromise = this.vscodeProposed.window.showInformationMessage( - message, - { - useCustom: true, - modal: true, - detail: `You must log in to access ${workspaceName}. If you've already logged in, you may close this dialog.`, - }, - "Log In", + const disposables: vscode.Disposable[] = []; + + try { + // Create OAuth session manager for this remote deployment + const remoteOAuthManager = await OAuthSessionManager.create( + baseUrlRaw, + parts.label, + this.serviceContainer, + this.extensionContext.extension.id, ); + disposables.push(remoteOAuthManager); + + const promptForLoginAndRetry = async (message: string, url: string) => { + const result = await this.loginCoordinator.promptForLoginWithDialog({ + url: url, + label: parts.label, + message, + detailPrefix: `You must log in to access ${workspaceName}.`, + oauthSessionManager: remoteOAuthManager, + }); - // Race between dialog and login detection - const result = await Promise.race([ - this.loginDetectedPromise.then(() => ({ type: "login" as const })), - dialogPromise.then((userChoice) => ({ - type: "dialog" as const, - userChoice, - })), - ]); - - if (result.type === "login") { - return this.setup(remoteAuthority, firstConnect, remoteSshExtensionId); - } else if (result.userChoice === "Log In") { - // Log in then try again. - await this.commands.login({ url: baseUrlRaw, label: parts.label }); - return this.setup(remoteAuthority, firstConnect, remoteSshExtensionId); - } else { - // User declined to log in. - await this.closeRemote(); - return; - } - }; + if (result.success) { + // Login successful, retry setup + return this.setup( + remoteAuthority, + firstConnect, + remoteSshExtensionId, + ); + } else { + // User cancelled or login failed + await this.closeRemote(); + } + }; - // It could be that the cli config was deleted. If so, ask for the url. - if ( - !baseUrlRaw || - (!token && needToken(vscode.workspace.getConfiguration())) - ) { - return showLoginDialog("You are not logged in..."); - } + // It could be that the cli config was deleted. If so, ask for the url. + if ( + !baseUrlRaw || + (!token && needToken(vscode.workspace.getConfiguration())) + ) { + const mementoManager = this.serviceContainer.getMementoManager(); + const newUrl = await maybeAskUrl( + mementoManager, + baseUrlRaw || parts.label, // TODO can we assume that "https://" is always valid? + ); + if (!newUrl) { + throw new Error("URL must be provided"); + } - this.logger.info("Using deployment URL", baseUrlRaw); - this.logger.info("Using deployment label", parts.label || "n/a"); - - // We could use the plugin client, but it is possible for the user to log - // out or log into a different deployment while still connected, which would - // break this connection. We could force close the remote session or - // disallow logging out/in altogether, but for now just use a separate - // client to remain unaffected by whatever the plugin is doing. - const workspaceClient = CoderApi.create(baseUrlRaw, token, this.logger); - // Store for use in commands. - this.commands.workspaceRestClient = workspaceClient; - - let binaryPath: string | undefined; - if (this.mode === vscode.ExtensionMode.Production) { - binaryPath = await this.cliManager.fetchBinary( - workspaceClient, - parts.label, + return promptForLoginAndRetry("You are not logged in...", baseUrlRaw); + } + + this.logger.info("Using deployment URL", baseUrlRaw); + this.logger.info("Using deployment label", parts.label || "n/a"); + + // We could use the plugin client, but it is possible for the user to log + // out or log into a different deployment while still connected, which would + // break this connection. We could force close the remote session or + // disallow logging out/in altogether, but for now just use a separate + // client to remain unaffected by whatever the plugin is doing. + const workspaceClient = CoderApi.create( + baseUrlRaw, + token, + this.logger, + remoteOAuthManager, ); - } else { - try { - // In development, try to use `/tmp/coder` as the binary path. - // This is useful for debugging with a custom bin! - binaryPath = path.join(os.tmpdir(), "coder"); - await fs.stat(binaryPath); - } catch { + // Store for use in commands. + this.commands.workspaceRestClient = workspaceClient; + + // Listen for token changes for this deployment and update CLI config + const authChangeDisposable = + this.secretsManager.onDidChangeDeploymentAuth( + parts.label, + async (auth) => { + const newToken = auth?.sessionToken; + if (newToken) { + // Update the client's token without breaking the connection + workspaceClient.setSessionToken(newToken); + + // Update CLI config with new token + if (auth?.url) { + await this.cliManager.configure( + parts.label, + auth.url, + newToken, + ); + this.logger.info( + "Updated CLI config with new token for remote deployment", + ); + } + } + }, + ); + disposables.push(authChangeDisposable); + + let binaryPath: string | undefined; + if ( + this.extensionContext.extensionMode === vscode.ExtensionMode.Production + ) { binaryPath = await this.cliManager.fetchBinary( workspaceClient, parts.label, ); + } else { + try { + // In development, try to use `/tmp/coder` as the binary path. + // This is useful for debugging with a custom bin! + binaryPath = path.join(os.tmpdir(), "coder"); + await fs.stat(binaryPath); + } catch { + binaryPath = await this.cliManager.fetchBinary( + workspaceClient, + parts.label, + ); + } } - } - // First thing is to check the version. - const buildInfo = await workspaceClient.getBuildInfo(); + // First thing is to check the version. + const buildInfo = await workspaceClient.getBuildInfo(); - let version: semver.SemVer | null = null; - try { - version = semver.parse(await cliUtils.version(binaryPath)); - } catch { - version = semver.parse(buildInfo.version); - } + let version: semver.SemVer | null = null; + try { + version = semver.parse(await cliUtils.version(binaryPath)); + } catch { + version = semver.parse(buildInfo.version); + } - const featureSet = featureSetForVersion(version); + const featureSet = featureSetForVersion(version); - // Server versions before v0.14.1 don't support the vscodessh command! - if (!featureSet.vscodessh) { - await this.vscodeProposed.window.showErrorMessage( - "Incompatible Server", - { - detail: - "Your Coder server is too old to support the Coder extension! Please upgrade to v0.14.1 or newer.", - modal: true, - useCustom: true, - }, - "Close Remote", - ); - await this.closeRemote(); - return; - } - - // Next is to find the workspace from the URI scheme provided. - let workspace: Workspace; - try { - this.logger.info(`Looking for workspace ${workspaceName}...`); - workspace = await workspaceClient.getWorkspaceByOwnerAndName( - parts.username, - parts.workspace, - ); - this.logger.info( - `Found workspace ${workspaceName} with status`, - workspace.latest_build.status, - ); - this.commands.workspace = workspace; - } catch (error) { - if (!isAxiosError(error)) { - throw error; + // Server versions before v0.14.1 don't support the vscodessh command! + if (!featureSet.vscodessh) { + await this.vscodeProposed.window.showErrorMessage( + "Incompatible Server", + { + detail: + "Your Coder server is too old to support the Coder extension! Please upgrade to v0.14.1 or newer.", + modal: true, + useCustom: true, + }, + "Close Remote", + ); + await this.closeRemote(); + return; } - switch (error.response?.status) { - case 404: { - const result = - await this.vscodeProposed.window.showInformationMessage( - `That workspace doesn't exist!`, - { - modal: true, - detail: `${workspaceName} cannot be found on ${baseUrlRaw}. Maybe it was deleted...`, - useCustom: true, - }, - "Open Workspace", + + // Next is to find the workspace from the URI scheme provided. + let workspace: Workspace; + try { + this.logger.info(`Looking for workspace ${workspaceName}...`); + workspace = await workspaceClient.getWorkspaceByOwnerAndName( + parts.username, + parts.workspace, + ); + this.logger.info( + `Found workspace ${workspaceName} with status`, + workspace.latest_build.status, + ); + this.commands.workspace = workspace; + } catch (error) { + if (!isAxiosError(error)) { + throw error; + } + switch (error.response?.status) { + case 404: { + const result = + await this.vscodeProposed.window.showInformationMessage( + `That workspace doesn't exist!`, + { + modal: true, + detail: `${workspaceName} cannot be found on ${baseUrlRaw}. Maybe it was deleted...`, + useCustom: true, + }, + "Open Workspace", + ); + if (!result) { + await this.closeRemote(); + } + await vscode.commands.executeCommand("coder.open"); + return; + } + case 401: { + return promptForLoginAndRetry( + "Your session expired...", + baseUrlRaw, ); - if (!result) { - await this.closeRemote(); } - await vscode.commands.executeCommand("coder.open"); - return; - } - case 401: { - return showLoginDialog("Your session expired..."); + default: + throw error; } - default: - throw error; } - } - const disposables: vscode.Disposable[] = []; - try { // Register before connection so the label still displays! let labelFormatterDisposable = this.registerLabelFormatter( remoteAuthority, @@ -529,6 +550,7 @@ export class Remote { // deployment in the sidebar. We use our own client in here for reasons // explained above. return { + label: parts.label, url: baseUrlRaw, token, dispose: () => { diff --git a/test/unit/core/cliManager.test.ts b/test/unit/core/cliManager.test.ts index 95755d31..6cdcdb07 100644 --- a/test/unit/core/cliManager.test.ts +++ b/test/unit/core/cliManager.test.ts @@ -165,52 +165,6 @@ describe("CliManager", () => { }); }); - describe("Read CLI Configuration", () => { - it("should read and trim stored configuration", async () => { - // Create directories and write files with whitespace - vol.mkdirSync("/path/base/deployment", { recursive: true }); - memfs.writeFileSync( - "/path/base/deployment/url", - " https://coder.example.com \n", - ); - memfs.writeFileSync( - "/path/base/deployment/session", - "\t test-token \r\n", - ); - - const result = await manager.readConfig("deployment"); - - expect(result).toEqual({ - url: "https://coder.example.com", - token: "test-token", - }); - }); - - it("should return empty strings for missing files", async () => { - const result = await manager.readConfig("deployment"); - - expect(result).toEqual({ - url: "", - token: "", - }); - }); - - it("should handle partial configuration", async () => { - vol.mkdirSync("/path/base/deployment", { recursive: true }); - memfs.writeFileSync( - "/path/base/deployment/url", - "https://coder.example.com", - ); - - const result = await manager.readConfig("deployment"); - - expect(result).toEqual({ - url: "https://coder.example.com", - token: "", - }); - }); - }); - describe("Binary Version Validation", () => { it("rejects invalid server versions", async () => { mockApi.getBuildInfo = vi.fn().mockResolvedValue({ version: "invalid" }); diff --git a/test/unit/core/secretsManager.test.ts b/test/unit/core/secretsManager.test.ts index bfe8c713..159aab0b 100644 --- a/test/unit/core/secretsManager.test.ts +++ b/test/unit/core/secretsManager.test.ts @@ -15,67 +15,95 @@ describe("SecretsManager", () => { describe("session token", () => { it("should store and retrieve tokens", async () => { - await secretsManager.setSessionToken("test-token"); - expect(await secretsManager.getSessionToken()).toBe("test-token"); + await secretsManager.setSessionToken("example-com", { + url: "https://example.com", + sessionToken: "test-token", + }); + expect(await secretsManager.getSessionToken("example-com")).toBe( + "test-token", + ); - await secretsManager.setSessionToken("new-token"); - expect(await secretsManager.getSessionToken()).toBe("new-token"); + await secretsManager.setSessionToken("example-com", { + url: "https://example.com", + sessionToken: "new-token", + }); + expect(await secretsManager.getSessionToken("example-com")).toBe( + "new-token", + ); }); - it("should delete token when empty or undefined", async () => { - await secretsManager.setSessionToken("test-token"); - await secretsManager.setSessionToken(""); - expect(await secretsManager.getSessionToken()).toBeUndefined(); - - await secretsManager.setSessionToken("test-token"); - await secretsManager.setSessionToken(undefined); - expect(await secretsManager.getSessionToken()).toBeUndefined(); + it("should delete token when undefined", async () => { + await secretsManager.setSessionToken("example-com", { + url: "https://example.com", + sessionToken: "test-token", + }); + await secretsManager.setSessionToken("example-com", undefined); + expect( + await secretsManager.getSessionToken("example-com"), + ).toBeUndefined(); }); it("should return undefined for corrupted storage", async () => { - await secretStorage.store("sessionToken", "valid-token"); + await secretStorage.store( + "coder.sessionAuthMap", + JSON.stringify({ + "example-com": { + url: "https://example.com", + sessionToken: "valid-token", + }, + }), + ); secretStorage.corruptStorage(); - expect(await secretsManager.getSessionToken()).toBeUndefined(); + expect( + await secretsManager.getSessionToken("example-com"), + ).toBeUndefined(); }); }); describe("login state", () => { it("should trigger login events", async () => { - const events: Array = []; - secretsManager.onDidChangeLoginState((state) => { - events.push(state); + const events: Array<{ state: AuthAction; label: string }> = []; + secretsManager.onDidChangeLoginState((state, label) => { + events.push({ state, label }); return Promise.resolve(); }); - await secretsManager.triggerLoginStateChange("login"); - expect(events).toEqual([AuthAction.LOGIN]); + await secretsManager.triggerLoginStateChange("example-com", "login"); + expect(events).toEqual([ + { state: AuthAction.LOGIN, label: "example-com" }, + ]); }); it("should trigger logout events", async () => { - const events: Array = []; - secretsManager.onDidChangeLoginState((state) => { - events.push(state); + const events: Array<{ state: AuthAction; label: string }> = []; + secretsManager.onDidChangeLoginState((state, label) => { + events.push({ state, label }); return Promise.resolve(); }); - await secretsManager.triggerLoginStateChange("logout"); - expect(events).toEqual([AuthAction.LOGOUT]); + await secretsManager.triggerLoginStateChange("example-com", "logout"); + expect(events).toEqual([ + { state: AuthAction.LOGOUT, label: "example-com" }, + ]); }); it("should fire same event twice in a row", async () => { vi.useFakeTimers(); - const events: Array = []; - secretsManager.onDidChangeLoginState((state) => { - events.push(state); + const events: Array<{ state: AuthAction; label: string }> = []; + secretsManager.onDidChangeLoginState((state, label) => { + events.push({ state, label }); return Promise.resolve(); }); - await secretsManager.triggerLoginStateChange("login"); + await secretsManager.triggerLoginStateChange("example-com", "login"); vi.advanceTimersByTime(5); - await secretsManager.triggerLoginStateChange("login"); + await secretsManager.triggerLoginStateChange("example-com", "login"); - expect(events).toEqual([AuthAction.LOGIN, AuthAction.LOGIN]); + expect(events).toEqual([ + { state: AuthAction.LOGIN, label: "example-com" }, + { state: AuthAction.LOGIN, label: "example-com" }, + ]); vi.useRealTimers(); }); }); From 7f9cc80cc46923ed7a3aed6e06de32816965bf53 Mon Sep 17 00:00:00 2001 From: Ehab Younes Date: Tue, 25 Nov 2025 11:39:37 +0300 Subject: [PATCH 15/20] Rely on the secrets API + Refactorings --- src/api/coderApi.ts | 128 +------- src/api/oauthInterceptors.ts | 122 ++++++++ src/commands.ts | 13 +- src/core/cliManager.ts | 40 ++- src/core/container.ts | 5 +- src/core/deployment.ts | 9 + src/core/secretsManager.ts | 425 ++++++++++---------------- src/extension.ts | 53 ++-- src/login/loginCoordinator.ts | 74 ++--- src/oauth/sessionManager.ts | 163 +++++----- src/promptUtils.ts | 8 +- src/remote/remote.ts | 80 ++--- test/unit/core/secretsManager.test.ts | 68 +++-- 13 files changed, 561 insertions(+), 627 deletions(-) create mode 100644 src/api/oauthInterceptors.ts create mode 100644 src/core/deployment.ts diff --git a/src/api/coderApi.ts b/src/api/coderApi.ts index f543667a..663bf716 100644 --- a/src/api/coderApi.ts +++ b/src/api/coderApi.ts @@ -3,8 +3,6 @@ import { type AxiosInstance, type AxiosHeaders, type AxiosResponseTransformer, - isAxiosError, - type AxiosError, } from "axios"; import { Api } from "coder/site/src/api/api"; import { @@ -33,8 +31,6 @@ import { HttpClientLogLevel, } from "../logging/types"; import { sizeOf } from "../logging/utils"; -import { parseOAuthError, requiresReAuthentication } from "../oauth/errors"; -import { type OAuthSessionManager } from "../oauth/sessionManager"; import { HttpStatusCode } from "../websocket/codes"; import { type UnidirectionalStream, @@ -76,7 +72,6 @@ export class CoderApi extends Api { baseUrl: string, token: string | undefined, output: Logger, - oauthSessionManager?: OAuthSessionManager, ): CoderApi { const client = new CoderApi(output); client.setHost(baseUrl); @@ -84,7 +79,7 @@ export class CoderApi extends Api { client.setSessionToken(token); } - setupInterceptors(client, output, oauthSessionManager); + setupInterceptors(client, output); return client; } @@ -395,11 +390,7 @@ export class CoderApi extends Api { /** * Set up logging and request interceptors for the CoderApi instance. */ -function setupInterceptors( - client: CoderApi, - output: Logger, - oauthSessionManager?: OAuthSessionManager, -): void { +function setupInterceptors(client: CoderApi, output: Logger): void { addLoggingInterceptors(client.getAxiosInstance(), output); client.getAxiosInstance().interceptors.request.use(async (config) => { @@ -437,11 +428,6 @@ function setupInterceptors( } }, ); - - // OAuth token refresh interceptors - if (oauthSessionManager) { - addOAuthInterceptors(client, output, oauthSessionManager); - } } function addLoggingInterceptors(client: AxiosInstance, logger: Logger) { @@ -487,116 +473,6 @@ function addLoggingInterceptors(client: AxiosInstance, logger: Logger) { ); } -/** - * Add OAuth token refresh interceptors. - * Success interceptor: proactively refreshes token when approaching expiry. - * Error interceptor: reactively refreshes token on 401 responses. - */ -function addOAuthInterceptors( - client: CoderApi, - logger: Logger, - oauthSessionManager: OAuthSessionManager, -) { - client.getAxiosInstance().interceptors.response.use( - // Success response interceptor: proactive token refresh - (response) => { - if (oauthSessionManager.shouldRefreshToken()) { - logger.debug( - "Token approaching expiry, triggering proactive refresh in background", - ); - - // Fire-and-forget: don't await, don't block response - oauthSessionManager.refreshToken().catch((error) => { - logger.warn("Background token refresh failed:", error); - }); - } - - return response; - }, - // Error response interceptor: reactive token refresh on 401 - async (error: unknown) => { - if (!isAxiosError(error)) { - throw error; - } - - if (error.config) { - const config = error.config as { - _oauthRetryAttempted?: boolean; - }; - if (config._oauthRetryAttempted) { - throw error; - } - } - - const status = error.response?.status; - - // These could indicate permanent auth failures that won't be fixed by token refresh - if (status === 400 || status === 403) { - handlePossibleOAuthError(error, logger, oauthSessionManager); - throw error; - } else if (status === 401) { - return handle401Error(error, client, logger, oauthSessionManager); - } - - throw error; - }, - ); -} - -function handlePossibleOAuthError( - error: unknown, - logger: Logger, - oauthSessionManager: OAuthSessionManager, -): void { - const oauthError = parseOAuthError(error); - if (oauthError && requiresReAuthentication(oauthError)) { - logger.error( - `OAuth error requires re-authentication: ${oauthError.errorCode}`, - ); - - oauthSessionManager.showReAuthenticationModal(oauthError).catch((err) => { - logger.error("Failed to show re-auth modal:", err); - }); - } -} - -async function handle401Error( - error: AxiosError, - client: CoderApi, - logger: Logger, - oauthSessionManager: OAuthSessionManager, -): Promise { - if (!oauthSessionManager.isLoggedInWithOAuth()) { - throw error; - } - - logger.info("Received 401 response, attempting token refresh"); - - try { - const newTokens = await oauthSessionManager.refreshToken(); - client.setSessionToken(newTokens.access_token); - - logger.info("Token refresh successful, retrying request"); - - // Retry the original request with the new token - if (error.config) { - const config = error.config as RequestConfigWithMeta & { - _oauthRetryAttempted?: boolean; - }; - config._oauthRetryAttempted = true; - config.headers[coderSessionTokenHeader] = newTokens.access_token; - return client.getAxiosInstance().request(config); - } - - throw error; - } catch (refreshError) { - logger.error("Token refresh failed:", refreshError); - - handlePossibleOAuthError(refreshError, logger, oauthSessionManager); - throw error; - } -} - function wrapRequestTransform( transformer: AxiosResponseTransformer | AxiosResponseTransformer[], config: RequestConfigWithMeta, diff --git a/src/api/oauthInterceptors.ts b/src/api/oauthInterceptors.ts new file mode 100644 index 00000000..adfe0efe --- /dev/null +++ b/src/api/oauthInterceptors.ts @@ -0,0 +1,122 @@ +import { type AxiosError, isAxiosError } from "axios"; + +import { type Logger } from "../logging/logger"; +import { type RequestConfigWithMeta } from "../logging/types"; +import { parseOAuthError, requiresReAuthentication } from "../oauth/errors"; +import { type OAuthSessionManager } from "../oauth/sessionManager"; + +import { type CoderApi } from "./coderApi"; + +const coderSessionTokenHeader = "Coder-Session-Token"; + +/** + * Attach OAuth token refresh interceptors to a CoderApi instance. + * This should be called after creating the CoderApi when OAuth authentication is being used. + * + * Success interceptor: proactively refreshes token when approaching expiry. + * Error interceptor: reactively refreshes token on 401 responses. + */ +export function attachOAuthInterceptors( + client: CoderApi, + logger: Logger, + oauthSessionManager: OAuthSessionManager, +): void { + client.getAxiosInstance().interceptors.response.use( + // Success response interceptor: proactive token refresh + (response) => { + if (oauthSessionManager.shouldRefreshToken()) { + logger.debug( + "Token approaching expiry, triggering proactive refresh in background", + ); + + // Fire-and-forget: don't await, don't block response + oauthSessionManager.refreshToken().catch((error) => { + logger.warn("Background token refresh failed:", error); + }); + } + + return response; + }, + // Error response interceptor: reactive token refresh on 401 + async (error: unknown) => { + if (!isAxiosError(error)) { + throw error; + } + + if (error.config) { + const config = error.config as { + _oauthRetryAttempted?: boolean; + }; + if (config._oauthRetryAttempted) { + throw error; + } + } + + const status = error.response?.status; + + // These could indicate permanent auth failures that won't be fixed by token refresh + if (status === 400 || status === 403) { + handlePossibleOAuthError(error, logger, oauthSessionManager); + throw error; + } else if (status === 401) { + return handle401Error(error, client, logger, oauthSessionManager); + } + + throw error; + }, + ); +} + +function handlePossibleOAuthError( + error: unknown, + logger: Logger, + oauthSessionManager: OAuthSessionManager, +): void { + const oauthError = parseOAuthError(error); + if (oauthError && requiresReAuthentication(oauthError)) { + logger.error( + `OAuth error requires re-authentication: ${oauthError.errorCode}`, + ); + + oauthSessionManager.showReAuthenticationModal(oauthError).catch((err) => { + logger.error("Failed to show re-auth modal:", err); + }); + } +} + +async function handle401Error( + error: AxiosError, + client: CoderApi, + logger: Logger, + oauthSessionManager: OAuthSessionManager, +): Promise { + if (!oauthSessionManager.isLoggedInWithOAuth()) { + throw error; + } + + logger.info("Received 401 response, attempting token refresh"); + + try { + const newTokens = await oauthSessionManager.refreshToken(); + client.setSessionToken(newTokens.access_token); + + logger.info("Token refresh successful, retrying request"); + + // Retry the original request with the new token + if (error.config) { + const config = error.config as RequestConfigWithMeta & { + _oauthRetryAttempted?: boolean; + }; + config._oauthRetryAttempted = true; + config.headers[coderSessionTokenHeader] = newTokens.access_token; + return client.getAxiosInstance().request(config); + } + + throw error; + } catch (refreshError) { + logger.error("Token refresh failed:", refreshError); + + handlePossibleOAuthError(refreshError, logger, oauthSessionManager); + throw error; + } +} diff --git a/src/commands.ts b/src/commands.ts index 04450600..79282296 100644 --- a/src/commands.ts +++ b/src/commands.ts @@ -68,7 +68,6 @@ export class Commands { */ public async login(args?: { url?: string; - token?: string; label?: string; autoLogin?: boolean; }): Promise { @@ -89,8 +88,7 @@ export class Commands { this.logger.info("Using deployment label", label); const result = await this.loginCoordinator.promptForLogin({ - url, - label, + deployment: { url, label }, autoLogin: args?.autoLogin, oauthSessionManager: this.oauthSessionManager, }); @@ -105,14 +103,11 @@ export class Commands { // Store for later sessions await this.mementoManager.setUrl(url); - await this.secretsManager.setSessionToken(label, { + await this.secretsManager.setSessionAuth(label, { url, - sessionToken: result.token, + token: result.token, }); - // Store on disk for CLI - await this.cliManager.configure(label, url, result.token); - // Update contexts this.contextManager.set("coder.authenticated", true); if (result.user.roles.some((role) => role.name === "owner")) { @@ -195,7 +190,7 @@ export class Commands { // Clear from memory. await this.mementoManager.setUrl(undefined); - await this.secretsManager.setSessionToken(label, undefined); + await this.secretsManager.clearAllAuthData(label); this.contextManager.set("coder.authenticated", false); vscode.window diff --git a/src/core/cliManager.ts b/src/core/cliManager.ts index 5e0b3d26..64ebc99e 100644 --- a/src/core/cliManager.ts +++ b/src/core/cliManager.ts @@ -721,8 +721,7 @@ export class CliManager { ): Promise { if (url) { const urlPath = this.pathResolver.getUrlPath(label); - await fs.mkdir(path.dirname(urlPath), { recursive: true }); - await fs.writeFile(urlPath, url); + await this.atomicWriteFile(urlPath, url); } } @@ -739,30 +738,27 @@ export class CliManager { ) { if (token !== null) { const tokenPath = this.pathResolver.getSessionTokenPath(label); - await fs.mkdir(path.dirname(tokenPath), { recursive: true }); - await fs.writeFile(tokenPath, token ?? ""); + await this.atomicWriteFile(tokenPath, token ?? ""); } } /** - * Read the CLI config for a deployment with the provided label. - * - * IF a config file does not exist, return an empty string. - * - * If the label is empty, read the old deployment-unaware config. + * Atomically write content to a file by writing to a temporary file first, + * then renaming it. */ - public async readConfig( - label: string, - ): Promise<{ url: string; token: string }> { - const urlPath = this.pathResolver.getUrlPath(label); - const tokenPath = this.pathResolver.getSessionTokenPath(label); - const [url, token] = await Promise.allSettled([ - fs.readFile(urlPath, "utf8"), - fs.readFile(tokenPath, "utf8"), - ]); - return { - url: url.status === "fulfilled" ? url.value.trim() : "", - token: token.status === "fulfilled" ? token.value.trim() : "", - }; + private async atomicWriteFile( + filePath: string, + content: string, + ): Promise { + await fs.mkdir(path.dirname(filePath), { recursive: true }); + const tempPath = + filePath + ".temp-" + Math.random().toString(36).substring(8); + try { + await fs.writeFile(tempPath, content); + await fs.rename(tempPath, filePath); + } catch (err) { + await fs.rm(tempPath, { force: true }).catch(() => {}); + throw err; + } } } diff --git a/src/core/container.ts b/src/core/container.ts index d83ae679..6c9d8b10 100644 --- a/src/core/container.ts +++ b/src/core/container.ts @@ -32,7 +32,10 @@ export class ServiceContainer implements vscode.Disposable { context.logUri.fsPath, ); this.mementoManager = new MementoManager(context.globalState); - this.secretsManager = new SecretsManager(context.secrets); + this.secretsManager = new SecretsManager( + context.secrets, + context.globalState, + ); this.cliManager = new CliManager( this.vscodeProposed, this.logger, diff --git a/src/core/deployment.ts b/src/core/deployment.ts new file mode 100644 index 00000000..a29c07ae --- /dev/null +++ b/src/core/deployment.ts @@ -0,0 +1,9 @@ +/** + * Represents a Coder deployment with its URL and label. + * The label is used as a unique identifier for storing credentials and configuration. + * It may be derived from the URL hostname (via toSafeHost) or come from SSH host parsing. + */ +export interface Deployment { + readonly url: string; + readonly label: string; +} diff --git a/src/core/secretsManager.ts b/src/core/secretsManager.ts index f4abdb45..e949e8e8 100644 --- a/src/core/secretsManager.ts +++ b/src/core/secretsManager.ts @@ -2,18 +2,19 @@ import { type TokenResponse, type ClientRegistrationResponse, } from "../oauth/types"; -import { toSafeHost } from "../util"; -import type { SecretStorage, Disposable } from "vscode"; +import type { Memento, SecretStorage, Disposable } from "vscode"; -const SESSION_TOKEN_KEY = "sessionToken"; - -const LOGIN_STATE_KEY = "loginState"; +const SESSION_KEY_PREFIX = "coder.session."; +const OAUTH_TOKENS_PREFIX = "coder.oauth.tokens."; +const OAUTH_CLIENT_PREFIX = "coder.oauth.client."; +const LOGIN_STATE_KEY = "coder.loginState"; const OAUTH_CALLBACK_KEY = "coder.oauthCallback"; -const SESSION_AUTH_MAP_KEY = "coder.sessionAuthMap"; -const OAUTH_DATA_MAP_KEY = "coder.oauthDataMap"; +const KNOWN_LABELS_KEY = "coder.knownLabels"; + +const LEGACY_SESSION_TOKEN_KEY = "sessionToken"; export type StoredOAuthTokens = Omit & { expiry_timestamp: number; @@ -22,17 +23,9 @@ export type StoredOAuthTokens = Omit & { export interface SessionAuth { url: string; - sessionToken: string; -} - -export interface OAuthData { - oauthClientRegistration?: ClientRegistrationResponse; - oauthTokens?: StoredOAuthTokens; + token: string; } -export type SessionAuthMap = Record; -export type OAuthDataMap = Record; - interface OAuthCallbackData { state: string; code: string | null; @@ -46,28 +39,13 @@ export enum AuthAction { } export class SecretsManager { - /** - * Track previous session tokens to detect actual changes. - * Maps label -> previous sessionToken value. - */ - private readonly previousSessionTokens = new Map< - string, - string | undefined - >(); - - constructor(private readonly secrets: SecretStorage) { - // Initialize previous session tokens - this.getSessionAuthMap().then((map) => { - for (const [label, auth] of Object.entries(map)) { - this.previousSessionTokens.set(label, auth.sessionToken); - } - }); - } + constructor( + private readonly secrets: SecretStorage, + private readonly memento: Memento, + ) {} /** * Triggers a login/logout event that propagates across all VS Code windows. - * Uses the secrets storage onDidChange event as a cross-window communication mechanism. - * Stores JSON with action, label, and timestamp to ensure the value always changes. */ public async triggerLoginStateChange( label: string, @@ -83,38 +61,38 @@ export class SecretsManager { /** * Listens for login/logout events from any VS Code window. - * The secrets storage onDidChange event fires across all windows, enabling cross-window sync. - * Parses JSON to extract action and label. */ public onDidChangeLoginState( listener: (state: AuthAction, label: string) => Promise, ): Disposable { return this.secrets.onDidChange(async (e) => { - if (e.key === LOGIN_STATE_KEY) { - const stateStr = await this.secrets.get(LOGIN_STATE_KEY); - if (!stateStr) { - await listener(AuthAction.INVALID, ""); - return; - } + if (e.key !== LOGIN_STATE_KEY) { + return; + } + + const stateStr = await this.secrets.get(LOGIN_STATE_KEY); + if (!stateStr) { + await listener(AuthAction.INVALID, ""); + return; + } - try { - const parsed = JSON.parse(stateStr) as { - action: string; - label: string; - timestamp: string; - }; - - if (parsed.action === "login") { - await listener(AuthAction.LOGIN, parsed.label); - } else if (parsed.action === "logout") { - await listener(AuthAction.LOGOUT, parsed.label); - } else { - await listener(AuthAction.INVALID, parsed.label); - } - } catch { - // Invalid JSON, treat as invalid state - await listener(AuthAction.INVALID, ""); + try { + const parsed = JSON.parse(stateStr) as { + action: string; + label: string; + timestamp: string; + }; + + if (parsed.action === "login") { + await listener(AuthAction.LOGIN, parsed.label); + } else if (parsed.action === "logout") { + await listener(AuthAction.LOGOUT, parsed.label); + } else { + await listener(AuthAction.INVALID, parsed.label); } + } catch { + // Invalid JSON, treat as invalid state + await listener(AuthAction.INVALID, ""); } }); } @@ -153,272 +131,179 @@ export class SecretsManager { /** * Listen for changes to a specific deployment's session auth. - * Only fires when the session token actually changes for this deployment. - * OAuth token/registration changes will NOT trigger this listener. */ public onDidChangeDeploymentAuth( label: string, listener: (auth: SessionAuth | undefined) => void | Promise, ): Disposable { - return this.onDidChangeSessionAuthMap(async (map) => { - const auth = map[label]; - const newToken = auth?.sessionToken ?? ""; - const previousToken = this.previousSessionTokens.get(label) ?? ""; - - // Only fire listener if session token actually changed - if (newToken !== previousToken) { - this.previousSessionTokens.set(label, newToken); - await listener(auth); - } - }); - } - - /** - * Listen for changes to the session auth map across all deployments. - * Fires whenever any deployment's session auth is updated. - */ - private onDidChangeSessionAuthMap( - listener: (map: SessionAuthMap) => void | Promise, - ): Disposable { + const key = `${SESSION_KEY_PREFIX}${label}`; return this.secrets.onDidChange(async (e) => { - if (e.key !== SESSION_AUTH_MAP_KEY) { + if (e.key !== key) { return; } + const auth = await this.getSessionAuth(label); + await listener(auth); + }); + } - try { - const map = await this.getSessionAuthMap(); - await listener(map); - } catch { - // Ignore errors in listener + public async getSessionAuth(label: string): Promise { + try { + const data = await this.secrets.get(`${SESSION_KEY_PREFIX}${label}`); + if (!data) { + return undefined; } - }); + return JSON.parse(data) as SessionAuth; + } catch { + return undefined; + } } - /** - * Get session token for a specific deployment. - */ public async getSessionToken(label: string): Promise { - const map = await this.getSessionAuthMap(); - return map[label]?.sessionToken; + const auth = await this.getSessionAuth(label); + return auth?.token; } - /** - * Set session token for a specific deployment. - */ - public async setSessionToken( - label: string, - auth: { url: string; sessionToken: string } | undefined, - ): Promise { - await this.updateSessionAuthMap((map) => { - if (auth === undefined) { - const newMap = { ...map }; - delete newMap[label]; - return newMap; - } + public async getUrl(label: string): Promise { + const auth = await this.getSessionAuth(label); + return auth?.url; + } - return { - ...map, - [label]: auth, - }; - }); + public async setSessionAuth(label: string, auth: SessionAuth): Promise { + await this.secrets.store( + `${SESSION_KEY_PREFIX}${label}`, + JSON.stringify(auth), + ); + await this.addKnownLabel(label); + } + + public async clearSessionAuth(label: string): Promise { + await this.secrets.delete(`${SESSION_KEY_PREFIX}${label}`); } - /** - * Get OAuth tokens for a specific deployment. - */ public async getOAuthTokens( label: string, ): Promise { - const map = await this.getOAuthDataMap(); - return map[label]?.oauthTokens; + try { + const data = await this.secrets.get(`${OAUTH_TOKENS_PREFIX}${label}`); + if (!data) { + return undefined; + } + return JSON.parse(data) as StoredOAuthTokens; + } catch { + return undefined; + } } - /** - * Set OAuth tokens for a specific deployment. - */ public async setOAuthTokens( label: string, - tokens: StoredOAuthTokens | undefined, + tokens: StoredOAuthTokens, ): Promise { - await this.updateOAuthDataMap((map) => { - const existing = map[label] || {}; - return { - ...map, - [label]: { - ...existing, - oauthTokens: tokens, - }, - }; - }); + await this.secrets.store( + `${OAUTH_TOKENS_PREFIX}${label}`, + JSON.stringify(tokens), + ); + await this.addKnownLabel(label); + } + + public async clearOAuthTokens(label: string): Promise { + await this.secrets.delete(`${OAUTH_TOKENS_PREFIX}${label}`); } - /** - * Get OAuth client registration for a specific deployment. - */ public async getOAuthClientRegistration( label: string, ): Promise { - const map = await this.getOAuthDataMap(); - return map[label]?.oauthClientRegistration; + try { + const data = await this.secrets.get(`${OAUTH_CLIENT_PREFIX}${label}`); + if (!data) { + return undefined; + } + return JSON.parse(data) as ClientRegistrationResponse; + } catch { + return undefined; + } } - /** - * Set OAuth client registration for a specific deployment. - * Creates the OAuth data entry if it doesn't exist. - */ public async setOAuthClientRegistration( label: string, - registration: ClientRegistrationResponse | undefined, + registration: ClientRegistrationResponse, ): Promise { - await this.updateOAuthDataMap((map) => { - const existing = map[label] || {}; - return { - ...map, - [label]: { - ...existing, - oauthClientRegistration: registration, - }, - }; - }); + await this.secrets.store( + `${OAUTH_CLIENT_PREFIX}${label}`, + JSON.stringify(registration), + ); + await this.addKnownLabel(label); } - public async clearOAuthData(label: string): Promise { - await this.updateOAuthDataMap((map) => { - const newMap = { ...map }; - delete newMap[label]; - return newMap; - }); + public async clearOAuthClientRegistration(label: string): Promise { + await this.secrets.delete(`${OAUTH_CLIENT_PREFIX}${label}`); } - /** - * Get the session auth map for all deployments. - */ - private async getSessionAuthMap(): Promise { - try { - const data = await this.secrets.get(SESSION_AUTH_MAP_KEY); - if (data) { - return JSON.parse(data) as SessionAuthMap; - } - } catch { - // Ignore parse errors - } - return {}; - } - - /** - * Set the session auth map for all deployments. - */ - private async setSessionAuthMap(map: SessionAuthMap): Promise { - await this.secrets.store(SESSION_AUTH_MAP_KEY, JSON.stringify(map)); + public async clearOAuthData(label: string): Promise { + await Promise.all([ + this.clearOAuthTokens(label), + this.clearOAuthClientRegistration(label), + ]); } /** - * Get the OAuth data map for all deployments. + * TODO currently it might be used wrong because we can be connected to a remote deployment + * and we log out from the sidebar causing the session to be removed and the auto-refresh disabled. + * + * Potential solutions: + * 1. Keep the last 10 auths and possibly remove entries not used in a while instead. + * We do not remove entries on logout! + * 2. Show the user a warning that their remote deployment might be disconnected. + * + * Update all usages of this after arriving at a decision! */ - private async getOAuthDataMap(): Promise { - try { - const data = await this.secrets.get(OAUTH_DATA_MAP_KEY); - if (data) { - return JSON.parse(data) as OAuthDataMap; - } - } catch { - // Ignore parse errors - } - return {}; + public async clearAllAuthData(label: string): Promise { + await Promise.all([ + this.clearSessionAuth(label), + this.clearOAuthData(label), + ]); + await this.removeKnownLabel(label); } - /** - * Set the OAuth data map for all deployments. - */ - private async setOAuthDataMap(map: OAuthDataMap): Promise { - await this.secrets.store(OAUTH_DATA_MAP_KEY, JSON.stringify(map)); + public getKnownLabels(): string[] { + return this.memento.get(KNOWN_LABELS_KEY) ?? []; } - /** - * Promise used for synchronizing session auth map updates. - */ - private sessionAuthUpdatePromise: Promise = Promise.resolve(); - - /** - * Promise used for synchronizing OAuth data map updates. - */ - private oauthDataUpdatePromise: Promise = Promise.resolve(); - - /** - * Atomically update the session auth map using a synchronized updater function. - * All write operations should go through this method to prevent race conditions. - */ - private async updateSessionAuthMap( - updater: (map: SessionAuthMap) => SessionAuthMap, - ): Promise { - this.sessionAuthUpdatePromise = this.sessionAuthUpdatePromise.then( - async () => { - const currentMap = await this.getSessionAuthMap(); - const newMap = updater(currentMap); - await this.setSessionAuthMap(newMap); - }, - ); - - return this.sessionAuthUpdatePromise; + private async addKnownLabel(label: string): Promise { + const labels = new Set(this.getKnownLabels()); + if (!labels.has(label)) { + labels.add(label); + await this.memento.update(KNOWN_LABELS_KEY, Array.from(labels)); + } } - /** - * Atomically update the OAuth data map using a synchronized updater function. - * All write operations should go through this method to prevent race conditions. - */ - private async updateOAuthDataMap( - updater: (map: OAuthDataMap) => OAuthDataMap, - ): Promise { - this.oauthDataUpdatePromise = this.oauthDataUpdatePromise.then(async () => { - const currentMap = await this.getOAuthDataMap(); - const newMap = updater(currentMap); - await this.setOAuthDataMap(newMap); - }); - - return this.oauthDataUpdatePromise; + private async removeKnownLabel(label: string): Promise { + const labels = new Set(this.getKnownLabels()); + if (labels.has(label)) { + labels.delete(label); + await this.memento.update(KNOWN_LABELS_KEY, Array.from(labels)); + } } /** - * Migrate from old flat storage format to new label-based format. - * This is a one-time operation that runs on extension activation. - * - * @param url The deployment URL to use for generating the label - * @returns true if migration was performed, false if already migrated or nothing to migrate + * Migrate from legacy flat sessionToken storage to new format. */ - public async migrateFromLegacyStorage(url: string): Promise { - try { - // Check if already migrated (new map exists and has data) - const existingMap = await this.getSessionAuthMap(); - if (Object.keys(existingMap).length > 0) { - return false; // Already migrated - } - - // Directly access old session token from flat storage - const oldToken = await this.secrets.get(SESSION_TOKEN_KEY); - if (!oldToken) { - return false; // Nothing to migrate - } - - // Generate label from URL - const label = toSafeHost(url); - - // Create new session auth map with migrated token - const sessionAuthMap: SessionAuthMap = { - [label]: { - url: url, - sessionToken: oldToken, - }, - }; + public async migrateFromLegacyStorage( + url: string, + label: string, + ): Promise { + const existing = await this.getSessionAuth(label); + if (existing) { + return false; + } - // Write new map to secrets - await this.setSessionAuthMap(sessionAuthMap); + const oldToken = await this.secrets.get(LEGACY_SESSION_TOKEN_KEY); + if (!oldToken) { + return false; + } - // Delete old session token key - await this.secrets.delete(SESSION_TOKEN_KEY); + await this.setSessionAuth(label, { url, token: oldToken }); + await this.secrets.delete(LEGACY_SESSION_TOKEN_KEY); - return true; // Migration successful - } catch (error) { - throw new Error(`Auth storage migration failed: ${error}`); - } + return true; } } diff --git a/src/extension.ts b/src/extension.ts index 38475a35..cbf67583 100644 --- a/src/extension.ts +++ b/src/extension.ts @@ -8,6 +8,7 @@ import * as vscode from "vscode"; import { errToStr } from "./api/api-helper"; import { CoderApi } from "./api/coderApi"; +import { attachOAuthInterceptors } from "./api/oauthInterceptors"; import { needToken } from "./api/utils"; import { Commands } from "./commands"; import { ServiceContainer } from "./core/container"; @@ -70,11 +71,11 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { const url = mementoManager.getUrl(); const label = url ? toSafeHost(url) : ""; + const deployment = url ? { url, label } : undefined; // Create OAuth session manager with login coordinator const oauthSessionManager = await OAuthSessionManager.create( - url || "", - label, + deployment, serviceContainer, ctx.extension.id, ); @@ -85,10 +86,10 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { // in commands that operate on the current login. const client = CoderApi.create( url || "", - (await secretsManager.getSessionToken(label)) || "", + await secretsManager.getSessionToken(label), output, - oauthSessionManager, ); + attachOAuthInterceptors(client, output, oauthSessionManager); const myWorkspacesProvider = new WorkspaceProvider( WorkspaceQuery.Mine, @@ -146,20 +147,12 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { output.debug("Registering auth listener for deployment", deploymentLabel); authChangeDisposable = secretsManager.onDidChangeDeploymentAuth( deploymentLabel, - async (auth) => { - const token = auth?.sessionToken ?? ""; - client.setSessionToken(token); + (auth) => { + client.setHost(auth?.url); + client.setSessionToken(auth?.token ?? ""); // Update authentication context for current deployment contextManager.set("coder.authenticated", auth !== undefined); - - output.debug("Session token changed, syncing state"); - - if (auth?.url) { - const cliManager = serviceContainer.getCliManager(); - await cliManager.configure(deploymentLabel, auth.url, token); - output.debug("Updated CLI config with new token"); - } }, ); }; @@ -180,7 +173,6 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { return; } - const cliManager = serviceContainer.getCliManager(); if (uri.path === "/open") { const owner = params.get("owner"); const workspace = params.get("workspace"); @@ -225,15 +217,9 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { const label = toSafeHost(url); if (token) { client.setSessionToken(token); - await secretsManager.setSessionToken(label, { - url, - sessionToken: token, - }); + await secretsManager.setSessionAuth(label, { url, token }); } - // Store on disk to be used by the cli. - await cliManager.configure(toSafeHost(url), url, token); - vscode.commands.executeCommand( "coder.open", owner, @@ -306,8 +292,10 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { ? params.get("token") : (params.get("token") ?? ""); - // Store on disk to be used by the cli. - await cliManager.configure(toSafeHost(url), url, token); + if (token) { + client.setSessionToken(token); + await secretsManager.setSessionAuth(label, { url, token }); + } vscode.commands.executeCommand( "coder.openDevContainer", @@ -390,16 +378,15 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { switch (state) { case AuthAction.LOGIN: { const url = mementoManager.getUrl(); - const token = await secretsManager.getSessionToken(label); // Should login the user directly if the URL+Token are valid - await commands.login({ url, token }); + await commands.login({ url }); // Re-register auth listener for the new deployment registerAuthListener(label); // Update OAuth session manager to match the new deployment if (url) { - await oauthSessionManager.setDeployment(label, url); + await oauthSessionManager.setDeployment({ label, url }); } break; } @@ -439,7 +426,10 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { // workspaces belonging to this deployment. client.setHost(details.url); client.setSessionToken(details.token); - await oauthSessionManager.setDeployment(details.label, details.url); + await oauthSessionManager.setDeployment({ + label: details.label, + url: details.url, + }); } } catch (ex) { if (ex instanceof CertificateError) { @@ -548,11 +538,12 @@ async function migrateAuthStorage( } // Perform migration using SecretsManager method - const migrated = await secretsManager.migrateFromLegacyStorage(url); + const label = toSafeHost(url); + const migrated = await secretsManager.migrateFromLegacyStorage(url, label); if (migrated) { output.info( - `Successfully migrated auth storage to label-based format (label: ${toSafeHost(url)})`, + `Successfully migrated auth storage to label-based format (label: ${label})`, ); } } catch (error) { diff --git a/src/login/loginCoordinator.ts b/src/login/loginCoordinator.ts index a923b12c..06a876d5 100644 --- a/src/login/loginCoordinator.ts +++ b/src/login/loginCoordinator.ts @@ -1,6 +1,8 @@ import { getErrorMessage } from "coder/site/src/api/errors"; import * as vscode from "vscode"; +import { type Deployment } from "src/core/deployment"; + import { CoderApi } from "../api/coderApi"; import { needToken } from "../api/utils"; import { type SecretsManager } from "../core/secretsManager"; @@ -19,12 +21,11 @@ export interface LoginResult { } export interface LoginOptions { - url: string; - label: string; + deployment: Deployment; oauthSessionManager: OAuthSessionManager; + autoLogin?: boolean; message?: string; detailPrefix?: string; - autoLogin?: boolean; } /** @@ -45,13 +46,10 @@ export class LoginCoordinator { public async promptForLogin( options: Omit, ): Promise { - return this.executeWithGuard(options.label, async () => { - const { url, label, oauthSessionManager } = options; - const existingToken = await this.secretsManager.getSessionToken(label); + const { deployment, oauthSessionManager } = options; + return this.executeWithGuard(options.deployment.label, async () => { return this.attemptLogin( - label, - url, - existingToken, + deployment, options.autoLogin ?? false, oauthSessionManager, ); @@ -64,15 +62,8 @@ export class LoginCoordinator { public async promptForLoginWithDialog( options: LoginOptions, ): Promise { - return this.executeWithGuard(options.label, () => { - const { - url, - label, - detailPrefix: reason, - message, - oauthSessionManager, - } = options; - + const { deployment, detailPrefix, message, oauthSessionManager } = options; + return this.executeWithGuard(deployment.label, () => { // Show dialog promise const dialogPromise = this.vscodeProposed.window .showErrorMessage( @@ -81,23 +72,29 @@ export class LoginCoordinator { modal: true, useCustom: true, detail: - (reason || `Authentication needed for ${label}`) + - " If you've already logged in, you may close this dialog.", + (detailPrefix || + `Authentication needed for ${deployment.label}.`) + + "\n\nIf you've already logged in, you may close this dialog.", }, "Login", ) .then(async (action) => { if (action === "Login") { // User clicked login - proceed with login flow - const existingToken = - await this.secretsManager.getSessionToken(label); - return this.attemptLogin( - label, - url, - existingToken, + const result = await this.attemptLogin( + deployment, false, oauthSessionManager, ); + + if (result.success && result.token) { + await this.secretsManager.setSessionAuth(deployment.label, { + url: deployment.url, + token: result.token, + }); + } + + return result; } else { // User cancelled return { success: false }; @@ -105,7 +102,10 @@ export class LoginCoordinator { }); // Race between user clicking login and cross-window detection - return Promise.race([dialogPromise, this.waitForCrossWindowLogin(label)]); + return Promise.race([ + dialogPromise, + this.waitForCrossWindowLogin(deployment.label), + ]); }); } @@ -139,9 +139,9 @@ export class LoginCoordinator { const disposable = this.secretsManager.onDidChangeDeploymentAuth( label, (auth) => { - if (auth?.sessionToken) { + if (auth?.token) { disposable.dispose(); - resolve({ success: true }); + resolve({ success: true, token: auth.token }); } }, ); @@ -155,13 +155,12 @@ export class LoginCoordinator { * failed (in which case an error notification will have been displayed). */ private async attemptLogin( - label: string, - url: string, - token: string | undefined, + deployment: Deployment, isAutoLogin: boolean, oauthSessionManager: OAuthSessionManager, ): Promise { - const client = CoderApi.create(url, token, this.logger); + const token = await this.secretsManager.getSessionToken(deployment.label); + const client = CoderApi.create(deployment.url, token, this.logger); const needsToken = needToken(vscode.workspace.getConfiguration()); if (!needsToken || token) { try { @@ -192,10 +191,11 @@ export class LoginCoordinator { const authMethod = await maybeAskAuthMethod(client); switch (authMethod) { case "oauth": - return this.loginWithOAuth(client, oauthSessionManager, label); + return this.loginWithOAuth(client, oauthSessionManager, deployment); case "legacy": { const initialToken = - token || (await this.secretsManager.getSessionToken(label)); + token || + (await this.secretsManager.getSessionToken(deployment.label)); return this.loginWithToken(client, initialToken); } case undefined: @@ -270,7 +270,7 @@ export class LoginCoordinator { private async loginWithOAuth( client: CoderApi, oauthSessionManager: OAuthSessionManager, - label: string, + deployment: Deployment, ): Promise { try { this.logger.info("Starting OAuth authentication"); @@ -282,7 +282,7 @@ export class LoginCoordinator { cancellable: false, }, async (progress) => - await oauthSessionManager.login(client, label, progress), + await oauthSessionManager.login(client, deployment, progress), ); // Validate token by fetching user diff --git a/src/oauth/sessionManager.ts b/src/oauth/sessionManager.ts index 8bfb7f56..a533850a 100644 --- a/src/oauth/sessionManager.ts +++ b/src/oauth/sessionManager.ts @@ -1,9 +1,9 @@ import { type AxiosInstance } from "axios"; import * as vscode from "vscode"; -import { type ServiceContainer } from "src/core/container"; - import { CoderApi } from "../api/coderApi"; +import { type ServiceContainer } from "../core/container"; +import { type Deployment } from "../core/deployment"; import { type LoginCoordinator } from "../login/loginCoordinator"; import { OAuthMetadataClient } from "./metadataClient"; @@ -76,14 +76,12 @@ export class OAuthSessionManager implements vscode.Disposable { * Create and initialize a new OAuth session manager. */ public static async create( - deploymentUrl: string, - label: string, + deployment: Deployment | undefined, container: ServiceContainer, extensionId: string, ): Promise { const manager = new OAuthSessionManager( - deploymentUrl, - label, + deployment, container.getSecretsManager(), container.getLogger(), container.getLoginCoordinator(), @@ -94,33 +92,49 @@ export class OAuthSessionManager implements vscode.Disposable { } private constructor( - private deploymentUrl: string, - private label: string, + private deployment: Deployment | undefined, private readonly secretsManager: SecretsManager, private readonly logger: Logger, private readonly loginCoordinator: LoginCoordinator, private readonly extensionId: string, ) {} + /** + * Get current deployment, throwing if not set. + * Use this in methods that require a deployment to be configured. + */ + private requireDeployment(): Deployment { + if (!this.deployment) { + throw new Error("No deployment configured for OAuth session manager"); + } + return this.deployment; + } + /** * Load stored tokens from storage. + * No-op if deployment is not set. * Validates that tokens belong to the current deployment URL. */ private async loadTokens(): Promise { - const tokens = await this.secretsManager.getOAuthTokens(this.label); + if (!this.deployment) { + return; + } + + const tokens = await this.secretsManager.getOAuthTokens( + this.deployment.label, + ); if (!tokens) { return; } - if (this.deploymentUrl && tokens.deployment_url !== this.deploymentUrl) { + if (tokens.deployment_url !== this.deployment.url) { this.logger.warn("Stored tokens for different deployment, clearing", { stored: tokens.deployment_url, - current: this.deploymentUrl, + current: this.deployment.url, }); - await this.clearTokenState(); + await this.clearOAuthState(); return; } - this.deploymentUrl = tokens.deployment_url; if (!this.hasRequiredScopes(tokens.scope)) { this.logger.warn( @@ -130,7 +144,7 @@ export class OAuthSessionManager implements vscode.Disposable { required_scopes: DEFAULT_OAUTH_SCOPES, }, ); - await this.secretsManager.setOAuthTokens(this.label, undefined); + await this.secretsManager.clearOAuthTokens(this.deployment.label); return; } @@ -138,10 +152,11 @@ export class OAuthSessionManager implements vscode.Disposable { this.logger.info(`Loaded stored OAuth tokens for ${tokens.deployment_url}`); } - private async clearTokenState(): Promise { + private async clearOAuthState(): Promise { this.clearInMemoryTokens(); - await this.secretsManager.setOAuthTokens(this.label, undefined); - await this.secretsManager.setOAuthClientRegistration(this.label, undefined); + if (this.deployment) { + await this.secretsManager.clearOAuthData(this.deployment.label); + } } private clearInMemoryTokens(): void { @@ -192,25 +207,23 @@ export class OAuthSessionManager implements vscode.Disposable { } /** - * Prepare common OAuth operation setup: CoderApi, metadata, and registration. + * Prepare common OAuth operation setup: client, metadata, and registration. * Used by refresh and revoke operations to reduce duplication. */ - private async prepareOAuthOperation( - deploymentUrl: string, - token?: string, - ): Promise<{ + private async prepareOAuthOperation(token?: string): Promise<{ axiosInstance: AxiosInstance; metadata: OAuthServerMetadata; registration: ClientRegistrationResponse; }> { - const client = CoderApi.create(deploymentUrl, token, this.logger); + const deployment = this.requireDeployment(); + const client = CoderApi.create(deployment.url, token, this.logger); const axiosInstance = client.getAxiosInstance(); const metadataClient = new OAuthMetadataClient(axiosInstance, this.logger); const metadata = await metadataClient.getMetadata(); const registration = await this.secretsManager.getOAuthClientRegistration( - this.label, + deployment.label, ); if (!registration) { throw new Error("No client registration found"); @@ -227,10 +240,11 @@ export class OAuthSessionManager implements vscode.Disposable { axiosInstance: AxiosInstance, metadata: OAuthServerMetadata, ): Promise { + const deployment = this.requireDeployment(); const redirectUri = this.getRedirectUri(); const existing = await this.secretsManager.getOAuthClientRegistration( - this.label, + deployment.label, ); if (existing?.client_id) { if (existing.redirect_uris.includes(redirectUri)) { @@ -262,7 +276,7 @@ export class OAuthSessionManager implements vscode.Disposable { ); await this.secretsManager.setOAuthClientRegistration( - this.label, + deployment.label, response.data, ); this.logger.info( @@ -273,44 +287,59 @@ export class OAuthSessionManager implements vscode.Disposable { return response.data; } - public async setDeployment(label: string, url: string): Promise { - this.logger.debug("Switching OAuth deployment", { label, url }); - this.label = label; - this.deploymentUrl = url; + public async setDeployment(deployment: Deployment): Promise { + if ( + this.deployment && + deployment.label === this.deployment.label && + deployment.url === this.deployment.url + ) { + return; + } + this.logger.debug("Switching OAuth deployment", deployment); + this.deployment = deployment; this.clearInMemoryTokens(); await this.loadTokens(); } + public clearDeployment(): void { + this.logger.debug("Clearing OAuth deployment state"); + this.deployment = undefined; + this.clearInMemoryTokens(); + } + /** - * Simplified OAuth login flow that handles the entire process. + * OAuth login flow that handles the entire process. * Fetches metadata, registers client, starts authorization, and exchanges tokens. * * @returns TokenResponse containing access token and optional refresh token */ public async login( client: CoderApi, - label: string, + deployment: Deployment, progress: vscode.Progress<{ message?: string; increment?: number }>, ): Promise { const baseUrl = client.getAxiosInstance().defaults.baseURL; if (!baseUrl) { - throw new Error("CoderApi instance has no base URL set"); + throw new Error("Client has no base URL set"); } - if (this.deploymentUrl && this.deploymentUrl !== baseUrl) { - this.logger.info("Deployment URL changed, clearing cached state", { - old: this.deploymentUrl, - new: baseUrl, - }); - this.clearInMemoryTokens(); - this.deploymentUrl = baseUrl; + if (baseUrl !== deployment.url) { + throw new Error( + `Client base URL (${baseUrl}) does not match deployment URL (${deployment.url})`, + ); } - if (this.label && this.label !== label) { - this.logger.info("Deployment label changed, clearing cached state", { - old: this.label, - new: label, + + // Update deployment if changed + if ( + !this.deployment || + this.deployment.url !== deployment.url || + this.deployment.label !== deployment.label + ) { + this.logger.info("Deployment changed, clearing cached state", { + old: this.deployment, + new: deployment, }); this.clearInMemoryTokens(); - this.label = label; + this.deployment = deployment; } const axiosInstance = client.getAxiosInstance(); @@ -546,7 +575,7 @@ export class OAuthSessionManager implements vscode.Disposable { this.refreshPromise = (async () => { try { const { axiosInstance, metadata, registration } = - await this.prepareOAuthOperation(this.deploymentUrl, accessToken); + await this.prepareOAuthOperation(accessToken); this.logger.debug("Refreshing access token"); @@ -587,26 +616,27 @@ export class OAuthSessionManager implements vscode.Disposable { * Also triggers event via secretsManager to update global client. */ private async saveTokens(tokenResponse: TokenResponse): Promise { + const deployment = this.requireDeployment(); const expiryTimestamp = tokenResponse.expires_in ? Date.now() + tokenResponse.expires_in * 1000 : Date.now() + ACCESS_TOKEN_DEFAULT_EXPIRY_MS; const tokens: StoredOAuthTokens = { ...tokenResponse, - deployment_url: this.deploymentUrl, + deployment_url: deployment.url, expiry_timestamp: expiryTimestamp, }; this.storedTokens = tokens; - await this.secretsManager.setOAuthTokens(this.label, tokens); - await this.secretsManager.setSessionToken(this.label, { - url: this.deploymentUrl, - sessionToken: tokenResponse.access_token, + await this.secretsManager.setOAuthTokens(deployment.label, tokens); + await this.secretsManager.setSessionAuth(deployment.label, { + url: deployment.url, + token: tokenResponse.access_token, }); this.logger.info("Tokens saved", { expires_at: new Date(expiryTimestamp).toISOString(), - deployment: this.deploymentUrl, + deployment: deployment.url, }); } @@ -643,10 +673,7 @@ export class OAuthSessionManager implements vscode.Disposable { tokenTypeHint: "access_token" | "refresh_token" = "refresh_token", ): Promise { const { axiosInstance, metadata, registration } = - await this.prepareOAuthOperation( - this.deploymentUrl, - this.storedTokens?.access_token, - ); + await this.prepareOAuthOperation(this.storedTokens?.access_token); const revocationEndpoint = metadata.revocation_endpoint || `${metadata.issuer}/oauth2/revoke`; @@ -693,7 +720,8 @@ export class OAuthSessionManager implements vscode.Disposable { } } - await this.clearTokenState(); + await this.clearOAuthState(); + this.deployment = undefined; this.logger.info("OAuth logout complete"); } @@ -711,27 +739,20 @@ export class OAuthSessionManager implements vscode.Disposable { * Clears tokens directly and lets listeners handle updates. */ public async showReAuthenticationModal(error: OAuthError): Promise { + const deployment = this.requireDeployment(); const errorMessage = error.description || "Your session is no longer valid. This could be due to token expiration or revocation."; // Clear invalid tokens - listeners will handle updates automatically - await this.clearTokenState(); - await this.secretsManager.setSessionToken(this.label, undefined); + this.clearInMemoryTokens(); + await this.secretsManager.clearAllAuthData(deployment.label); - const result = await this.loginCoordinator.promptForLoginWithDialog({ - url: this.deploymentUrl, - label: this.label, + await this.loginCoordinator.promptForLoginWithDialog({ + deployment, detailPrefix: errorMessage, oauthSessionManager: this, }); - - if (result.token) { - await this.secretsManager.setSessionToken(this.label, { - url: this.deploymentUrl, - sessionToken: result.token, - }); - } } /** @@ -742,9 +763,7 @@ export class OAuthSessionManager implements vscode.Disposable { this.pendingAuthReject(new Error("OAuth session manager disposed")); } this.pendingAuthReject = undefined; - this.storedTokens = undefined; - this.refreshPromise = null; - this.lastRefreshAttempt = 0; + this.clearInMemoryTokens(); this.logger.debug("OAuth session manager disposed"); } diff --git a/src/promptUtils.ts b/src/promptUtils.ts index 8f8369f5..35a8b838 100644 --- a/src/promptUtils.ts +++ b/src/promptUtils.ts @@ -65,14 +65,17 @@ export async function maybeAskAgent( */ async function askURL( mementoManager: MementoManager, + selection: string | undefined, ): Promise { + const lastUsedUrl = mementoManager.getUrl(); const defaultURL = vscode.workspace .getConfiguration() .get("coder.defaultUrl") ?.trim(); const quickPick = vscode.window.createQuickPick(); quickPick.value = - mementoManager.getUrl() || + selection || + lastUsedUrl || defaultURL || process.env.CODER_URL?.trim() || ""; @@ -117,8 +120,9 @@ async function askURL( export async function maybeAskUrl( mementoManager: MementoManager, providedUrl: string | undefined | null, + lastUsedUrl?: string, ): Promise { - let url = providedUrl || (await askURL(mementoManager)); + let url = providedUrl || (await askURL(mementoManager, lastUsedUrl)); if (!url) { // User aborted. return undefined; diff --git a/src/remote/remote.ts b/src/remote/remote.ts index 2655ff8d..fce41359 100644 --- a/src/remote/remote.ts +++ b/src/remote/remote.ts @@ -19,6 +19,7 @@ import { } from "../api/agentMetadataHelper"; import { extractAgents } from "../api/api-helper"; import { CoderApi } from "../api/coderApi"; +import { attachOAuthInterceptors } from "../api/oauthInterceptors"; import { needToken } from "../api/utils"; import { type Commands } from "../commands"; import { type CliManager } from "../core/cliManager"; @@ -99,17 +100,38 @@ export class Remote { await this.migrateSessionToken(parts.label); // Get the URL and token belonging to this host. - const { url: baseUrlRaw, token } = await this.cliManager.readConfig( - parts.label, - ); + const baseUrlRaw = (await this.secretsManager.getUrl(parts.label)) ?? ""; + const token = + (await this.secretsManager.getSessionToken(parts.label)) ?? ""; + if (baseUrlRaw && token) { + await this.cliManager.configure(parts.label, baseUrlRaw, token); + } const disposables: vscode.Disposable[] = []; try { + disposables.push( + this.secretsManager.onDidChangeDeploymentAuth( + parts.label, + async (auth) => { + if (auth?.token && auth.url) { + // Update CLI config with new token + await this.cliManager.configure( + parts.label, + auth.url, + auth.token, + ); + this.logger.info( + "Updated CLI config with new token for remote deployment", + ); + } + }, + ), + ); + // Create OAuth session manager for this remote deployment const remoteOAuthManager = await OAuthSessionManager.create( - baseUrlRaw, - parts.label, + { url: baseUrlRaw, label: parts.label }, this.serviceContainer, this.extensionContext.extension.id, ); @@ -117,8 +139,7 @@ export class Remote { const promptForLoginAndRetry = async (message: string, url: string) => { const result = await this.loginCoordinator.promptForLoginWithDialog({ - url: url, - label: parts.label, + deployment: { url, label: parts.label }, message, detailPrefix: `You must log in to access ${workspaceName}.`, oauthSessionManager: remoteOAuthManager, @@ -145,13 +166,14 @@ export class Remote { const mementoManager = this.serviceContainer.getMementoManager(); const newUrl = await maybeAskUrl( mementoManager, - baseUrlRaw || parts.label, // TODO can we assume that "https://" is always valid? + baseUrlRaw, // TODO can we assume that "https://" is always valid? + parts.label, ); if (!newUrl) { throw new Error("URL must be provided"); } - return promptForLoginAndRetry("You are not logged in...", baseUrlRaw); + return promptForLoginAndRetry("You are not logged in...", newUrl); } this.logger.info("Using deployment URL", baseUrlRaw); @@ -162,40 +184,18 @@ export class Remote { // break this connection. We could force close the remote session or // disallow logging out/in altogether, but for now just use a separate // client to remain unaffected by whatever the plugin is doing. - const workspaceClient = CoderApi.create( - baseUrlRaw, - token, - this.logger, - remoteOAuthManager, - ); + const workspaceClient = CoderApi.create(baseUrlRaw, token, this.logger); + attachOAuthInterceptors(workspaceClient, this.logger, remoteOAuthManager); // Store for use in commands. this.commands.workspaceRestClient = workspaceClient; - // Listen for token changes for this deployment and update CLI config - const authChangeDisposable = - this.secretsManager.onDidChangeDeploymentAuth( - parts.label, - async (auth) => { - const newToken = auth?.sessionToken; - if (newToken) { - // Update the client's token without breaking the connection - workspaceClient.setSessionToken(newToken); - - // Update CLI config with new token - if (auth?.url) { - await this.cliManager.configure( - parts.label, - auth.url, - newToken, - ); - this.logger.info( - "Updated CLI config with new token for remote deployment", - ); - } - } - }, - ); - disposables.push(authChangeDisposable); + // Listen for token changes for this deployment + disposables.push( + this.secretsManager.onDidChangeDeploymentAuth(parts.label, (auth) => { + workspaceClient.setHost(auth?.url); + workspaceClient.setSessionToken(auth?.token ?? ""); + }), + ); let binaryPath: string | undefined; if ( diff --git a/test/unit/core/secretsManager.test.ts b/test/unit/core/secretsManager.test.ts index 159aab0b..9de34c83 100644 --- a/test/unit/core/secretsManager.test.ts +++ b/test/unit/core/secretsManager.test.ts @@ -2,42 +2,50 @@ import { beforeEach, describe, expect, it, vi } from "vitest"; import { AuthAction, SecretsManager } from "@/core/secretsManager"; -import { InMemorySecretStorage } from "../../mocks/testHelpers"; +import { + InMemoryMemento, + InMemorySecretStorage, +} from "../../mocks/testHelpers"; describe("SecretsManager", () => { let secretStorage: InMemorySecretStorage; + let memento: InMemoryMemento; let secretsManager: SecretsManager; beforeEach(() => { secretStorage = new InMemorySecretStorage(); - secretsManager = new SecretsManager(secretStorage); + memento = new InMemoryMemento(); + secretsManager = new SecretsManager(secretStorage, memento); }); - describe("session token", () => { - it("should store and retrieve tokens", async () => { - await secretsManager.setSessionToken("example-com", { + describe("session auth", () => { + it("should store and retrieve session auth", async () => { + await secretsManager.setSessionAuth("example-com", { url: "https://example.com", - sessionToken: "test-token", + token: "test-token", }); expect(await secretsManager.getSessionToken("example-com")).toBe( "test-token", ); + expect(await secretsManager.getUrl("example-com")).toBe( + "https://example.com", + ); - await secretsManager.setSessionToken("example-com", { + await secretsManager.setSessionAuth("example-com", { url: "https://example.com", - sessionToken: "new-token", + token: "new-token", }); expect(await secretsManager.getSessionToken("example-com")).toBe( "new-token", ); }); - it("should delete token when undefined", async () => { - await secretsManager.setSessionToken("example-com", { + it("should clear session auth", async () => { + await secretsManager.setSessionAuth("example-com", { url: "https://example.com", - sessionToken: "test-token", + token: "test-token", }); - await secretsManager.setSessionToken("example-com", undefined); + await secretsManager.clearSessionAuth("example-com"); expect( await secretsManager.getSessionToken("example-com"), ).toBeUndefined(); @@ -45,12 +53,10 @@ describe("SecretsManager", () => { it("should return undefined for corrupted storage", async () => { await secretStorage.store( - "coder.sessionAuthMap", + "coder.session.example-com", JSON.stringify({ - "example-com": { - url: "https://example.com", - sessionToken: "valid-token", - }, + url: "https://example.com", + token: "valid-token", }), ); secretStorage.corruptStorage(); @@ -59,6 +65,34 @@ describe("SecretsManager", () => { await secretsManager.getSessionToken("example-com"), ).toBeUndefined(); }); + + it("should track known labels", async () => { + expect(secretsManager.getKnownLabels()).toEqual([]); + + await secretsManager.setSessionAuth("example-com", { + url: "https://example.com", + token: "test-token", + }); + expect(secretsManager.getKnownLabels()).toContain("example-com"); + + await secretsManager.setSessionAuth("other-com", { + url: "https://other.com", + token: "other-token", + }); + expect(secretsManager.getKnownLabels()).toContain("example-com"); + expect(secretsManager.getKnownLabels()).toContain("other-com"); + }); + + it("should remove label on clearAllAuthData", async () => { + await secretsManager.setSessionAuth("example-com", { + url: "https://example.com", + token: "test-token", + }); + expect(secretsManager.getKnownLabels()).toContain("example-com"); + + await secretsManager.clearAllAuthData("example-com"); + expect(secretsManager.getKnownLabels()).not.toContain("example-com"); + }); }); describe("login state", () => { From 13213f0b8b0fae7c7c8aa84e2e904c12085f90e1 Mon Sep 17 00:00:00 2001 From: Ehab Younes Date: Wed, 26 Nov 2025 16:54:43 +0300 Subject: [PATCH 16/20] Improve login experience --- src/commands.ts | 73 +++++---- src/core/container.ts | 1 + src/core/mementoManager.ts | 17 +- src/core/secretsManager.ts | 117 ++++++++------ src/extension.ts | 215 +++++++++++++------------- src/login/loginCoordinator.ts | 82 ++++++---- src/oauth/sessionManager.ts | 5 +- src/promptUtils.ts | 14 +- src/remote/remote.ts | 25 ++- test/unit/core/mementoManager.test.ts | 20 +-- test/unit/core/secretsManager.test.ts | 113 ++++++++------ 11 files changed, 358 insertions(+), 324 deletions(-) diff --git a/src/commands.ts b/src/commands.ts index 79282296..8e1593c8 100644 --- a/src/commands.ts +++ b/src/commands.ts @@ -9,6 +9,7 @@ import { createWorkspaceIdentifier, extractAgents } from "./api/api-helper"; import { type CliManager } from "./core/cliManager"; import { type ServiceContainer } from "./core/container"; import { type ContextManager } from "./core/contextManager"; +import { type Deployment } from "./core/deployment"; import { type MementoManager } from "./core/mementoManager"; import { type PathResolver } from "./core/pathResolver"; import { type SecretsManager } from "./core/secretsManager"; @@ -61,6 +62,17 @@ export class Commands { this.loginCoordinator = serviceContainer.getLoginCoordinator(); } + /** + * Get the current deployment, throwing if not logged in. + */ + private async requireDeployment(): Promise { + const deployment = await this.secretsManager.getCurrentDeployment(); + if (!deployment) { + throw new Error("You are not logged in"); + } + return deployment; + } + /** * 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 @@ -76,7 +88,12 @@ export class Commands { } this.logger.info("Logging in"); - const url = await maybeAskUrl(this.mementoManager, args?.url); + const currentDeployment = await this.secretsManager.getCurrentDeployment(); + const url = await maybeAskUrl( + this.mementoManager, + args?.url, + currentDeployment?.url, + ); if (!url) { return; } @@ -88,7 +105,8 @@ export class Commands { this.logger.info("Using deployment label", label); const result = await this.loginCoordinator.promptForLogin({ - deployment: { url, label }, + label, + url, autoLogin: args?.autoLogin, oauthSessionManager: this.oauthSessionManager, }); @@ -97,16 +115,13 @@ export class Commands { return; } - // Authorize the global client + // Set client immediately so subsequent operations in this function have the correct host/token. + // The cross-window listener will also update the client, but that's async. this.restClient.setHost(url); this.restClient.setSessionToken(result.token); - // Store for later sessions - await this.mementoManager.setUrl(url); - await this.secretsManager.setSessionAuth(label, { - url, - token: result.token, - }); + // Set as current deployment (triggers cross-window sync). + await this.secretsManager.setCurrentDeployment({ url, label }); // Update contexts this.contextManager.set("coder.authenticated", true); @@ -129,7 +144,6 @@ export class Commands { } }); - await this.secretsManager.triggerLoginStateChange(label, "login"); vscode.commands.executeCommand("coder.refreshWorkspaces"); } @@ -162,13 +176,8 @@ export class Commands { * Log out from the currently logged-in deployment. */ public async logout(): Promise { - const url = this.mementoManager.getUrl(); - if (!url) { - // Sanity check; command should not be available if no url. - throw new Error("You are not logged in"); - } - - await this.forceLogout(toSafeHost(url)); + const deployment = await this.requireDeployment(); + await this.forceLogout(deployment.label); } public async forceLogout(label: string): Promise { @@ -177,8 +186,7 @@ export class Commands { } this.logger.info(`Logging out of deployment: ${label}`); - // Only clear REST client and UI context if logging out of current deployment - // Fire and forget + // Fire and forget OAuth logout this.oauthSessionManager.logout().catch((error) => { this.logger.warn("OAuth logout failed, continuing with cleanup:", error); }); @@ -188,8 +196,10 @@ export class Commands { this.restClient.setHost(""); this.restClient.setSessionToken(""); - // Clear from memory. - await this.mementoManager.setUrl(undefined); + // Clear current deployment (triggers cross-window sync) + await this.secretsManager.setCurrentDeployment(undefined); + + // Clear all auth data for this deployment await this.secretsManager.clearAllAuthData(label); this.contextManager.set("coder.authenticated", false); @@ -203,8 +213,6 @@ export class Commands { // This will result in clearing the workspace list. vscode.commands.executeCommand("coder.refreshWorkspaces"); - - await this.secretsManager.triggerLoginStateChange(label, "logout"); } /** @@ -213,7 +221,8 @@ export class Commands { * Must only be called if currently logged in. */ public async createWorkspace(): Promise { - const uri = this.mementoManager.getUrl() + "/templates"; + const deployment = await this.requireDeployment(); + const uri = deployment.url + "/templates"; await vscode.commands.executeCommand("vscode.open", uri); } @@ -227,8 +236,9 @@ export class Commands { */ public async navigateToWorkspace(item: OpenableTreeItem) { if (item) { + const deployment = await this.requireDeployment(); const workspaceId = createWorkspaceIdentifier(item.workspace); - const uri = this.mementoManager.getUrl() + `/@${workspaceId}`; + const uri = deployment.url + `/@${workspaceId}`; await vscode.commands.executeCommand("vscode.open", uri); } else if (this.workspace && this.workspaceRestClient) { const baseUrl = @@ -250,8 +260,9 @@ export class Commands { */ public async navigateToWorkspaceSettings(item: OpenableTreeItem) { if (item) { + const deployment = await this.requireDeployment(); const workspaceId = createWorkspaceIdentifier(item.workspace); - const uri = this.mementoManager.getUrl() + `/@${workspaceId}/settings`; + const uri = deployment.url + `/@${workspaceId}/settings`; await vscode.commands.executeCommand("vscode.open", uri); } else if (this.workspace && this.workspaceRestClient) { const baseUrl = @@ -328,18 +339,14 @@ export class Commands { const terminal = vscode.window.createTerminal(app.name); // If workspace_name is provided, run coder ssh before the command - - const url = this.mementoManager.getUrl(); - if (!url) { - throw new Error("No coder url found for sidebar"); - } + const deployment = await this.requireDeployment(); const binary = await this.cliManager.fetchBinary( this.restClient, - toSafeHost(url), + deployment.label, ); const configDir = this.pathResolver.getGlobalConfigDir( - toSafeHost(url), + deployment.label, ); const globalFlags = getGlobalFlags( vscode.workspace.getConfiguration(), diff --git a/src/core/container.ts b/src/core/container.ts index 6c9d8b10..10cbf162 100644 --- a/src/core/container.ts +++ b/src/core/container.ts @@ -44,6 +44,7 @@ export class ServiceContainer implements vscode.Disposable { this.contextManager = new ContextManager(); this.loginCoordinator = new LoginCoordinator( this.secretsManager, + this.mementoManager, this.vscodeProposed, this.logger, ); diff --git a/src/core/mementoManager.ts b/src/core/mementoManager.ts index a317ffe5..3cf4478e 100644 --- a/src/core/mementoManager.ts +++ b/src/core/mementoManager.ts @@ -7,27 +7,16 @@ export class MementoManager { constructor(private readonly memento: Memento) {} /** - * Add the URL to the list of recently accessed URLs in global storage, then - * set it as the last used URL. - * - * If the URL is falsey, then remove it as the last used URL and do not touch - * the history. + * Add a URL to the history of recently accessed URLs. + * Used by the URL picker to show recent deployments. */ - public async setUrl(url: string | undefined): Promise { - await this.memento.update("url", url); + public async addToUrlHistory(url: string): Promise { if (url) { const history = this.withUrlHistory(url); await this.memento.update("urlHistory", history); } } - /** - * Get the last used URL. - */ - public getUrl(): string | undefined { - return this.memento.get("url"); - } - /** * Get the most recently accessed URLs (oldest to newest) with the provided * values appended. Duplicates will be removed. diff --git a/src/core/secretsManager.ts b/src/core/secretsManager.ts index e949e8e8..6069439c 100644 --- a/src/core/secretsManager.ts +++ b/src/core/secretsManager.ts @@ -1,15 +1,16 @@ -import { - type TokenResponse, - type ClientRegistrationResponse, -} from "../oauth/types"; +import { toSafeHost } from "../util"; import type { Memento, SecretStorage, Disposable } from "vscode"; +import type { TokenResponse, ClientRegistrationResponse } from "../oauth/types"; + +import type { Deployment } from "./deployment"; + const SESSION_KEY_PREFIX = "coder.session."; const OAUTH_TOKENS_PREFIX = "coder.oauth.tokens."; const OAUTH_CLIENT_PREFIX = "coder.oauth.client."; -const LOGIN_STATE_KEY = "coder.loginState"; +const CURRENT_DEPLOYMENT_KEY = "coder.currentDeployment"; const OAUTH_CALLBACK_KEY = "coder.oauthCallback"; const KNOWN_LABELS_KEY = "coder.knownLabels"; @@ -32,10 +33,8 @@ interface OAuthCallbackData { error: string | null; } -export enum AuthAction { - LOGIN, - LOGOUT, - INVALID, +export interface CurrentDeploymentState { + deployment: Deployment | null; } export class SecretsManager { @@ -45,54 +44,57 @@ export class SecretsManager { ) {} /** - * Triggers a login/logout event that propagates across all VS Code windows. + * Sets the current deployment and triggers a cross-window sync event. + * This is the single source of truth for which deployment is currently active. */ - public async triggerLoginStateChange( - label: string, - action: "login" | "logout", + public async setCurrentDeployment( + deployment: Deployment | undefined, ): Promise { - const loginState = { - action, - label, + const state = { + deployment: deployment ?? null, timestamp: new Date().toISOString(), }; - await this.secrets.store(LOGIN_STATE_KEY, JSON.stringify(loginState)); + await this.secrets.store(CURRENT_DEPLOYMENT_KEY, JSON.stringify(state)); } /** - * Listens for login/logout events from any VS Code window. + * Gets the current deployment from storage. */ - public onDidChangeLoginState( - listener: (state: AuthAction, label: string) => Promise, - ): Disposable { - return this.secrets.onDidChange(async (e) => { - if (e.key !== LOGIN_STATE_KEY) { - return; + public async getCurrentDeployment(): Promise { + try { + const data = await this.secrets.get(CURRENT_DEPLOYMENT_KEY); + if (!data) { + return undefined; } + const parsed = JSON.parse(data) as { deployment: Deployment | null }; + return parsed.deployment ?? undefined; + } catch { + return undefined; + } + } - const stateStr = await this.secrets.get(LOGIN_STATE_KEY); - if (!stateStr) { - await listener(AuthAction.INVALID, ""); + /** + * Listens for deployment changes from any VS Code window. + * Fires when login, logout, or deployment switch occurs. + */ + public onDidChangeCurrentDeployment( + listener: (state: CurrentDeploymentState) => void | Promise, + ): Disposable { + return this.secrets.onDidChange(async (e) => { + if (e.key !== CURRENT_DEPLOYMENT_KEY) { return; } try { - const parsed = JSON.parse(stateStr) as { - action: string; - label: string; - timestamp: string; - }; - - if (parsed.action === "login") { - await listener(AuthAction.LOGIN, parsed.label); - } else if (parsed.action === "logout") { - await listener(AuthAction.LOGOUT, parsed.label); - } else { - await listener(AuthAction.INVALID, parsed.label); + 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 { - // Invalid JSON, treat as invalid state - await listener(AuthAction.INVALID, ""); + // Ignore parse errors } }); } @@ -132,7 +134,7 @@ export class SecretsManager { /** * Listen for changes to a specific deployment's session auth. */ - public onDidChangeDeploymentAuth( + public onDidChangeSessionAuth( label: string, listener: (auth: SessionAuth | undefined) => void | Promise, ): Disposable { @@ -147,6 +149,10 @@ export class SecretsManager { } public async getSessionAuth(label: string): Promise { + if (!label) { + return undefined; + } + try { const data = await this.secrets.get(`${SESSION_KEY_PREFIX}${label}`); if (!data) { @@ -286,24 +292,35 @@ export class SecretsManager { /** * Migrate from legacy flat sessionToken storage to new format. + * Also sets the current deployment if none exists. */ - public async migrateFromLegacyStorage( - url: string, - label: string, - ): Promise { + public async migrateFromLegacyStorage(): Promise { + const legacyUrl = this.memento.get("url"); + if (!legacyUrl) { + return undefined; + } + + const label = toSafeHost(legacyUrl); + const existing = await this.getSessionAuth(label); if (existing) { - return false; + return undefined; } const oldToken = await this.secrets.get(LEGACY_SESSION_TOKEN_KEY); if (!oldToken) { - return false; + return undefined; } - await this.setSessionAuth(label, { url, token: oldToken }); + await this.setSessionAuth(label, { url: legacyUrl, token: oldToken }); await this.secrets.delete(LEGACY_SESSION_TOKEN_KEY); - return true; + // Also set as current deployment if none exists + const currentDeployment = await this.getCurrentDeployment(); + if (!currentDeployment) { + await this.setCurrentDeployment({ url: legacyUrl, label }); + } + + return label; } } diff --git a/src/extension.ts b/src/extension.ts index cbf67583..7f513f1c 100644 --- a/src/extension.ts +++ b/src/extension.ts @@ -12,7 +12,6 @@ import { attachOAuthInterceptors } from "./api/oauthInterceptors"; import { needToken } from "./api/utils"; import { Commands } from "./commands"; import { ServiceContainer } from "./core/container"; -import { AuthAction } from "./core/secretsManager"; import { CertificateError, getErrorDetail } from "./error"; import { OAuthSessionManager } from "./oauth/sessionManager"; import { CALLBACK_PATH } from "./oauth/utils"; @@ -69,9 +68,7 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { // Try to clear this flag ASAP const isFirstConnect = await mementoManager.getAndClearFirstConnect(); - const url = mementoManager.getUrl(); - const label = url ? toSafeHost(url) : ""; - const deployment = url ? { url, label } : undefined; + const deployment = await secretsManager.getCurrentDeployment(); // Create OAuth session manager with login coordinator const oauthSessionManager = await OAuthSessionManager.create( @@ -85,8 +82,8 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { // the plugin to poll workspaces for the current login, as well as being used // in commands that operate on the current login. const client = CoderApi.create( - url || "", - await secretsManager.getSessionToken(label), + deployment?.url || "", + await secretsManager.getSessionToken(deployment?.label ?? ""), output, ); attachOAuthInterceptors(client, output, oauthSessionManager); @@ -134,10 +131,10 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { ctx.subscriptions, ); - // Listen for deployment auth changes and sync state across all components + // Listen for deployment auth changes (token updates) for the current deployment // This listener is re-registered when the user logs into a different deployment let authChangeDisposable: vscode.Disposable | undefined; - const registerAuthListener = (deploymentLabel: string) => { + const registerAuthListener = (deploymentLabel: string | undefined) => { authChangeDisposable?.dispose(); if (!deploymentLabel) { @@ -145,7 +142,7 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { } output.debug("Registering auth listener for deployment", deploymentLabel); - authChangeDisposable = secretsManager.onDidChangeDeploymentAuth( + authChangeDisposable = secretsManager.onDidChangeSessionAuth( deploymentLabel, (auth) => { client.setHost(auth?.url); @@ -157,7 +154,8 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { ); }; - registerAuthListener(label); + // Initialize auth listener for current deployment + registerAuthListener(deployment?.label); ctx.subscriptions.push({ dispose: () => authChangeDisposable?.dispose() }); // Handle vscode:// URIs. @@ -189,37 +187,17 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { throw new Error("workspace must be specified as a query parameter"); } - // We are not guaranteed that the URL we currently have is for the URL - // this workspace belongs to, or that we even have a URL at all (the - // queries will default to localhost) so ask for it if missing. - // Pre-populate in case we do have the right URL so the user can just - // hit enter and move on. - const url = await maybeAskUrl(mementoManager, params.get("url")); - if (url) { - client.setHost(url); - await mementoManager.setUrl(url); - } else { + const deployment = await setupDeploymentFromUri( + params, + client, + serviceContainer, + ); + if (!deployment) { throw new Error( "url must be provided or specified as a query parameter", ); } - // If the token is missing we will get a 401 later and the user will be - // prompted to sign in again, so we do not need to ensure it is set now. - // For non-token auth, we write a blank token since the `vscodessh` - // command currently always requires a token file. However, if there is - // a query parameter for non-token auth go ahead and use it anyway; all - // that really matters is the file is created. - const token = needToken(vscode.workspace.getConfiguration()) - ? params.get("token") - : (params.get("token") ?? ""); - - const label = toSafeHost(url); - if (token) { - client.setSessionToken(token); - await secretsManager.setSessionAuth(label, { url, token }); - } - vscode.commands.executeCommand( "coder.open", owner, @@ -267,36 +245,17 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { ); } - // We are not guaranteed that the URL we currently have is for the URL - // this workspace belongs to, or that we even have a URL at all (the - // queries will default to localhost) so ask for it if missing. - // Pre-populate in case we do have the right URL so the user can just - // hit enter and move on. - const url = await maybeAskUrl(mementoManager, params.get("url")); - if (url) { - client.setHost(url); - await mementoManager.setUrl(url); - } else { + const deployment = await setupDeploymentFromUri( + params, + client, + serviceContainer, + ); + if (!deployment) { throw new Error( "url must be provided or specified as a query parameter", ); } - // If the token is missing we will get a 401 later and the user will be - // prompted to sign in again, so we do not need to ensure it is set now. - // For non-token auth, we write a blank token since the `vscodessh` - // command currently always requires a token file. However, if there is - // a query parameter for non-token auth go ahead and use it anyway; all - // that really matters is the file is created. - const token = needToken(vscode.workspace.getConfiguration()) - ? params.get("token") - : (params.get("token") ?? ""); - - if (token) { - client.setSessionToken(token); - await secretsManager.setSessionAuth(label, { url, token }); - } - vscode.commands.executeCommand( "coder.openDevContainer", workspaceOwner, @@ -373,33 +332,29 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { const remote = new Remote(serviceContainer, commands, ctx); + // Listen for deployment changes from other windows (cross-window sync) ctx.subscriptions.push( - secretsManager.onDidChangeLoginState(async (state, label) => { - switch (state) { - case AuthAction.LOGIN: { - const url = mementoManager.getUrl(); - // Should login the user directly if the URL+Token are valid - await commands.login({ url }); - - // Re-register auth listener for the new deployment - registerAuthListener(label); - - // Update OAuth session manager to match the new deployment - if (url) { - await oauthSessionManager.setDeployment({ label, url }); - } - break; - } - case AuthAction.LOGOUT: - await commands.forceLogout(label); - - // Dispose auth listener when logged out - authChangeDisposable?.dispose(); - authChangeDisposable = undefined; - break; - case AuthAction.INVALID: - break; + secretsManager.onDidChangeCurrentDeployment(async ({ deployment }) => { + output.info("Deployment changed from another window"); + + // Update client + client.setHost(deployment?.url); + if (deployment) { + const token = await secretsManager.getSessionToken(deployment.label); + client.setSessionToken(token ?? ""); + await oauthSessionManager.setDeployment(deployment); + } else { + client.setSessionToken(""); + oauthSessionManager.clearDeployment(); } + registerAuthListener(deployment?.label); + + // Update context + contextManager.set("coder.authenticated", deployment !== undefined); + + // Refresh workspaces + myWorkspacesProvider.fetchAndRefresh(); + allWorkspacesProvider.fetchAndRefresh(); }), ); @@ -419,17 +374,19 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { isFirstConnect, remoteSshExtension.id, ); - // TODO this is weird, why do we change the host here?? if (details) { ctx.subscriptions.push(details); - // Authenticate the plugin client which is used in the sidebar to display - // workspaces belonging to this deployment. + + // Set client host/token immediately for subsequent operations. client.setHost(details.url); client.setSessionToken(details.token); - await oauthSessionManager.setDeployment({ - label: details.label, + + // Persist and sync deployment across windows + await secretsManager.setCurrentDeployment({ url: details.url, + label: details.label, }); + await mementoManager.addToUrlHistory(details.url); } } catch (ex) { if (ex instanceof CertificateError) { @@ -526,31 +483,20 @@ async function migrateAuthStorage( serviceContainer: ServiceContainer, ): Promise { const secretsManager = serviceContainer.getSecretsManager(); - const mementoManager = serviceContainer.getMementoManager(); const output = serviceContainer.getLogger(); try { - // Get deployment URL from memento - const url = mementoManager.getUrl(); - if (!url) { - output.info("No URL configured, skipping migration"); - return; - } - - // Perform migration using SecretsManager method - const label = toSafeHost(url); - const migrated = await secretsManager.migrateFromLegacyStorage(url, label); + const migratedLabel = await secretsManager.migrateFromLegacyStorage(); - if (migrated) { + if (migratedLabel) { output.info( - `Successfully migrated auth storage to label-based format (label: ${label})`, + `Successfully migrated auth storage to label-based format (label: ${migratedLabel})`, ); } } catch (error) { output.error( `Auth storage migration failed: ${error}. You may need to log in again.`, ); - // Don't throw - allow extension to continue } } @@ -558,3 +504,62 @@ async function showTreeViewSearch(id: string): Promise { await vscode.commands.executeCommand(`${id}.focus`); await vscode.commands.executeCommand("list.find"); } + +/** + * Sets up deployment from URI parameters. Handles URL prompting, client setup, + * and token storage. Throws if user cancels URL input. + * + * Sets client host/token immediately for subsequent operations. + * Other updates (auth listener, OAuth manager, context, workspaces) are handled + * asynchronously by the onDidChangeCurrentDeployment listener. + */ +async function setupDeploymentFromUri( + params: URLSearchParams, + client: CoderApi, + serviceContainer: ServiceContainer, +): Promise<{ url: string; label: string }> { + const secretsManager = serviceContainer.getSecretsManager(); + const mementoManager = serviceContainer.getMementoManager(); + const currentDeployment = await secretsManager.getCurrentDeployment(); + + // We are not guaranteed that the URL we currently have is for the URL + // this workspace belongs to, or that we even have a URL at all (the + // queries will default to localhost) so ask for it if missing. + // Pre-populate in case we do have the right URL so the user can just + // hit enter and move on. + const url = await maybeAskUrl( + mementoManager, + params.get("url"), + currentDeployment?.url, + ); + if (!url) { + throw new Error("url must be provided or specified as a query parameter"); + } + + const label = toSafeHost(url); + + // Set client host immediately for subsequent operations. + client.setHost(url); + + // If the token is missing we will get a 401 later and the user will be + // prompted to sign in again, so we do not need to ensure it is set now. + // For non-token auth, we write a blank token since the `vscodessh` + // command currently always requires a token file. However, if there is + // a query parameter for non-token auth go ahead and use it anyway; all + // that really matters is the file is created. + const token = needToken(vscode.workspace.getConfiguration()) + ? params.get("token") + : (params.get("token") ?? ""); + + if (token) { + // Set token immediately for subsequent operations + client.setSessionToken(token); + await secretsManager.setSessionAuth(label, { url, token }); + } + + // Persist and sync deployment across windows + await secretsManager.setCurrentDeployment({ url, label }); + await mementoManager.addToUrlHistory(url); + + return { url, label }; +} diff --git a/src/login/loginCoordinator.ts b/src/login/loginCoordinator.ts index 06a876d5..a2cd84cb 100644 --- a/src/login/loginCoordinator.ts +++ b/src/login/loginCoordinator.ts @@ -1,31 +1,30 @@ import { getErrorMessage } from "coder/site/src/api/errors"; import * as vscode from "vscode"; -import { type Deployment } from "src/core/deployment"; - import { CoderApi } from "../api/coderApi"; import { needToken } from "../api/utils"; +import { type Deployment } from "../core/deployment"; +import { type MementoManager } from "../core/mementoManager"; import { type SecretsManager } from "../core/secretsManager"; import { CertificateError } from "../error"; import { type Logger } from "../logging/logger"; -import { maybeAskAuthMethod } from "../promptUtils"; +import { maybeAskAuthMethod, maybeAskUrl } from "../promptUtils"; import type { User } from "coder/site/src/api/typesGenerated"; import type { OAuthSessionManager } from "../oauth/sessionManager"; -export interface LoginResult { +interface LoginResult { success: boolean; user?: User; token?: string; } -export interface LoginOptions { - deployment: Deployment; +interface LoginOptions { + label: string; + url: string | undefined; oauthSessionManager: OAuthSessionManager; autoLogin?: boolean; - message?: string; - detailPrefix?: string; } /** @@ -36,23 +35,29 @@ export class LoginCoordinator { constructor( private readonly secretsManager: SecretsManager, + private readonly mementoManager: MementoManager, private readonly vscodeProposed: typeof vscode, private readonly logger: Logger, ) {} /** * Direct login - for user-initiated login via commands. + * Stores session auth and URL history on success. */ public async promptForLogin( - options: Omit, + options: LoginOptions & { url: string }, ): Promise { - const { deployment, oauthSessionManager } = options; - return this.executeWithGuard(options.deployment.label, async () => { - return this.attemptLogin( - deployment, + const { label, url, oauthSessionManager } = options; + return this.executeWithGuard(label, async () => { + const result = await this.attemptLogin( + { label, url }, options.autoLogin ?? false, oauthSessionManager, ); + + await this.persistSessionAuth(result, label, url); + + return result; }); } @@ -60,10 +65,10 @@ export class LoginCoordinator { * Shows dialog then login - for system-initiated auth (remote, OAuth refresh). */ public async promptForLoginWithDialog( - options: LoginOptions, + options: LoginOptions & { message?: string; detailPrefix?: string }, ): Promise { - const { deployment, detailPrefix, message, oauthSessionManager } = options; - return this.executeWithGuard(deployment.label, () => { + const { label, url, detailPrefix, message, oauthSessionManager } = options; + return this.executeWithGuard(label, () => { // Show dialog promise const dialogPromise = this.vscodeProposed.window .showErrorMessage( @@ -72,27 +77,31 @@ export class LoginCoordinator { modal: true, useCustom: true, detail: - (detailPrefix || - `Authentication needed for ${deployment.label}.`) + + (detailPrefix || `Authentication needed for ${label}.`) + "\n\nIf you've already logged in, you may close this dialog.", }, "Login", ) .then(async (action) => { if (action === "Login") { - // User clicked login - proceed with login flow + // Proceed with the login flow, handling logging in from another window + const storedUrl = await this.secretsManager.getUrl(label); + const newUrl = await maybeAskUrl( + this.mementoManager, + url, + storedUrl, + ); + if (!newUrl) { + throw new Error("URL must be provided"); + } + const result = await this.attemptLogin( - deployment, + { url: newUrl, label }, false, oauthSessionManager, ); - if (result.success && result.token) { - await this.secretsManager.setSessionAuth(deployment.label, { - url: deployment.url, - token: result.token, - }); - } + await this.persistSessionAuth(result, label, newUrl); return result; } else { @@ -102,13 +111,24 @@ export class LoginCoordinator { }); // Race between user clicking login and cross-window detection - return Promise.race([ - dialogPromise, - this.waitForCrossWindowLogin(deployment.label), - ]); + return Promise.race([dialogPromise, this.waitForCrossWindowLogin(label)]); }); } + private async persistSessionAuth( + result: LoginResult, + label: string, + url: string, + ): Promise { + if (result.success && result.token) { + await this.secretsManager.setSessionAuth(label, { + url, + token: result.token, + }); + await this.mementoManager.addToUrlHistory(url); + } + } + /** * Same-window guard wrapper. */ @@ -136,7 +156,7 @@ export class LoginCoordinator { */ private async waitForCrossWindowLogin(label: string): Promise { return new Promise((resolve) => { - const disposable = this.secretsManager.onDidChangeDeploymentAuth( + const disposable = this.secretsManager.onDidChangeSessionAuth( label, (auth) => { if (auth?.token) { diff --git a/src/oauth/sessionManager.ts b/src/oauth/sessionManager.ts index a533850a..3e667c22 100644 --- a/src/oauth/sessionManager.ts +++ b/src/oauth/sessionManager.ts @@ -149,7 +149,7 @@ export class OAuthSessionManager implements vscode.Disposable { } this.storedTokens = tokens; - this.logger.info(`Loaded stored OAuth tokens for ${tokens.deployment_url}`); + this.logger.info(`Loaded stored OAuth tokens for ${this.deployment.label}`); } private async clearOAuthState(): Promise { @@ -749,7 +749,8 @@ export class OAuthSessionManager implements vscode.Disposable { await this.secretsManager.clearAllAuthData(deployment.label); await this.loginCoordinator.promptForLoginWithDialog({ - deployment, + label: deployment.label, + url: deployment.url, detailPrefix: errorMessage, oauthSessionManager: this, }); diff --git a/src/promptUtils.ts b/src/promptUtils.ts index 35a8b838..9e3d8895 100644 --- a/src/promptUtils.ts +++ b/src/promptUtils.ts @@ -65,20 +65,16 @@ export async function maybeAskAgent( */ async function askURL( mementoManager: MementoManager, - selection: string | undefined, + prePopulateUrl: string | undefined, ): Promise { - const lastUsedUrl = mementoManager.getUrl(); const defaultURL = vscode.workspace .getConfiguration() .get("coder.defaultUrl") ?.trim(); const quickPick = vscode.window.createQuickPick(); + quickPick.ignoreFocusOut = true; quickPick.value = - selection || - lastUsedUrl || - defaultURL || - process.env.CODER_URL?.trim() || - ""; + prePopulateUrl || defaultURL || process.env.CODER_URL?.trim() || ""; quickPick.placeholder = "https://example.coder.com"; quickPick.title = "Enter the URL of your Coder deployment."; @@ -120,9 +116,9 @@ async function askURL( export async function maybeAskUrl( mementoManager: MementoManager, providedUrl: string | undefined | null, - lastUsedUrl?: string, + prePopulateUrl?: string, ): Promise { - let url = providedUrl || (await askURL(mementoManager, lastUsedUrl)); + let url = providedUrl || (await askURL(mementoManager, prePopulateUrl)); if (!url) { // User aborted. return undefined; diff --git a/src/remote/remote.ts b/src/remote/remote.ts index fce41359..bdad051f 100644 --- a/src/remote/remote.ts +++ b/src/remote/remote.ts @@ -34,7 +34,6 @@ import { Inbox } from "../inbox"; import { type Logger } from "../logging/logger"; import { type LoginCoordinator } from "../login/loginCoordinator"; import { OAuthSessionManager } from "../oauth/sessionManager"; -import { maybeAskUrl } from "../promptUtils"; import { AuthorityPrefix, escapeCommandArg, @@ -111,7 +110,7 @@ export class Remote { try { disposables.push( - this.secretsManager.onDidChangeDeploymentAuth( + this.secretsManager.onDidChangeSessionAuth( parts.label, async (auth) => { if (auth?.token && auth.url) { @@ -137,9 +136,13 @@ export class Remote { ); disposables.push(remoteOAuthManager); - const promptForLoginAndRetry = async (message: string, url: string) => { + const promptForLoginAndRetry = async ( + message: string, + url: string | undefined, + ) => { const result = await this.loginCoordinator.promptForLoginWithDialog({ - deployment: { url, label: parts.label }, + label: parts.label, + url, message, detailPrefix: `You must log in to access ${workspaceName}.`, oauthSessionManager: remoteOAuthManager, @@ -163,17 +166,7 @@ export class Remote { !baseUrlRaw || (!token && needToken(vscode.workspace.getConfiguration())) ) { - const mementoManager = this.serviceContainer.getMementoManager(); - const newUrl = await maybeAskUrl( - mementoManager, - baseUrlRaw, // TODO can we assume that "https://" is always valid? - parts.label, - ); - if (!newUrl) { - throw new Error("URL must be provided"); - } - - return promptForLoginAndRetry("You are not logged in...", newUrl); + return promptForLoginAndRetry("You are not logged in...", baseUrlRaw); } this.logger.info("Using deployment URL", baseUrlRaw); @@ -191,7 +184,7 @@ export class Remote { // Listen for token changes for this deployment disposables.push( - this.secretsManager.onDidChangeDeploymentAuth(parts.label, (auth) => { + this.secretsManager.onDidChangeSessionAuth(parts.label, (auth) => { workspaceClient.setHost(auth?.url); workspaceClient.setSessionToken(auth?.token ?? ""); }), diff --git a/test/unit/core/mementoManager.test.ts b/test/unit/core/mementoManager.test.ts index 54289a65..791f7602 100644 --- a/test/unit/core/mementoManager.test.ts +++ b/test/unit/core/mementoManager.test.ts @@ -13,28 +13,22 @@ describe("MementoManager", () => { mementoManager = new MementoManager(memento); }); - describe("setUrl", () => { - it("should store URL and add to history", async () => { - await mementoManager.setUrl("https://coder.example.com"); + describe("addToUrlHistory", () => { + it("should add URL to history", async () => { + await mementoManager.addToUrlHistory("https://coder.example.com"); - expect(mementoManager.getUrl()).toBe("https://coder.example.com"); expect(memento.get("urlHistory")).toEqual(["https://coder.example.com"]); }); it("should not update history for falsy values", async () => { - await mementoManager.setUrl(undefined); - expect(mementoManager.getUrl()).toBeUndefined(); - expect(memento.get("urlHistory")).toBeUndefined(); - - await mementoManager.setUrl(""); - expect(mementoManager.getUrl()).toBe(""); + await mementoManager.addToUrlHistory(""); expect(memento.get("urlHistory")).toBeUndefined(); }); it("should deduplicate URLs in history", async () => { - await mementoManager.setUrl("url1"); - await mementoManager.setUrl("url2"); - await mementoManager.setUrl("url1"); // Re-add first URL + await mementoManager.addToUrlHistory("url1"); + await mementoManager.addToUrlHistory("url2"); + await mementoManager.addToUrlHistory("url1"); // Re-add first URL expect(memento.get("urlHistory")).toEqual(["url2", "url1"]); }); diff --git a/test/unit/core/secretsManager.test.ts b/test/unit/core/secretsManager.test.ts index 9de34c83..5aac3425 100644 --- a/test/unit/core/secretsManager.test.ts +++ b/test/unit/core/secretsManager.test.ts @@ -1,6 +1,9 @@ import { beforeEach, describe, expect, it, vi } from "vitest"; -import { AuthAction, SecretsManager } from "@/core/secretsManager"; +import { + type CurrentDeploymentState, + SecretsManager, +} from "@/core/secretsManager"; import { InMemoryMemento, @@ -13,6 +16,7 @@ describe("SecretsManager", () => { let secretsManager: SecretsManager; beforeEach(() => { + vi.useRealTimers(); secretStorage = new InMemorySecretStorage(); memento = new InMemoryMemento(); secretsManager = new SecretsManager(secretStorage, memento); @@ -20,40 +24,40 @@ describe("SecretsManager", () => { describe("session auth", () => { it("should store and retrieve session auth", async () => { - await secretsManager.setSessionAuth("example-com", { + await secretsManager.setSessionAuth("example.com", { url: "https://example.com", token: "test-token", }); - expect(await secretsManager.getSessionToken("example-com")).toBe( + expect(await secretsManager.getSessionToken("example.com")).toBe( "test-token", ); - expect(await secretsManager.getUrl("example-com")).toBe( + expect(await secretsManager.getUrl("example.com")).toBe( "https://example.com", ); - await secretsManager.setSessionAuth("example-com", { + await secretsManager.setSessionAuth("example.com", { url: "https://example.com", token: "new-token", }); - expect(await secretsManager.getSessionToken("example-com")).toBe( + expect(await secretsManager.getSessionToken("example.com")).toBe( "new-token", ); }); it("should clear session auth", async () => { - await secretsManager.setSessionAuth("example-com", { + await secretsManager.setSessionAuth("example.com", { url: "https://example.com", token: "test-token", }); - await secretsManager.clearSessionAuth("example-com"); + await secretsManager.clearSessionAuth("example.com"); expect( - await secretsManager.getSessionToken("example-com"), + await secretsManager.getSessionToken("example.com"), ).toBeUndefined(); }); it("should return undefined for corrupted storage", async () => { await secretStorage.store( - "coder.session.example-com", + "coder.session.example.com", JSON.stringify({ url: "https://example.com", token: "valid-token", @@ -62,83 +66,90 @@ describe("SecretsManager", () => { secretStorage.corruptStorage(); expect( - await secretsManager.getSessionToken("example-com"), + await secretsManager.getSessionToken("example.com"), ).toBeUndefined(); }); it("should track known labels", async () => { expect(secretsManager.getKnownLabels()).toEqual([]); - await secretsManager.setSessionAuth("example-com", { + await secretsManager.setSessionAuth("example.com", { url: "https://example.com", token: "test-token", }); - expect(secretsManager.getKnownLabels()).toContain("example-com"); + expect(secretsManager.getKnownLabels()).toContain("example.com"); await secretsManager.setSessionAuth("other-com", { url: "https://other.com", token: "other-token", }); - expect(secretsManager.getKnownLabels()).toContain("example-com"); + expect(secretsManager.getKnownLabels()).toContain("example.com"); expect(secretsManager.getKnownLabels()).toContain("other-com"); }); it("should remove label on clearAllAuthData", async () => { - await secretsManager.setSessionAuth("example-com", { + await secretsManager.setSessionAuth("example.com", { url: "https://example.com", token: "test-token", }); - expect(secretsManager.getKnownLabels()).toContain("example-com"); + expect(secretsManager.getKnownLabels()).toContain("example.com"); - await secretsManager.clearAllAuthData("example-com"); - expect(secretsManager.getKnownLabels()).not.toContain("example-com"); + await secretsManager.clearAllAuthData("example.com"); + expect(secretsManager.getKnownLabels()).not.toContain("example.com"); }); }); - describe("login state", () => { - it("should trigger login events", async () => { - const events: Array<{ state: AuthAction; label: string }> = []; - secretsManager.onDidChangeLoginState((state, label) => { - events.push({ state, label }); - return Promise.resolve(); - }); + describe("current deployment", () => { + it("should store and retrieve current deployment", async () => { + const deployment = { url: "https://example.com", label: "example.com" }; + await secretsManager.setCurrentDeployment(deployment); - await secretsManager.triggerLoginStateChange("example-com", "login"); - expect(events).toEqual([ - { state: AuthAction.LOGIN, label: "example-com" }, - ]); + const result = await secretsManager.getCurrentDeployment(); + expect(result).toEqual(deployment); }); - it("should trigger logout events", async () => { - const events: Array<{ state: AuthAction; label: string }> = []; - secretsManager.onDidChangeLoginState((state, label) => { - events.push({ state, label }); - return Promise.resolve(); - }); + it("should clear current deployment with undefined", async () => { + const deployment = { url: "https://example.com", label: "example.com" }; + await secretsManager.setCurrentDeployment(deployment); + await secretsManager.setCurrentDeployment(undefined); + + const result = await secretsManager.getCurrentDeployment(); + expect(result).toBeUndefined(); + }); - await secretsManager.triggerLoginStateChange("example-com", "logout"); - expect(events).toEqual([ - { state: AuthAction.LOGOUT, label: "example-com" }, - ]); + it("should return undefined when no deployment set", async () => { + const result = await secretsManager.getCurrentDeployment(); + expect(result).toBeUndefined(); }); - it("should fire same event twice in a row", async () => { + it("should notify listeners on deployment change", async () => { vi.useFakeTimers(); - const events: Array<{ state: AuthAction; label: string }> = []; - secretsManager.onDidChangeLoginState((state, label) => { - events.push({ state, label }); - return Promise.resolve(); + const events: Array = []; + secretsManager.onDidChangeCurrentDeployment((state) => { + events.push(state); }); - await secretsManager.triggerLoginStateChange("example-com", "login"); + const deployments = [ + { url: "https://example.com", label: "example.com" }, + { url: "https://another.org", label: "another.org" }, + { url: "https://another.org", label: "another.org" }, + ]; + await secretsManager.setCurrentDeployment(deployments[0]); vi.advanceTimersByTime(5); - await secretsManager.triggerLoginStateChange("example-com", "login"); + await secretsManager.setCurrentDeployment(deployments[1]); + vi.advanceTimersByTime(5); + await secretsManager.setCurrentDeployment(deployments[2]); + vi.advanceTimersByTime(5); + + // Trigger an event even if the deployment did not change + expect(events).toEqual(deployments.map((deployment) => ({ deployment }))); + }); + + it("should handle corrupted storage gracefully", async () => { + await secretStorage.store("coder.currentDeployment", "invalid-json{"); - expect(events).toEqual([ - { state: AuthAction.LOGIN, label: "example-com" }, - { state: AuthAction.LOGIN, label: "example-com" }, - ]); - vi.useRealTimers(); + const result = await secretsManager.getCurrentDeployment(); + expect(result).toBeUndefined(); }); }); }); From 131a6e9fa1c85bb83261f63481a4f3214ebf294f Mon Sep 17 00:00:00 2001 From: Ehab Younes Date: Thu, 27 Nov 2025 13:44:48 +0300 Subject: [PATCH 17/20] Add WebSocket suspension --- src/api/coderApi.ts | 326 ++++++++++-------- src/commands.ts | 20 +- src/extension.ts | 69 ++-- src/remote/remote.ts | 4 +- src/websocket/codes.ts | 6 +- src/websocket/reconnectingWebSocket.ts | 74 +++- test/unit/api/coderApi.test.ts | 144 +++++++- .../websocket/reconnectingWebSocket.test.ts | 230 ++++++++++-- 8 files changed, 623 insertions(+), 250 deletions(-) diff --git a/src/api/coderApi.ts b/src/api/coderApi.ts index 663bf716..0f52f585 100644 --- a/src/api/coderApi.ts +++ b/src/api/coderApi.ts @@ -31,7 +31,7 @@ import { HttpClientLogLevel, } from "../logging/types"; import { sizeOf } from "../logging/utils"; -import { HttpStatusCode } from "../websocket/codes"; +import { HttpStatusCode, WebSocketCloseCode } from "../websocket/codes"; import { type UnidirectionalStream, type CloseEvent, @@ -55,7 +55,7 @@ const coderSessionTokenHeader = "Coder-Session-Token"; * Unified API class that includes both REST API methods from the base Api class * and WebSocket methods for real-time functionality. */ -export class CoderApi extends Api { +export class CoderApi extends Api implements vscode.Disposable { private readonly reconnectingSockets = new Set< ReconnectingWebSocket >(); @@ -74,38 +74,63 @@ export class CoderApi extends Api { output: Logger, ): CoderApi { const client = new CoderApi(output); - client.setHost(baseUrl); - if (token) { - client.setSessionToken(token); - } + client.setCredentials(baseUrl, token); setupInterceptors(client, output); return client; } - setSessionToken = (token: string): void => { - const defaultHeaders = this.getAxiosInstance().defaults.headers.common; - const currentToken = defaultHeaders[coderSessionTokenHeader]; - defaultHeaders[coderSessionTokenHeader] = token; + /** + * Set both host and token together. Useful for login/logout/switch to + * avoid triggering multiple reconnection events. + */ + setCredentials = ( + host: string | undefined, + token: string | undefined, + ): void => { + const defaults = this.getAxiosInstance().defaults; + const currentHost = defaults.baseURL; + const currentToken = defaults.headers.common[coderSessionTokenHeader]; - if (currentToken !== token) { + defaults.baseURL = host; + defaults.headers.common[coderSessionTokenHeader] = token; + + const hostChanged = currentHost !== host; + const tokenChanged = currentToken !== token; + + if (hostChanged || tokenChanged) { for (const socket of this.reconnectingSockets) { - socket.reconnect(); + if (host) { + socket.reconnect(); + } else { + socket.suspend(WebSocketCloseCode.NORMAL, "Host cleared"); + } } } }; + setSessionToken = (token: string): void => { + const currentHost = this.getAxiosInstance().defaults.baseURL; + this.setCredentials(currentHost, token); + }; + setHost = (host: string | undefined): void => { - const defaults = this.getAxiosInstance().defaults; - const currentHost = defaults.baseURL; - defaults.baseURL = host; + const currentToken = this.getAxiosInstance().defaults.headers.common[ + coderSessionTokenHeader + ] as string | undefined; + this.setCredentials(host, currentToken); + }; - if (currentHost !== host) { - for (const socket of this.reconnectingSockets) { - socket.reconnect(); - } + /** + * Permanently dispose all WebSocket connections. + * This clears handlers and prevents reconnection. + */ + dispose(): void { + for (const socket of this.reconnectingSockets) { + socket.close(); } - }; + this.reconnectingSockets.clear(); + } watchInboxNotifications = async ( watchTemplates: string[], @@ -125,7 +150,7 @@ export class CoderApi extends Api { }; watchWorkspace = async (workspace: Workspace, options?: ClientOptions) => { - return this.createWebSocketWithFallback({ + return this.createWebSocketWithFallback({ apiRoute: `/api/v2/workspaces/${workspace.id}/watch-ws`, fallbackApiRoute: `/api/v2/workspaces/${workspace.id}/watch`, options, @@ -137,7 +162,7 @@ export class CoderApi extends Api { agentId: WorkspaceAgent["id"], options?: ClientOptions, ) => { - return this.createWebSocketWithFallback({ + return this.createWebSocketWithFallback({ apiRoute: `/api/v2/workspaceagents/${agentId}/watch-metadata-ws`, fallbackApiRoute: `/api/v2/workspaceagents/${agentId}/watch-metadata`, options, @@ -198,68 +223,62 @@ export class CoderApi extends Api { throw new Error("No base URL set on REST client"); } - const baseUrl = new URL(baseUrlRaw); - const token = this.getAxiosInstance().defaults.headers.common[ - coderSessionTokenHeader - ] as string | undefined; + return this.createOneWayWebSocket(socketConfigs); + }; - const headersFromCommand = await getHeaders( - baseUrlRaw, - getHeaderCommand(vscode.workspace.getConfiguration()), - this.output, - ); + if (enableRetry) { + return this.createReconnectingSocket(socketFactory, configs.apiRoute); + } + return socketFactory(); + } - const httpAgent = await createHttpAgent( - vscode.workspace.getConfiguration(), - ); + private async createOneWayWebSocket( + configs: Omit, + ): Promise> { + const baseUrlRaw = this.getAxiosInstance().defaults.baseURL; + if (!baseUrlRaw) { + throw new Error("No base URL set on REST client"); + } + const token = this.getAxiosInstance().defaults.headers.common[ + coderSessionTokenHeader + ] as string | undefined; - /** - * Similar to the REST client, we want to prioritize headers in this order (highest to lowest): - * 1. Headers from the header command - * 2. Any headers passed directly to this function - * 3. Coder session token from the Api client (if set) - */ - const headers = { - ...(token ? { [coderSessionTokenHeader]: token } : {}), - ...configs.options?.headers, - ...headersFromCommand, - }; + const headersFromCommand = await getHeaders( + baseUrlRaw, + getHeaderCommand(vscode.workspace.getConfiguration()), + this.output, + ); - const webSocket = new OneWayWebSocket({ - location: baseUrl, - ...socketConfigs, - options: { - ...configs.options, - agent: httpAgent, - followRedirects: true, - headers, - }, - }); + const httpAgent = await createHttpAgent( + vscode.workspace.getConfiguration(), + ); - this.attachStreamLogger(webSocket); - return webSocket; + /** + * Similar to the REST client, we want to prioritize headers in this order (highest to lowest): + * 1. Headers from the header command + * 2. Any headers passed directly to this function + * 3. Coder session token from the Api client (if set) + */ + const headers = { + ...(token ? { [coderSessionTokenHeader]: token } : {}), + ...configs.options?.headers, + ...headersFromCommand, }; - if (enableRetry) { - const reconnectingSocket = await ReconnectingWebSocket.create( - socketFactory, - this.output, - configs.apiRoute, - undefined, - () => - this.reconnectingSockets.delete( - reconnectingSocket as ReconnectingWebSocket, - ), - ); - - this.reconnectingSockets.add( - reconnectingSocket as ReconnectingWebSocket, - ); + const baseUrl = new URL(baseUrlRaw); + const ws = new OneWayWebSocket({ + location: baseUrl, + ...configs, + options: { + ...configs.options, + agent: httpAgent, + followRedirects: true, + headers, + }, + }); - return reconnectingSocket; - } else { - return socketFactory(); - } + this.attachStreamLogger(ws); + return ws; } private attachStreamLogger( @@ -288,44 +307,79 @@ export class CoderApi extends Api { /** * Create a WebSocket connection with SSE fallback on 404. * + * The factory tries WS first, falls back to SSE on 404. Since the factory + * is called on every reconnect. + * * Note: The fallback on SSE ignores all passed client options except the headers. */ - private async createWebSocketWithFallback(configs: { - apiRoute: string; - fallbackApiRoute: string; - searchParams?: Record | URLSearchParams; - options?: ClientOptions; - enableRetry?: boolean; - }): Promise> { - let webSocket: UnidirectionalStream; - try { - webSocket = await this.createWebSocket({ - apiRoute: configs.apiRoute, - searchParams: configs.searchParams, - options: configs.options, - enableRetry: configs.enableRetry, - }); - } catch { - // Failed to create WebSocket, use SSE fallback - return this.createSseFallback( - configs.fallbackApiRoute, - configs.searchParams, - configs.options?.headers, + private async createWebSocketWithFallback( + configs: Omit & { + fallbackApiRoute: string; + enableRetry?: boolean; + }, + ): Promise> { + const { fallbackApiRoute, enableRetry, ...socketConfigs } = configs; + const socketFactory: SocketFactory = async () => { + try { + const ws = + await this.createOneWayWebSocket(socketConfigs); + return await this.waitForOpen(ws); + } catch (error) { + if (this.is404Error(error)) { + this.output.warn( + `WebSocket failed, using SSE fallback: ${socketConfigs.apiRoute}`, + ); + const sse = this.createSseConnection( + fallbackApiRoute, + socketConfigs.searchParams, + socketConfigs.options?.headers, + ); + return await this.waitForOpen(sse); + } + throw error; + } + }; + + if (enableRetry) { + return this.createReconnectingSocket( + socketFactory, + socketConfigs.apiRoute, ); } + return socketFactory(); + } - return this.waitForConnection(webSocket, () => - this.createSseFallback( - configs.fallbackApiRoute, - configs.searchParams, - configs.options?.headers, - ), - ); + /** + * Create an SSE connection without waiting for connection. + */ + private createSseConnection( + apiRoute: string, + searchParams?: Record | URLSearchParams, + optionsHeaders?: Record, + ): SseConnection { + const baseUrlRaw = this.getAxiosInstance().defaults.baseURL; + if (!baseUrlRaw) { + throw new Error("No base URL set on REST client"); + } + const url = new URL(baseUrlRaw); + const sse = new SseConnection({ + location: url, + apiRoute, + searchParams, + axiosInstance: this.getAxiosInstance(), + optionsHeaders, + logger: this.output, + }); + + this.attachStreamLogger(sse); + return sse; } - private waitForConnection( + /** + * Wait for a connection to open. Rejects on error. + */ + private waitForOpen( connection: UnidirectionalStream, - onNotFound?: () => Promise>, ): Promise> { return new Promise((resolve, reject) => { const cleanup = () => { @@ -340,16 +394,8 @@ export class CoderApi extends Api { const handleError = (event: ErrorEvent) => { cleanup(); - const is404 = - event.message?.includes(String(HttpStatusCode.NOT_FOUND)) || - event.error?.message?.includes(String(HttpStatusCode.NOT_FOUND)); - - if (is404 && onNotFound) { - connection.close(); - onNotFound().then(resolve).catch(reject); - } else { - reject(event.error || new Error(event.message)); - } + connection.close(); + reject(event.error || new Error(event.message)); }; connection.addEventListener("open", handleOpen); @@ -358,32 +404,36 @@ export class CoderApi extends Api { } /** - * Create SSE fallback connection + * Check if an error is a 404 Not Found error. */ - private async createSseFallback( - apiRoute: string, - searchParams?: Record | URLSearchParams, - optionsHeaders?: Record, - ): Promise> { - this.output.warn(`WebSocket failed, using SSE fallback: ${apiRoute}`); - - const baseUrlRaw = this.getAxiosInstance().defaults.baseURL; - if (!baseUrlRaw) { - throw new Error("No base URL set on REST client"); - } + private is404Error(error: unknown): boolean { + const msg = error instanceof Error ? error.message : String(error); + return msg.includes(String(HttpStatusCode.NOT_FOUND)); + } - const baseUrl = new URL(baseUrlRaw); - const sseConnection = new SseConnection({ - location: baseUrl, + /** + * Create a ReconnectingWebSocket and track it for lifecycle management. + */ + private async createReconnectingSocket( + socketFactory: SocketFactory, + apiRoute: string, + ): Promise> { + const reconnectingSocket = await ReconnectingWebSocket.create( + socketFactory, + this.output, apiRoute, - searchParams, - axiosInstance: this.getAxiosInstance(), - optionsHeaders: optionsHeaders, - logger: this.output, - }); + undefined, + () => + this.reconnectingSockets.delete( + reconnectingSocket as ReconnectingWebSocket, + ), + ); + + this.reconnectingSockets.add( + reconnectingSocket as ReconnectingWebSocket, + ); - this.attachStreamLogger(sseConnection); - return this.waitForConnection(sseConnection); + return reconnectingSocket; } } diff --git a/src/commands.ts b/src/commands.ts index 8e1593c8..62f323ea 100644 --- a/src/commands.ts +++ b/src/commands.ts @@ -1,4 +1,3 @@ -import { type Api } from "coder/site/src/api/api"; import { type Workspace, type WorkspaceAgent, @@ -6,6 +5,7 @@ import { import * as vscode from "vscode"; import { createWorkspaceIdentifier, extractAgents } from "./api/api-helper"; +import { type CoderApi } from "./api/coderApi"; import { type CliManager } from "./core/cliManager"; import { type ServiceContainer } from "./core/container"; import { type ContextManager } from "./core/contextManager"; @@ -45,11 +45,11 @@ export class Commands { // if you use multiple deployments). public workspace?: Workspace; public workspaceLogPath?: string; - public workspaceRestClient?: Api; + public workspaceRestClient?: CoderApi; public constructor( serviceContainer: ServiceContainer, - private readonly restClient: Api, + private readonly restClient: CoderApi, private readonly oauthSessionManager: OAuthSessionManager, ) { this.vscodeProposed = serviceContainer.getVsCodeProposed(); @@ -117,8 +117,7 @@ export class Commands { // Set client immediately so subsequent operations in this function have the correct host/token. // The cross-window listener will also update the client, but that's async. - this.restClient.setHost(url); - this.restClient.setSessionToken(result.token); + this.restClient.setCredentials(url, result.token); // Set as current deployment (triggers cross-window sync). await this.secretsManager.setCurrentDeployment({ url, label }); @@ -143,8 +142,6 @@ export class Commands { vscode.commands.executeCommand("coder.open"); } }); - - vscode.commands.executeCommand("coder.refreshWorkspaces"); } /** @@ -193,15 +190,11 @@ export class Commands { // Clear from the REST client. An empty url will indicate to other parts of // the code that we are logged out. - this.restClient.setHost(""); - this.restClient.setSessionToken(""); + this.restClient.setCredentials(undefined, undefined); // Clear current deployment (triggers cross-window sync) await this.secretsManager.setCurrentDeployment(undefined); - // Clear all auth data for this deployment - await this.secretsManager.clearAllAuthData(label); - this.contextManager.set("coder.authenticated", false); vscode.window .showInformationMessage("You've been logged out of Coder!", "Login") @@ -210,9 +203,6 @@ export class Commands { this.login(); } }); - - // This will result in clearing the workspace list. - vscode.commands.executeCommand("coder.refreshWorkspaces"); } /** diff --git a/src/extension.ts b/src/extension.ts index 7f513f1c..77bf89c5 100644 --- a/src/extension.ts +++ b/src/extension.ts @@ -12,6 +12,8 @@ import { attachOAuthInterceptors } from "./api/oauthInterceptors"; import { needToken } from "./api/utils"; import { Commands } from "./commands"; import { ServiceContainer } from "./core/container"; +import { type Deployment } from "./core/deployment"; +import { type SecretsManager } from "./core/secretsManager"; import { CertificateError, getErrorDetail } from "./error"; import { OAuthSessionManager } from "./oauth/sessionManager"; import { CALLBACK_PATH } from "./oauth/utils"; @@ -86,6 +88,7 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { await secretsManager.getSessionToken(deployment?.label ?? ""), output, ); + ctx.subscriptions.push(client); attachOAuthInterceptors(client, output, oauthSessionManager); const myWorkspacesProvider = new WorkspaceProvider( @@ -145,8 +148,7 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { authChangeDisposable = secretsManager.onDidChangeSessionAuth( deploymentLabel, (auth) => { - client.setHost(auth?.url); - client.setSessionToken(auth?.token ?? ""); + client.setCredentials(auth?.url, auth?.token); // Update authentication context for current deployment contextManager.set("coder.authenticated", auth !== undefined); @@ -187,16 +189,7 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { throw new Error("workspace must be specified as a query parameter"); } - const deployment = await setupDeploymentFromUri( - params, - client, - serviceContainer, - ); - if (!deployment) { - throw new Error( - "url must be provided or specified as a query parameter", - ); - } + await setupDeploymentFromUri(params, client, serviceContainer); vscode.commands.executeCommand( "coder.open", @@ -245,16 +238,7 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { ); } - const deployment = await setupDeploymentFromUri( - params, - client, - serviceContainer, - ); - if (!deployment) { - throw new Error( - "url must be provided or specified as a query parameter", - ); - } + await setupDeploymentFromUri(params, client, serviceContainer); vscode.commands.executeCommand( "coder.openDevContainer", @@ -338,13 +322,12 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { output.info("Deployment changed from another window"); // Update client - client.setHost(deployment?.url); if (deployment) { const token = await secretsManager.getSessionToken(deployment.label); - client.setSessionToken(token ?? ""); + client.setCredentials(deployment.url, token); await oauthSessionManager.setDeployment(deployment); } else { - client.setSessionToken(""); + client.setCredentials(undefined, undefined); oauthSessionManager.clearDeployment(); } registerAuthListener(deployment?.label); @@ -378,8 +361,7 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { ctx.subscriptions.push(details); // Set client host/token immediately for subsequent operations. - client.setHost(details.url); - client.setSessionToken(details.token); + client.setCredentials(details.url, details.token); // Persist and sync deployment across windows await secretsManager.setCurrentDeployment({ @@ -517,7 +499,7 @@ async function setupDeploymentFromUri( params: URLSearchParams, client: CoderApi, serviceContainer: ServiceContainer, -): Promise<{ url: string; label: string }> { +): Promise { const secretsManager = serviceContainer.getSecretsManager(); const mementoManager = serviceContainer.getMementoManager(); const currentDeployment = await secretsManager.getCurrentDeployment(); @@ -538,22 +520,14 @@ async function setupDeploymentFromUri( const label = toSafeHost(url); - // Set client host immediately for subsequent operations. - client.setHost(url); - // If the token is missing we will get a 401 later and the user will be // prompted to sign in again, so we do not need to ensure it is set now. // For non-token auth, we write a blank token since the `vscodessh` // command currently always requires a token file. However, if there is - // a query parameter for non-token auth go ahead and use it anyway; all - // that really matters is the file is created. - const token = needToken(vscode.workspace.getConfiguration()) - ? params.get("token") - : (params.get("token") ?? ""); - + // a query parameter for non-token auth go ahead and use it anyway; + const token = await getToken(params, label, secretsManager); + client.setCredentials(url, token); if (token) { - // Set token immediately for subsequent operations - client.setSessionToken(token); await secretsManager.setSessionAuth(label, { url, token }); } @@ -563,3 +537,20 @@ async function setupDeploymentFromUri( return { url, label }; } + +async function getToken( + params: URLSearchParams, + label: string, + secretsManager: SecretsManager, +): Promise { + const paramsToken = params.get("token"); + if (paramsToken !== null) { + // Always prefer the passed token if set + return paramsToken; + } + + if (needToken(vscode.workspace.getConfiguration())) { + return await secretsManager.getSessionToken(label); + } + return ""; +} diff --git a/src/remote/remote.ts b/src/remote/remote.ts index bdad051f..e8c11060 100644 --- a/src/remote/remote.ts +++ b/src/remote/remote.ts @@ -178,6 +178,7 @@ export class Remote { // disallow logging out/in altogether, but for now just use a separate // client to remain unaffected by whatever the plugin is doing. const workspaceClient = CoderApi.create(baseUrlRaw, token, this.logger); + disposables.push(workspaceClient); attachOAuthInterceptors(workspaceClient, this.logger, remoteOAuthManager); // Store for use in commands. this.commands.workspaceRestClient = workspaceClient; @@ -185,8 +186,7 @@ export class Remote { // Listen for token changes for this deployment disposables.push( this.secretsManager.onDidChangeSessionAuth(parts.label, (auth) => { - workspaceClient.setHost(auth?.url); - workspaceClient.setSessionToken(auth?.token ?? ""); + workspaceClient.setCredentials(auth?.url, auth?.token); }), ); diff --git a/src/websocket/codes.ts b/src/websocket/codes.ts index ac8eccf7..f3fd95cd 100644 --- a/src/websocket/codes.ts +++ b/src/websocket/codes.ts @@ -19,7 +19,9 @@ export const WebSocketCloseCode = { /** HTTP status codes used for socket creation and connection logic */ export const HttpStatusCode = { - /** Authentication or permission denied */ + /** Authentication required */ + UNAUTHORIZED: 401, + /** Permission denied */ FORBIDDEN: 403, /** Endpoint not found */ NOT_FOUND: 404, @@ -43,7 +45,9 @@ export const UNRECOVERABLE_WS_CLOSE_CODES = new Set([ * These appear during socket creation and should stop reconnection attempts. */ export const UNRECOVERABLE_HTTP_CODES = new Set([ + HttpStatusCode.UNAUTHORIZED, HttpStatusCode.FORBIDDEN, + HttpStatusCode.NOT_FOUND, HttpStatusCode.GONE, HttpStatusCode.UPGRADE_REQUIRED, ]); diff --git a/src/websocket/reconnectingWebSocket.ts b/src/websocket/reconnectingWebSocket.ts index 2ced9351..55f72988 100644 --- a/src/websocket/reconnectingWebSocket.ts +++ b/src/websocket/reconnectingWebSocket.ts @@ -41,7 +41,8 @@ export class ReconnectingWebSocket #currentSocket: UnidirectionalStream | null = null; #backoffMs: number; #reconnectTimeoutId: NodeJS.Timeout | null = null; - #isDisposed = false; + #isSuspended = false; // Temporary pause, can be resumed via reconnect() + #isDisposed = false; // Permanent disposal, cannot be resumed #isConnecting = false; #pendingReconnect = false; readonly #onDispose?: () => void; @@ -102,6 +103,11 @@ export class ReconnectingWebSocket } reconnect(): void { + if (this.#isSuspended) { + this.#isSuspended = false; + this.#backoffMs = this.#options.initialBackoffMs; + } + if (this.#isDisposed) { return; } @@ -121,6 +127,18 @@ export class ReconnectingWebSocket this.connect().catch((error) => this.handleConnectionError(error)); } + /** + * Temporarily suspend the socket. Can be resumed via reconnect(). + */ + suspend(code?: number, reason?: string): void { + if (this.#isDisposed || this.#isSuspended) { + return; + } + + this.#isSuspended = true; + this.clearCurrentSocket(code, reason); + } + close(code?: number, reason?: string): void { if (this.#isDisposed) { return; @@ -139,7 +157,7 @@ export class ReconnectingWebSocket } private async connect(): Promise { - if (this.#isDisposed || this.#isConnecting) { + if (this.#isDisposed || this.#isSuspended || this.#isConnecting) { return; } @@ -168,10 +186,21 @@ export class ReconnectingWebSocket socket.addEventListener("error", (event) => { this.executeHandlers("error", event); + + // Check for unrecoverable HTTP errors in the error event + // HTTP errors during handshake fire 'error' then 'close' with 1006 + // We need to suspend here to prevent infinite reconnect loops + const errorMessage = event.error?.message ?? event.message ?? ""; + if (this.isUnrecoverableHttpError(errorMessage)) { + this.#logger.error( + `Unrecoverable HTTP error for ${this.#apiRoute}: ${errorMessage}`, + ); + this.suspend(); + } }); socket.addEventListener("close", (event) => { - if (this.#isDisposed) { + if (this.#isDisposed || this.#isSuspended) { return; } @@ -181,7 +210,8 @@ export class ReconnectingWebSocket this.#logger.error( `WebSocket connection closed with unrecoverable error code ${event.code}`, ); - this.dispose(); + // Suspend instead of dispose - allows recovery when credentials change + this.suspend(); return; } @@ -204,7 +234,11 @@ export class ReconnectingWebSocket } private scheduleReconnect(): void { - if (this.#isDisposed || this.#reconnectTimeoutId !== null) { + if ( + this.#isDisposed || + this.#isSuspended || + this.#reconnectTimeoutId !== null + ) { return; } @@ -241,11 +275,11 @@ export class ReconnectingWebSocket } /** - * Checks if the error is unrecoverable and disposes the connection, + * Checks if the error is unrecoverable and suspends the connection, * otherwise schedules a reconnect. */ private handleConnectionError(error: unknown): void { - if (this.#isDisposed) { + if (this.#isDisposed || this.#isSuspended) { return; } @@ -254,7 +288,7 @@ export class ReconnectingWebSocket `Unrecoverable HTTP error during connection for ${this.#apiRoute}`, error, ); - this.dispose(); + this.suspend(); return; } @@ -266,12 +300,12 @@ export class ReconnectingWebSocket } /** - * Check if an error contains an unrecoverable HTTP status code. + * Check if an error message contains an unrecoverable HTTP status code. */ private isUnrecoverableHttpError(error: unknown): boolean { - const errorMessage = error instanceof Error ? error.message : String(error); + const message = (error as { message?: string }).message || String(error); for (const code of UNRECOVERABLE_HTTP_CODES) { - if (errorMessage.includes(String(code))) { + if (message.includes(String(code))) { return true; } } @@ -284,6 +318,18 @@ export class ReconnectingWebSocket } this.#isDisposed = true; + this.clearCurrentSocket(code, reason); + + for (const set of Object.values(this.#eventHandlers)) { + set.clear(); + } + + this.#onDispose?.(); + } + + private clearCurrentSocket(code?: number, reason?: string): void { + // Clear pending reconnect to prevent resume + this.#pendingReconnect = false; if (this.#reconnectTimeoutId !== null) { clearTimeout(this.#reconnectTimeoutId); @@ -294,11 +340,5 @@ export class ReconnectingWebSocket this.#currentSocket.close(code, reason); this.#currentSocket = null; } - - for (const set of Object.values(this.#eventHandlers)) { - set.clear(); - } - - this.#onDispose?.(); } } diff --git a/test/unit/api/coderApi.test.ts b/test/unit/api/coderApi.test.ts index 4f90f33e..7a7758c0 100644 --- a/test/unit/api/coderApi.test.ts +++ b/test/unit/api/coderApi.test.ts @@ -11,7 +11,6 @@ import { CertificateError } from "@/error"; import { getHeaders } from "@/headers"; import { type RequestConfigWithMeta } from "@/logging/types"; import { ReconnectingWebSocket } from "@/websocket/reconnectingWebSocket"; -import { SseConnection } from "@/websocket/sseConnection"; import { createMockLogger, @@ -336,18 +335,20 @@ describe("CoderApi", () => { expect(EventSource).not.toHaveBeenCalled(); }); - it("falls back to SSE when WebSocket creation fails", async () => { + it("falls back to SSE when WebSocket creation fails with 404", async () => { + // Only 404 errors trigger SSE fallback - other errors are thrown vi.mocked(Ws).mockImplementation(() => { - throw new Error("WebSocket creation failed"); + throw new Error("Unexpected server response: 404"); }); const connection = await api.watchAgentMetadata(AGENT_ID); - expect(connection).toBeInstanceOf(SseConnection); + // Returns ReconnectingWebSocket (which wraps SSE internally) + expect(connection).toBeInstanceOf(ReconnectingWebSocket); expect(EventSource).toHaveBeenCalled(); }); - it("falls back to SSE on 404 error from WebSocket", async () => { + it("falls back to SSE on 404 error from WebSocket open", async () => { const mockWs = createMockWebSocket( `wss://${CODER_URL.replace("https://", "")}/api/v2/test`, { @@ -368,9 +369,64 @@ describe("CoderApi", () => { const connection = await api.watchAgentMetadata(AGENT_ID); - expect(connection).toBeInstanceOf(SseConnection); + // Returns ReconnectingWebSocket (which wraps SSE internally after WS 404) + expect(connection).toBeInstanceOf(ReconnectingWebSocket); expect(EventSource).toHaveBeenCalled(); }); + + it("throws non-404 errors without SSE fallback", async () => { + vi.mocked(Ws).mockImplementation(() => { + throw new Error("Network error"); + }); + + await expect(api.watchAgentMetadata(AGENT_ID)).rejects.toThrow( + "Network error", + ); + expect(EventSource).not.toHaveBeenCalled(); + }); + + describe("reconnection after fallback", () => { + beforeEach(() => vi.useFakeTimers({ shouldAdvanceTime: true })); + afterEach(() => vi.useRealTimers()); + + it("reconnects after SSE fallback and retries WS on each reconnect", async () => { + let wsAttempts = 0; + const mockEventSources: MockEventSource[] = []; + + vi.mocked(Ws).mockImplementation(() => { + wsAttempts++; + const mockWs = createMockWebSocket("wss://test", { + on: vi.fn((event: string, handler: (e: unknown) => void) => { + if (event === "error") { + setImmediate(() => + handler({ error: new Error("Something 404") }), + ); + } + return mockWs as Ws; + }), + }); + return mockWs as Ws; + }); + + vi.mocked(EventSource).mockImplementation(() => { + const es = createMockEventSource(`${CODER_URL}/api/v2/test`); + mockEventSources.push(es); + return es as unknown as EventSource; + }); + + const connection = await api.watchAgentMetadata(AGENT_ID); + expect(wsAttempts).toBe(1); + expect(EventSource).toHaveBeenCalledTimes(1); + + mockEventSources[0].fireError(); + await vi.advanceTimersByTimeAsync(300); + + expect(wsAttempts).toBe(2); + expect(EventSource).toHaveBeenCalledTimes(2); + + connection.close(); + }); + }); }); describe("Reconnection on Host/Token Changes", () => { @@ -413,6 +469,7 @@ describe("CoderApi", () => { expect(wsWrap.url).toContain(CODER_URL.replace("http", "ws")); api.setHost("https://new-coder.example.com"); + // Wait for the async reconnect to complete (factory is async) await new Promise((resolve) => setImmediate(resolve)); expect(sockets[0].close).toHaveBeenCalledWith( @@ -420,7 +477,8 @@ describe("CoderApi", () => { "Replacing connection", ); expect(sockets).toHaveLength(2); - expect(wsWrap.url).toContain("wss://new-coder.example.com"); + // Verify the new socket was created with the correct URL + expect(sockets[1].url).toContain("wss://new-coder.example.com"); }); it("does not reconnect when token or host are unchanged", async () => { @@ -435,6 +493,58 @@ describe("CoderApi", () => { expect(sockets[0].close).not.toHaveBeenCalled(); expect(sockets).toHaveLength(1); }); + + it("suspends sockets when host is set to empty string (logout)", async () => { + const sockets = setupAutoOpeningWebSocket(); + api = createApi(CODER_URL, AXIOS_TOKEN); + await api.watchAgentMetadata(AGENT_ID); + + // Setting host to empty string (logout) should suspend (not permanently close) + api.setHost(""); + await new Promise((resolve) => setImmediate(resolve)); + + expect(sockets[0].close).toHaveBeenCalledWith(1000, "Host cleared"); + expect(sockets).toHaveLength(1); + }); + + it("does not reconnect when setting token after clearing host", async () => { + const sockets = setupAutoOpeningWebSocket(); + api = createApi(CODER_URL, AXIOS_TOKEN); + await api.watchAgentMetadata(AGENT_ID); + + api.setHost(""); + api.setSessionToken("new-token"); + await new Promise((resolve) => setImmediate(resolve)); + + // Should only have the initial socket - no reconnection after token change + expect(sockets).toHaveLength(1); + expect(sockets[0].close).toHaveBeenCalledWith(1000, "Host cleared"); + }); + + it("setCredentials sets both host and token together", async () => { + const sockets = setupAutoOpeningWebSocket(); + api = createApi(CODER_URL, AXIOS_TOKEN); + await api.watchAgentMetadata(AGENT_ID); + + api.setCredentials("https://new-coder.example.com", "new-token"); + await new Promise((resolve) => setImmediate(resolve)); + + // Should reconnect only once despite both values changing + expect(sockets).toHaveLength(2); + expect(sockets[1].url).toContain("wss://new-coder.example.com"); + }); + + it("setCredentials suspends when host is cleared", async () => { + const sockets = setupAutoOpeningWebSocket(); + api = createApi(CODER_URL, AXIOS_TOKEN); + await api.watchAgentMetadata(AGENT_ID); + + api.setCredentials(undefined, undefined); + await new Promise((resolve) => setImmediate(resolve)); + + expect(sockets).toHaveLength(1); + expect(sockets[0].close).toHaveBeenCalledWith(1000, "Host cleared"); + }); }); describe("Error Handling", () => { @@ -472,18 +582,32 @@ function createMockWebSocket( }; } -function createMockEventSource(url: string): Partial { - return { +type MockEventSource = Partial & { + readyState: number; + fireOpen: () => void; + fireError: () => void; +}; + +function createMockEventSource(url: string): MockEventSource { + const handlers: Record void) | undefined> = {}; + const mock: MockEventSource = { url, readyState: EventSource.CONNECTING, - addEventListener: vi.fn((event, handler) => { + addEventListener: vi.fn((event: string, handler: (e: Event) => void) => { + handlers[event] = handler; if (event === "open") { setImmediate(() => handler(new Event("open"))); } }), removeEventListener: vi.fn(), close: vi.fn(), + fireOpen: () => handlers.open?.(new Event("open")), + fireError: () => { + mock.readyState = EventSource.CLOSED; + handlers.error?.(new Event("error")); + }, }; + return mock; } function setupWebSocketMock(ws: Partial): void { diff --git a/test/unit/websocket/reconnectingWebSocket.test.ts b/test/unit/websocket/reconnectingWebSocket.test.ts index cdf08949..1434b6a6 100644 --- a/test/unit/websocket/reconnectingWebSocket.test.ts +++ b/test/unit/websocket/reconnectingWebSocket.test.ts @@ -104,6 +104,32 @@ describe("ReconnectingWebSocket", () => { }, ); + it.each([ + HttpStatusCode.UNAUTHORIZED, + HttpStatusCode.FORBIDDEN, + HttpStatusCode.GONE, + ])( + "does not reconnect on unrecoverable HTTP error via error event: %i", + async (statusCode) => { + // HTTP errors during handshake fire 'error' event, then 'close' with 1006 + const { ws, sockets } = await createReconnectingWebSocket(); + + sockets[0].fireError( + new Error(`Unexpected server response: ${statusCode}`), + ); + sockets[0].fireClose({ + code: WebSocketCloseCode.ABNORMAL, + reason: "Connection failed", + }); + + // Should not reconnect - unrecoverable HTTP error + await vi.advanceTimersByTimeAsync(10000); + expect(sockets).toHaveLength(1); + + ws.close(); + }, + ); + it("reconnect() connects immediately and cancels pending reconnections", async () => { const { ws, sockets } = await createReconnectingWebSocket(); @@ -125,26 +151,8 @@ describe("ReconnectingWebSocket", () => { }); it("queues reconnect() calls made during connection", async () => { - const sockets: MockSocket[] = []; - let pendingResolve: ((socket: MockSocket) => void) | null = null; - - const factory = vi.fn(() => { - const socket = createMockSocket(); - sockets.push(socket); - - // First call resolves immediately, other calls wait for manual resolve - if (sockets.length === 1) { - return Promise.resolve(socket); - } else { - return new Promise((resolve) => { - pendingResolve = resolve; - }); - } - }); - - const ws = await fromFactory(factory); - sockets[0].fireOpen(); - expect(sockets).toHaveLength(1); + const { ws, sockets, completeConnection } = + await createBlockingReconnectingWebSocket(); // Start first reconnect (will block on factory promise) ws.reconnect(); @@ -154,17 +162,33 @@ describe("ReconnectingWebSocket", () => { // Still only 2 sockets (queued reconnect hasn't started) expect(sockets).toHaveLength(2); - // Complete the first reconnect - pendingResolve!(sockets[1]); - sockets[1].fireOpen(); - - // Wait a tick for the queued reconnect to execute + completeConnection(); await Promise.resolve(); // Now queued reconnect should have executed, creating third socket expect(sockets).toHaveLength(3); ws.close(); }); + + it("suspend() cancels pending reconnect queued during connection", async () => { + const { ws, sockets, failConnection } = + await createBlockingReconnectingWebSocket(); + + ws.reconnect(); + ws.reconnect(); // queued + expect(sockets).toHaveLength(2); + + // This should cancel the queued request + ws.suspend(); + failConnection(new Error("No base URL")); + await Promise.resolve(); + + expect(sockets).toHaveLength(2); + await vi.advanceTimersByTimeAsync(10000); + expect(sockets).toHaveLength(2); + + ws.close(); + }); }); describe("Event Handlers", () => { @@ -216,6 +240,48 @@ describe("ReconnectingWebSocket", () => { ws.close(); }); + + it("preserves event handlers after suspend() and reconnect()", async () => { + const { ws, sockets } = await createReconnectingWebSocket(); + sockets[0].fireOpen(); + + const handler = vi.fn(); + ws.addEventListener("message", handler); + sockets[0].fireMessage({ test: 1 }); + expect(handler).toHaveBeenCalledTimes(1); + + // Suspend the socket + ws.suspend(); + + // Reconnect (async operation) + ws.reconnect(); + await Promise.resolve(); // Wait for async connect() + expect(sockets).toHaveLength(2); + sockets[1].fireOpen(); + + // Handler should still work after suspend/reconnect + sockets[1].fireMessage({ test: 2 }); + expect(handler).toHaveBeenCalledTimes(2); + + ws.close(); + }); + + it("clears event handlers after close()", async () => { + const { ws, sockets } = await createReconnectingWebSocket(); + sockets[0].fireOpen(); + + const handler = vi.fn(); + ws.addEventListener("message", handler); + sockets[0].fireMessage({ test: 1 }); + expect(handler).toHaveBeenCalledTimes(1); + + // Close permanently + ws.close(); + + // Even if we could reconnect (we can't), handlers would be cleared + // Verify handler was removed by checking it's no longer in the set + // We can't easily test this without exposing internals, but close() clears handlers + }); }); describe("close() and Disposal", () => { @@ -258,9 +324,9 @@ describe("ReconnectingWebSocket", () => { expect(disposeCount).toBe(1); }); - it("calls onDispose callback on unrecoverable WebSocket close code", async () => { + it("suspends (not disposes) on unrecoverable WebSocket close code", async () => { let disposeCount = 0; - const { sockets } = await createReconnectingWebSocket( + const { ws, sockets } = await createReconnectingWebSocket( () => ++disposeCount, ); @@ -270,7 +336,14 @@ describe("ReconnectingWebSocket", () => { reason: "Protocol error", }); - expect(disposeCount).toBe(1); + // Should suspend, not dispose - allows recovery when credentials change + expect(disposeCount).toBe(0); + + // Should be able to reconnect after suspension + ws.reconnect(); + expect(sockets).toHaveLength(2); + + ws.close(); }); it("does not call onDispose callback during reconnection", async () => { @@ -291,6 +364,41 @@ describe("ReconnectingWebSocket", () => { ws.close(); expect(disposeCount).toBe(1); }); + + it("reconnect() resumes suspended socket after HTTP 403 error", async () => { + const { ws, sockets, setFactoryError } = + await createReconnectingWebSocketWithErrorControl(); + sockets[0].fireOpen(); + + // Trigger reconnect that will fail with 403 + setFactoryError( + new Error(`Unexpected server response: ${HttpStatusCode.FORBIDDEN}`), + ); + ws.reconnect(); + await Promise.resolve(); + + // Socket should be suspended - no automatic reconnection + await vi.advanceTimersByTimeAsync(10000); + expect(sockets).toHaveLength(1); + + // reconnect() should resume the suspended socket + setFactoryError(null); + ws.reconnect(); + await Promise.resolve(); + expect(sockets).toHaveLength(2); + + ws.close(); + }); + + it("reconnect() does nothing after close()", async () => { + const { ws, sockets } = await createReconnectingWebSocket(); + sockets[0].fireOpen(); + + ws.close(); + ws.reconnect(); + + expect(sockets).toHaveLength(1); + }); }); describe("Backoff Strategy", () => { @@ -454,6 +562,35 @@ async function createReconnectingWebSocket(onDispose?: () => void): Promise<{ return { ws, sockets }; } +async function createReconnectingWebSocketWithErrorControl(): Promise<{ + ws: ReconnectingWebSocket; + sockets: MockSocket[]; + setFactoryError: (error: Error | null) => void; +}> { + const sockets: MockSocket[] = []; + let factoryError: Error | null = null; + + const factory = vi.fn(() => { + if (factoryError) { + return Promise.reject(factoryError); + } + const socket = createMockSocket(); + sockets.push(socket); + return Promise.resolve(socket); + }); + + const ws = await fromFactory(factory); + expect(sockets).toHaveLength(1); + + return { + ws, + sockets, + setFactoryError: (error: Error | null) => { + factoryError = error; + }, + }; +} + async function fromFactory( factory: SocketFactory, onDispose?: () => void, @@ -466,3 +603,40 @@ async function fromFactory( onDispose, ); } + +async function createBlockingReconnectingWebSocket(): Promise<{ + ws: ReconnectingWebSocket; + sockets: MockSocket[]; + completeConnection: () => void; + failConnection: (error: Error) => void; +}> { + const sockets: MockSocket[] = []; + let pendingResolve: ((socket: MockSocket) => void) | null = null; + let pendingReject: ((error: Error) => void) | null = null; + + const factory = vi.fn(() => { + const socket = createMockSocket(); + sockets.push(socket); + if (sockets.length === 1) { + return Promise.resolve(socket); + } + return new Promise((resolve, reject) => { + pendingResolve = resolve; + pendingReject = reject; + }); + }); + + const ws = await fromFactory(factory); + sockets[0].fireOpen(); + + return { + ws, + sockets, + completeConnection: () => { + const socket = sockets.at(-1)!; + pendingResolve?.(socket); + socket.fireOpen(); + }, + failConnection: (error: Error) => pendingReject?.(error), + }; +} From 1f0284ba58e73157bd1ee7c631cbf52e345cb3be Mon Sep 17 00:00:00 2001 From: Ehab Younes Date: Fri, 28 Nov 2025 13:16:21 +0300 Subject: [PATCH 18/20] Prune old login records + Add debug flag for deployments --- package.json | 6 +++ src/core/container.ts | 2 +- src/core/contextManager.ts | 11 +++-- src/core/secretsManager.ts | 70 ++++++++++++++++----------- src/extension.ts | 25 ++++++++++ test/unit/core/secretsManager.test.ts | 39 +++++++++++++++ 6 files changed, 120 insertions(+), 33 deletions(-) diff --git a/package.json b/package.json index bd60a54c..a3b585ab 100644 --- a/package.json +++ b/package.json @@ -255,6 +255,12 @@ "title": "Search", "category": "Coder", "icon": "$(search)" + }, + { + "command": "coder.debug.listDeployments", + "title": "List Stored Deployments", + "category": "Coder Debug", + "when": "coder.devMode" } ], "menus": { diff --git a/src/core/container.ts b/src/core/container.ts index 10cbf162..f140f628 100644 --- a/src/core/container.ts +++ b/src/core/container.ts @@ -41,7 +41,7 @@ export class ServiceContainer implements vscode.Disposable { this.logger, this.pathResolver, ); - this.contextManager = new ContextManager(); + this.contextManager = new ContextManager(context); this.loginCoordinator = new LoginCoordinator( this.secretsManager, this.mementoManager, diff --git a/src/core/contextManager.ts b/src/core/contextManager.ts index a5a18397..9a0f3d00 100644 --- a/src/core/contextManager.ts +++ b/src/core/contextManager.ts @@ -5,6 +5,7 @@ const CONTEXT_DEFAULTS = { "coder.isOwner": false, "coder.loaded": false, "coder.workspace.updatable": false, + "coder.devMode": false, } as const; type CoderContext = keyof typeof CONTEXT_DEFAULTS; @@ -12,10 +13,14 @@ type CoderContext = keyof typeof CONTEXT_DEFAULTS; export class ContextManager implements vscode.Disposable { private readonly context = new Map(); - public constructor() { - (Object.keys(CONTEXT_DEFAULTS) as CoderContext[]).forEach((key) => { + public constructor(extensionContext: vscode.ExtensionContext) { + for (const key of Object.keys(CONTEXT_DEFAULTS) as CoderContext[]) { this.set(key, CONTEXT_DEFAULTS[key]); - }); + } + this.set( + "coder.devMode", + extensionContext.extensionMode === vscode.ExtensionMode.Development, + ); } public set(key: CoderContext, value: boolean): void { diff --git a/src/core/secretsManager.ts b/src/core/secretsManager.ts index 6069439c..e1e9411b 100644 --- a/src/core/secretsManager.ts +++ b/src/core/secretsManager.ts @@ -13,10 +13,16 @@ const OAUTH_CLIENT_PREFIX = "coder.oauth.client."; const CURRENT_DEPLOYMENT_KEY = "coder.currentDeployment"; const OAUTH_CALLBACK_KEY = "coder.oauthCallback"; -const KNOWN_LABELS_KEY = "coder.knownLabels"; +const DEPLOYMENT_USAGE_KEY = "coder.deploymentUsage"; +const DEFAULT_MAX_DEPLOYMENTS = 10; const LEGACY_SESSION_TOKEN_KEY = "sessionToken"; +export interface DeploymentUsage { + label: string; + lastAccessedAt: string; +} + export type StoredOAuthTokens = Omit & { expiry_timestamp: number; deployment_url: string; @@ -179,7 +185,7 @@ export class SecretsManager { `${SESSION_KEY_PREFIX}${label}`, JSON.stringify(auth), ); - await this.addKnownLabel(label); + await this.recordDeploymentAccess(label); } public async clearSessionAuth(label: string): Promise { @@ -208,7 +214,7 @@ export class SecretsManager { `${OAUTH_TOKENS_PREFIX}${label}`, JSON.stringify(tokens), ); - await this.addKnownLabel(label); + await this.recordDeploymentAccess(label); } public async clearOAuthTokens(label: string): Promise { @@ -237,7 +243,7 @@ export class SecretsManager { `${OAUTH_CLIENT_PREFIX}${label}`, JSON.stringify(registration), ); - await this.addKnownLabel(label); + await this.recordDeploymentAccess(label); } public async clearOAuthClientRegistration(label: string): Promise { @@ -252,42 +258,48 @@ export class SecretsManager { } /** - * TODO currently it might be used wrong because we can be connected to a remote deployment - * and we log out from the sidebar causing the session to be removed and the auto-refresh disabled. - * - * Potential solutions: - * 1. Keep the last 10 auths and possibly remove entries not used in a while instead. - * We do not remove entries on logout! - * 2. Show the user a warning that their remote deployment might be disconnected. - * - * Update all usages of this after arriving at a decision! + * Record that a deployment was accessed, moving it to the front of the LRU list. + * Prunes deployments beyond maxCount, clearing their auth data. + */ + public async recordDeploymentAccess( + label: string, + maxCount = DEFAULT_MAX_DEPLOYMENTS, + ): Promise { + const usage = this.getDeploymentUsage(); + const filtered = usage.filter((u) => u.label !== label); + filtered.unshift({ label, lastAccessedAt: new Date().toISOString() }); + + const toKeep = filtered.slice(0, maxCount); + const toRemove = filtered.slice(maxCount); + + await Promise.all(toRemove.map((u) => this.clearAllAuthData(u.label))); + await this.memento.update(DEPLOYMENT_USAGE_KEY, toKeep); + } + + /** + * Clear all auth data for a deployment and remove it from the usage list. */ public async clearAllAuthData(label: string): Promise { await Promise.all([ this.clearSessionAuth(label), this.clearOAuthData(label), ]); - await this.removeKnownLabel(label); + const usage = this.getDeploymentUsage().filter((u) => u.label !== label); + await this.memento.update(DEPLOYMENT_USAGE_KEY, usage); } + /** + * Get all known deployment labels, ordered by most recently accessed. + */ public getKnownLabels(): string[] { - return this.memento.get(KNOWN_LABELS_KEY) ?? []; + return this.getDeploymentUsage().map((u) => u.label); } - private async addKnownLabel(label: string): Promise { - const labels = new Set(this.getKnownLabels()); - if (!labels.has(label)) { - labels.add(label); - await this.memento.update(KNOWN_LABELS_KEY, Array.from(labels)); - } - } - - private async removeKnownLabel(label: string): Promise { - const labels = new Set(this.getKnownLabels()); - if (labels.has(label)) { - labels.delete(label); - await this.memento.update(KNOWN_LABELS_KEY, Array.from(labels)); - } + /** + * Get the full deployment usage list with access timestamps. + */ + private getDeploymentUsage(): DeploymentUsage[] { + return this.memento.get(DEPLOYMENT_USAGE_KEY) ?? []; } /** diff --git a/src/extension.ts b/src/extension.ts index 77bf89c5..2e017b5f 100644 --- a/src/extension.ts +++ b/src/extension.ts @@ -312,6 +312,9 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { vscode.commands.registerCommand("coder.searchAllWorkspaces", async () => showTreeViewSearch(ALL_WORKSPACES_TREE_ID), ), + vscode.commands.registerCommand("coder.debug.listDeployments", () => + listStoredDeployments(secretsManager), + ), ); const remote = new Remote(serviceContainer, commands, ctx); @@ -554,3 +557,25 @@ async function getToken( } return ""; } + +async function listStoredDeployments( + secretsManager: SecretsManager, +): Promise { + const labels = secretsManager.getKnownLabels(); + if (labels.length === 0) { + vscode.window.showInformationMessage("No deployments stored."); + return; + } + + const selected = await vscode.window.showQuickPick( + labels.map((label) => ({ label, description: "Click to forget" })), + { placeHolder: "Select a deployment to forget" }, + ); + + if (selected) { + await secretsManager.clearAllAuthData(selected.label); + vscode.window.showInformationMessage( + `Cleared auth data for ${selected.label}`, + ); + } +} diff --git a/test/unit/core/secretsManager.test.ts b/test/unit/core/secretsManager.test.ts index 5aac3425..f4456fa5 100644 --- a/test/unit/core/secretsManager.test.ts +++ b/test/unit/core/secretsManager.test.ts @@ -97,6 +97,45 @@ describe("SecretsManager", () => { await secretsManager.clearAllAuthData("example.com"); expect(secretsManager.getKnownLabels()).not.toContain("example.com"); }); + + it("should order labels by most recently accessed", async () => { + await secretsManager.setSessionAuth("first.com", { + url: "https://first.com", + token: "token1", + }); + await secretsManager.setSessionAuth("second.com", { + url: "https://second.com", + token: "token2", + }); + await secretsManager.setSessionAuth("first.com", { + url: "https://first.com", + token: "token1-updated", + }); + + expect(secretsManager.getKnownLabels()).toEqual([ + "first.com", + "second.com", + ]); + }); + + it("should prune old deployments when exceeding maxCount", async () => { + for (let i = 1; i <= 5; i++) { + await secretsManager.setSessionAuth(`host${i}.com`, { + url: `https://host${i}.com`, + token: `token${i}`, + }); + } + + await secretsManager.recordDeploymentAccess("new.com", 3); + + expect(secretsManager.getKnownLabels()).toEqual([ + "new.com", + "host5.com", + "host4.com", + ]); + expect(await secretsManager.getSessionToken("host1.com")).toBeUndefined(); + expect(await secretsManager.getSessionToken("host2.com")).toBeUndefined(); + }); }); describe("current deployment", () => { From 04ec9e61068d4aa8f35a1e9e7c9ead7108c4329b Mon Sep 17 00:00:00 2001 From: Ehab Younes Date: Fri, 28 Nov 2025 15:50:25 +0300 Subject: [PATCH 19/20] More login synchronization issues --- src/extension.ts | 2 +- src/login/loginCoordinator.ts | 59 ++++++++++++++++++++--------------- src/oauth/sessionManager.ts | 15 ++++----- 3 files changed, 40 insertions(+), 36 deletions(-) diff --git a/src/extension.ts b/src/extension.ts index 2e017b5f..fdf423cc 100644 --- a/src/extension.ts +++ b/src/extension.ts @@ -336,7 +336,7 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { registerAuthListener(deployment?.label); // Update context - contextManager.set("coder.authenticated", deployment !== undefined); + contextManager.set("coder.authenticated", Boolean(deployment)); // Refresh workspaces myWorkspacesProvider.fetchAndRefresh(); diff --git a/src/login/loginCoordinator.ts b/src/login/loginCoordinator.ts index a2cd84cb..f170baeb 100644 --- a/src/login/loginCoordinator.ts +++ b/src/login/loginCoordinator.ts @@ -179,17 +179,28 @@ export class LoginCoordinator { isAutoLogin: boolean, oauthSessionManager: OAuthSessionManager, ): Promise { - const token = await this.secretsManager.getSessionToken(deployment.label); - const client = CoderApi.create(deployment.url, token, this.logger); const needsToken = needToken(vscode.workspace.getConfiguration()); - if (!needsToken || token) { - 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 { success: true, token: needsToken && token ? token : "", user }; - } catch (err) { + const client = CoderApi.create(deployment.url, "", this.logger); + + let storedToken: string | undefined; + if (needsToken) { + storedToken = await this.secretsManager.getSessionToken(deployment.label); + if (storedToken) { + client.setSessionToken(storedToken); + } + } + + // Attempt authentication with current credentials (token or mTLS) + try { + const user = await client.getAuthenticatedUser(); + // Return the token that was used (empty string for mTLS since + // the `vscodessh` command currently always requires a token file) + return { success: true, token: storedToken ?? "", user }; + } catch (err) { + if (needsToken) { + // For token auth: silently continue to prompt for new credentials + } else { + // For mTLS: show error and abort (no credentials to prompt for) const message = getErrorMessage(err, "no response from the server"); if (isAutoLogin) { this.logger.warn("Failed to log in to Coder server:", message); @@ -203,7 +214,6 @@ export class LoginCoordinator { }, ); } - // Invalid certificate, most likely. return { success: false }; } } @@ -212,12 +222,8 @@ export class LoginCoordinator { switch (authMethod) { case "oauth": return this.loginWithOAuth(client, oauthSessionManager, deployment); - case "legacy": { - const initialToken = - token || - (await this.secretsManager.getSessionToken(deployment.label)); - return this.loginWithToken(client, initialToken); - } + case "legacy": + return this.loginWithToken(client); case undefined: return { success: false }; // User aborted } @@ -226,10 +232,7 @@ export class LoginCoordinator { /** * Session token authentication flow. */ - private async loginWithToken( - client: CoderApi, - initialToken: string | undefined, - ): Promise { + private async loginWithToken(client: CoderApi): Promise { const url = client.getAxiosInstance().defaults.baseURL; if (!url) { throw new Error("No base URL set on REST client"); @@ -246,7 +249,6 @@ export class LoginCoordinator { title: "Coder API Key", password: true, placeHolder: "Paste your API key.", - value: initialToken, ignoreFocusOut: true, validateInput: async (value) => { if (!value) { @@ -315,10 +317,15 @@ export class LoginCoordinator { user, }; } catch (error) { - this.logger.error("OAuth authentication failed:", error); - vscode.window.showErrorMessage( - `OAuth authentication failed: ${getErrorMessage(error, "Unknown error")}`, - ); + const title = "OAuth authentication failed"; + this.logger.error(title, error); + if (error instanceof CertificateError) { + error.showNotification(title); + } else { + vscode.window.showErrorMessage( + `${title}: ${getErrorMessage(error, "Unknown error")}`, + ); + } return { success: false }; } } diff --git a/src/oauth/sessionManager.ts b/src/oauth/sessionManager.ts index 3e667c22..bd0a5f3a 100644 --- a/src/oauth/sessionManager.ts +++ b/src/oauth/sessionManager.ts @@ -132,7 +132,8 @@ export class OAuthSessionManager implements vscode.Disposable { stored: tokens.deployment_url, current: this.deployment.url, }); - await this.clearOAuthState(); + this.clearInMemoryTokens(); + await this.secretsManager.clearOAuthData(this.deployment.label); return; } @@ -144,6 +145,7 @@ export class OAuthSessionManager implements vscode.Disposable { required_scopes: DEFAULT_OAUTH_SCOPES, }, ); + this.clearInMemoryTokens(); await this.secretsManager.clearOAuthTokens(this.deployment.label); return; } @@ -152,13 +154,6 @@ export class OAuthSessionManager implements vscode.Disposable { this.logger.info(`Loaded stored OAuth tokens for ${this.deployment.label}`); } - private async clearOAuthState(): Promise { - this.clearInMemoryTokens(); - if (this.deployment) { - await this.secretsManager.clearOAuthData(this.deployment.label); - } - } - private clearInMemoryTokens(): void { this.storedTokens = undefined; this.refreshPromise = null; @@ -714,13 +709,15 @@ export class OAuthSessionManager implements vscode.Disposable { // Revoke refresh token (which also invalidates access token per RFC 7009) if (this.storedTokens?.refresh_token) { try { + // TODO what if other windows are using this? + // We should only revoke if we are clearing the OAuth data await this.revokeToken(this.storedTokens.refresh_token); } catch (error) { this.logger.warn("Token revocation failed during logout:", error); } } - await this.clearOAuthState(); + this.clearInMemoryTokens(); this.deployment = undefined; this.logger.info("OAuth logout complete"); From 40f1d6380064e4caf9214761d33ae8c698328da5 Mon Sep 17 00:00:00 2001 From: Ehab Younes Date: Wed, 3 Dec 2025 13:02:39 +0300 Subject: [PATCH 20/20] Improve login synchronization + Add background oauth refresh --- src/api/oauthInterceptors.ts | 14 ++--- src/commands.ts | 74 +++++++++++----------- src/extension.ts | 119 +++++++++++++++++++++-------------- src/oauth/sessionManager.ts | 46 +++++++++++++- src/remote/remote.ts | 2 +- 5 files changed, 155 insertions(+), 100 deletions(-) diff --git a/src/api/oauthInterceptors.ts b/src/api/oauthInterceptors.ts index adfe0efe..b80e1d96 100644 --- a/src/api/oauthInterceptors.ts +++ b/src/api/oauthInterceptors.ts @@ -24,16 +24,10 @@ export function attachOAuthInterceptors( client.getAxiosInstance().interceptors.response.use( // Success response interceptor: proactive token refresh (response) => { - if (oauthSessionManager.shouldRefreshToken()) { - logger.debug( - "Token approaching expiry, triggering proactive refresh in background", - ); - - // Fire-and-forget: don't await, don't block response - oauthSessionManager.refreshToken().catch((error) => { - logger.warn("Background token refresh failed:", error); - }); - } + // Fire-and-forget: don't await, don't block response + oauthSessionManager.refreshIfAlmostExpired().catch((error) => { + logger.warn("Proactive background token refresh failed:", error); + }); return response; }, diff --git a/src/commands.ts b/src/commands.ts index 62f323ea..00cb2ee0 100644 --- a/src/commands.ts +++ b/src/commands.ts @@ -9,7 +9,6 @@ import { type CoderApi } from "./api/coderApi"; import { type CliManager } from "./core/cliManager"; import { type ServiceContainer } from "./core/container"; import { type ContextManager } from "./core/contextManager"; -import { type Deployment } from "./core/deployment"; import { type MementoManager } from "./core/mementoManager"; import { type PathResolver } from "./core/pathResolver"; import { type SecretsManager } from "./core/secretsManager"; @@ -45,11 +44,11 @@ export class Commands { // if you use multiple deployments). public workspace?: Workspace; public workspaceLogPath?: string; - public workspaceRestClient?: CoderApi; + public remoteWorkspaceClient?: CoderApi; public constructor( serviceContainer: ServiceContainer, - private readonly restClient: CoderApi, + private readonly extensionClient: CoderApi, private readonly oauthSessionManager: OAuthSessionManager, ) { this.vscodeProposed = serviceContainer.getVsCodeProposed(); @@ -65,12 +64,12 @@ export class Commands { /** * Get the current deployment, throwing if not logged in. */ - private async requireDeployment(): Promise { - const deployment = await this.secretsManager.getCurrentDeployment(); - if (!deployment) { + private requireExtensionBaseUrl(): string { + const url = this.extensionClient.getAxiosInstance().defaults.baseURL; + if (!url) { throw new Error("You are not logged in"); } - return deployment; + return url; } /** @@ -117,9 +116,9 @@ export class Commands { // Set client immediately so subsequent operations in this function have the correct host/token. // The cross-window listener will also update the client, but that's async. - this.restClient.setCredentials(url, result.token); + this.extensionClient.setCredentials(url, result.token); - // Set as current deployment (triggers cross-window sync). + // Set as current deployment await this.secretsManager.setCurrentDeployment({ url, label }); // Update contexts @@ -173,14 +172,12 @@ export class Commands { * Log out from the currently logged-in deployment. */ public async logout(): Promise { - const deployment = await this.requireDeployment(); - await this.forceLogout(deployment.label); - } - - public async forceLogout(label: string): Promise { + const baseUrl = this.requireExtensionBaseUrl(); if (!this.contextManager.get("coder.authenticated")) { return; } + + const label = toSafeHost(baseUrl); this.logger.info(`Logging out of deployment: ${label}`); // Fire and forget OAuth logout @@ -190,7 +187,7 @@ export class Commands { // Clear from the REST client. An empty url will indicate to other parts of // the code that we are logged out. - this.restClient.setCredentials(undefined, undefined); + this.extensionClient.setCredentials(undefined, undefined); // Clear current deployment (triggers cross-window sync) await this.secretsManager.setCurrentDeployment(undefined); @@ -211,8 +208,8 @@ export class Commands { * Must only be called if currently logged in. */ public async createWorkspace(): Promise { - const deployment = await this.requireDeployment(); - const uri = deployment.url + "/templates"; + const baseUrl = this.requireExtensionBaseUrl(); + const uri = baseUrl + "/templates"; await vscode.commands.executeCommand("vscode.open", uri); } @@ -226,13 +223,13 @@ export class Commands { */ public async navigateToWorkspace(item: OpenableTreeItem) { if (item) { - const deployment = await this.requireDeployment(); + const baseUrl = this.requireExtensionBaseUrl(); const workspaceId = createWorkspaceIdentifier(item.workspace); - const uri = deployment.url + `/@${workspaceId}`; + const uri = baseUrl + `/@${workspaceId}`; await vscode.commands.executeCommand("vscode.open", uri); - } else if (this.workspace && this.workspaceRestClient) { + } else if (this.workspace && this.remoteWorkspaceClient) { const baseUrl = - this.workspaceRestClient.getAxiosInstance().defaults.baseURL; + this.remoteWorkspaceClient.getAxiosInstance().defaults.baseURL; const uri = `${baseUrl}/@${createWorkspaceIdentifier(this.workspace)}`; await vscode.commands.executeCommand("vscode.open", uri); } else { @@ -250,13 +247,13 @@ export class Commands { */ public async navigateToWorkspaceSettings(item: OpenableTreeItem) { if (item) { - const deployment = await this.requireDeployment(); + const baseUrl = this.requireExtensionBaseUrl(); const workspaceId = createWorkspaceIdentifier(item.workspace); - const uri = deployment.url + `/@${workspaceId}/settings`; + const uri = baseUrl + `/@${workspaceId}/settings`; await vscode.commands.executeCommand("vscode.open", uri); - } else if (this.workspace && this.workspaceRestClient) { + } else if (this.workspace && this.remoteWorkspaceClient) { const baseUrl = - this.workspaceRestClient.getAxiosInstance().defaults.baseURL; + this.remoteWorkspaceClient.getAxiosInstance().defaults.baseURL; const uri = `${baseUrl}/@${createWorkspaceIdentifier(this.workspace)}/settings`; await vscode.commands.executeCommand("vscode.open", uri); } else { @@ -274,7 +271,7 @@ export class Commands { */ public async openFromSidebar(item: OpenableTreeItem) { if (item) { - const baseUrl = this.restClient.getAxiosInstance().defaults.baseURL; + const baseUrl = this.extensionClient.getAxiosInstance().defaults.baseURL; if (!baseUrl) { throw new Error("You are not logged in"); } @@ -329,15 +326,14 @@ export class Commands { const terminal = vscode.window.createTerminal(app.name); // If workspace_name is provided, run coder ssh before the command - const deployment = await this.requireDeployment(); + const baseUrl = this.requireExtensionBaseUrl(); + const label = toSafeHost(baseUrl); const binary = await this.cliManager.fetchBinary( - this.restClient, - deployment.label, + this.extensionClient, + label, ); - const configDir = this.pathResolver.getGlobalConfigDir( - deployment.label, - ); + const configDir = this.pathResolver.getGlobalConfigDir(label); const globalFlags = getGlobalFlags( vscode.workspace.getConfiguration(), configDir, @@ -374,14 +370,14 @@ export class Commands { folderPath?: string, openRecent?: boolean, ): Promise { - const baseUrl = this.restClient.getAxiosInstance().defaults.baseURL; + const baseUrl = this.extensionClient.getAxiosInstance().defaults.baseURL; if (!baseUrl) { throw new Error("You are not logged in"); } let workspace: Workspace | undefined; if (workspaceOwner && workspaceName) { - workspace = await this.restClient.getWorkspaceByOwnerAndName( + workspace = await this.extensionClient.getWorkspaceByOwnerAndName( workspaceOwner, workspaceName, ); @@ -417,7 +413,7 @@ export class Commands { localWorkspaceFolder: string = "", localConfigFile: string = "", ): Promise { - const baseUrl = this.restClient.getAxiosInstance().defaults.baseURL; + const baseUrl = this.extensionClient.getAxiosInstance().defaults.baseURL; if (!baseUrl) { throw new Error("You are not logged in"); } @@ -473,7 +469,7 @@ export class Commands { * this is a no-op. */ public async updateWorkspace(): Promise { - if (!this.workspace || !this.workspaceRestClient) { + if (!this.workspace || !this.remoteWorkspaceClient) { return; } const action = await this.vscodeProposed.window.showWarningMessage( @@ -486,7 +482,7 @@ export class Commands { "Update", ); if (action === "Update") { - await this.workspaceRestClient.updateWorkspaceVersion(this.workspace); + await this.remoteWorkspaceClient.updateWorkspaceVersion(this.workspace); } } @@ -501,7 +497,7 @@ export class Commands { let lastWorkspaces: readonly Workspace[]; quickPick.onDidChangeValue((value) => { quickPick.busy = true; - this.restClient + this.extensionClient .getWorkspaces({ q: value, }) @@ -564,7 +560,7 @@ export class Commands { // we need to fetch the agents through the resources API, as the // workspaces query does not include agents when off. this.logger.info("Fetching agents from template version"); - const resources = await this.restClient.getTemplateVersionResources( + const resources = await this.extensionClient.getTemplateVersionResources( workspace.latest_build.template_version_id, ); return extractAgents(resources); diff --git a/src/extension.ts b/src/extension.ts index fdf423cc..093fc5e7 100644 --- a/src/extension.ts +++ b/src/extension.ts @@ -151,6 +151,7 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { client.setCredentials(auth?.url, auth?.token); // Update authentication context for current deployment + // TODO(ehab) this might never even happen :thinking: contextManager.set("coder.authenticated", auth !== undefined); }, ); @@ -160,6 +161,55 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { registerAuthListener(deployment?.label); ctx.subscriptions.push({ dispose: () => authChangeDisposable?.dispose() }); + const changeDeployment = async ( + deployment: Deployment | null, + sessionToken?: string, + ) => { + // Update client + if (deployment) { + const token = + sessionToken || + (await secretsManager.getSessionToken(deployment.label)); + client.setCredentials(deployment.url, token); + await oauthSessionManager.setDeployment(deployment); + } else { + client.setCredentials(undefined, undefined); + oauthSessionManager.clearDeployment(); + } + registerAuthListener(deployment?.label); + + // Update context + contextManager.set("coder.authenticated", Boolean(deployment)); + + // Refresh workspaces + myWorkspacesProvider.fetchAndRefresh(); + allWorkspacesProvider.fetchAndRefresh(); + }; + + const changeDeploymentAndPersist = async ( + deployment: Deployment | null, + sessionToken?: string, + ) => { + await changeDeployment(deployment, sessionToken); + // Persist and sync deployment across windows + await secretsManager.setCurrentDeployment(deployment ?? undefined); + await mementoManager.addToUrlHistory(deployment?.url ?? ""); + }; + + // Listen for deployment changes from other windows (cross-window sync) + ctx.subscriptions.push( + secretsManager.onDidChangeCurrentDeployment(async ({ deployment }) => { + const isLoggedIn = contextManager.get("coder.authenticated"); + if (isLoggedIn) { + // We keep whatever deployment we have if we're logged in + return; + } + + output.info("Deployment changed from another window"); + return changeDeployment(deployment); + }), + ); + // Handle vscode:// URIs. const uriHandler = vscode.window.registerUriHandler({ handleUri: async (uri) => { @@ -189,7 +239,11 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { throw new Error("workspace must be specified as a query parameter"); } - await setupDeploymentFromUri(params, client, serviceContainer); + await setupDeploymentFromUri( + params, + serviceContainer, + changeDeploymentAndPersist, + ); vscode.commands.executeCommand( "coder.open", @@ -238,7 +292,11 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { ); } - await setupDeploymentFromUri(params, client, serviceContainer); + await setupDeploymentFromUri( + params, + serviceContainer, + changeDeploymentAndPersist, + ); vscode.commands.executeCommand( "coder.openDevContainer", @@ -319,31 +377,6 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { const remote = new Remote(serviceContainer, commands, ctx); - // Listen for deployment changes from other windows (cross-window sync) - ctx.subscriptions.push( - secretsManager.onDidChangeCurrentDeployment(async ({ deployment }) => { - output.info("Deployment changed from another window"); - - // Update client - if (deployment) { - const token = await secretsManager.getSessionToken(deployment.label); - client.setCredentials(deployment.url, token); - await oauthSessionManager.setDeployment(deployment); - } else { - client.setCredentials(undefined, undefined); - oauthSessionManager.clearDeployment(); - } - registerAuthListener(deployment?.label); - - // Update context - contextManager.set("coder.authenticated", Boolean(deployment)); - - // Refresh workspaces - myWorkspacesProvider.fetchAndRefresh(); - allWorkspacesProvider.fetchAndRefresh(); - }), - ); - // 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. @@ -363,15 +396,10 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { if (details) { ctx.subscriptions.push(details); - // Set client host/token immediately for subsequent operations. - client.setCredentials(details.url, details.token); - - // Persist and sync deployment across windows - await secretsManager.setCurrentDeployment({ - url: details.url, - label: details.label, - }); - await mementoManager.addToUrlHistory(details.url); + await changeDeploymentAndPersist( + { label: details.label, url: details.url }, + details.token, + ); } } catch (ex) { if (ex instanceof CertificateError) { @@ -494,15 +522,17 @@ async function showTreeViewSearch(id: string): Promise { * Sets up deployment from URI parameters. Handles URL prompting, client setup, * and token storage. Throws if user cancels URL input. * - * Sets client host/token immediately for subsequent operations. - * Other updates (auth listener, OAuth manager, context, workspaces) are handled - * asynchronously by the onDidChangeCurrentDeployment listener. + * Updates the client host/token, auth listener, OAuth manager, context, etc. + * through the `changeDeploymentAndPersist` callback. */ async function setupDeploymentFromUri( params: URLSearchParams, - client: CoderApi, serviceContainer: ServiceContainer, -): Promise { + changeDeploymentAndPersist: ( + deployment: Deployment | null, + sessionToken?: string, + ) => Promise, +): Promise { const secretsManager = serviceContainer.getSecretsManager(); const mementoManager = serviceContainer.getMementoManager(); const currentDeployment = await secretsManager.getCurrentDeployment(); @@ -529,16 +559,11 @@ async function setupDeploymentFromUri( // command currently always requires a token file. However, if there is // a query parameter for non-token auth go ahead and use it anyway; const token = await getToken(params, label, secretsManager); - client.setCredentials(url, token); if (token) { await secretsManager.setSessionAuth(label, { url, token }); } - // Persist and sync deployment across windows - await secretsManager.setCurrentDeployment({ url, label }); - await mementoManager.addToUrlHistory(url); - - return { url, label }; + await changeDeploymentAndPersist({ label, url }, token); } async function getToken( diff --git a/src/oauth/sessionManager.ts b/src/oauth/sessionManager.ts index bd0a5f3a..482065af 100644 --- a/src/oauth/sessionManager.ts +++ b/src/oauth/sessionManager.ts @@ -34,7 +34,7 @@ const RESPONSE_TYPE = "code" as const; const PKCE_CHALLENGE_METHOD = "S256" as const; /** - * Token refresh threshold: refresh when token expires in less than this time + * Token refresh threshold: refresh when token expires in less than this time. */ const TOKEN_REFRESH_THRESHOLD_MS = 10 * 60 * 1000; @@ -44,10 +44,15 @@ const TOKEN_REFRESH_THRESHOLD_MS = 10 * 60 * 1000; const ACCESS_TOKEN_DEFAULT_EXPIRY_MS = 60 * 60 * 1000; /** - * Minimum time between refresh attempts to prevent thrashing + * Minimum time between refresh attempts to prevent thrashing. */ const REFRESH_THROTTLE_MS = 30 * 1000; +/** + * Background token refresh check interval. + */ +const BACKGROUND_REFRESH_INTERVAL_MS = 5 * 60 * 1000; + /** * Minimal scopes required by the VS Code extension. */ @@ -69,6 +74,7 @@ export class OAuthSessionManager implements vscode.Disposable { private storedTokens: StoredOAuthTokens | undefined; private refreshPromise: Promise | null = null; private lastRefreshAttempt = 0; + private refreshTimer: NodeJS.Timeout | undefined; private pendingAuthReject: ((reason: Error) => void) | undefined; @@ -88,6 +94,7 @@ export class OAuthSessionManager implements vscode.Disposable { extensionId, ); await manager.loadTokens(); + manager.scheduleBackgroundRefresh(); return manager; } @@ -160,6 +167,25 @@ export class OAuthSessionManager implements vscode.Disposable { this.lastRefreshAttempt = 0; } + /** + * Schedule the next background token refresh check. + * Only schedules the next check after the current one completes. + */ + private scheduleBackgroundRefresh(): void { + if (this.refreshTimer) { + clearTimeout(this.refreshTimer); + } + + this.refreshTimer = setTimeout(async () => { + try { + await this.refreshIfAlmostExpired(); + } catch (error) { + this.logger.warn("Background token refresh failed:", error); + } + this.scheduleBackgroundRefresh(); + }, BACKGROUND_REFRESH_INTERVAL_MS); + } + /** * Check if granted scopes cover all required scopes. * Supports wildcard scopes like "workspace:*". @@ -635,6 +661,16 @@ export class OAuthSessionManager implements vscode.Disposable { }); } + /** + * Refreshes the token if it is approaching expiry. + */ + public async refreshIfAlmostExpired(): Promise { + if (this.shouldRefreshToken()) { + this.logger.debug("Token approaching expiry, triggering refresh"); + await this.refreshToken(); + } + } + /** * Check if token should be refreshed. * Returns true if: @@ -642,7 +678,7 @@ export class OAuthSessionManager implements vscode.Disposable { * 2. Last refresh attempt was more than REFRESH_THROTTLE_MS ago * 3. No refresh is currently in progress */ - public shouldRefreshToken(): boolean { + private shouldRefreshToken(): boolean { if ( !this.isLoggedInWithOAuth() || !this.storedTokens?.refresh_token || @@ -757,6 +793,10 @@ export class OAuthSessionManager implements vscode.Disposable { * Clears all in-memory state. */ public dispose(): void { + if (this.refreshTimer) { + clearTimeout(this.refreshTimer); + this.refreshTimer = undefined; + } if (this.pendingAuthReject) { this.pendingAuthReject(new Error("OAuth session manager disposed")); } diff --git a/src/remote/remote.ts b/src/remote/remote.ts index e8c11060..f992c18e 100644 --- a/src/remote/remote.ts +++ b/src/remote/remote.ts @@ -181,7 +181,7 @@ export class Remote { disposables.push(workspaceClient); attachOAuthInterceptors(workspaceClient, this.logger, remoteOAuthManager); // Store for use in commands. - this.commands.workspaceRestClient = workspaceClient; + this.commands.remoteWorkspaceClient = workspaceClient; // Listen for token changes for this deployment disposables.push(