fix(api): validate API responses at the boundary

src/api/coderApi.ts

@@ -49,15 +49,22 @@ import {
 import { SseConnection } from "../websocket/sseConnection";
 
 import { getRefreshCommand, refreshCertificates } from "./certificateRefresh";
+import {
+	parseApiResponse,
+	VALIDATED_RESPONSES,
+	type ValidatedMethods,
+} from "./responseValidation";
 import { createHttpAgent } from "./utils";
 
 import type {
 	GetInboxNotificationResponse,
+	ProvisionerJob,
 	ProvisionerJobLog,
 	ServerSentEvent,
 	Workspace,
 	WorkspaceAgent,
 	WorkspaceAgentLog,
+	WorkspaceBuild,
 } from "coder/site/src/api/typesGenerated";
 import type { ClientOptions } from "ws";
 
@@ -114,6 +121,7 @@ export class CoderApi extends Api implements vscode.Disposable {
 		private readonly authConfigTracker: AuthConfigTracker,
 	) {
 		super();
+		wrapWithValidation(this);
 		this.configWatcher = this.watchConfigChanges();
 	}
 
@@ -149,6 +157,30 @@ export class CoderApi extends Api implements vscode.Disposable {
 		return this.getAxiosInstance().defaults.baseURL;
 	}
 
+	/**
+	 * Reimplemented because the SDK version polls inside a voided IIFE that
+	 * swallows errors, hanging callers forever if a poll throws (e.g. on
+	 * failed response validation).
+	 */
+	override waitForBuild = async (
+		build: WorkspaceBuild,
+	): Promise<ProvisionerJob | undefined> => {
+		while (true) {
+			const { job } = await this.getWorkspaceBuildByNumber(
+				build.workspace_owner_name,
+				build.workspace_name,
+				build.build_number,
+			);
+			if (job.status === "failed") {
+				throw new Error(`Build ${build.build_number} failed`);
+			}
+			if (job.status === "succeeded" || job.status === "canceled") {
+				return job;
+			}
+			await new Promise((resolve) => setTimeout(resolve, 1000));
+		}
+	};
+
 	hasAuthConfigChangedSince(version: number | undefined): boolean {
 		return this.authConfigTracker.hasChangedSince(version);
 	}
@@ -747,6 +779,21 @@ function wrapResponseTransform(
 	];
 }
 
+/**
+ * Validate the fields the extension reads on each response, since the SDK
+ * casts bodies to the generated types with no runtime check. The methods
+ * are instance arrow properties, so wrapping is by reassignment;
+ * `override` fields would depend on declaration order.
+ */
+function wrapWithValidation(api: CoderApi): void {
+	const methods: ValidatedMethods = api;
+	for (const [name, schema] of VALIDATED_RESPONSES) {
+		const method = methods[name];
+		methods[name] = async (...args) =>
+			parseApiResponse(schema, await method(...args), name, api.getHost());
+	}
+}
+
 function getSize(headers: AxiosHeaders, data: unknown): number | undefined {
 	const contentLength = headers["content-length"] as unknown;
 	if (typeof contentLength === "string") {

src/api/responseValidation.ts

@@ -0,0 +1,122 @@
+import { z } from "zod";
+
+import type { CoderApi } from "./coderApi";
+
+/**
+ * Thrown when a 2xx response body does not match the shape the extension
+ * needs, which almost always means the URL does not point at a Coder
+ * deployment (a proxy error page, a different service, a partial body).
+ */
+export class InvalidApiResponseError extends Error {
+	constructor(
+		public readonly endpoint: string,
+		url: string | undefined,
+		options?: { cause?: unknown },
+	) {
+		super(
+			`${url ?? "The deployment"} did not return a valid Coder API response ` +
+				`for ${endpoint}. Check that the URL points to a Coder deployment.`,
+			options,
+		);
+		this.name = "InvalidApiResponseError";
+	}
+}
+
+/**
+ * Validate a response body, returning the original value with the caller's
+ * type.
+ *
+ * Schemas list only the fields the extension reads and use looseObject so
+ * unknown fields pass through. Require a field only if every deployment back
+ * to Coder 0.25 sends it; anything newer must be .optional().
+ *
+ * @throws {InvalidApiResponseError} naming the endpoint when validation fails.
+ */
+export function parseApiResponse<T>(
+	schema: z.ZodType<unknown>,
+	data: T,
+	endpoint: string,
+	url?: string,
+): T {
+	const result = schema.safeParse(data);
+	if (!result.success) {
+		throw new InvalidApiResponseError(endpoint, url, { cause: result.error });
+	}
+	return data;
+}
+
+export const UserSchema = z.looseObject({
+	id: z.string(),
+	username: z.string(),
+	roles: z.array(z.looseObject({ name: z.string() })),
+});
+
+const WorkspaceAgentSchema = z.looseObject({
+	id: z.string(),
+	name: z.string(),
+	status: z.string(),
+	operating_system: z.string(),
+});
+
+const WorkspaceResourceSchema = z.looseObject({
+	agents: z.array(WorkspaceAgentSchema).nullable().optional(),
+});
+
+export const WorkspaceSchema = z.looseObject({
+	id: z.string(),
+	name: z.string(),
+	owner_name: z.string(),
+	template_id: z.string(),
+	latest_build: z.looseObject({
+		id: z.string(),
+		status: z.string(),
+		template_version_id: z.string(),
+		resources: z.array(WorkspaceResourceSchema),
+	}),
+});
+
+/** waitForBuild reads the identifiers to poll and the job status to stop. */
+export const WorkspaceBuildSchema = z.looseObject({
+	workspace_owner_name: z.string(),
+	workspace_name: z.string(),
+	build_number: z.number(),
+	job: z.looseObject({ status: z.string() }),
+});
+
+export const TemplateSchema = z.looseObject({
+	active_version_id: z.string(),
+});
+
+export const WorkspaceResourcesSchema = z.array(WorkspaceResourceSchema);
+
+export const SSHConfigResponseSchema = z.looseObject({
+	ssh_config_options: z.record(z.string(), z.string()),
+});
+
+/**
+ * The schema each SDK method's response must match, applied by CoderApi. Add a
+ * pair to validate another method; the name doubles as the endpoint in the
+ * error. Pairs, not an object, so iterating keeps the names as literal types.
+ * OAuth endpoints are plain axios calls and pass their schema at the call site.
+ */
+export const VALIDATED_RESPONSES = [
+	["getAuthenticatedUser", UserSchema],
+	["getDeploymentSSHConfig", SSHConfigResponseSchema],
+	["getTemplate", TemplateSchema],
+	["getTemplateVersionResources", WorkspaceResourcesSchema],
+	["getWorkspace", WorkspaceSchema],
+	["getWorkspaceByOwnerAndName", WorkspaceSchema],
+	["getWorkspaceBuildByNumber", WorkspaceBuildSchema],
+	["startWorkspace", WorkspaceBuildSchema],
+	["stopWorkspace", WorkspaceBuildSchema],
+] as const satisfies ReadonlyArray<readonly [keyof CoderApi, z.ZodType]>;
+
+/**
+ * The methods above, reduced to what the wrapper needs. CoderApi satisfies this
+ * with no assertion: `never` parameters accept any signature, and one uniform
+ * value type is what allows assigning by a name held in a variable.
+ */
+export type ValidatedMethods = Record<
+	(typeof VALIDATED_RESPONSES)[number][0],
+	(...args: never[]) => Promise<unknown>
+>;

src/oauth/authorizer.ts

@@ -17,6 +17,11 @@ import {
 	generateState,
 	toUrlSearchParams,
 } from "./utils";
+import {
+	OAuth2ClientRegistrationResponseSchema,
+	OAuth2TokenResponseSchema,
+	parseOAuthResponse,
+} from "./validation";
 
 import type { AxiosInstance } from "axios";
 import type {
@@ -173,16 +178,22 @@ export class OAuthAuthorizer implements vscode.Disposable {
 			registrationRequest,
 		);
 
+		const registrationResponse = parseOAuthResponse(
+			OAuth2ClientRegistrationResponseSchema,
+			response.data,
+			metadata.registration_endpoint,
+		);
+
 		await this.secretsManager.setOAuthClientRegistration(
 			deployment.safeHostname,
-			response.data,
+			registrationResponse,
 		);
 		this.logger.debug(
 			"Saved OAuth client registration:",
-			response.data.client_id,
+			registrationResponse.client_id,
 		);
 
-		return response.data;
+		return registrationResponse;
 	}
 
 	/**
@@ -360,7 +371,11 @@ export class OAuthAuthorizer implements vscode.Disposable {
 
 		this.logger.debug("Token exchange successful");
 
-		return response.data;
+		return parseOAuthResponse(
+			OAuth2TokenResponseSchema,
+			response.data,
+			metadata.token_endpoint,
+		);
 	}
 
 	public dispose(): void {

src/oauth/metadataClient.ts

@@ -1,10 +1,13 @@
+import { parseApiResponse } from "../api/responseValidation";
+
 import {
 	AUTH_GRANT_TYPE,
 	PKCE_CHALLENGE_METHOD,
 	REFRESH_GRANT_TYPE,
 	RESPONSE_TYPE,
 	TOKEN_ENDPOINT_AUTH_METHOD,
 } from "./constants";
+import { OAuth2AuthorizationServerMetadataSchema } from "./validation";
 
 import type { AxiosInstance } from "axios";
 import type {
@@ -69,9 +72,13 @@ export class OAuthMetadataClient {
 				OAUTH_DISCOVERY_ENDPOINT,
 			);
 
-		const metadata = response.data;
+		const metadata = parseApiResponse(
+			OAuth2AuthorizationServerMetadataSchema,
+			response.data,
+			OAUTH_DISCOVERY_ENDPOINT,
+			this.axiosInstance.defaults.baseURL,
+		);
 
-		this.validateRequiredEndpoints(metadata);
 		this.validateGrantTypes(metadata);
 		this.validateResponseTypes(metadata);
 		this.validateAuthMethods(metadata);
@@ -87,21 +94,6 @@ export class OAuthMetadataClient {
 		return metadata;
 	}
 
-	private validateRequiredEndpoints(
-		metadata: OAuth2AuthorizationServerMetadata,
-	): 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: OAuth2AuthorizationServerMetadata,
 	): void {

src/oauth/sessionManager.ts

@@ -8,6 +8,7 @@ import { DEFAULT_OAUTH_SCOPES, REFRESH_GRANT_TYPE } from "./constants";
 import { OAuthError, parseOAuthError } from "./errors";
 import { OAuthMetadataClient } from "./metadataClient";
 import { buildOAuthTokenData, toUrlSearchParams } from "./utils";
+import { OAuth2TokenResponseSchema, parseOAuthResponse } from "./validation";
 
 import type { AxiosInstance } from "axios";
 import type {
@@ -421,17 +422,23 @@ export class OAuthSessionManager implements vscode.Disposable {
 
 					this.logger.debug("Token refresh successful");
 
+					const tokenResponse = parseOAuthResponse(
+						OAuth2TokenResponseSchema,
+						response.data,
+						metadata.token_endpoint,
+					);
+
 					await this.secretsManager.setSessionAuth(deployment.safeHostname, {
 						url: deployment.url,
-						token: response.data.access_token,
+						token: tokenResponse.access_token,
 						username: await this.fetchUsername(
 							deployment,
-							response.data.access_token,
+							tokenResponse.access_token,
 						),
-						oauth: buildOAuthTokenData(response.data),
+						oauth: buildOAuthTokenData(tokenResponse),
 					});
 
-					return response.data;
+					return tokenResponse;
 				},
 			);
 		} catch (error) {

src/oauth/validation.ts

@@ -0,0 +1,53 @@
+import { z } from "zod";
+
+import { parseApiResponse } from "../api/responseValidation";
+
+/**
+ * parseApiResponse for OAuth endpoints, whose absolute URLs come from server
+ * metadata and may live on a different origin than the deployment. The schemas
+ * below follow the same rules.
+ */
+export function parseOAuthResponse<T>(
+	schema: z.ZodType<unknown>,
+	data: T,
+	endpoint: string,
+): T {
+	const { origin, pathname } = new URL(endpoint);
+	return parseApiResponse(schema, data, pathname, origin);
+}
+
+/** An empty endpoint or identifier is as unusable as a missing one. */
+const REQUIRED_STRING = z.string().min(1);
+
+/**
+ * Plain strings rather than the generated enums, so a server adding a value
+ * does not fail validation. Absent means the RFC 8414 default applies.
+ */
+const CAPABILITIES = z.array(z.string()).optional();
+
+export const OAuth2AuthorizationServerMetadataSchema = z.looseObject({
+	issuer: REQUIRED_STRING,
+	authorization_endpoint: REQUIRED_STRING,
+	token_endpoint: REQUIRED_STRING,
+	// Callers report these as unsupported when absent, so no .min(1) here.
+	registration_endpoint: z.string().optional(),
+	revocation_endpoint: z.string().optional(),
+	grant_types_supported: CAPABILITIES,
+	response_types_supported: CAPABILITIES,
+	token_endpoint_auth_methods_supported: CAPABILITIES,
+	code_challenge_methods_supported: CAPABILITIES,
+	scopes_supported: CAPABILITIES,
+});
+
+export const OAuth2ClientRegistrationResponseSchema = z.looseObject({
+	client_id: REQUIRED_STRING,
+	client_secret: z.string().optional(),
+	redirect_uris: z.array(z.string()).optional(),
+});
+
+export const OAuth2TokenResponseSchema = z.looseObject({
+	access_token: REQUIRED_STRING,
+	token_type: z.string(),
+	refresh_token: z.string().optional(),
+	expires_in: z.number().optional(),
+});