Skip to content

Commit ab50022

Browse files
committed
Ensure login prompt before opening workspace from URI
When opening a workspace via URI without a token, ensure we prompt for login and set the deployment before attempting to open. Extract URI handler into dedicated module with route-based dispatch.
1 parent 9078076 commit ab50022

File tree

5 files changed

+508
-140
lines changed

5 files changed

+508
-140
lines changed

src/deployment/deploymentManager.ts

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -83,8 +83,13 @@ export class DeploymentManager implements vscode.Disposable {
8383
* Returns true if deployment was changed, false otherwise.
8484
*/
8585
public async setDeploymentIfValid(
86-
deployment: Deployment & { token?: string },
86+
deployment: DeploymentWithAuth,
8787
): Promise<boolean> {
88+
if (deployment.user) {
89+
await this.setDeployment({ ...deployment, user: deployment.user });
90+
return true;
91+
}
92+
8893
const auth = await this.secretsManager.getSessionAuth(
8994
deployment.safeHostname,
9095
);

src/extension.ts

Lines changed: 8 additions & 139 deletions
Original file line numberDiff line numberDiff line change
@@ -13,10 +13,9 @@ import { ServiceContainer } from "./core/container";
1313
import { type SecretsManager } from "./core/secretsManager";
1414
import { DeploymentManager } from "./deployment/deploymentManager";
1515
import { CertificateError, getErrorDetail } from "./error";
16-
import { maybeAskUrl } from "./promptUtils";
1716
import { Remote } from "./remote/remote";
1817
import { getRemoteSshExtension } from "./remote/sshExtension";
19-
import { toSafeHost } from "./util";
18+
import { registerUriHandler } from "./uri/uriHandler";
2019
import {
2120
WorkspaceProvider,
2221
WorkspaceQuery,
@@ -129,103 +128,17 @@ export async function activate(ctx: vscode.ExtensionContext): Promise<void> {
129128
]);
130129
ctx.subscriptions.push(deploymentManager);
131130

132-
// Handle vscode:// URIs.
133-
const uriHandler = vscode.window.registerUriHandler({
134-
handleUri: async (uri) => {
135-
const params = new URLSearchParams(uri.query);
136-
137-
if (uri.path === "/open") {
138-
const owner = params.get("owner");
139-
const workspace = params.get("workspace");
140-
const agent = params.get("agent");
141-
const folder = params.get("folder");
142-
const openRecent =
143-
params.has("openRecent") &&
144-
(!params.get("openRecent") || params.get("openRecent") === "true");
145-
146-
if (!owner) {
147-
throw new Error("owner must be specified as a query parameter");
148-
}
149-
if (!workspace) {
150-
throw new Error("workspace must be specified as a query parameter");
151-
}
152-
153-
await setupDeploymentFromUri(params, serviceContainer);
154-
155-
await commands.open(
156-
owner,
157-
workspace,
158-
agent ?? undefined,
159-
folder ?? undefined,
160-
openRecent,
161-
);
162-
} else if (uri.path === "/openDevContainer") {
163-
const workspaceOwner = params.get("owner");
164-
const workspaceName = params.get("workspace");
165-
const workspaceAgent = params.get("agent");
166-
const devContainerName = params.get("devContainerName");
167-
const devContainerFolder = params.get("devContainerFolder");
168-
const localWorkspaceFolder = params.get("localWorkspaceFolder");
169-
const localConfigFile = params.get("localConfigFile");
170-
171-
if (!workspaceOwner) {
172-
throw new Error(
173-
"workspace owner must be specified as a query parameter",
174-
);
175-
}
176-
177-
if (!workspaceName) {
178-
throw new Error(
179-
"workspace name must be specified as a query parameter",
180-
);
181-
}
182-
183-
if (!workspaceAgent) {
184-
throw new Error(
185-
"workspace agent must be specified as a query parameter",
186-
);
187-
}
188-
189-
if (!devContainerName) {
190-
throw new Error(
191-
"dev container name must be specified as a query parameter",
192-
);
193-
}
194-
195-
if (!devContainerFolder) {
196-
throw new Error(
197-
"dev container folder must be specified as a query parameter",
198-
);
199-
}
200-
201-
if (localConfigFile && !localWorkspaceFolder) {
202-
throw new Error(
203-
"local workspace folder must be specified as a query parameter if local config file is provided",
204-
);
205-
}
206-
207-
await setupDeploymentFromUri(params, serviceContainer);
208-
209-
await commands.openDevContainer(
210-
workspaceOwner,
211-
workspaceName,
212-
workspaceAgent,
213-
devContainerName,
214-
devContainerFolder,
215-
localWorkspaceFolder ?? "",
216-
localConfigFile ?? "",
217-
);
218-
} else {
219-
throw new Error(`Unknown path ${uri.path}`);
220-
}
221-
},
222-
});
223-
ctx.subscriptions.push(uriHandler);
224-
225131
// Register globally available commands. Many of these have visibility
226132
// controlled by contexts, see `when` in the package.json.
227133
const commands = new Commands(serviceContainer, client, deploymentManager);
134+
228135
ctx.subscriptions.push(
136+
registerUriHandler(
137+
serviceContainer,
138+
deploymentManager,
139+
commands,
140+
vscodeProposed,
141+
),
229142
vscode.commands.registerCommand(
230143
"coder.login",
231144
commands.login.bind(commands),
@@ -418,50 +331,6 @@ async function showTreeViewSearch(id: string): Promise<void> {
418331
await vscode.commands.executeCommand("list.find");
419332
}
420333

421-
/**
422-
* Sets up deployment from URI parameters. Handles URL prompting, client setup,
423-
* and token storage. Throws if user cancels URL input.
424-
*/
425-
async function setupDeploymentFromUri(
426-
params: URLSearchParams,
427-
serviceContainer: ServiceContainer,
428-
): Promise<void> {
429-
const secretsManager = serviceContainer.getSecretsManager();
430-
const mementoManager = serviceContainer.getMementoManager();
431-
const currentDeployment = await secretsManager.getCurrentDeployment();
432-
433-
// We are not guaranteed that the URL we currently have is for the URL
434-
// this workspace belongs to, or that we even have a URL at all (the
435-
// queries will default to localhost) so ask for it if missing.
436-
// Pre-populate in case we do have the right URL so the user can just
437-
// hit enter and move on.
438-
const url = await maybeAskUrl(
439-
mementoManager,
440-
params.get("url"),
441-
currentDeployment?.url,
442-
);
443-
if (!url) {
444-
throw new Error("url must be provided or specified as a query parameter");
445-
}
446-
447-
const safeHost = toSafeHost(url);
448-
449-
// If the token is missing we will get a 401 later and the user will be
450-
// prompted to sign in again, so we do not need to ensure it is set now.
451-
const token: string | null = params.get("token");
452-
if (token === null) {
453-
// We need to ensure there is at least an entry for this in storage
454-
// so that we know what URL to prompt the user with when logging in.
455-
const auth = await secretsManager.getSessionAuth(safeHost);
456-
if (!auth) {
457-
// Racy, we could accidentally overwrite the token that is written in the meantime.
458-
await secretsManager.setSessionAuth(safeHost, { url, token: "" });
459-
}
460-
} else {
461-
await secretsManager.setSessionAuth(safeHost, { url, token });
462-
}
463-
}
464-
465334
async function listStoredDeployments(
466335
secretsManager: SecretsManager,
467336
): Promise<void> {

src/uri/uriHandler.ts

Lines changed: 182 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,182 @@
1+
import * as vscode from "vscode";
2+
3+
import { errToStr } from "../api/api-helper";
4+
import { type Commands } from "../commands";
5+
import { type ServiceContainer } from "../core/container";
6+
import { type DeploymentManager } from "../deployment/deploymentManager";
7+
import { maybeAskUrl } from "../promptUtils";
8+
import { toSafeHost } from "../util";
9+
10+
interface UriRouteContext {
11+
params: URLSearchParams;
12+
serviceContainer: ServiceContainer;
13+
deploymentManager: DeploymentManager;
14+
commands: Commands;
15+
}
16+
17+
type UriRouteHandler = (ctx: UriRouteContext) => Promise<void>;
18+
19+
const routes: Record<string, UriRouteHandler> = {
20+
"/open": handleOpen,
21+
"/openDevContainer": handleOpenDevContainer,
22+
};
23+
24+
/**
25+
* Registers the URI handler for `{vscode.env.uriScheme}://coder.coder-remote`... URIs.
26+
*/
27+
export function registerUriHandler(
28+
serviceContainer: ServiceContainer,
29+
deploymentManager: DeploymentManager,
30+
commands: Commands,
31+
vscodeProposed: typeof vscode,
32+
): vscode.Disposable {
33+
const output = serviceContainer.getLogger();
34+
35+
return vscode.window.registerUriHandler({
36+
handleUri: async (uri) => {
37+
try {
38+
await routeUri(uri, serviceContainer, deploymentManager, commands);
39+
} catch (error) {
40+
const message = errToStr(error, "No error message was provided");
41+
output.warn(`Failed to handle URI ${uri.toString()}: ${message}`);
42+
vscodeProposed.window.showErrorMessage("Failed to handle URI", {
43+
detail: message,
44+
modal: true,
45+
useCustom: true,
46+
});
47+
}
48+
},
49+
});
50+
}
51+
52+
async function routeUri(
53+
uri: vscode.Uri,
54+
serviceContainer: ServiceContainer,
55+
deploymentManager: DeploymentManager,
56+
commands: Commands,
57+
): Promise<void> {
58+
const handler = routes[uri.path];
59+
if (!handler) {
60+
throw new Error(`Unknown path ${uri.path}`);
61+
}
62+
63+
await handler({
64+
params: new URLSearchParams(uri.query),
65+
serviceContainer,
66+
deploymentManager,
67+
commands,
68+
});
69+
}
70+
71+
function getRequiredParam(params: URLSearchParams, name: string): string {
72+
const value = params.get(name);
73+
if (!value) {
74+
throw new Error(`${name} must be specified as a query parameter`);
75+
}
76+
return value;
77+
}
78+
79+
async function handleOpen(ctx: UriRouteContext): Promise<void> {
80+
const { params, serviceContainer, deploymentManager, commands } = ctx;
81+
82+
const owner = getRequiredParam(params, "owner");
83+
const workspace = getRequiredParam(params, "workspace");
84+
const agent = params.get("agent");
85+
const folder = params.get("folder");
86+
const openRecent =
87+
params.has("openRecent") &&
88+
(!params.get("openRecent") || params.get("openRecent") === "true");
89+
90+
await setupDeployment(params, serviceContainer, deploymentManager);
91+
92+
await commands.open(
93+
owner,
94+
workspace,
95+
agent ?? undefined,
96+
folder ?? undefined,
97+
openRecent,
98+
);
99+
}
100+
101+
async function handleOpenDevContainer(ctx: UriRouteContext): Promise<void> {
102+
const { params, serviceContainer, deploymentManager, commands } = ctx;
103+
104+
const owner = getRequiredParam(params, "owner");
105+
const workspace = getRequiredParam(params, "workspace");
106+
const agent = getRequiredParam(params, "agent");
107+
const devContainerName = getRequiredParam(params, "devContainerName");
108+
const devContainerFolder = getRequiredParam(params, "devContainerFolder");
109+
const localWorkspaceFolder = params.get("localWorkspaceFolder");
110+
const localConfigFile = params.get("localConfigFile");
111+
112+
if (localConfigFile && !localWorkspaceFolder) {
113+
throw new Error(
114+
"localWorkspaceFolder must be specified as a query parameter if localConfigFile is provided",
115+
);
116+
}
117+
118+
await setupDeployment(params, serviceContainer, deploymentManager);
119+
120+
await commands.openDevContainer(
121+
owner,
122+
workspace,
123+
agent,
124+
devContainerName,
125+
devContainerFolder,
126+
localWorkspaceFolder ?? "",
127+
localConfigFile ?? "",
128+
);
129+
}
130+
131+
/**
132+
* Sets up deployment from URI parameters. Handles URL prompting, client setup,
133+
* and token storage. Throws if user cancels URL input or login fails.
134+
*/
135+
async function setupDeployment(
136+
params: URLSearchParams,
137+
serviceContainer: ServiceContainer,
138+
deploymentManager: DeploymentManager,
139+
): Promise<void> {
140+
const secretsManager = serviceContainer.getSecretsManager();
141+
const mementoManager = serviceContainer.getMementoManager();
142+
const loginCoordinator = serviceContainer.getLoginCoordinator();
143+
144+
const currentDeployment = await secretsManager.getCurrentDeployment();
145+
146+
// We are not guaranteed that the URL we currently have is for the URL
147+
// this workspace belongs to, or that we even have a URL at all (the
148+
// queries will default to localhost) so ask for it if missing.
149+
// Pre-populate in case we do have the right URL so the user can just
150+
// hit enter and move on.
151+
const url = await maybeAskUrl(
152+
mementoManager,
153+
params.get("url"),
154+
currentDeployment?.url,
155+
);
156+
if (!url) {
157+
throw new Error("url must be provided or specified as a query parameter");
158+
}
159+
160+
const safeHostname = toSafeHost(url);
161+
162+
const token: string | null = params.get("token");
163+
if (token !== null) {
164+
await secretsManager.setSessionAuth(safeHostname, { url, token });
165+
}
166+
167+
const result = await loginCoordinator.ensureLoggedIn({
168+
safeHostname,
169+
url,
170+
});
171+
172+
if (!result.success) {
173+
throw new Error("Failed to login to deployment from URI");
174+
}
175+
176+
await deploymentManager.setDeploymentIfValid({
177+
safeHostname,
178+
url,
179+
token: result.token,
180+
user: result.user,
181+
});
182+
}

test/mocks/vscode.runtime.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -99,6 +99,7 @@ export const window = {
9999
clear: vi.fn(),
100100
})),
101101
createStatusBarItem: vi.fn(),
102+
registerUriHandler: vi.fn(() => ({ dispose: vi.fn() })),
102103
};
103104

104105
export const commands = {

0 commit comments

Comments
 (0)