Skip to content

Commit b4526cf

Browse files
EhabYEhabY
committed
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
1 parent 7a44e86 commit b4526cf

File tree

13 files changed

+1737
-16
lines changed

13 files changed

+1737
-16
lines changed

‎src/api/oauthInterceptors.ts‎

Lines changed: 116 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,116 @@
1+
import { type AxiosError, isAxiosError } from "axios";
2+
3+
import { type Logger } from "../logging/logger";
4+
import { type RequestConfigWithMeta } from "../logging/types";
5+
import { parseOAuthError, requiresReAuthentication } from "../oauth/errors";
6+
import { type OAuthSessionManager } from "../oauth/sessionManager";
7+
8+
import { type CoderApi } from "./coderApi";
9+
10+
const coderSessionTokenHeader = "Coder-Session-Token";
11+
12+
/**
13+
* Attach OAuth token refresh interceptors to a CoderApi instance.
14+
* This should be called after creating the CoderApi when OAuth authentication is being used.
15+
*
16+
* Success interceptor: proactively refreshes token when approaching expiry.
17+
* Error interceptor: reactively refreshes token on 401 responses.
18+
*/
19+
export function attachOAuthInterceptors(
20+
client: CoderApi,
21+
logger: Logger,
22+
oauthSessionManager: OAuthSessionManager,
23+
): void {
24+
client.getAxiosInstance().interceptors.response.use(
25+
// Success response interceptor: proactive token refresh
26+
(response) => {
27+
// Fire-and-forget: don't await, don't block response
28+
oauthSessionManager.refreshIfAlmostExpired().catch((error) => {
29+
logger.warn("Proactive background token refresh failed:", error);
30+
});
31+
32+
return response;
33+
},
34+
// Error response interceptor: reactive token refresh on 401
35+
async (error: unknown) => {
36+
if (!isAxiosError(error)) {
37+
throw error;
38+
}
39+
40+
if (error.config) {
41+
const config = error.config as {
42+
_oauthRetryAttempted?: boolean;
43+
};
44+
if (config._oauthRetryAttempted) {
45+
throw error;
46+
}
47+
}
48+
49+
const status = error.response?.status;
50+
51+
// These could indicate permanent auth failures that won't be fixed by token refresh
52+
if (status === 400 || status === 403) {
53+
handlePossibleOAuthError(error, logger, oauthSessionManager);
54+
throw error;
55+
} else if (status === 401) {
56+
return handle401Error(error, client, logger, oauthSessionManager);
57+
}
58+
59+
throw error;
60+
},
61+
);
62+
}
63+
64+
function handlePossibleOAuthError(
65+
error: unknown,
66+
logger: Logger,
67+
oauthSessionManager: OAuthSessionManager,
68+
): void {
69+
const oauthError = parseOAuthError(error);
70+
if (oauthError && requiresReAuthentication(oauthError)) {
71+
logger.error(
72+
`OAuth error requires re-authentication: ${oauthError.errorCode}`,
73+
);
74+
75+
oauthSessionManager.showReAuthenticationModal(oauthError).catch((err) => {
76+
logger.error("Failed to show re-auth modal:", err);
77+
});
78+
}
79+
}
80+
81+
async function handle401Error(
82+
error: AxiosError,
83+
client: CoderApi,
84+
logger: Logger,
85+
oauthSessionManager: OAuthSessionManager,
86+
): Promise<void> {
87+
if (!oauthSessionManager.isLoggedInWithOAuth()) {
88+
throw error;
89+
}
90+
91+
logger.info("Received 401 response, attempting token refresh");
92+
93+
try {
94+
const newTokens = await oauthSessionManager.refreshToken();
95+
client.setSessionToken(newTokens.access_token);
96+
97+
logger.info("Token refresh successful, retrying request");
98+
99+
// Retry the original request with the new token
100+
if (error.config) {
101+
const config = error.config as RequestConfigWithMeta & {
102+
_oauthRetryAttempted?: boolean;
103+
};
104+
config._oauthRetryAttempted = true;
105+
config.headers[coderSessionTokenHeader] = newTokens.access_token;
106+
return client.getAxiosInstance().request(config);
107+
}
108+
109+
throw error;
110+
} catch (refreshError) {
111+
logger.error("Token refresh failed:", refreshError);
112+
113+
handlePossibleOAuthError(refreshError, logger, oauthSessionManager);
114+
throw error;
115+
}
116+
}

‎src/commands.ts‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ import { CertificateError } from "./error";
1717
import { getGlobalFlags } from "./globalFlags";
1818
import { type Logger } from "./logging/logger";
1919
import { type LoginCoordinator } from "./login/loginCoordinator";
20+
import { type OAuthSessionManager } from "./oauth/sessionManager";
2021
import { maybeAskAgent, maybeAskUrl } from "./promptUtils";
2122
import { escapeCommandArg, toRemoteAuthority, toSafeHost } from "./util";
2223
import {
@@ -49,6 +50,7 @@ export class Commands {
4950
public constructor(
5051
serviceContainer: ServiceContainer,
5152
private readonly extensionClient: CoderApi,
53+
private readonly oauthSessionManager: OAuthSessionManager,
5254
private readonly deploymentManager: DeploymentManager,
5355
) {
5456
this.vscodeProposed = serviceContainer.getVsCodeProposed();
@@ -107,6 +109,7 @@ export class Commands {
107109
label,
108110
url,
109111
autoLogin: args?.autoLogin,
112+
oauthSessionManager: this.oauthSessionManager,
110113
});
111114

112115
if (!result.success || !result.user || result.token === undefined) {

‎src/core/secretsManager.ts‎

Lines changed: 116 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,10 +3,14 @@ import { toSafeHost } from "../util";
33
import type { Memento, SecretStorage, Disposable } from "vscode";
44

55
import type { Deployment } from "../deployment";
6+
import type { TokenResponse, ClientRegistrationResponse } from "../oauth/types";
67

78
const SESSION_KEY_PREFIX = "coder.session.";
9+
const OAUTH_TOKENS_PREFIX = "coder.oauth.tokens.";
10+
const OAUTH_CLIENT_PREFIX = "coder.oauth.client.";
811

912
const CURRENT_DEPLOYMENT_KEY = "coder.currentDeployment";
13+
const OAUTH_CALLBACK_KEY = "coder.oauthCallback";
1014

1115
const DEPLOYMENT_USAGE_KEY = "coder.deploymentUsage";
1216
const DEFAULT_MAX_DEPLOYMENTS = 10;
@@ -18,11 +22,22 @@ export interface DeploymentUsage {
1822
lastAccessedAt: string;
1923
}
2024

25+
export type StoredOAuthTokens = Omit<TokenResponse, "expires_in"> & {
26+
expiry_timestamp: number;
27+
deployment_url: string;
28+
};
29+
2130
export interface SessionAuth {
2231
url: string;
2332
token: string;
2433
}
2534

35+
interface OAuthCallbackData {
36+
state: string;
37+
code: string | null;
38+
error: string | null;
39+
}
40+
2641
export interface CurrentDeploymentState {
2742
deployment: Deployment | null;
2843
}
@@ -89,6 +104,38 @@ export class SecretsManager {
89104
});
90105
}
91106

107+
/**
108+
* Write an OAuth callback result to secrets storage.
109+
* Used for cross-window communication when OAuth callback arrives in a different window.
110+
*/
111+
public async setOAuthCallback(data: OAuthCallbackData): Promise<void> {
112+
await this.secrets.store(OAUTH_CALLBACK_KEY, JSON.stringify(data));
113+
}
114+
115+
/**
116+
* Listen for OAuth callback results from any VS Code window.
117+
* The listener receives the state parameter, code (if success), and error (if failed).
118+
*/
119+
public onDidChangeOAuthCallback(
120+
listener: (data: OAuthCallbackData) => void,
121+
): Disposable {
122+
return this.secrets.onDidChange(async (e) => {
123+
if (e.key !== OAUTH_CALLBACK_KEY) {
124+
return;
125+
}
126+
127+
try {
128+
const data = await this.secrets.get(OAUTH_CALLBACK_KEY);
129+
if (data) {
130+
const parsed = JSON.parse(data) as OAuthCallbackData;
131+
listener(parsed);
132+
}
133+
} catch {
134+
// Ignore parse errors
135+
}
136+
});
137+
}
138+
92139
/**
93140
* Listen for changes to a specific deployment's session auth.
94141
*/
@@ -144,6 +191,71 @@ export class SecretsManager {
144191
await this.secrets.delete(`${SESSION_KEY_PREFIX}${label}`);
145192
}
146193

194+
public async getOAuthTokens(
195+
label: string,
196+
): Promise<StoredOAuthTokens | undefined> {
197+
try {
198+
const data = await this.secrets.get(`${OAUTH_TOKENS_PREFIX}${label}`);
199+
if (!data) {
200+
return undefined;
201+
}
202+
return JSON.parse(data) as StoredOAuthTokens;
203+
} catch {
204+
return undefined;
205+
}
206+
}
207+
208+
public async setOAuthTokens(
209+
label: string,
210+
tokens: StoredOAuthTokens,
211+
): Promise<void> {
212+
await this.secrets.store(
213+
`${OAUTH_TOKENS_PREFIX}${label}`,
214+
JSON.stringify(tokens),
215+
);
216+
await this.recordDeploymentAccess(label);
217+
}
218+
219+
public async clearOAuthTokens(label: string): Promise<void> {
220+
await this.secrets.delete(`${OAUTH_TOKENS_PREFIX}${label}`);
221+
}
222+
223+
public async getOAuthClientRegistration(
224+
label: string,
225+
): Promise<ClientRegistrationResponse | undefined> {
226+
try {
227+
const data = await this.secrets.get(`${OAUTH_CLIENT_PREFIX}${label}`);
228+
if (!data) {
229+
return undefined;
230+
}
231+
return JSON.parse(data) as ClientRegistrationResponse;
232+
} catch {
233+
return undefined;
234+
}
235+
}
236+
237+
public async setOAuthClientRegistration(
238+
label: string,
239+
registration: ClientRegistrationResponse,
240+
): Promise<void> {
241+
await this.secrets.store(
242+
`${OAUTH_CLIENT_PREFIX}${label}`,
243+
JSON.stringify(registration),
244+
);
245+
await this.recordDeploymentAccess(label);
246+
}
247+
248+
public async clearOAuthClientRegistration(label: string): Promise<void> {
249+
await this.secrets.delete(`${OAUTH_CLIENT_PREFIX}${label}`);
250+
}
251+
252+
public async clearOAuthData(label: string): Promise<void> {
253+
await Promise.all([
254+
this.clearOAuthTokens(label),
255+
this.clearOAuthClientRegistration(label),
256+
]);
257+
}
258+
147259
/**
148260
* Record that a deployment was accessed, moving it to the front of the LRU list.
149261
* Prunes deployments beyond maxCount, clearing their auth data.
@@ -167,7 +279,10 @@ export class SecretsManager {
167279
* Clear all auth data for a deployment and remove it from the usage list.
168280
*/
169281
public async clearAllAuthData(label: string): Promise<void> {
170-
await this.clearSessionAuth(label);
282+
await Promise.all([
283+
this.clearSessionAuth(label),
284+
this.clearOAuthData(label),
285+
]);
171286
const usage = this.getDeploymentUsage().filter((u) => u.label !== label);
172287
await this.memento.update(DEPLOYMENT_USAGE_KEY, usage);
173288
}

‎src/deployment/deploymentManager.ts‎

Lines changed: 24 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import { type ContextManager } from "../core/contextManager";
44
import { type MementoManager } from "../core/mementoManager";
55
import { type SecretsManager } from "../core/secretsManager";
66
import { type Logger } from "../logging/logger";
7+
import { type OAuthSessionManager } from "../oauth/sessionManager";
78
import { type WorkspaceProvider } from "../workspace/workspacesProvider";
89

910
import {
@@ -21,6 +22,7 @@ import type * as vscode from "vscode";
2122
* Centralizes:
2223
* - In-memory deployment state (url, label, token, user)
2324
* - Client credential updates
25+
* - OAuth session management
2426
* - Auth listener registration
2527
* - Context updates (coder.authenticated, coder.isOwner)
2628
* - Workspace provider refresh
@@ -39,6 +41,7 @@ export class DeploymentManager implements vscode.Disposable {
3941
private constructor(
4042
serviceContainer: ServiceContainer,
4143
private readonly client: CoderApi,
44+
private readonly oauthSessionManager: OAuthSessionManager,
4245
private readonly workspaceProviders: WorkspaceProvider[],
4346
) {
4447
this.secretsManager = serviceContainer.getSecretsManager();
@@ -50,11 +53,13 @@ export class DeploymentManager implements vscode.Disposable {
5053
public static create(
5154
serviceContainer: ServiceContainer,
5255
client: CoderApi,
56+
oauthSessionManager: OAuthSessionManager,
5357
workspaceProviders: WorkspaceProvider[],
5458
): DeploymentManager {
5559
const manager = new DeploymentManager(
5660
serviceContainer,
5761
client,
62+
oauthSessionManager,
5863
workspaceProviders,
5964
);
6065
manager.subscribeToCrossWindowChanges();
@@ -82,7 +87,7 @@ export class DeploymentManager implements vscode.Disposable {
8287
public async changeDeployment(
8388
deployment: DeploymentWithAuth & { user: User },
8489
): Promise<void> {
85-
this.setDeploymentCore(deployment);
90+
await this.setDeploymentCore(deployment);
8691

8792
this.refreshWorkspaces();
8893
await this.persistDeployment(deployment);
@@ -96,7 +101,12 @@ export class DeploymentManager implements vscode.Disposable {
96101
public async setDeploymentWithoutAuth(
97102
deployment: Deployment & { token?: string },
98103
): Promise<void> {
99-
this.setDeploymentCore({ ...deployment });
104+
await this.setDeploymentCore({ ...deployment });
105+
106+
// This can be async since we'll get an auth event if the token is refreshed
107+
this.oauthSessionManager.refreshIfAlmostExpired().catch((error) => {
108+
this.logger.warn("Setup token refresh failed:", error);
109+
});
100110

101111
await this.tryFetchAndUpgradeUser();
102112
}
@@ -106,6 +116,7 @@ export class DeploymentManager implements vscode.Disposable {
106116
*/
107117
public async logout(): Promise<void> {
108118
this.client.setCredentials(undefined, undefined);
119+
this.oauthSessionManager.clearDeployment();
109120

110121
this.authListenerDisposable?.dispose();
111122
this.authListenerDisposable = undefined;
@@ -116,12 +127,22 @@ export class DeploymentManager implements vscode.Disposable {
116127
await this.secretsManager.setCurrentDeployment(undefined);
117128
}
118129

119-
private setDeploymentCore(deployment: DeploymentWithAuth): void {
130+
/**
131+
* Clear OAuth state for a deployment when switching to token auth.
132+
*/
133+
public async clearOAuthState(label: string): Promise<void> {
134+
await this.oauthSessionManager.clearOAuthState(label);
135+
}
136+
137+
private async setDeploymentCore(
138+
deployment: DeploymentWithAuth,
139+
): Promise<void> {
120140
if (deployment.token === undefined) {
121141
this.client.setHost(deployment.url);
122142
} else {
123143
this.client.setCredentials(deployment.url, deployment.token);
124144
}
145+
await this.oauthSessionManager.setDeployment(deployment);
125146
this.registerAuthListener(deployment.label);
126147
this.currentDeployment = { ...deployment };
127148
this.updateAuthContexts(deployment.user);

0 commit comments

Comments
 (0)