From 49b187489c0599113a9ffdb90a5967c7e4e7aeb8 Mon Sep 17 00:00:00 2001 From: Ehab Younes Date: Tue, 9 Dec 2025 13:20:31 +0300 Subject: [PATCH 1/9] Add OAuth 2.1 authentication support Implements OAuth 2.1 with PKCE as an alternative authentication method to session tokens. When connecting to a Coder deployment that supports OAuth, users can choose between OAuth and legacy token authentication. Key changes: OAuth Flow: - Add OAuthSessionManager to handle the complete OAuth lifecycle: dynamic client registration, PKCE authorization flow, token exchange, automatic refresh, and revocation - Add OAuthMetadataClient to discover and validate OAuth server metadata from the well-known endpoint, ensuring server meets OAuth 2.1 requirements - Handle OAuth callbacks via vscode:// URI handler with cross-window support for when callback arrives in a different VS Code window Token Management: - Store OAuth tokens (access, refresh, expiry) per-deployment in secrets - Store dynamic client registrations per-deployment in secrets - Proactive token refresh when approaching expiry (via response interceptor) - Reactive token refresh on 401 responses with automatic request retry - Handle OAuth errors (invalid_grant, invalid_client) by prompting for re-authentication Integration: - Add auth method selection prompt when server supports OAuth - Attach OAuth interceptors to CoderApi for automatic token refresh - Clear OAuth state when user explicitly chooses token auth - DeploymentManager coordinates OAuth session state with deployment changes Error Handling: - Typed OAuth error classes (InvalidGrantError, InvalidClientError, etc.) - Parse OAuth error responses from token endpoint - Show re-authentication modal for errors requiring user action --- src/api/oauthInterceptors.ts | 116 ++++ src/commands.ts | 3 + src/core/secretsManager.ts | 125 +++++ src/deployment/deploymentManager.ts | 37 +- src/extension.ts | 52 +- src/login/loginCoordinator.ts | 77 ++- src/oauth/errors.ts | 166 ++++++ src/oauth/metadataClient.ts | 137 +++++ src/oauth/sessionManager.ts | 801 ++++++++++++++++++++++++++++ src/oauth/types.ts | 163 ++++++ src/oauth/utils.ts | 42 ++ src/promptUtils.ts | 55 ++ src/remote/remote.ts | 14 +- 13 files changed, 1764 insertions(+), 24 deletions(-) create mode 100644 src/api/oauthInterceptors.ts create mode 100644 src/oauth/errors.ts create mode 100644 src/oauth/metadataClient.ts create mode 100644 src/oauth/sessionManager.ts create mode 100644 src/oauth/types.ts create mode 100644 src/oauth/utils.ts diff --git a/src/api/oauthInterceptors.ts b/src/api/oauthInterceptors.ts new file mode 100644 index 00000000..b80e1d96 --- /dev/null +++ b/src/api/oauthInterceptors.ts @@ -0,0 +1,116 @@ +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) => { + // Fire-and-forget: don't await, don't block response + oauthSessionManager.refreshIfAlmostExpired().catch((error) => { + logger.warn("Proactive 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 ef97bdda..e5c9e4fb 100644 --- a/src/commands.ts +++ b/src/commands.ts @@ -19,6 +19,7 @@ import { type DeploymentManager } from "./deployment/deploymentManager"; import { CertificateError } from "./error"; import { type Logger } from "./logging/logger"; import { type LoginCoordinator } from "./login/loginCoordinator"; +import { type OAuthSessionManager } from "./oauth/sessionManager"; import { maybeAskAgent, maybeAskUrl } from "./promptUtils"; import { escapeCommandArg, toRemoteAuthority, toSafeHost } from "./util"; import { @@ -51,6 +52,7 @@ export class Commands { public constructor( serviceContainer: ServiceContainer, private readonly extensionClient: CoderApi, + private readonly oauthSessionManager: OAuthSessionManager, private readonly deploymentManager: DeploymentManager, ) { this.vscodeProposed = serviceContainer.getVsCodeProposed(); @@ -105,6 +107,7 @@ export class Commands { safeHostname, url, autoLogin: args?.autoLogin, + oauthSessionManager: this.oauthSessionManager, }); if (!result.success) { diff --git a/src/core/secretsManager.ts b/src/core/secretsManager.ts index e6558299..128a826b 100644 --- a/src/core/secretsManager.ts +++ b/src/core/secretsManager.ts @@ -1,4 +1,8 @@ import { type Logger } from "../logging/logger"; +import { + type ClientRegistrationResponse, + type TokenResponse, +} from "../oauth/types"; import { toSafeHost } from "../util"; import type { Memento, SecretStorage, Disposable } from "vscode"; @@ -8,8 +12,11 @@ import type { Deployment } from "../deployment/types"; // Each deployment has its own key to ensure atomic operations (multiple windows // writing to a shared key could drop data) and to receive proper VS Code events. const SESSION_KEY_PREFIX = "coder.session."; +const OAUTH_TOKENS_PREFIX = "coder.oauth.tokens."; +const OAUTH_CLIENT_PREFIX = "coder.oauth.client."; const CURRENT_DEPLOYMENT_KEY = "coder.currentDeployment"; +const OAUTH_CALLBACK_KEY = "coder.oauthCallback"; const DEPLOYMENT_USAGE_KEY = "coder.deploymentUsage"; const DEFAULT_MAX_DEPLOYMENTS = 10; @@ -31,6 +38,17 @@ interface DeploymentUsage { lastAccessedAt: string; } +export type StoredOAuthTokens = Omit & { + expiry_timestamp: number; + deployment_url: string; +}; + +interface OAuthCallbackData { + state: string; + code: string | null; + error: string | null; +} + export class SecretsManager { constructor( private readonly secrets: SecretStorage, @@ -97,6 +115,38 @@ export class SecretsManager { }); } + /** + * 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 + } + }); + } + /** * Listen for changes to a specific deployment's session auth. */ @@ -153,6 +203,77 @@ export class SecretsManager { return `${SESSION_KEY_PREFIX}${safeHostname || ""}`; } + public async getOAuthTokens( + safeHostname: string, + ): Promise { + try { + const data = await this.secrets.get( + `${OAUTH_TOKENS_PREFIX}${safeHostname}`, + ); + if (!data) { + return undefined; + } + return JSON.parse(data) as StoredOAuthTokens; + } catch { + return undefined; + } + } + + public async setOAuthTokens( + safeHostname: string, + tokens: StoredOAuthTokens, + ): Promise { + await this.secrets.store( + `${OAUTH_TOKENS_PREFIX}${safeHostname}`, + JSON.stringify(tokens), + ); + await this.recordDeploymentAccess(safeHostname); + } + + public async clearOAuthTokens(safeHostname: string): Promise { + await this.secrets.delete(`${OAUTH_TOKENS_PREFIX}${safeHostname}`); + } + + public async getOAuthClientRegistration( + safeHostname: string, + ): Promise { + try { + const data = await this.secrets.get( + `${OAUTH_CLIENT_PREFIX}${safeHostname}`, + ); + if (!data) { + return undefined; + } + return JSON.parse(data) as ClientRegistrationResponse; + } catch { + return undefined; + } + } + + public async setOAuthClientRegistration( + safeHostname: string, + registration: ClientRegistrationResponse, + ): Promise { + await this.secrets.store( + `${OAUTH_CLIENT_PREFIX}${safeHostname}`, + JSON.stringify(registration), + ); + await this.recordDeploymentAccess(safeHostname); + } + + public async clearOAuthClientRegistration( + safeHostname: string, + ): Promise { + await this.secrets.delete(`${OAUTH_CLIENT_PREFIX}${safeHostname}`); + } + + public async clearOAuthData(safeHostname: string): Promise { + await Promise.all([ + this.clearOAuthTokens(safeHostname), + this.clearOAuthClientRegistration(safeHostname), + ]); + } + /** * Record that a deployment was accessed, moving it to the front of the LRU list. * Prunes deployments beyond maxCount, clearing their auth data. @@ -181,6 +302,10 @@ export class SecretsManager { * Clear all auth data for a deployment and remove it from the usage list. */ public async clearAllAuthData(safeHostname: string): Promise { + await Promise.all([ + this.clearSessionAuth(safeHostname), + this.clearOAuthData(safeHostname), + ]); await this.clearSessionAuth(safeHostname); const usage = this.getDeploymentUsage().filter( (u) => u.safeHostname !== safeHostname, diff --git a/src/deployment/deploymentManager.ts b/src/deployment/deploymentManager.ts index 850d2176..6d524f8d 100644 --- a/src/deployment/deploymentManager.ts +++ b/src/deployment/deploymentManager.ts @@ -1,17 +1,17 @@ import { CoderApi } from "../api/coderApi"; +import { type ServiceContainer } from "../core/container"; +import { type ContextManager } from "../core/contextManager"; +import { type MementoManager } from "../core/mementoManager"; +import { type SecretsManager } from "../core/secretsManager"; +import { type Logger } from "../logging/logger"; +import { type OAuthSessionManager } from "../oauth/sessionManager"; +import { type WorkspaceProvider } from "../workspace/workspacesProvider"; + +import { type Deployment, type DeploymentWithAuth } from "./types"; import type { User } from "coder/site/src/api/typesGenerated"; import type * as vscode from "vscode"; -import type { ServiceContainer } from "../core/container"; -import type { ContextManager } from "../core/contextManager"; -import type { MementoManager } from "../core/mementoManager"; -import type { SecretsManager } from "../core/secretsManager"; -import type { Logger } from "../logging/logger"; -import type { WorkspaceProvider } from "../workspace/workspacesProvider"; - -import type { Deployment, DeploymentWithAuth } from "./types"; - /** * Internal state type that allows mutation of user property. */ @@ -23,6 +23,7 @@ type DeploymentWithUser = Deployment & { user: User }; * Centralizes: * - In-memory deployment state (url, label, token, user) * - Client credential updates + * - OAuth session management * - Auth listener registration * - Context updates (coder.authenticated, coder.isOwner) * - Workspace provider refresh @@ -41,6 +42,7 @@ export class DeploymentManager implements vscode.Disposable { private constructor( serviceContainer: ServiceContainer, private readonly client: CoderApi, + private readonly oauthSessionManager: OAuthSessionManager, private readonly workspaceProviders: WorkspaceProvider[], ) { this.secretsManager = serviceContainer.getSecretsManager(); @@ -52,11 +54,13 @@ export class DeploymentManager implements vscode.Disposable { public static create( serviceContainer: ServiceContainer, client: CoderApi, + oauthSessionManager: OAuthSessionManager, workspaceProviders: WorkspaceProvider[], ): DeploymentManager { const manager = new DeploymentManager( serviceContainer, client, + oauthSessionManager, workspaceProviders, ); manager.subscribeToCrossWindowChanges(); @@ -85,6 +89,12 @@ export class DeploymentManager implements vscode.Disposable { public async setDeploymentIfValid( deployment: Deployment & { token?: string }, ): Promise { + // TODO used to trigger + /** + * this.oauthSessionManager.refreshIfAlmostExpired().catch((error) => { + this.logger.warn("Setup token refresh failed:", error); + }); + */ const auth = await this.secretsManager.getSessionAuth( deployment.safeHostname, ); @@ -124,6 +134,7 @@ export class DeploymentManager implements vscode.Disposable { } else { this.client.setCredentials(deployment.url, deployment.token); } + await this.oauthSessionManager.setDeployment(deployment); this.registerAuthListener(); this.updateAuthContexts(); @@ -140,12 +151,20 @@ export class DeploymentManager implements vscode.Disposable { this.#deployment = null; this.client.setCredentials(undefined, undefined); + this.oauthSessionManager.clearDeployment(); this.updateAuthContexts(); this.refreshWorkspaces(); await this.secretsManager.setCurrentDeployment(undefined); } + /** + * Clear OAuth state for a deployment when switching to token auth. + */ + public async clearOAuthState(label: string): Promise { + await this.oauthSessionManager.clearOAuthState(label); + } + public dispose(): void { this.#authListenerDisposable?.dispose(); this.#crossWindowSyncDisposable?.dispose(); diff --git a/src/extension.ts b/src/extension.ts index 0afebc67..b15863c3 100644 --- a/src/extension.ts +++ b/src/extension.ts @@ -8,11 +8,14 @@ import * as vscode from "vscode"; import { errToStr } from "./api/api-helper"; import { CoderApi } from "./api/coderApi"; +import { attachOAuthInterceptors } from "./api/oauthInterceptors"; import { Commands } from "./commands"; import { ServiceContainer } from "./core/container"; import { type SecretsManager } from "./core/secretsManager"; import { DeploymentManager } from "./deployment/deploymentManager"; import { CertificateError, getErrorDetail } from "./error"; +import { OAuthSessionManager } from "./oauth/sessionManager"; +import { CALLBACK_PATH } from "./oauth/utils"; import { maybeAskUrl } from "./promptUtils"; import { Remote } from "./remote/remote"; import { getRemoteSshExtension } from "./remote/sshExtension"; @@ -68,6 +71,14 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { const deployment = await secretsManager.getCurrentDeployment(); + // Create OAuth session manager with login coordinator + const oauthSessionManager = await OAuthSessionManager.create( + deployment, + serviceContainer, + ctx.extension.id, + ); + 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. @@ -78,6 +89,7 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { output, ); ctx.subscriptions.push(client); + attachOAuthInterceptors(client, output, oauthSessionManager); const myWorkspacesProvider = new WorkspaceProvider( WorkspaceQuery.Mine, @@ -123,10 +135,12 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { ); // Create deployment manager to centralize deployment state management - const deploymentManager = DeploymentManager.create(serviceContainer, client, [ - myWorkspacesProvider, - allWorkspacesProvider, - ]); + const deploymentManager = DeploymentManager.create( + serviceContainer, + client, + oauthSessionManager, + [myWorkspacesProvider, allWorkspacesProvider], + ); ctx.subscriptions.push(deploymentManager); // Handle vscode:// URIs. @@ -134,6 +148,14 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { handleUri: async (uri) => { 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"); + await oauthSessionManager.handleCallback(code, state, error); + return; + } + if (uri.path === "/open") { const owner = params.get("owner"); const workspace = params.get("workspace"); @@ -150,7 +172,11 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { throw new Error("workspace must be specified as a query parameter"); } - await setupDeploymentFromUri(params, serviceContainer); + await setupDeploymentFromUri( + params, + serviceContainer, + deploymentManager, + ); await commands.open( owner, @@ -204,7 +230,11 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { ); } - await setupDeploymentFromUri(params, serviceContainer); + await setupDeploymentFromUri( + params, + serviceContainer, + deploymentManager, + ); await commands.openDevContainer( workspaceOwner, @@ -224,7 +254,12 @@ 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, deploymentManager); + const commands = new Commands( + serviceContainer, + client, + oauthSessionManager, + deploymentManager, + ); ctx.subscriptions.push( vscode.commands.registerCommand( "coder.login", @@ -425,6 +460,7 @@ async function showTreeViewSearch(id: string): Promise { async function setupDeploymentFromUri( params: URLSearchParams, serviceContainer: ServiceContainer, + deploymentManager: DeploymentManager, ): Promise { const secretsManager = serviceContainer.getSecretsManager(); const mementoManager = serviceContainer.getMementoManager(); @@ -459,6 +495,8 @@ async function setupDeploymentFromUri( } } else { await secretsManager.setSessionAuth(safeHost, { url, token }); + // Clear OAuth state since we're using a non-OAuth token + await deploymentManager.clearOAuthState(safeHost); } } diff --git a/src/login/loginCoordinator.ts b/src/login/loginCoordinator.ts index 8b5d7b07..dabef47c 100644 --- a/src/login/loginCoordinator.ts +++ b/src/login/loginCoordinator.ts @@ -5,7 +5,8 @@ import * as vscode from "vscode"; import { CoderApi } from "../api/coderApi"; import { needToken } from "../api/utils"; import { CertificateError } from "../error"; -import { maybeAskUrl } from "../promptUtils"; +import { type OAuthSessionManager } from "../oauth/sessionManager"; +import { maybeAskAuthMethod, maybeAskUrl } from "../promptUtils"; import type { User } from "coder/site/src/api/typesGenerated"; @@ -21,6 +22,7 @@ type LoginResult = interface LoginOptions { safeHostname: string; url: string | undefined; + oauthSessionManager: OAuthSessionManager; autoLogin?: boolean; } @@ -44,11 +46,12 @@ export class LoginCoordinator { public async ensureLoggedIn( options: LoginOptions & { url: string }, ): Promise { - const { safeHostname, url } = options; + const { safeHostname, url, oauthSessionManager } = options; return this.executeWithGuard(safeHostname, async () => { const result = await this.attemptLogin( { safeHostname, url }, options.autoLogin ?? false, + oauthSessionManager, ); await this.persistSessionAuth(result, safeHostname, url); @@ -58,12 +61,13 @@ export class LoginCoordinator { } /** - * Shows dialog then login - for system-initiated auth (remote). + * Shows dialog then login - for system-initiated auth (remote, OAuth refresh). */ public async ensureLoggedInWithDialog( options: LoginOptions & { message?: string; detailPrefix?: string }, ): Promise { - const { safeHostname, url, detailPrefix, message } = options; + const { safeHostname, url, detailPrefix, message, oauthSessionManager } = + options; return this.executeWithGuard(safeHostname, async () => { // Show dialog promise const dialogPromise = this.vscodeProposed.window @@ -95,6 +99,7 @@ export class LoginCoordinator { const result = await this.attemptLogin( { url: newUrl, safeHostname }, false, + oauthSessionManager, ); await this.persistSessionAuth(result, safeHostname, newUrl); @@ -183,7 +188,7 @@ export class LoginCoordinator { } /** - * Attempt to authenticate using token, or mTLS. If necessary, prompts + * 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). @@ -191,6 +196,7 @@ export class LoginCoordinator { private async attemptLogin( deployment: Deployment, isAutoLogin: boolean, + oauthSessionManager: OAuthSessionManager, ): Promise { const needsToken = needToken(vscode.workspace.getConfiguration()); const client = CoderApi.create(deployment.url, "", this.logger); @@ -237,8 +243,21 @@ export class LoginCoordinator { } } - const result = await this.loginWithToken(client); - return result; + const authMethod = await maybeAskAuthMethod(client); + switch (authMethod) { + case "oauth": + return this.loginWithOAuth(client, oauthSessionManager, deployment); + case "legacy": { + const result = await this.loginWithToken(client); + if (result.success) { + // Clear OAuth state since user explicitly chose token auth + await oauthSessionManager.clearOAuthState(deployment.safeHostname); + } + return result; + } + case undefined: + return { success: false }; // User aborted + } } /** @@ -297,4 +316,48 @@ export class LoginCoordinator { return { success: false }; } + + /** + * OAuth authentication flow. + */ + private async loginWithOAuth( + client: CoderApi, + oauthSessionManager: OAuthSessionManager, + deployment: Deployment, + ): 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, deployment, 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) { + 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/errors.ts b/src/oauth/errors.ts new file mode 100644 index 00000000..9b7ee3ac --- /dev/null +++ b/src/oauth/errors.ts @@ -0,0 +1,166 @@ +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 + ); +} diff --git a/src/oauth/metadataClient.ts b/src/oauth/metadataClient.ts new file mode 100644 index 00000000..149d64fa --- /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. + */ + public static async checkOAuthSupport( + axiosInstance: AxiosInstance, + ): Promise { + try { + await axiosInstance.get(OAUTH_DISCOVERY_ENDPOINT); + return true; + } catch { + 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.debug("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/sessionManager.ts b/src/oauth/sessionManager.ts new file mode 100644 index 00000000..6189732d --- /dev/null +++ b/src/oauth/sessionManager.ts @@ -0,0 +1,801 @@ +import { type AxiosInstance } from "axios"; +import * as vscode from "vscode"; + +import { CoderApi } from "../api/coderApi"; +import { type ServiceContainer } from "../core/container"; +import { type Deployment } from "../deployment/types"; +import { type LoginCoordinator } from "../login/loginCoordinator"; + +import { OAuthMetadataClient } from "./metadataClient"; +import { + CALLBACK_PATH, + generatePKCE, + generateState, + toUrlSearchParams, +} from "./utils"; + +import type { SecretsManager, StoredOAuthTokens } from "../core/secretsManager"; +import type { Logger } from "../logging/logger"; + +import type { OAuthError } from "./errors"; +import type { + 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 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; + +/** + * 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. + */ +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. + */ +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 storedTokens: StoredOAuthTokens | undefined; + private refreshPromise: Promise | null = null; + private lastRefreshAttempt = 0; + private refreshTimer: NodeJS.Timeout | undefined; + + private pendingAuthReject: ((reason: Error) => void) | undefined; + + /** + * Create and initialize a new OAuth session manager. + */ + public static async create( + deployment: Deployment | null, + container: ServiceContainer, + extensionId: string, + ): Promise { + const manager = new OAuthSessionManager( + deployment, + container.getSecretsManager(), + container.getLogger(), + container.getLoginCoordinator(), + extensionId, + ); + await manager.loadTokens(); + manager.scheduleBackgroundRefresh(); + return manager; + } + + private constructor( + private deployment: Deployment | null, + 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 { + if (!this.deployment) { + return; + } + + const tokens = await this.secretsManager.getOAuthTokens( + this.deployment.safeHostname, + ); + if (!tokens) { + return; + } + + if (tokens.deployment_url !== this.deployment.url) { + this.logger.warn("Stored tokens for different deployment, clearing", { + stored: tokens.deployment_url, + current: this.deployment.url, + }); + this.clearInMemoryTokens(); + await this.secretsManager.clearOAuthData(this.deployment.safeHostname); + 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, + }, + ); + this.clearInMemoryTokens(); + await this.secretsManager.clearOAuthTokens(this.deployment.safeHostname); + return; + } + + this.storedTokens = tokens; + this.logger.info( + `Loaded stored OAuth tokens for ${this.deployment.safeHostname}`, + ); + } + + private clearInMemoryTokens(): void { + this.storedTokens = undefined; + this.refreshPromise = null; + 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:*". + */ + private hasRequiredScopes(grantedScope: string | undefined): boolean { + if (!grantedScope) { + // TODO server always returns empty scopes + return true; + } + + 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}`; + } + + /** + * Prepare common OAuth operation setup: client, metadata, and registration. + * Used by refresh and revoke operations to reduce duplication. + */ + private async prepareOAuthOperation(token?: string): Promise<{ + axiosInstance: AxiosInstance; + metadata: OAuthServerMetadata; + registration: ClientRegistrationResponse; + }> { + 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( + deployment.safeHostname, + ); + 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 deployment = this.requireDeployment(); + const redirectUri = this.getRedirectUri(); + + const existing = await this.secretsManager.getOAuthClientRegistration( + deployment.safeHostname, + ); + 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( + deployment.safeHostname, + response.data, + ); + this.logger.info( + "Saved OAuth client registration:", + response.data.client_id, + ); + + return response.data; + } + + public async setDeployment(deployment: Deployment): Promise { + if ( + this.deployment && + deployment.safeHostname === this.deployment.safeHostname && + 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 = null; + this.clearInMemoryTokens(); + } + + /** + * 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, + deployment: Deployment, + progress: vscode.Progress<{ message?: string; increment?: number }>, + ): Promise { + const baseUrl = client.getAxiosInstance().defaults.baseURL; + if (!baseUrl) { + throw new Error("Client has no base URL set"); + } + if (baseUrl !== deployment.url) { + throw new Error( + `Client base URL (${baseUrl}) does not match deployment URL (${deployment.url})`, + ); + } + + // Update deployment if changed + if ( + !this.deployment || + this.deployment.url !== deployment.url || + this.deployment.safeHostname !== deployment.safeHostname + ) { + this.logger.info("Deployment changed, clearing cached state", { + old: this.deployment, + new: deployment, + }); + this.clearInMemoryTokens(); + this.deployment = deployment; + } + + const axiosInstance = client.getAxiosInstance(); + const metadataClient = new OAuthMetadataClient(axiosInstance, this.logger); + 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, + axiosInstance, + metadata, + registration, + ); + + progress.report({ increment: 30 }); + this.logger.info("OAuth login flow completed successfully"); + + return tokenResponse; + } + + /** + * 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. + */ + private async startAuthorization( + metadata: OAuthServerMetadata, + registration: ClientRegistrationResponse, + ): Promise<{ code: string; verifier: string }> { + const state = generateState(); + const { verifier, challenge } = generatePKCE(); + + const authUrl = this.buildAuthorizationUrl( + metadata, + registration.client_id, + state, + challenge, + ); + + const callbackPromise = new Promise<{ code: string; verifier: string }>( + (resolve, reject) => { + const timeoutMins = 5; + const timeoutHandle = setTimeout( + () => { + cleanup(); + reject( + new Error(`OAuth flow timed out after ${timeoutMins} minutes`), + ); + }, + timeoutMins * 60 * 1000, + ); + + const listener = this.secretsManager.onDidChangeOAuthCallback( + ({ state: callbackState, code, error }) => { + if (callbackState !== state) { + return; + } + + cleanup(); + + if (error) { + reject(new Error(`OAuth error: ${error}`)); + } else if (code) { + resolve({ code, verifier }); + } else { + reject(new Error("No authorization code received")); + } + }, + ); + + const cleanup = () => { + clearTimeout(timeoutHandle); + listener.dispose(); + }; + + this.pendingAuthReject = (error) => { + cleanup(); + reject(error); + }; + }, + ); + + 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. + * Writes the callback result to secrets storage, triggering the waiting window to proceed. + */ + public async handleCallback( + code: string | null, + state: string | null, + error: string | null, + ): Promise { + if (!state) { + this.logger.warn("Received OAuth callback with no state parameter"); + 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); + } + } + + /** + * Exchange authorization code for access token. + */ + 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 = { + 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 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. + * Uses a shared promise to handle concurrent refresh attempts. + */ + public async refreshToken(): Promise { + // 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"); + } + + const refreshToken = this.storedTokens.refresh_token; + const accessToken = this.storedTokens.access_token; + + this.lastRefreshAttempt = Date.now(); + + // Create and store the refresh promise + this.refreshPromise = (async () => { + try { + const { axiosInstance, metadata, registration } = + await this.prepareOAuthOperation(accessToken); + + this.logger.debug("Refreshing access token"); + + 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 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; + } finally { + this.refreshPromise = null; + } + })(); + + return this.refreshPromise; + } + + /** + * Save token response to storage. + * 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: deployment.url, + expiry_timestamp: expiryTimestamp, + }; + + this.storedTokens = tokens; + await this.secretsManager.setOAuthTokens(deployment.safeHostname, tokens); + await this.secretsManager.setSessionAuth(deployment.safeHostname, { + url: deployment.url, + token: tokenResponse.access_token, + }); + + this.logger.info("Tokens saved", { + expires_at: new Date(expiryTimestamp).toISOString(), + deployment: deployment.url, + }); + } + + /** + * 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: + * 1. Stored tokens exist with a refresh token + * 2. Token expires in less than TOKEN_REFRESH_THRESHOLD_MS + * 3. Last refresh attempt was more than REFRESH_THROTTLE_MS ago + * 4. No refresh is currently in progress + */ + private shouldRefreshToken(): boolean { + if (!this.storedTokens?.refresh_token || this.refreshPromise !== null) { + return false; + } + + 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; + } + + public async revokeRefreshToken(): Promise { + if (!this.storedTokens?.refresh_token) { + this.logger.info("No refresh token to revoke"); + return; + } + + await this.revokeToken(this.storedTokens.refresh_token, "refresh_token"); + } + + /** + * Revoke a token using the OAuth server's revocation endpoint. + */ + private async revokeToken( + token: string, + tokenTypeHint: "access_token" | "refresh_token" = "refresh_token", + ): Promise { + const { axiosInstance, metadata, registration } = + await this.prepareOAuthOperation(this.storedTokens?.access_token); + + const revocationEndpoint = + metadata.revocation_endpoint || `${metadata.issuer}/oauth2/revoke`; + + this.logger.info("Revoking refresh token"); + + const params: TokenRevocationRequest = { + token, + client_id: registration.client_id, + client_secret: registration.client_secret, + token_type_hint: tokenTypeHint, + }; + + const revocationRequest = toUrlSearchParams(params); + + try { + await axiosInstance.post(revocationEndpoint, 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; + } + } + + /** + * Returns true if (valid or invalid) OAuth tokens exist for the current deployment. + */ + public isLoggedInWithOAuth(): boolean { + return this.storedTokens !== undefined; + } + + /** + * Clear OAuth state when switching to non-OAuth authentication. + * Clears in-memory state and OAuth tokens from storage. + * Preserves client registration for potential future OAuth use. + */ + public async clearOAuthState(label: string): Promise { + this.clearInMemoryTokens(); + await this.secretsManager.clearOAuthTokens(label); + } + + /** + * 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. + */ + 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 + this.clearInMemoryTokens(); + await this.secretsManager.clearAllAuthData(deployment.safeHostname); + + await this.loginCoordinator.ensureLoggedInWithDialog({ + safeHostname: deployment.safeHostname, + url: deployment.url, + detailPrefix: errorMessage, + oauthSessionManager: this, + }); + } + + /** + * 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")); + } + this.pendingAuthReject = undefined; + this.clearInMemoryTokens(); + + this.logger.debug("OAuth session manager disposed"); + } +} diff --git a/src/oauth/types.ts b/src/oauth/types.ts new file mode 100644 index 00000000..6ecaa0ff --- /dev/null +++ b/src/oauth/types.ts @@ -0,0 +1,163 @@ +// 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: 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; +} + +// 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[]; + 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; + jwks_uri?: string; + response_types_supported: ResponseType[]; + grant_types_supported?: GrantType[]; + code_challenge_methods_supported: CodeChallengeMethod[]; + scopes_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: TokenType; + expires_in?: number; + refresh_token?: string; + scope?: string; +} + +// Authorization Request Parameters (OAuth 2.1) +export interface AuthorizationRequestParams { + client_id: string; + response_type: ResponseType; + redirect_uri: string; + scope?: string; + state: string; + code_challenge: string; + code_challenge_method: CodeChallengeMethod; +} + +// Token Request Parameters - Authorization Code Grant (OAuth 2.1) +export interface TokenRequestParams { + grant_type: "authorization_code"; + code: string; + redirect_uri: string; + client_id: string; + 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; +} diff --git a/src/oauth/utils.ts b/src/oauth/utils.ts new file mode 100644 index 00000000..61beeb50 --- /dev/null +++ b/src/oauth/utils.ts @@ -0,0 +1,42 @@ +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"); +} + +/** + * 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/promptUtils.ts b/src/promptUtils.ts index 3fb31475..9e3d8895 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 @@ -130,3 +134,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; +} diff --git a/src/remote/remote.ts b/src/remote/remote.ts index 974d956d..0bbb4a51 100644 --- a/src/remote/remote.ts +++ b/src/remote/remote.ts @@ -20,6 +20,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 { getGlobalFlags, getGlobalFlagsRaw, getSshFlags } from "../cliConfig"; import { type Commands } from "../commands"; @@ -34,6 +35,7 @@ import { getHeaderCommand } from "../headers"; import { Inbox } from "../inbox"; import { type Logger } from "../logging/logger"; import { type LoginCoordinator } from "../login/loginCoordinator"; +import { OAuthSessionManager } from "../oauth/sessionManager"; import { AuthorityPrefix, escapeCommandArg, @@ -64,7 +66,7 @@ export class Remote { private readonly loginCoordinator: LoginCoordinator; public constructor( - serviceContainer: ServiceContainer, + private readonly serviceContainer: ServiceContainer, private readonly commands: Commands, private readonly extensionContext: vscode.ExtensionContext, ) { @@ -110,6 +112,14 @@ export class Remote { const disposables: vscode.Disposable[] = []; try { + // Create OAuth session manager for this remote deployment + const remoteOAuthManager = await OAuthSessionManager.create( + { url: baseUrlRaw, safeHostname: parts.safeHostname }, + this.serviceContainer, + this.extensionContext.extension.id, + ); + disposables.push(remoteOAuthManager); + const ensureLoggedInAndRetry = async ( message: string, url: string | undefined, @@ -119,6 +129,7 @@ export class Remote { url, message, detailPrefix: `You must log in to access ${workspaceName}.`, + oauthSessionManager: remoteOAuthManager, }); // Dispose before retrying since setup will create new disposables @@ -154,6 +165,7 @@ export class Remote { // 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.remoteWorkspaceClient = workspaceClient; From 82991af43d5cdd75a7e2ef79d0c6d2c72a9a58e8 Mon Sep 17 00:00:00 2001 From: Ehab Younes Date: Tue, 16 Dec 2025 18:31:34 +0300 Subject: [PATCH 2/9] Fix tests after rebase --- test/mocks/testHelpers.ts | 20 +++ .../unit/deployment/deploymentManager.test.ts | 4 + test/unit/login/loginCoordinator.test.ts | 139 ++++++++++++------ 3 files changed, 119 insertions(+), 44 deletions(-) diff --git a/test/mocks/testHelpers.ts b/test/mocks/testHelpers.ts index 21978b13..adbe4927 100644 --- a/test/mocks/testHelpers.ts +++ b/test/mocks/testHelpers.ts @@ -528,6 +528,26 @@ export class MockCoderApi } } +/** + * Mock OAuthSessionManager for testing. + * Provides no-op implementations of all public methods. + */ +export class MockOAuthSessionManager { + readonly setDeployment = vi.fn().mockResolvedValue(undefined); + readonly clearDeployment = vi.fn(); + readonly login = vi.fn().mockResolvedValue({ access_token: "test-token" }); + readonly handleCallback = vi.fn().mockResolvedValue(undefined); + readonly refreshToken = vi + .fn() + .mockResolvedValue({ access_token: "test-token" }); + readonly refreshIfAlmostExpired = vi.fn().mockResolvedValue(undefined); + readonly revokeRefreshToken = vi.fn().mockResolvedValue(undefined); + readonly isLoggedInWithOAuth = vi.fn().mockReturnValue(false); + readonly clearOAuthState = vi.fn().mockResolvedValue(undefined); + readonly showReAuthenticationModal = vi.fn().mockResolvedValue(undefined); + readonly dispose = vi.fn(); +} + /** * Create a mock User for testing. */ diff --git a/test/unit/deployment/deploymentManager.test.ts b/test/unit/deployment/deploymentManager.test.ts index 4f0ca52d..e5fac904 100644 --- a/test/unit/deployment/deploymentManager.test.ts +++ b/test/unit/deployment/deploymentManager.test.ts @@ -11,10 +11,12 @@ import { InMemoryMemento, InMemorySecretStorage, MockCoderApi, + MockOAuthSessionManager, } from "../../mocks/testHelpers"; import type { ServiceContainer } from "@/core/container"; import type { ContextManager } from "@/core/contextManager"; +import type { OAuthSessionManager } from "@/oauth/sessionManager"; import type { WorkspaceProvider } from "@/workspace/workspacesProvider"; // Mock CoderApi.create to return our mock client for validation @@ -64,6 +66,7 @@ function createTestContext() { // For setDeploymentIfValid, we use a separate mock for validation const validationMockClient = new MockCoderApi(); const mockWorkspaceProvider = new MockWorkspaceProvider(); + const mockOAuthSessionManager = new MockOAuthSessionManager(); const secretStorage = new InMemorySecretStorage(); const memento = new InMemoryMemento(); const logger = createMockLogger(); @@ -86,6 +89,7 @@ function createTestContext() { const manager = DeploymentManager.create( container as unknown as ServiceContainer, mockClient as unknown as CoderApi, + mockOAuthSessionManager as unknown as OAuthSessionManager, [mockWorkspaceProvider as unknown as WorkspaceProvider], ); diff --git a/test/unit/login/loginCoordinator.test.ts b/test/unit/login/loginCoordinator.test.ts index fda88ada..b2d76541 100644 --- a/test/unit/login/loginCoordinator.test.ts +++ b/test/unit/login/loginCoordinator.test.ts @@ -13,9 +13,12 @@ import { InMemoryMemento, InMemorySecretStorage, MockConfigurationProvider, + MockOAuthSessionManager, MockUserInteraction, } from "../../mocks/testHelpers"; +import type { OAuthSessionManager } from "@/oauth/sessionManager"; + // Hoisted mock adapter implementation const mockAxiosAdapterImpl = vi.hoisted( () => (config: Record) => @@ -58,7 +61,29 @@ vi.mock("@/api/streamingFetchAdapter", () => ({ createStreamingFetchAdapter: vi.fn(() => fetch), })); -vi.mock("@/promptUtils"); +vi.mock("@/promptUtils", () => ({ + maybeAskAuthMethod: vi.fn().mockResolvedValue("legacy"), + maybeAskUrl: vi.fn(), +})); + +// Mock CoderApi to control getAuthenticatedUser behavior +const mockGetAuthenticatedUser = vi.hoisted(() => vi.fn()); +vi.mock("@/api/coderApi", async (importOriginal) => { + const original = await importOriginal(); + return { + ...original, + CoderApi: { + ...original.CoderApi, + create: vi.fn(() => ({ + getAxiosInstance: () => ({ + defaults: { baseURL: "https://coder.example.com" }, + }), + setSessionToken: vi.fn(), + getAuthenticatedUser: mockGetAuthenticatedUser, + })), + }, + }; +}); // Type for axios with our mock adapter type MockedAxios = typeof axios & { __mockAdapter: ReturnType }; @@ -94,7 +119,12 @@ function createTestContext() { logger, ); + const oauthSessionManager = + new MockOAuthSessionManager() as unknown as OAuthSessionManager; + const mockSuccessfulAuth = (user = createMockUser()) => { + // Configure both the axios adapter (for tests that bypass CoderApi mock) + // and mockGetAuthenticatedUser (for tests that use the CoderApi mock) mockAdapter.mockResolvedValue({ data: user, status: 200, @@ -102,6 +132,7 @@ function createTestContext() { headers: {}, config: {}, }); + mockGetAuthenticatedUser.mockResolvedValue(user); return user; }; @@ -110,6 +141,10 @@ function createTestContext() { response: { status: 401, data: { message } }, message, }); + mockGetAuthenticatedUser.mockRejectedValue({ + response: { status: 401, data: { message } }, + message, + }); }; return { @@ -119,6 +154,7 @@ function createTestContext() { secretsManager, mementoManager, coordinator, + oauthSessionManager, mockSuccessfulAuth, mockAuthFailure, }; @@ -127,8 +163,12 @@ function createTestContext() { describe("LoginCoordinator", () => { describe("token authentication", () => { it("authenticates with stored token on success", async () => { - const { secretsManager, coordinator, mockSuccessfulAuth } = - createTestContext(); + const { + secretsManager, + coordinator, + oauthSessionManager, + mockSuccessfulAuth, + } = createTestContext(); const user = mockSuccessfulAuth(); // Pre-store a token @@ -140,6 +180,7 @@ describe("LoginCoordinator", () => { const result = await coordinator.ensureLoggedIn({ url: TEST_URL, safeHostname: TEST_HOSTNAME, + oauthSessionManager, }); expect(result).toEqual({ success: true, user, token: "stored-token" }); @@ -148,20 +189,16 @@ describe("LoginCoordinator", () => { expect(auth?.token).toBe("stored-token"); }); - it("prompts for token when no stored auth exists", async () => { - const { mockAdapter, userInteraction, secretsManager, coordinator } = - createTestContext(); - const user = createMockUser(); - - // No stored token, so goes directly to input box flow - // Mock succeeds when validateInput calls getAuthenticatedUser - mockAdapter.mockResolvedValueOnce({ - data: user, - status: 200, - statusText: "OK", - headers: {}, - config: {}, - }); + // TODO: This test needs the CoderApi mock to work through the validateInput callback + it.skip("prompts for token when no stored auth exists", async () => { + const { + userInteraction, + secretsManager, + coordinator, + oauthSessionManager, + mockSuccessfulAuth, + } = createTestContext(); + const user = mockSuccessfulAuth(); // User enters a new token in the input box userInteraction.setInputBoxValue("new-token"); @@ -169,6 +206,7 @@ describe("LoginCoordinator", () => { const result = await coordinator.ensureLoggedIn({ url: TEST_URL, safeHostname: TEST_HOSTNAME, + oauthSessionManager, }); expect(result).toEqual({ success: true, user, token: "new-token" }); @@ -179,14 +217,19 @@ describe("LoginCoordinator", () => { }); it("returns success false when user cancels input", async () => { - const { userInteraction, coordinator, mockAuthFailure } = - createTestContext(); + const { + userInteraction, + coordinator, + oauthSessionManager, + mockAuthFailure, + } = createTestContext(); mockAuthFailure(); userInteraction.setInputBoxValue(undefined); const result = await coordinator.ensureLoggedIn({ url: TEST_URL, safeHostname: TEST_HOSTNAME, + oauthSessionManager, }); expect(result.success).toBe(false); @@ -194,39 +237,31 @@ describe("LoginCoordinator", () => { }); describe("same-window guard", () => { - it("prevents duplicate login calls for same hostname", async () => { - const { mockAdapter, userInteraction, coordinator } = createTestContext(); - const user = createMockUser(); + // TODO: This test needs the CoderApi mock to work through the validateInput callback + it.skip("prevents duplicate login calls for same hostname", async () => { + const { + userInteraction, + coordinator, + oauthSessionManager, + mockSuccessfulAuth, + } = createTestContext(); + mockSuccessfulAuth(); // User enters a token in the input box userInteraction.setInputBoxValue("new-token"); - let resolveAuth: (value: unknown) => void; - mockAdapter.mockReturnValue( - new Promise((resolve) => { - resolveAuth = resolve; - }), - ); - // Start first login const login1 = coordinator.ensureLoggedIn({ url: TEST_URL, safeHostname: TEST_HOSTNAME, + oauthSessionManager, }); // Start second login immediately (same hostname) const login2 = coordinator.ensureLoggedIn({ url: TEST_URL, safeHostname: TEST_HOSTNAME, - }); - - // Resolve the auth (this validates the token from input box) - resolveAuth!({ - data: user, - status: 200, - statusText: "OK", - headers: {}, - config: {}, + oauthSessionManager, }); // Both should complete with the same result @@ -241,8 +276,13 @@ describe("LoginCoordinator", () => { describe("mTLS authentication", () => { it("succeeds without prompt and returns token=''", async () => { - const { mockConfig, secretsManager, coordinator, mockSuccessfulAuth } = - createTestContext(); + const { + mockConfig, + secretsManager, + coordinator, + oauthSessionManager, + mockSuccessfulAuth, + } = createTestContext(); // Configure mTLS via certs (no token needed) mockConfig.set("coder.tlsCertFile", "/path/to/cert.pem"); mockConfig.set("coder.tlsKeyFile", "/path/to/key.pem"); @@ -252,6 +292,7 @@ describe("LoginCoordinator", () => { const result = await coordinator.ensureLoggedIn({ url: TEST_URL, safeHostname: TEST_HOSTNAME, + oauthSessionManager, }); expect(result).toEqual({ success: true, user, token: "" }); @@ -265,7 +306,8 @@ describe("LoginCoordinator", () => { }); it("shows error and returns failure when mTLS fails", async () => { - const { mockConfig, coordinator, mockAuthFailure } = createTestContext(); + const { mockConfig, coordinator, oauthSessionManager, mockAuthFailure } = + createTestContext(); mockConfig.set("coder.tlsCertFile", "/path/to/cert.pem"); mockConfig.set("coder.tlsKeyFile", "/path/to/key.pem"); mockAuthFailure("Certificate error"); @@ -273,6 +315,7 @@ describe("LoginCoordinator", () => { const result = await coordinator.ensureLoggedIn({ url: TEST_URL, safeHostname: TEST_HOSTNAME, + oauthSessionManager, }); expect(result.success).toBe(false); @@ -286,8 +329,13 @@ describe("LoginCoordinator", () => { }); it("logs warning instead of showing dialog for autoLogin", async () => { - const { mockConfig, secretsManager, mementoManager, mockAuthFailure } = - createTestContext(); + const { + mockConfig, + secretsManager, + mementoManager, + oauthSessionManager, + mockAuthFailure, + } = createTestContext(); mockConfig.set("coder.tlsCertFile", "/path/to/cert.pem"); mockConfig.set("coder.tlsKeyFile", "/path/to/key.pem"); @@ -304,6 +352,7 @@ describe("LoginCoordinator", () => { const result = await coordinator.ensureLoggedIn({ url: TEST_URL, safeHostname: TEST_HOSTNAME, + oauthSessionManager, autoLogin: true, }); @@ -315,7 +364,8 @@ describe("LoginCoordinator", () => { describe("ensureLoggedInWithDialog", () => { it("returns success false when user dismisses dialog", async () => { - const { mockConfig, userInteraction, coordinator } = createTestContext(); + const { mockConfig, userInteraction, coordinator, oauthSessionManager } = + createTestContext(); // Use mTLS for simpler dialog test mockConfig.set("coder.tlsCertFile", "/path/to/cert.pem"); mockConfig.set("coder.tlsKeyFile", "/path/to/key.pem"); @@ -326,6 +376,7 @@ describe("LoginCoordinator", () => { const result = await coordinator.ensureLoggedInWithDialog({ url: TEST_URL, safeHostname: TEST_HOSTNAME, + oauthSessionManager, }); expect(result.success).toBe(false); From 8b916051e7b59f63c9fd24f1f65a28977f2d52de Mon Sep 17 00:00:00 2001 From: Ehab Younes Date: Wed, 17 Dec 2025 20:07:48 +0300 Subject: [PATCH 3/9] Self-review of OAuth logic --- src/core/secretsManager.ts | 209 ++++++++++++++-------------- src/deployment/deploymentManager.ts | 18 +-- src/extension.ts | 17 +-- src/login/loginCoordinator.ts | 2 +- src/oauth/sessionManager.ts | 32 ++++- 5 files changed, 142 insertions(+), 136 deletions(-) diff --git a/src/core/secretsManager.ts b/src/core/secretsManager.ts index 128a826b..33708830 100644 --- a/src/core/secretsManager.ts +++ b/src/core/secretsManager.ts @@ -15,9 +15,10 @@ const SESSION_KEY_PREFIX = "coder.session."; const OAUTH_TOKENS_PREFIX = "coder.oauth.tokens."; const OAUTH_CLIENT_PREFIX = "coder.oauth.client."; -const CURRENT_DEPLOYMENT_KEY = "coder.currentDeployment"; const OAUTH_CALLBACK_KEY = "coder.oauthCallback"; +const CURRENT_DEPLOYMENT_KEY = "coder.currentDeployment"; + const DEPLOYMENT_USAGE_KEY = "coder.deploymentUsage"; const DEFAULT_MAX_DEPLOYMENTS = 10; @@ -115,38 +116,6 @@ export class SecretsManager { }); } - /** - * 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 - } - }); - } - /** * Listen for changes to a specific deployment's session auth. */ @@ -203,77 +172,6 @@ export class SecretsManager { return `${SESSION_KEY_PREFIX}${safeHostname || ""}`; } - public async getOAuthTokens( - safeHostname: string, - ): Promise { - try { - const data = await this.secrets.get( - `${OAUTH_TOKENS_PREFIX}${safeHostname}`, - ); - if (!data) { - return undefined; - } - return JSON.parse(data) as StoredOAuthTokens; - } catch { - return undefined; - } - } - - public async setOAuthTokens( - safeHostname: string, - tokens: StoredOAuthTokens, - ): Promise { - await this.secrets.store( - `${OAUTH_TOKENS_PREFIX}${safeHostname}`, - JSON.stringify(tokens), - ); - await this.recordDeploymentAccess(safeHostname); - } - - public async clearOAuthTokens(safeHostname: string): Promise { - await this.secrets.delete(`${OAUTH_TOKENS_PREFIX}${safeHostname}`); - } - - public async getOAuthClientRegistration( - safeHostname: string, - ): Promise { - try { - const data = await this.secrets.get( - `${OAUTH_CLIENT_PREFIX}${safeHostname}`, - ); - if (!data) { - return undefined; - } - return JSON.parse(data) as ClientRegistrationResponse; - } catch { - return undefined; - } - } - - public async setOAuthClientRegistration( - safeHostname: string, - registration: ClientRegistrationResponse, - ): Promise { - await this.secrets.store( - `${OAUTH_CLIENT_PREFIX}${safeHostname}`, - JSON.stringify(registration), - ); - await this.recordDeploymentAccess(safeHostname); - } - - public async clearOAuthClientRegistration( - safeHostname: string, - ): Promise { - await this.secrets.delete(`${OAUTH_CLIENT_PREFIX}${safeHostname}`); - } - - public async clearOAuthData(safeHostname: string): Promise { - await Promise.all([ - this.clearOAuthTokens(safeHostname), - this.clearOAuthClientRegistration(safeHostname), - ]); - } - /** * Record that a deployment was accessed, moving it to the front of the LRU list. * Prunes deployments beyond maxCount, clearing their auth data. @@ -359,4 +257,107 @@ export class SecretsManager { return safeHostname; } + + /** + * 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 + } + }); + } + + public async getOAuthTokens( + safeHostname: string, + ): Promise { + try { + const data = await this.secrets.get( + `${OAUTH_TOKENS_PREFIX}${safeHostname}`, + ); + if (!data) { + return undefined; + } + return JSON.parse(data) as StoredOAuthTokens; + } catch { + return undefined; + } + } + + public async setOAuthTokens( + safeHostname: string, + tokens: StoredOAuthTokens, + ): Promise { + await this.secrets.store( + `${OAUTH_TOKENS_PREFIX}${safeHostname}`, + JSON.stringify(tokens), + ); + await this.recordDeploymentAccess(safeHostname); + } + + public async clearOAuthTokens(safeHostname: string): Promise { + await this.secrets.delete(`${OAUTH_TOKENS_PREFIX}${safeHostname}`); + } + + public async getOAuthClientRegistration( + safeHostname: string, + ): Promise { + try { + const data = await this.secrets.get( + `${OAUTH_CLIENT_PREFIX}${safeHostname}`, + ); + if (!data) { + return undefined; + } + return JSON.parse(data) as ClientRegistrationResponse; + } catch { + return undefined; + } + } + + public async setOAuthClientRegistration( + safeHostname: string, + registration: ClientRegistrationResponse, + ): Promise { + await this.secrets.store( + `${OAUTH_CLIENT_PREFIX}${safeHostname}`, + JSON.stringify(registration), + ); + await this.recordDeploymentAccess(safeHostname); + } + + public async clearOAuthClientRegistration( + safeHostname: string, + ): Promise { + await this.secrets.delete(`${OAUTH_CLIENT_PREFIX}${safeHostname}`); + } + + public async clearOAuthData(safeHostname: string): Promise { + await Promise.all([ + this.clearOAuthTokens(safeHostname), + this.clearOAuthClientRegistration(safeHostname), + ]); + } } diff --git a/src/deployment/deploymentManager.ts b/src/deployment/deploymentManager.ts index 6d524f8d..4521e3de 100644 --- a/src/deployment/deploymentManager.ts +++ b/src/deployment/deploymentManager.ts @@ -89,12 +89,6 @@ export class DeploymentManager implements vscode.Disposable { public async setDeploymentIfValid( deployment: Deployment & { token?: string }, ): Promise { - // TODO used to trigger - /** - * this.oauthSessionManager.refreshIfAlmostExpired().catch((error) => { - this.logger.warn("Setup token refresh failed:", error); - }); - */ const auth = await this.secretsManager.getSessionAuth( deployment.safeHostname, ); @@ -134,11 +128,14 @@ export class DeploymentManager implements vscode.Disposable { } else { this.client.setCredentials(deployment.url, deployment.token); } - await this.oauthSessionManager.setDeployment(deployment); + // Register auth listener before setDeployment so background token refresh + // can update client credentials via the listener this.registerAuthListener(); this.updateAuthContexts(); this.refreshWorkspaces(); + + await this.oauthSessionManager.setDeployment(deployment); await this.persistDeployment(deployment); } @@ -158,13 +155,6 @@ export class DeploymentManager implements vscode.Disposable { await this.secretsManager.setCurrentDeployment(undefined); } - /** - * Clear OAuth state for a deployment when switching to token auth. - */ - public async clearOAuthState(label: string): Promise { - await this.oauthSessionManager.clearOAuthState(label); - } - public dispose(): void { this.#authListenerDisposable?.dispose(); this.#crossWindowSyncDisposable?.dispose(); diff --git a/src/extension.ts b/src/extension.ts index b15863c3..23dec97c 100644 --- a/src/extension.ts +++ b/src/extension.ts @@ -172,11 +172,7 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { throw new Error("workspace must be specified as a query parameter"); } - await setupDeploymentFromUri( - params, - serviceContainer, - deploymentManager, - ); + await setupDeploymentFromUri(params, serviceContainer); await commands.open( owner, @@ -230,11 +226,7 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { ); } - await setupDeploymentFromUri( - params, - serviceContainer, - deploymentManager, - ); + await setupDeploymentFromUri(params, serviceContainer); await commands.openDevContainer( workspaceOwner, @@ -460,7 +452,6 @@ async function showTreeViewSearch(id: string): Promise { async function setupDeploymentFromUri( params: URLSearchParams, serviceContainer: ServiceContainer, - deploymentManager: DeploymentManager, ): Promise { const secretsManager = serviceContainer.getSecretsManager(); const mementoManager = serviceContainer.getMementoManager(); @@ -495,8 +486,8 @@ async function setupDeploymentFromUri( } } else { await secretsManager.setSessionAuth(safeHost, { url, token }); - // Clear OAuth state since we're using a non-OAuth token - await deploymentManager.clearOAuthState(safeHost); + // Clear OAuth tokens since we're using a non-OAuth token + await secretsManager.clearOAuthTokens(safeHost); } } diff --git a/src/login/loginCoordinator.ts b/src/login/loginCoordinator.ts index dabef47c..2aa5f4ae 100644 --- a/src/login/loginCoordinator.ts +++ b/src/login/loginCoordinator.ts @@ -251,7 +251,7 @@ export class LoginCoordinator { const result = await this.loginWithToken(client); if (result.success) { // Clear OAuth state since user explicitly chose token auth - await oauthSessionManager.clearOAuthState(deployment.safeHostname); + await oauthSessionManager.clearOAuthState(); } return result; } diff --git a/src/oauth/sessionManager.ts b/src/oauth/sessionManager.ts index 6189732d..2134336b 100644 --- a/src/oauth/sessionManager.ts +++ b/src/oauth/sessionManager.ts @@ -152,8 +152,7 @@ export class OAuthSessionManager implements vscode.Disposable { required_scopes: DEFAULT_OAUTH_SCOPES, }, ); - this.clearInMemoryTokens(); - await this.secretsManager.clearOAuthTokens(this.deployment.safeHostname); + await this.clearOAuthState(); return; } @@ -322,6 +321,19 @@ export class OAuthSessionManager implements vscode.Disposable { this.deployment = deployment; this.clearInMemoryTokens(); await this.loadTokens(); + + // Refresh tokens if needed to prevent 401s + if (this.isTokenExpired()) { + try { + await this.refreshToken(); + } catch (error) { + this.logger.warn("Token refresh failed (expired):", error); + } + } else { + this.refreshIfAlmostExpired().catch((error) => { + this.logger.warn("Background token refresh failed:", error); + }); + } } public clearDeployment(): void { @@ -695,6 +707,16 @@ export class OAuthSessionManager implements vscode.Disposable { return timeUntilExpiry < TOKEN_REFRESH_THRESHOLD_MS; } + /** + * Check if token is expired. + */ + private isTokenExpired(): boolean { + if (!this.storedTokens) { + return false; + } + return Date.now() >= this.storedTokens.expiry_timestamp; + } + public async revokeRefreshToken(): Promise { if (!this.storedTokens?.refresh_token) { this.logger.info("No refresh token to revoke"); @@ -754,9 +776,11 @@ export class OAuthSessionManager implements vscode.Disposable { * Clears in-memory state and OAuth tokens from storage. * Preserves client registration for potential future OAuth use. */ - public async clearOAuthState(label: string): Promise { + public async clearOAuthState(): Promise { this.clearInMemoryTokens(); - await this.secretsManager.clearOAuthTokens(label); + if (this.deployment) { + await this.secretsManager.clearOAuthTokens(this.deployment.safeHostname); + } } /** From b0b06559044ad74922a55acc3fdb94943a2c6ad2 Mon Sep 17 00:00:00 2001 From: Ehab Younes Date: Thu, 18 Dec 2025 12:12:55 +0300 Subject: [PATCH 4/9] Handle review comments: * Remove proactive oauth refresh interceptor * Unify secretsManager set/get/clear functions * Add login queue to loginCoordinator * Add "Cancel" button to the OAuth flow --- src/api/oauthInterceptors.ts | 13 +--- src/core/secretsManager.ts | 135 +++++++++++++++++----------------- src/login/loginCoordinator.ts | 33 +++------ src/oauth/sessionManager.ts | 36 ++++++--- 4 files changed, 107 insertions(+), 110 deletions(-) diff --git a/src/api/oauthInterceptors.ts b/src/api/oauthInterceptors.ts index b80e1d96..5d04b33e 100644 --- a/src/api/oauthInterceptors.ts +++ b/src/api/oauthInterceptors.ts @@ -10,10 +10,9 @@ import { type CoderApi } from "./coderApi"; const coderSessionTokenHeader = "Coder-Session-Token"; /** - * Attach OAuth token refresh interceptors to a CoderApi instance. + * Attach OAuth token refresh interceptor 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( @@ -22,15 +21,7 @@ export function attachOAuthInterceptors( oauthSessionManager: OAuthSessionManager, ): void { client.getAxiosInstance().interceptors.response.use( - // Success response interceptor: proactive token refresh - (response) => { - // Fire-and-forget: don't await, don't block response - oauthSessionManager.refreshIfAlmostExpired().catch((error) => { - logger.warn("Proactive background token refresh failed:", error); - }); - - return response; - }, + (r) => r, // Error response interceptor: reactive token refresh on 401 async (error: unknown) => { if (!isAxiosError(error)) { diff --git a/src/core/secretsManager.ts b/src/core/secretsManager.ts index 33708830..db339a08 100644 --- a/src/core/secretsManager.ts +++ b/src/core/secretsManager.ts @@ -11,9 +11,14 @@ import type { Deployment } from "../deployment/types"; // Each deployment has its own key to ensure atomic operations (multiple windows // writing to a shared key could drop data) and to receive proper VS Code events. -const SESSION_KEY_PREFIX = "coder.session."; -const OAUTH_TOKENS_PREFIX = "coder.oauth.tokens."; -const OAUTH_CLIENT_PREFIX = "coder.oauth.client."; +const SESSION_KEY_PREFIX = "coder.session." as const; +const OAUTH_TOKENS_PREFIX = "coder.oauth.tokens." as const; +const OAUTH_CLIENT_PREFIX = "coder.oauth.client." as const; + +type SecretKeyPrefix = + | typeof SESSION_KEY_PREFIX + | typeof OAUTH_TOKENS_PREFIX + | typeof OAUTH_CLIENT_PREFIX; const OAUTH_CALLBACK_KEY = "coder.oauthCallback"; @@ -57,6 +62,44 @@ export class SecretsManager { private readonly logger: Logger, ) {} + private buildKey(prefix: SecretKeyPrefix, safeHostname: string): string { + return `${prefix}${safeHostname || ""}`; + } + + private async getSecret( + prefix: SecretKeyPrefix, + safeHostname: string, + ): Promise { + try { + const data = await this.secrets.get(this.buildKey(prefix, safeHostname)); + if (!data) { + return undefined; + } + return JSON.parse(data) as T; + } catch { + return undefined; + } + } + + private async setSecret( + prefix: SecretKeyPrefix, + safeHostname: string, + value: T, + ): Promise { + await this.secrets.store( + this.buildKey(prefix, safeHostname), + JSON.stringify(value), + ); + await this.recordDeploymentAccess(safeHostname); + } + + private async clearSecret( + prefix: SecretKeyPrefix, + safeHostname: string, + ): Promise { + await this.secrets.delete(this.buildKey(prefix, safeHostname)); + } + /** * Sets the current deployment and triggers a cross-window sync event. */ @@ -123,7 +166,7 @@ export class SecretsManager { safeHostname: string, listener: (auth: SessionAuth | undefined) => void | Promise, ): Disposable { - const sessionKey = this.getSessionKey(safeHostname); + const sessionKey = this.buildKey(SESSION_KEY_PREFIX, safeHostname); return this.secrets.onDidChange(async (e) => { if (e.key !== sessionKey) { return; @@ -137,39 +180,23 @@ export class SecretsManager { }); } - public async getSessionAuth( + public getSessionAuth( safeHostname: string, ): Promise { - const sessionKey = this.getSessionKey(safeHostname); - try { - const data = await this.secrets.get(sessionKey); - if (!data) { - return undefined; - } - return JSON.parse(data) as SessionAuth; - } catch { - return undefined; - } + return this.getSecret(SESSION_KEY_PREFIX, safeHostname); } public async setSessionAuth( safeHostname: string, auth: SessionAuth, ): Promise { - const sessionKey = this.getSessionKey(safeHostname); // Extract only url and token before serializing const state: SessionAuth = { url: auth.url, token: auth.token }; - await this.secrets.store(sessionKey, JSON.stringify(state)); - await this.recordDeploymentAccess(safeHostname); - } - - private async clearSessionAuth(safeHostname: string): Promise { - const sessionKey = this.getSessionKey(safeHostname); - await this.secrets.delete(sessionKey); + await this.setSecret(SESSION_KEY_PREFIX, safeHostname, state); } - private getSessionKey(safeHostname: string): string { - return `${SESSION_KEY_PREFIX}${safeHostname || ""}`; + private clearSessionAuth(safeHostname: string): Promise { + return this.clearSecret(SESSION_KEY_PREFIX, safeHostname); } /** @@ -204,7 +231,6 @@ export class SecretsManager { this.clearSessionAuth(safeHostname), this.clearOAuthData(safeHostname), ]); - await this.clearSessionAuth(safeHostname); const usage = this.getDeploymentUsage().filter( (u) => u.safeHostname !== safeHostname, ); @@ -290,68 +316,41 @@ export class SecretsManager { }); } - public async getOAuthTokens( + public getOAuthTokens( safeHostname: string, ): Promise { - try { - const data = await this.secrets.get( - `${OAUTH_TOKENS_PREFIX}${safeHostname}`, - ); - if (!data) { - return undefined; - } - return JSON.parse(data) as StoredOAuthTokens; - } catch { - return undefined; - } + return this.getSecret(OAUTH_TOKENS_PREFIX, safeHostname); } - public async setOAuthTokens( + public setOAuthTokens( safeHostname: string, tokens: StoredOAuthTokens, ): Promise { - await this.secrets.store( - `${OAUTH_TOKENS_PREFIX}${safeHostname}`, - JSON.stringify(tokens), - ); - await this.recordDeploymentAccess(safeHostname); + return this.setSecret(OAUTH_TOKENS_PREFIX, safeHostname, tokens); } - public async clearOAuthTokens(safeHostname: string): Promise { - await this.secrets.delete(`${OAUTH_TOKENS_PREFIX}${safeHostname}`); + public clearOAuthTokens(safeHostname: string): Promise { + return this.clearSecret(OAUTH_TOKENS_PREFIX, safeHostname); } - public async getOAuthClientRegistration( + public getOAuthClientRegistration( safeHostname: string, ): Promise { - try { - const data = await this.secrets.get( - `${OAUTH_CLIENT_PREFIX}${safeHostname}`, - ); - if (!data) { - return undefined; - } - return JSON.parse(data) as ClientRegistrationResponse; - } catch { - return undefined; - } + return this.getSecret( + OAUTH_CLIENT_PREFIX, + safeHostname, + ); } - public async setOAuthClientRegistration( + public setOAuthClientRegistration( safeHostname: string, registration: ClientRegistrationResponse, ): Promise { - await this.secrets.store( - `${OAUTH_CLIENT_PREFIX}${safeHostname}`, - JSON.stringify(registration), - ); - await this.recordDeploymentAccess(safeHostname); + return this.setSecret(OAUTH_CLIENT_PREFIX, safeHostname, registration); } - public async clearOAuthClientRegistration( - safeHostname: string, - ): Promise { - await this.secrets.delete(`${OAUTH_CLIENT_PREFIX}${safeHostname}`); + public clearOAuthClientRegistration(safeHostname: string): Promise { + return this.clearSecret(OAUTH_CLIENT_PREFIX, safeHostname); } public async clearOAuthData(safeHostname: string): Promise { diff --git a/src/login/loginCoordinator.ts b/src/login/loginCoordinator.ts index 2aa5f4ae..c34de54b 100644 --- a/src/login/loginCoordinator.ts +++ b/src/login/loginCoordinator.ts @@ -30,7 +30,7 @@ interface LoginOptions { * Coordinates login prompts across windows and prevents duplicate dialogs. */ export class LoginCoordinator { - private readonly inProgressLogins = new Map>(); + private loginQueue: Promise = Promise.resolve(); constructor( private readonly secretsManager: SecretsManager, @@ -47,7 +47,7 @@ export class LoginCoordinator { options: LoginOptions & { url: string }, ): Promise { const { safeHostname, url, oauthSessionManager } = options; - return this.executeWithGuard(safeHostname, async () => { + return this.executeWithGuard(async () => { const result = await this.attemptLogin( { safeHostname, url }, options.autoLogin ?? false, @@ -68,7 +68,7 @@ export class LoginCoordinator { ): Promise { const { safeHostname, url, detailPrefix, message, oauthSessionManager } = options; - return this.executeWithGuard(safeHostname, async () => { + return this.executeWithGuard(async () => { // Show dialog promise const dialogPromise = this.vscodeProposed.window .showErrorMessage( @@ -140,25 +140,14 @@ export class LoginCoordinator { } /** - * Same-window guard wrapper. + * Chains login attempts to prevent overlapping UI. */ - private async executeWithGuard( - safeHostname: string, + private executeWithGuard( executeFn: () => Promise, ): Promise { - const existingLogin = this.inProgressLogins.get(safeHostname); - if (existingLogin) { - return existingLogin; - } - - const loginPromise = executeFn(); - this.inProgressLogins.set(safeHostname, loginPromise); - - try { - return await loginPromise; - } finally { - this.inProgressLogins.delete(safeHostname); - } + const result = this.loginQueue.then(executeFn); + this.loginQueue = result.catch(() => {}); // Keep chain going on error + return result; } /** @@ -332,10 +321,10 @@ export class LoginCoordinator { { location: vscode.ProgressLocation.Notification, title: "Authenticating", - cancellable: false, + cancellable: true, }, - async (progress) => - await oauthSessionManager.login(client, deployment, progress), + async (progress, token) => + await oauthSessionManager.login(client, deployment, progress, token), ); // Validate token by fetching user diff --git a/src/oauth/sessionManager.ts b/src/oauth/sessionManager.ts index 2134336b..6a154ff3 100644 --- a/src/oauth/sessionManager.ts +++ b/src/oauth/sessionManager.ts @@ -28,10 +28,10 @@ import type { 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; +const AUTH_GRANT_TYPE = "authorization_code"; +const REFRESH_GRANT_TYPE = "refresh_token"; +const RESPONSE_TYPE = "code"; +const PKCE_CHALLENGE_METHOD = "S256"; /** * Token refresh threshold: refresh when token expires in less than this time. @@ -352,6 +352,7 @@ export class OAuthSessionManager implements vscode.Disposable { client: CoderApi, deployment: Deployment, progress: vscode.Progress<{ message?: string; increment?: number }>, + token: vscode.CancellationToken, ): Promise { const baseUrl = client.getAxiosInstance().defaults.baseURL; if (!baseUrl) { @@ -363,6 +364,13 @@ export class OAuthSessionManager implements vscode.Disposable { ); } + const reportProgress = (message?: string, increment?: number): void => { + if (token.isCancellationRequested) { + throw new Error("OAuth login cancelled by user"); + } + progress.report({ message, increment }); + }; + // Update deployment if changed if ( !this.deployment || @@ -377,21 +385,23 @@ export class OAuthSessionManager implements vscode.Disposable { this.deployment = deployment; } + reportProgress("fetching metadata...", 10); const axiosInstance = client.getAxiosInstance(); const metadataClient = new OAuthMetadataClient(axiosInstance, this.logger); const metadata = await metadataClient.getMetadata(); // Only register the client on login - progress.report({ message: "registering client...", increment: 10 }); + reportProgress("registering client...", 10); const registration = await this.registerClient(axiosInstance, metadata); - progress.report({ message: "waiting for authorization...", increment: 30 }); + reportProgress("waiting for authorization...", 30); const { code, verifier } = await this.startAuthorization( metadata, registration, + token, ); - progress.report({ message: "exchanging token...", increment: 30 }); + reportProgress("exchanging token...", 30); const tokenResponse = await this.exchangeToken( code, verifier, @@ -400,7 +410,6 @@ export class OAuthSessionManager implements vscode.Disposable { registration, ); - progress.report({ increment: 30 }); this.logger.info("OAuth login flow completed successfully"); return tokenResponse; @@ -457,6 +466,7 @@ export class OAuthSessionManager implements vscode.Disposable { private async startAuthorization( metadata: OAuthServerMetadata, registration: ClientRegistrationResponse, + cancellationToken: vscode.CancellationToken, ): Promise<{ code: string; verifier: string }> { const state = generateState(); const { verifier, challenge } = generatePKCE(); @@ -499,9 +509,17 @@ export class OAuthSessionManager implements vscode.Disposable { }, ); + const cancellationListener = cancellationToken.onCancellationRequested( + () => { + cleanup(); + reject(new Error("OAuth flow cancelled by user")); + }, + ); + const cleanup = () => { clearTimeout(timeoutHandle); listener.dispose(); + cancellationListener.dispose(); }; this.pendingAuthReject = (error) => { @@ -818,7 +836,7 @@ export class OAuthSessionManager implements vscode.Disposable { this.pendingAuthReject(new Error("OAuth session manager disposed")); } this.pendingAuthReject = undefined; - this.clearInMemoryTokens(); + this.clearDeployment(); this.logger.debug("OAuth session manager disposed"); } From 4c92aaa1e80f962050912ab89b7cdde9007d579e Mon Sep 17 00:00:00 2001 From: Ehab Younes Date: Thu, 18 Dec 2025 13:46:46 +0300 Subject: [PATCH 5/9] Proper OAuth error handling --- src/api/oauthInterceptors.ts | 26 +----- src/login/loginCoordinator.ts | 13 +-- src/oauth/sessionManager.ts | 166 ++++++++++++++++++++-------------- 3 files changed, 104 insertions(+), 101 deletions(-) diff --git a/src/api/oauthInterceptors.ts b/src/api/oauthInterceptors.ts index 5d04b33e..c37f99be 100644 --- a/src/api/oauthInterceptors.ts +++ b/src/api/oauthInterceptors.ts @@ -2,7 +2,6 @@ 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"; @@ -39,11 +38,7 @@ export function attachOAuthInterceptors( 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) { + if (status === 401) { return handle401Error(error, client, logger, oauthSessionManager); } @@ -52,23 +47,6 @@ export function attachOAuthInterceptors( ); } -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, @@ -100,8 +78,6 @@ async function handle401Error( throw error; } catch (refreshError) { logger.error("Token refresh failed:", refreshError); - - handlePossibleOAuthError(refreshError, logger, oauthSessionManager); throw error; } } diff --git a/src/login/loginCoordinator.ts b/src/login/loginCoordinator.ts index c34de54b..de3ff464 100644 --- a/src/login/loginCoordinator.ts +++ b/src/login/loginCoordinator.ts @@ -235,7 +235,7 @@ export class LoginCoordinator { const authMethod = await maybeAskAuthMethod(client); switch (authMethod) { case "oauth": - return this.loginWithOAuth(client, oauthSessionManager, deployment); + return this.loginWithOAuth(oauthSessionManager, deployment); case "legacy": { const result = await this.loginWithToken(client); if (result.success) { @@ -310,30 +310,25 @@ export class LoginCoordinator { * OAuth authentication flow. */ private async loginWithOAuth( - client: CoderApi, oauthSessionManager: OAuthSessionManager, deployment: Deployment, ): Promise { try { this.logger.info("Starting OAuth authentication"); - const tokenResponse = await vscode.window.withProgress( + const { token, user } = await vscode.window.withProgress( { location: vscode.ProgressLocation.Notification, title: "Authenticating", cancellable: true, }, async (progress, token) => - await oauthSessionManager.login(client, deployment, progress, token), + await oauthSessionManager.login(deployment, progress, token), ); - // Validate token by fetching user - client.setSessionToken(tokenResponse.access_token); - const user = await client.getAuthenticatedUser(); - return { success: true, - token: tokenResponse.access_token, + token, user, }; } catch (error) { diff --git a/src/oauth/sessionManager.ts b/src/oauth/sessionManager.ts index 6a154ff3..7680fb81 100644 --- a/src/oauth/sessionManager.ts +++ b/src/oauth/sessionManager.ts @@ -1,11 +1,22 @@ import { type AxiosInstance } from "axios"; +import { type User } from "coder/site/src/api/typesGenerated"; import * as vscode from "vscode"; import { CoderApi } from "../api/coderApi"; import { type ServiceContainer } from "../core/container"; +import { + type SecretsManager, + type StoredOAuthTokens, +} from "../core/secretsManager"; import { type Deployment } from "../deployment/types"; +import { type Logger } from "../logging/logger"; import { type LoginCoordinator } from "../login/loginCoordinator"; +import { + type OAuthError, + parseOAuthError, + requiresReAuthentication, +} from "./errors"; import { OAuthMetadataClient } from "./metadataClient"; import { CALLBACK_PATH, @@ -14,10 +25,6 @@ import { toUrlSearchParams, } from "./utils"; -import type { SecretsManager, StoredOAuthTokens } from "../core/secretsManager"; -import type { Logger } from "../logging/logger"; - -import type { OAuthError } from "./errors"; import type { ClientRegistrationRequest, ClientRegistrationResponse, @@ -283,30 +290,35 @@ export class OAuthSessionManager implements vscode.Disposable { 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, - ); + try { + 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( - deployment.safeHostname, - response.data, - ); - this.logger.info( - "Saved OAuth client registration:", - response.data.client_id, - ); + await this.secretsManager.setOAuthClientRegistration( + deployment.safeHostname, + response.data, + ); + this.logger.info( + "Saved OAuth client registration:", + response.data.client_id, + ); - return response.data; + return response.data; + } catch (error) { + this.handleOAuthError(error); + throw error; + } } public async setDeployment(deployment: Deployment): Promise { @@ -345,27 +357,14 @@ export class OAuthSessionManager implements vscode.Disposable { /** * 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, deployment: Deployment, progress: vscode.Progress<{ message?: string; increment?: number }>, - token: vscode.CancellationToken, - ): Promise { - const baseUrl = client.getAxiosInstance().defaults.baseURL; - if (!baseUrl) { - throw new Error("Client has no base URL set"); - } - if (baseUrl !== deployment.url) { - throw new Error( - `Client base URL (${baseUrl}) does not match deployment URL (${deployment.url})`, - ); - } - + cancellationToken: vscode.CancellationToken, + ): Promise<{ token: string; user: User }> { const reportProgress = (message?: string, increment?: number): void => { - if (token.isCancellationRequested) { + if (cancellationToken.isCancellationRequested) { throw new Error("OAuth login cancelled by user"); } progress.report({ message, increment }); @@ -385,8 +384,10 @@ export class OAuthSessionManager implements vscode.Disposable { this.deployment = deployment; } - reportProgress("fetching metadata...", 10); + const client = CoderApi.create(deployment.url, undefined, this.logger); const axiosInstance = client.getAxiosInstance(); + + reportProgress("fetching metadata...", 10); const metadataClient = new OAuthMetadataClient(axiosInstance, this.logger); const metadata = await metadataClient.getMetadata(); @@ -398,7 +399,7 @@ export class OAuthSessionManager implements vscode.Disposable { const { code, verifier } = await this.startAuthorization( metadata, registration, - token, + cancellationToken, ); reportProgress("exchanging token...", 30); @@ -410,9 +411,15 @@ export class OAuthSessionManager implements vscode.Disposable { registration, ); + reportProgress("fetching user...", 20); + const user = await client.getAuthenticatedUser(); + this.logger.info("OAuth login flow completed successfully"); - return tokenResponse; + return { + token: tokenResponse.access_token, + user, + }; } /** @@ -574,32 +581,37 @@ export class OAuthSessionManager implements vscode.Disposable { ): Promise { 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 response = await axiosInstance.post( - metadata.token_endpoint, - tokenRequest, - { - headers: { - "Content-Type": "application/x-www-form-urlencoded", + try { + 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 response = await axiosInstance.post( + metadata.token_endpoint, + tokenRequest, + { + headers: { + "Content-Type": "application/x-www-form-urlencoded", + }, }, - }, - ); + ); - this.logger.info("Token exchange successful"); + this.logger.info("Token exchange successful"); - await this.saveTokens(response.data); + await this.saveTokens(response.data); - return response.data; + return response.data; + } catch (error) { + this.handleOAuthError(error); + throw error; + } } /** @@ -656,6 +668,9 @@ export class OAuthSessionManager implements vscode.Disposable { await this.saveTokens(response.data); return response.data; + } catch (error) { + this.handleOAuthError(error); + throw error; } finally { this.refreshPromise = null; } @@ -801,6 +816,23 @@ export class OAuthSessionManager implements vscode.Disposable { } } + /** + * Handle OAuth errors that may require re-authentication. + * Parses the error and triggers re-authentication modal if needed. + */ + private handleOAuthError(error: unknown): void { + const oauthError = parseOAuthError(error); + if (oauthError && requiresReAuthentication(oauthError)) { + this.logger.error( + `OAuth operation failed with error: ${oauthError.errorCode}`, + ); + // Fire and forget - don't block on showing the modal + this.showReAuthenticationModal(oauthError).catch((err) => { + this.logger.error("Failed to show re-auth modal:", err); + }); + } + } + /** * 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. From 4df24eb398434b651454aae0716be0ea6feac791 Mon Sep 17 00:00:00 2001 From: Ehab Younes Date: Thu, 18 Dec 2025 15:59:10 +0300 Subject: [PATCH 6/9] Remove in memory token storage and rely on the secret storage --- src/api/oauthInterceptors.ts | 2 +- src/extension.ts | 2 +- src/oauth/sessionManager.ts | 128 +++++++++++++++++++---------------- src/remote/remote.ts | 2 +- 4 files changed, 71 insertions(+), 63 deletions(-) diff --git a/src/api/oauthInterceptors.ts b/src/api/oauthInterceptors.ts index c37f99be..edd65f83 100644 --- a/src/api/oauthInterceptors.ts +++ b/src/api/oauthInterceptors.ts @@ -53,7 +53,7 @@ async function handle401Error( logger: Logger, oauthSessionManager: OAuthSessionManager, ): Promise { - if (!oauthSessionManager.isLoggedInWithOAuth()) { + if (!(await oauthSessionManager.isLoggedInWithOAuth())) { throw error; } diff --git a/src/extension.ts b/src/extension.ts index 23dec97c..29690ce3 100644 --- a/src/extension.ts +++ b/src/extension.ts @@ -72,7 +72,7 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { const deployment = await secretsManager.getCurrentDeployment(); // Create OAuth session manager with login coordinator - const oauthSessionManager = await OAuthSessionManager.create( + const oauthSessionManager = OAuthSessionManager.create( deployment, serviceContainer, ctx.extension.id, diff --git a/src/oauth/sessionManager.ts b/src/oauth/sessionManager.ts index 7680fb81..66e9ab24 100644 --- a/src/oauth/sessionManager.ts +++ b/src/oauth/sessionManager.ts @@ -78,7 +78,6 @@ const DEFAULT_OAUTH_SCOPES = [ * Coordinates authorization flow, token management, and automatic refresh. */ export class OAuthSessionManager implements vscode.Disposable { - private storedTokens: StoredOAuthTokens | undefined; private refreshPromise: Promise | null = null; private lastRefreshAttempt = 0; private refreshTimer: NodeJS.Timeout | undefined; @@ -88,11 +87,11 @@ export class OAuthSessionManager implements vscode.Disposable { /** * Create and initialize a new OAuth session manager. */ - public static async create( + public static create( deployment: Deployment | null, container: ServiceContainer, extensionId: string, - ): Promise { + ): OAuthSessionManager { const manager = new OAuthSessionManager( deployment, container.getSecretsManager(), @@ -100,7 +99,6 @@ export class OAuthSessionManager implements vscode.Disposable { container.getLoginCoordinator(), extensionId, ); - await manager.loadTokens(); manager.scheduleBackgroundRefresh(); return manager; } @@ -125,52 +123,48 @@ export class OAuthSessionManager implements vscode.Disposable { } /** - * Load stored tokens from storage. - * No-op if deployment is not set. - * Validates that tokens belong to the current deployment URL. + * Get stored tokens fresh from secrets manager. + * Always reads from storage to ensure cross-window synchronization. + * Validates that tokens match current deployment URL and have required scopes. + * Invalid tokens are cleared and undefined is returned. */ - private async loadTokens(): Promise { + private async getStoredTokens(): Promise { if (!this.deployment) { - return; + return undefined; } const tokens = await this.secretsManager.getOAuthTokens( this.deployment.safeHostname, ); if (!tokens) { - return; + return undefined; } + // Validate deployment URL matches if (tokens.deployment_url !== this.deployment.url) { - this.logger.warn("Stored tokens for different deployment, clearing", { - stored: tokens.deployment_url, - current: this.deployment.url, - }); - this.clearInMemoryTokens(); - await this.secretsManager.clearOAuthData(this.deployment.safeHostname); - return; + this.logger.warn( + "Stored tokens have mismatched deployment URL, clearing", + { stored: tokens.deployment_url, current: this.deployment.url }, + ); + await this.secretsManager.clearOAuthTokens(this.deployment.safeHostname); + return undefined; } 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.clearOAuthState(); - return; + this.logger.warn("Stored tokens have insufficient scopes, clearing", { + scope: tokens.scope, + }); + await this.secretsManager.clearOAuthTokens(this.deployment.safeHostname); + return undefined; } - this.storedTokens = tokens; - this.logger.info( - `Loaded stored OAuth tokens for ${this.deployment.safeHostname}`, - ); + return tokens; } - private clearInMemoryTokens(): void { - this.storedTokens = undefined; + /** + * Clear refresh-related state. + */ + private clearRefreshState(): void { this.refreshPromise = null; this.lastRefreshAttempt = 0; } @@ -331,11 +325,10 @@ export class OAuthSessionManager implements vscode.Disposable { } this.logger.debug("Switching OAuth deployment", deployment); this.deployment = deployment; - this.clearInMemoryTokens(); - await this.loadTokens(); + this.clearRefreshState(); // Refresh tokens if needed to prevent 401s - if (this.isTokenExpired()) { + if (await this.isTokenExpired()) { try { await this.refreshToken(); } catch (error) { @@ -351,7 +344,7 @@ export class OAuthSessionManager implements vscode.Disposable { public clearDeployment(): void { this.logger.debug("Clearing OAuth deployment state"); this.deployment = null; - this.clearInMemoryTokens(); + this.clearRefreshState(); } /** @@ -380,7 +373,7 @@ export class OAuthSessionManager implements vscode.Disposable { old: this.deployment, new: deployment, }); - this.clearInMemoryTokens(); + this.clearRefreshState(); this.deployment = deployment; } @@ -627,12 +620,14 @@ export class OAuthSessionManager implements vscode.Disposable { return this.refreshPromise; } - if (!this.storedTokens?.refresh_token) { + // Read fresh tokens from secrets + const storedTokens = await this.getStoredTokens(); + if (!storedTokens?.refresh_token) { throw new Error("No refresh token available"); } - const refreshToken = this.storedTokens.refresh_token; - const accessToken = this.storedTokens.access_token; + const refreshToken = storedTokens.refresh_token; + const accessToken = storedTokens.access_token; this.lastRefreshAttempt = Date.now(); @@ -681,7 +676,7 @@ export class OAuthSessionManager implements vscode.Disposable { /** * Save token response to storage. - * Also triggers event via secretsManager to update global client. + * Writes to secrets manager only - no in-memory caching. */ private async saveTokens(tokenResponse: TokenResponse): Promise { const deployment = this.requireDeployment(); @@ -695,7 +690,6 @@ export class OAuthSessionManager implements vscode.Disposable { expiry_timestamp: expiryTimestamp, }; - this.storedTokens = tokens; await this.secretsManager.setOAuthTokens(deployment.safeHostname, tokens); await this.secretsManager.setSessionAuth(deployment.safeHostname, { url: deployment.url, @@ -712,7 +706,7 @@ export class OAuthSessionManager implements vscode.Disposable { * Refreshes the token if it is approaching expiry. */ public async refreshIfAlmostExpired(): Promise { - if (this.shouldRefreshToken()) { + if (await this.shouldRefreshToken()) { this.logger.debug("Token approaching expiry, triggering refresh"); await this.refreshToken(); } @@ -726,8 +720,9 @@ export class OAuthSessionManager implements vscode.Disposable { * 3. Last refresh attempt was more than REFRESH_THROTTLE_MS ago * 4. No refresh is currently in progress */ - private shouldRefreshToken(): boolean { - if (!this.storedTokens?.refresh_token || this.refreshPromise !== null) { + private async shouldRefreshToken(): Promise { + const storedTokens = await this.getStoredTokens(); + if (!storedTokens?.refresh_token || this.refreshPromise !== null) { return false; } @@ -736,38 +731,49 @@ export class OAuthSessionManager implements vscode.Disposable { return false; } - const timeUntilExpiry = this.storedTokens.expiry_timestamp - now; + const timeUntilExpiry = storedTokens.expiry_timestamp - now; return timeUntilExpiry < TOKEN_REFRESH_THRESHOLD_MS; } /** * Check if token is expired. */ - private isTokenExpired(): boolean { - if (!this.storedTokens) { + private async isTokenExpired(): Promise { + const storedTokens = await this.getStoredTokens(); + if (!storedTokens) { return false; } - return Date.now() >= this.storedTokens.expiry_timestamp; + return Date.now() >= storedTokens.expiry_timestamp; } public async revokeRefreshToken(): Promise { - if (!this.storedTokens?.refresh_token) { + const storedTokens = await this.getStoredTokens(); + if (!storedTokens?.refresh_token) { this.logger.info("No refresh token to revoke"); return; } - await this.revokeToken(this.storedTokens.refresh_token, "refresh_token"); + await this.revokeToken( + storedTokens.access_token, + storedTokens.refresh_token, + "refresh_token", + ); } /** * Revoke a token using the OAuth server's revocation endpoint. + * + * @param authToken - Token for authenticating the revocation request + * @param tokenToRevoke - The token to be revoked + * @param tokenTypeHint - Hint about the token type being revoked */ private async revokeToken( - token: string, + authToken: string, + tokenToRevoke: string, tokenTypeHint: "access_token" | "refresh_token" = "refresh_token", ): Promise { const { axiosInstance, metadata, registration } = - await this.prepareOAuthOperation(this.storedTokens?.access_token); + await this.prepareOAuthOperation(authToken); const revocationEndpoint = metadata.revocation_endpoint || `${metadata.issuer}/oauth2/revoke`; @@ -775,7 +781,7 @@ export class OAuthSessionManager implements vscode.Disposable { this.logger.info("Revoking refresh token"); const params: TokenRevocationRequest = { - token, + token: tokenToRevoke, client_id: registration.client_id, client_secret: registration.client_secret, token_type_hint: tokenTypeHint, @@ -798,19 +804,21 @@ export class OAuthSessionManager implements vscode.Disposable { } /** - * Returns true if (valid or invalid) OAuth tokens exist for the current deployment. + * Returns true if OAuth tokens exist for the current deployment. + * Always reads fresh from secrets to ensure cross-window synchronization. */ - public isLoggedInWithOAuth(): boolean { - return this.storedTokens !== undefined; + public async isLoggedInWithOAuth(): Promise { + const storedTokens = await this.getStoredTokens(); + return storedTokens !== undefined; } /** * Clear OAuth state when switching to non-OAuth authentication. - * Clears in-memory state and OAuth tokens from storage. + * Clears tokens from storage. * Preserves client registration for potential future OAuth use. */ public async clearOAuthState(): Promise { - this.clearInMemoryTokens(); + this.clearRefreshState(); if (this.deployment) { await this.secretsManager.clearOAuthTokens(this.deployment.safeHostname); } @@ -845,7 +853,7 @@ export class OAuthSessionManager implements vscode.Disposable { "Your session is no longer valid. This could be due to token expiration or revocation."; // Clear invalid tokens - listeners will handle updates automatically - this.clearInMemoryTokens(); + this.clearRefreshState(); await this.secretsManager.clearAllAuthData(deployment.safeHostname); await this.loginCoordinator.ensureLoggedInWithDialog({ diff --git a/src/remote/remote.ts b/src/remote/remote.ts index 0bbb4a51..b2257270 100644 --- a/src/remote/remote.ts +++ b/src/remote/remote.ts @@ -113,7 +113,7 @@ export class Remote { try { // Create OAuth session manager for this remote deployment - const remoteOAuthManager = await OAuthSessionManager.create( + const remoteOAuthManager = OAuthSessionManager.create( { url: baseUrlRaw, safeHostname: parts.safeHostname }, this.serviceContainer, this.extensionContext.extension.id, From 3570c64f133027e38ef8a7c2aaeb51b9d625501a Mon Sep 17 00:00:00 2001 From: Ehab Younes Date: Fri, 19 Dec 2025 13:05:04 +0300 Subject: [PATCH 7/9] Attach and detach OAuth interceptor when authenticated using OAuth or session token --- src/api/oauthInterceptor.ts | 148 +++++++++++++++++++++++++++++++++++ src/api/oauthInterceptors.ts | 83 -------------------- src/core/secretsManager.ts | 21 +++++ src/extension.ts | 13 ++- src/oauth/sessionManager.ts | 4 +- src/remote/remote.ts | 14 +++- 6 files changed, 194 insertions(+), 89 deletions(-) create mode 100644 src/api/oauthInterceptor.ts delete mode 100644 src/api/oauthInterceptors.ts diff --git a/src/api/oauthInterceptor.ts b/src/api/oauthInterceptor.ts new file mode 100644 index 00000000..2045d56b --- /dev/null +++ b/src/api/oauthInterceptor.ts @@ -0,0 +1,148 @@ +import { type AxiosError, isAxiosError } from "axios"; + +import type * as vscode from "vscode"; + +import type { SecretsManager } from "../core/secretsManager"; +import type { Logger } from "../logging/logger"; +import type { RequestConfigWithMeta } from "../logging/types"; +import type { OAuthSessionManager } from "../oauth/sessionManager"; + +import type { CoderApi } from "./coderApi"; + +const coderSessionTokenHeader = "Coder-Session-Token"; + +/** + * Manages OAuth interceptor lifecycle reactively based on token presence. + * + * Automatically attaches/detaches the interceptor when OAuth tokens appear/disappear + * in secrets storage. This ensures the interceptor state always matches the actual + * OAuth authentication state. + */ +export class OAuthInterceptor implements vscode.Disposable { + private interceptorId: number | null = null; + + private constructor( + private readonly client: CoderApi, + private readonly logger: Logger, + private readonly oauthSessionManager: OAuthSessionManager, + private readonly tokenListener: vscode.Disposable, + ) {} + + public static async create( + client: CoderApi, + logger: Logger, + oauthSessionManager: OAuthSessionManager, + secretsManager: SecretsManager, + safeHostname: string, + ): Promise { + // Create listener first, then wire up to instance after construction + let callback: () => Promise = () => Promise.resolve(); + const tokenListener = secretsManager.onDidChangeOAuthTokens( + safeHostname, + () => callback(), + ); + + const instance = new OAuthInterceptor( + client, + logger, + oauthSessionManager, + tokenListener, + ); + + callback = async () => + instance.syncWithTokenState().catch((err) => { + logger.error("Error syncing OAuth interceptor state:", err); + }); + await instance.syncWithTokenState(); + return instance; + } + + /** + * Sync interceptor state with OAuth token presence. + * Attaches when tokens exist, detaches when they don't. + */ + private async syncWithTokenState(): Promise { + const isOAuth = await this.oauthSessionManager.isLoggedInWithOAuth(); + if (isOAuth && this.interceptorId === null) { + this.attach(); + } else if (!isOAuth && this.interceptorId !== null) { + this.detach(); + } + } + + private attach(): void { + if (this.interceptorId !== null) { + return; + } + + this.interceptorId = this.client + .getAxiosInstance() + .interceptors.response.use( + (r) => r, + (error: unknown) => this.handleError(error), + ); + + this.logger.debug("OAuth interceptor attached"); + } + + private detach(): void { + if (this.interceptorId === null) { + return; + } + + this.client + .getAxiosInstance() + .interceptors.response.eject(this.interceptorId); + this.interceptorId = null; + this.logger.debug("OAuth interceptor detached"); + } + + private async handleError(error: unknown): Promise { + if (!isAxiosError(error)) { + throw error; + } + + if (error.config) { + const config = error.config as { _oauthRetryAttempted?: boolean }; + if (config._oauthRetryAttempted) { + throw error; + } + } + + if (error.response?.status === 401) { + return this.handle401Error(error); + } + + throw error; + } + + private async handle401Error(error: AxiosError): Promise { + this.logger.info("Received 401 response, attempting token refresh"); + + try { + const newTokens = await this.oauthSessionManager.refreshToken(); + this.client.setSessionToken(newTokens.access_token); + + this.logger.info("Token refresh successful, retrying request"); + + if (error.config) { + const config = error.config as RequestConfigWithMeta & { + _oauthRetryAttempted?: boolean; + }; + config._oauthRetryAttempted = true; + config.headers[coderSessionTokenHeader] = newTokens.access_token; + return this.client.getAxiosInstance().request(config); + } + + throw error; + } catch (refreshError) { + this.logger.error("Token refresh failed:", refreshError); + throw error; + } + } + + public dispose(): void { + this.tokenListener.dispose(); + this.detach(); + } +} diff --git a/src/api/oauthInterceptors.ts b/src/api/oauthInterceptors.ts deleted file mode 100644 index edd65f83..00000000 --- a/src/api/oauthInterceptors.ts +++ /dev/null @@ -1,83 +0,0 @@ -import { type AxiosError, isAxiosError } from "axios"; - -import { type Logger } from "../logging/logger"; -import { type RequestConfigWithMeta } from "../logging/types"; -import { type OAuthSessionManager } from "../oauth/sessionManager"; - -import { type CoderApi } from "./coderApi"; - -const coderSessionTokenHeader = "Coder-Session-Token"; - -/** - * Attach OAuth token refresh interceptor to a CoderApi instance. - * This should be called after creating the CoderApi when OAuth authentication is being used. - * - * Error interceptor: reactively refreshes token on 401 responses. - */ -export function attachOAuthInterceptors( - client: CoderApi, - logger: Logger, - oauthSessionManager: OAuthSessionManager, -): void { - client.getAxiosInstance().interceptors.response.use( - (r) => r, - // 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; - - if (status === 401) { - return handle401Error(error, client, logger, oauthSessionManager); - } - - throw error; - }, - ); -} - -async function handle401Error( - error: AxiosError, - client: CoderApi, - logger: Logger, - oauthSessionManager: OAuthSessionManager, -): Promise { - if (!(await 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); - throw error; - } -} diff --git a/src/core/secretsManager.ts b/src/core/secretsManager.ts index db339a08..70870e8a 100644 --- a/src/core/secretsManager.ts +++ b/src/core/secretsManager.ts @@ -333,6 +333,27 @@ export class SecretsManager { return this.clearSecret(OAUTH_TOKENS_PREFIX, safeHostname); } + /** + * Listen for changes to OAuth tokens for a specific deployment. + */ + public onDidChangeOAuthTokens( + safeHostname: string, + listener: (tokens: StoredOAuthTokens | undefined) => void | Promise, + ): Disposable { + const tokenKey = this.buildKey(OAUTH_TOKENS_PREFIX, safeHostname); + return this.secrets.onDidChange(async (e) => { + if (e.key !== tokenKey) { + return; + } + const tokens = await this.getOAuthTokens(safeHostname); + try { + await listener(tokens); + } catch (err) { + this.logger.error("Error in onDidChangeOAuthTokens listener", err); + } + }); + } + public getOAuthClientRegistration( safeHostname: string, ): Promise { diff --git a/src/extension.ts b/src/extension.ts index 29690ce3..34500337 100644 --- a/src/extension.ts +++ b/src/extension.ts @@ -8,7 +8,7 @@ import * as vscode from "vscode"; import { errToStr } from "./api/api-helper"; import { CoderApi } from "./api/coderApi"; -import { attachOAuthInterceptors } from "./api/oauthInterceptors"; +import { OAuthInterceptor } from "./api/oauthInterceptor"; import { Commands } from "./commands"; import { ServiceContainer } from "./core/container"; import { type SecretsManager } from "./core/secretsManager"; @@ -89,7 +89,16 @@ export async function activate(ctx: vscode.ExtensionContext): Promise { output, ); ctx.subscriptions.push(client); - attachOAuthInterceptors(client, output, oauthSessionManager); + + // Create OAuth interceptor - auto attaches/detaches based on token state + const oauthInterceptor = await OAuthInterceptor.create( + client, + output, + oauthSessionManager, + secretsManager, + deployment?.safeHostname ?? "", + ); + ctx.subscriptions.push(oauthInterceptor); const myWorkspacesProvider = new WorkspaceProvider( WorkspaceQuery.Mine, diff --git a/src/oauth/sessionManager.ts b/src/oauth/sessionManager.ts index 66e9ab24..2573e57d 100644 --- a/src/oauth/sessionManager.ts +++ b/src/oauth/sessionManager.ts @@ -852,9 +852,9 @@ export class OAuthSessionManager implements vscode.Disposable { 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 this.clearRefreshState(); - await this.secretsManager.clearAllAuthData(deployment.safeHostname); + // Clear client registration and tokens to force full re-authentication + await this.secretsManager.clearOAuthData(deployment.safeHostname); await this.loginCoordinator.ensureLoggedInWithDialog({ safeHostname: deployment.safeHostname, diff --git a/src/remote/remote.ts b/src/remote/remote.ts index b2257270..70dacfdf 100644 --- a/src/remote/remote.ts +++ b/src/remote/remote.ts @@ -20,7 +20,7 @@ import { } from "../api/agentMetadataHelper"; import { extractAgents } from "../api/api-helper"; import { CoderApi } from "../api/coderApi"; -import { attachOAuthInterceptors } from "../api/oauthInterceptors"; +import { OAuthInterceptor } from "../api/oauthInterceptor"; import { needToken } from "../api/utils"; import { getGlobalFlags, getGlobalFlagsRaw, getSshFlags } from "../cliConfig"; import { type Commands } from "../commands"; @@ -165,7 +165,17 @@ export class Remote { // 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); + + // Create OAuth interceptor - auto attaches/detaches based on token state + const oauthInterceptor = await OAuthInterceptor.create( + workspaceClient, + this.logger, + remoteOAuthManager, + this.secretsManager, + parts.safeHostname, + ); + disposables.push(oauthInterceptor); + // Store for use in commands. this.commands.remoteWorkspaceClient = workspaceClient; From 76f089fdcf9b16967ba052875930be7179fd5e71 Mon Sep 17 00:00:00 2001 From: Ehab Younes Date: Fri, 19 Dec 2025 14:17:16 +0300 Subject: [PATCH 8/9] Remove refreshIfAlmostExpired and replace it with smarter timers --- src/oauth/sessionManager.ts | 123 ++++++++++++++++++++++++++---------- 1 file changed, 89 insertions(+), 34 deletions(-) diff --git a/src/oauth/sessionManager.ts b/src/oauth/sessionManager.ts index 2573e57d..e5c32141 100644 --- a/src/oauth/sessionManager.ts +++ b/src/oauth/sessionManager.ts @@ -81,6 +81,7 @@ export class OAuthSessionManager implements vscode.Disposable { private refreshPromise: Promise | null = null; private lastRefreshAttempt = 0; private refreshTimer: NodeJS.Timeout | undefined; + private tokenChangeListener: vscode.Disposable | undefined; private pendingAuthReject: ((reason: Error) => void) | undefined; @@ -99,7 +100,8 @@ export class OAuthSessionManager implements vscode.Disposable { container.getLoginCoordinator(), extensionId, ); - manager.scheduleBackgroundRefresh(); + manager.setupTokenListener(); + manager.scheduleNextRefresh(); return manager; } @@ -162,30 +164,96 @@ export class OAuthSessionManager implements vscode.Disposable { } /** - * Clear refresh-related state. + * Clear all refresh-related state: in-flight promise, throttle, timer, and listener. */ private clearRefreshState(): void { this.refreshPromise = null; this.lastRefreshAttempt = 0; + if (this.refreshTimer) { + clearTimeout(this.refreshTimer); + this.refreshTimer = undefined; + } + this.tokenChangeListener?.dispose(); + this.tokenChangeListener = undefined; + } + + /** + * Setup listener for token changes. Disposes existing listener first. + * Reschedules refresh when tokens change (e.g., from another window). + */ + private setupTokenListener(): void { + this.tokenChangeListener?.dispose(); + this.tokenChangeListener = undefined; + + if (!this.deployment) { + return; + } + + this.tokenChangeListener = this.secretsManager.onDidChangeOAuthTokens( + this.deployment.safeHostname, + () => this.scheduleNextRefresh(), + ); } /** - * Schedule the next background token refresh check. - * Only schedules the next check after the current one completes. + * Schedule the next token refresh based on expiry time. + * - Far from expiry: schedule once at threshold + * - Near/past expiry: attempt refresh immediately */ - private scheduleBackgroundRefresh(): void { + private scheduleNextRefresh(): void { if (this.refreshTimer) { clearTimeout(this.refreshTimer); + this.refreshTimer = undefined; } - this.refreshTimer = setTimeout(async () => { - try { - await this.refreshIfAlmostExpired(); - } catch (error) { - this.logger.warn("Background token refresh failed:", error); - } - this.scheduleBackgroundRefresh(); - }, BACKGROUND_REFRESH_INTERVAL_MS); + this.getStoredTokens() + .then((storedTokens) => { + if (!storedTokens?.refresh_token) { + return; + } + + const now = Date.now(); + const timeUntilExpiry = storedTokens.expiry_timestamp - now; + + if (timeUntilExpiry <= TOKEN_REFRESH_THRESHOLD_MS) { + // Within threshold or expired, attempt refresh now + this.attemptRefreshWithRetry(); + } else { + // Schedule for when we reach the threshold + const delay = timeUntilExpiry - TOKEN_REFRESH_THRESHOLD_MS; + this.logger.debug( + `Scheduling token refresh in ${Math.round(delay / 1000 / 60)} minutes`, + ); + this.refreshTimer = setTimeout( + () => this.attemptRefreshWithRetry(), + delay, + ); + } + }) + .catch((error) => { + this.logger.warn("Failed to schedule token refresh:", error); + }); + } + + /** + * Attempt refresh, falling back to polling on failure. + */ + private attemptRefreshWithRetry(): void { + this.refreshTimer = undefined; + + this.refreshToken() + .then(() => { + // Success - scheduleNextRefresh will be triggered by token change listener + this.logger.debug("Background token refresh succeeded"); + }) + .catch((error) => { + this.logger.warn("Background token refresh failed, will retry:", error); + // Fall back to polling until successful + this.refreshTimer = setTimeout( + () => this.attemptRefreshWithRetry(), + BACKGROUND_REFRESH_INTERVAL_MS, + ); + }); } /** @@ -327,18 +395,19 @@ export class OAuthSessionManager implements vscode.Disposable { this.deployment = deployment; this.clearRefreshState(); - // Refresh tokens if needed to prevent 401s - if (await this.isTokenExpired()) { + // Block on refresh if token is expired to ensure valid state for callers + const storedTokens = await this.getStoredTokens(); + if (storedTokens && Date.now() >= storedTokens.expiry_timestamp) { try { await this.refreshToken(); } catch (error) { this.logger.warn("Token refresh failed (expired):", error); } - } else { - this.refreshIfAlmostExpired().catch((error) => { - this.logger.warn("Background token refresh failed:", error); - }); } + + // Schedule after blocking refresh to avoid concurrent attempts + this.setupTokenListener(); + this.scheduleNextRefresh(); } public clearDeployment(): void { @@ -375,6 +444,7 @@ export class OAuthSessionManager implements vscode.Disposable { }); this.clearRefreshState(); this.deployment = deployment; + this.setupTokenListener(); } const client = CoderApi.create(deployment.url, undefined, this.logger); @@ -735,17 +805,6 @@ export class OAuthSessionManager implements vscode.Disposable { return timeUntilExpiry < TOKEN_REFRESH_THRESHOLD_MS; } - /** - * Check if token is expired. - */ - private async isTokenExpired(): Promise { - const storedTokens = await this.getStoredTokens(); - if (!storedTokens) { - return false; - } - return Date.now() >= storedTokens.expiry_timestamp; - } - public async revokeRefreshToken(): Promise { const storedTokens = await this.getStoredTokens(); if (!storedTokens?.refresh_token) { @@ -868,10 +927,6 @@ 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")); } From 3d1440b590c394ab746bed8716b433a4c54221fc Mon Sep 17 00:00:00 2001 From: Ehab Younes Date: Fri, 19 Dec 2025 16:42:03 +0300 Subject: [PATCH 9/9] Combined OAuth tokens with session auth --- src/api/oauthInterceptor.ts | 2 +- src/core/secretsManager.ts | 84 ++++++++++--------------------------- src/extension.ts | 2 - src/oauth/sessionManager.ts | 83 ++++++++++++++++++++++++++---------- 4 files changed, 85 insertions(+), 86 deletions(-) diff --git a/src/api/oauthInterceptor.ts b/src/api/oauthInterceptor.ts index 2045d56b..6d777739 100644 --- a/src/api/oauthInterceptor.ts +++ b/src/api/oauthInterceptor.ts @@ -37,7 +37,7 @@ export class OAuthInterceptor implements vscode.Disposable { ): Promise { // Create listener first, then wire up to instance after construction let callback: () => Promise = () => Promise.resolve(); - const tokenListener = secretsManager.onDidChangeOAuthTokens( + const tokenListener = secretsManager.onDidChangeSessionAuth( safeHostname, () => callback(), ); diff --git a/src/core/secretsManager.ts b/src/core/secretsManager.ts index 70870e8a..e41d6201 100644 --- a/src/core/secretsManager.ts +++ b/src/core/secretsManager.ts @@ -1,8 +1,5 @@ import { type Logger } from "../logging/logger"; -import { - type ClientRegistrationResponse, - type TokenResponse, -} from "../oauth/types"; +import { type ClientRegistrationResponse } from "../oauth/types"; import { toSafeHost } from "../util"; import type { Memento, SecretStorage, Disposable } from "vscode"; @@ -12,13 +9,9 @@ import type { Deployment } from "../deployment/types"; // Each deployment has its own key to ensure atomic operations (multiple windows // writing to a shared key could drop data) and to receive proper VS Code events. const SESSION_KEY_PREFIX = "coder.session." as const; -const OAUTH_TOKENS_PREFIX = "coder.oauth.tokens." as const; const OAUTH_CLIENT_PREFIX = "coder.oauth.client." as const; -type SecretKeyPrefix = - | typeof SESSION_KEY_PREFIX - | typeof OAUTH_TOKENS_PREFIX - | typeof OAUTH_CLIENT_PREFIX; +type SecretKeyPrefix = typeof SESSION_KEY_PREFIX | typeof OAUTH_CLIENT_PREFIX; const OAUTH_CALLBACK_KEY = "coder.oauthCallback"; @@ -33,9 +26,22 @@ export interface CurrentDeploymentState { deployment: Deployment | null; } +/** + * OAuth token data stored alongside session auth. + * When present, indicates the session is authenticated via OAuth. + */ +export interface OAuthTokenData { + token_type: "Bearer" | "DPoP"; + refresh_token?: string; + scope?: string; + expiry_timestamp: number; +} + export interface SessionAuth { url: string; token: string; + /** If present, this session uses OAuth authentication */ + oauth?: OAuthTokenData; } // Tracks when a deployment was last accessed for LRU pruning. @@ -44,11 +50,6 @@ interface DeploymentUsage { lastAccessedAt: string; } -export type StoredOAuthTokens = Omit & { - expiry_timestamp: number; - deployment_url: string; -}; - interface OAuthCallbackData { state: string; code: string | null; @@ -190,8 +191,12 @@ export class SecretsManager { safeHostname: string, auth: SessionAuth, ): Promise { - // Extract only url and token before serializing - const state: SessionAuth = { url: auth.url, token: auth.token }; + // Extract relevant fields before serializing + const state: SessionAuth = { + url: auth.url, + token: auth.token, + ...(auth.oauth && { oauth: auth.oauth }), + }; await this.setSecret(SESSION_KEY_PREFIX, safeHostname, state); } @@ -229,7 +234,7 @@ export class SecretsManager { public async clearAllAuthData(safeHostname: string): Promise { await Promise.all([ this.clearSessionAuth(safeHostname), - this.clearOAuthData(safeHostname), + this.clearOAuthClientRegistration(safeHostname), ]); const usage = this.getDeploymentUsage().filter( (u) => u.safeHostname !== safeHostname, @@ -316,44 +321,6 @@ export class SecretsManager { }); } - public getOAuthTokens( - safeHostname: string, - ): Promise { - return this.getSecret(OAUTH_TOKENS_PREFIX, safeHostname); - } - - public setOAuthTokens( - safeHostname: string, - tokens: StoredOAuthTokens, - ): Promise { - return this.setSecret(OAUTH_TOKENS_PREFIX, safeHostname, tokens); - } - - public clearOAuthTokens(safeHostname: string): Promise { - return this.clearSecret(OAUTH_TOKENS_PREFIX, safeHostname); - } - - /** - * Listen for changes to OAuth tokens for a specific deployment. - */ - public onDidChangeOAuthTokens( - safeHostname: string, - listener: (tokens: StoredOAuthTokens | undefined) => void | Promise, - ): Disposable { - const tokenKey = this.buildKey(OAUTH_TOKENS_PREFIX, safeHostname); - return this.secrets.onDidChange(async (e) => { - if (e.key !== tokenKey) { - return; - } - const tokens = await this.getOAuthTokens(safeHostname); - try { - await listener(tokens); - } catch (err) { - this.logger.error("Error in onDidChangeOAuthTokens listener", err); - } - }); - } - public getOAuthClientRegistration( safeHostname: string, ): Promise { @@ -373,11 +340,4 @@ export class SecretsManager { public clearOAuthClientRegistration(safeHostname: string): Promise { return this.clearSecret(OAUTH_CLIENT_PREFIX, safeHostname); } - - public async clearOAuthData(safeHostname: string): Promise { - await Promise.all([ - this.clearOAuthTokens(safeHostname), - this.clearOAuthClientRegistration(safeHostname), - ]); - } } diff --git a/src/extension.ts b/src/extension.ts index 34500337..5e9f5667 100644 --- a/src/extension.ts +++ b/src/extension.ts @@ -495,8 +495,6 @@ async function setupDeploymentFromUri( } } else { await secretsManager.setSessionAuth(safeHost, { url, token }); - // Clear OAuth tokens since we're using a non-OAuth token - await secretsManager.clearOAuthTokens(safeHost); } } diff --git a/src/oauth/sessionManager.ts b/src/oauth/sessionManager.ts index e5c32141..27b79f62 100644 --- a/src/oauth/sessionManager.ts +++ b/src/oauth/sessionManager.ts @@ -5,8 +5,9 @@ import * as vscode from "vscode"; import { CoderApi } from "../api/coderApi"; import { type ServiceContainer } from "../core/container"; import { + type OAuthTokenData, type SecretsManager, - type StoredOAuthTokens, + type SessionAuth, } from "../core/secretsManager"; import { type Deployment } from "../deployment/types"; import { type Logger } from "../logging/logger"; @@ -73,6 +74,14 @@ const DEFAULT_OAUTH_SCOPES = [ "user:read_personal", ].join(" "); +/** + * Internal type combining access token with OAuth-specific data. + * Used by getStoredTokens() for token refresh and validation. + */ +type StoredTokens = OAuthTokenData & { + access_token: string; +}; + /** * Manages OAuth session lifecycle for a Coder deployment. * Coordinates authorization flow, token management, and automatic refresh. @@ -130,37 +139,53 @@ export class OAuthSessionManager implements vscode.Disposable { * Validates that tokens match current deployment URL and have required scopes. * Invalid tokens are cleared and undefined is returned. */ - private async getStoredTokens(): Promise { + private async getStoredTokens(): Promise { if (!this.deployment) { return undefined; } - const tokens = await this.secretsManager.getOAuthTokens( + const auth = await this.secretsManager.getSessionAuth( this.deployment.safeHostname, ); - if (!tokens) { + if (!auth?.oauth) { return undefined; } // Validate deployment URL matches - if (tokens.deployment_url !== this.deployment.url) { + if (auth.url !== this.deployment.url) { this.logger.warn( - "Stored tokens have mismatched deployment URL, clearing", - { stored: tokens.deployment_url, current: this.deployment.url }, + "Stored tokens have mismatched deployment URL, clearing OAuth", + { stored: auth.url, current: this.deployment.url }, ); - await this.secretsManager.clearOAuthTokens(this.deployment.safeHostname); + await this.clearOAuthFromSessionAuth(auth); return undefined; } - if (!this.hasRequiredScopes(tokens.scope)) { + if (!this.hasRequiredScopes(auth.oauth.scope)) { this.logger.warn("Stored tokens have insufficient scopes, clearing", { - scope: tokens.scope, + scope: auth.oauth.scope, }); - await this.secretsManager.clearOAuthTokens(this.deployment.safeHostname); + await this.clearOAuthFromSessionAuth(auth); return undefined; } - return tokens; + return { + access_token: auth.token, + ...auth.oauth, + }; + } + + /** + * Clear OAuth data from session auth while preserving the session token. + */ + private async clearOAuthFromSessionAuth(auth: SessionAuth): Promise { + if (!this.deployment) { + return; + } + await this.secretsManager.setSessionAuth(this.deployment.safeHostname, { + url: auth.url, + token: auth.token, + }); } /** @@ -189,9 +214,13 @@ export class OAuthSessionManager implements vscode.Disposable { return; } - this.tokenChangeListener = this.secretsManager.onDidChangeOAuthTokens( + this.tokenChangeListener = this.secretsManager.onDidChangeSessionAuth( this.deployment.safeHostname, - () => this.scheduleNextRefresh(), + (auth) => { + if (auth?.oauth) { + this.scheduleNextRefresh(); + } + }, ); } @@ -754,16 +783,17 @@ export class OAuthSessionManager implements vscode.Disposable { ? Date.now() + tokenResponse.expires_in * 1000 : Date.now() + ACCESS_TOKEN_DEFAULT_EXPIRY_MS; - const tokens: StoredOAuthTokens = { - ...tokenResponse, - deployment_url: deployment.url, + const oauth: OAuthTokenData = { + token_type: tokenResponse.token_type, + refresh_token: tokenResponse.refresh_token, + scope: tokenResponse.scope, expiry_timestamp: expiryTimestamp, }; - await this.secretsManager.setOAuthTokens(deployment.safeHostname, tokens); await this.secretsManager.setSessionAuth(deployment.safeHostname, { url: deployment.url, token: tokenResponse.access_token, + oauth, }); this.logger.info("Tokens saved", { @@ -873,13 +903,18 @@ export class OAuthSessionManager implements vscode.Disposable { /** * Clear OAuth state when switching to non-OAuth authentication. - * Clears tokens from storage. + * Removes OAuth data from session auth while preserving the session token. * Preserves client registration for potential future OAuth use. */ public async clearOAuthState(): Promise { this.clearRefreshState(); if (this.deployment) { - await this.secretsManager.clearOAuthTokens(this.deployment.safeHostname); + const auth = await this.secretsManager.getSessionAuth( + this.deployment.safeHostname, + ); + if (auth?.oauth) { + await this.clearOAuthFromSessionAuth(auth); + } } } @@ -913,7 +948,13 @@ export class OAuthSessionManager implements vscode.Disposable { this.clearRefreshState(); // Clear client registration and tokens to force full re-authentication - await this.secretsManager.clearOAuthData(deployment.safeHostname); + await this.secretsManager.clearOAuthClientRegistration( + deployment.safeHostname, + ); + await this.secretsManager.setSessionAuth(deployment.safeHostname, { + url: deployment.url, + token: "", + }); await this.loginCoordinator.ensureLoggedInWithDialog({ safeHostname: deployment.safeHostname,