Skip to content

Commit cdea111

Browse files
mafredrimafredri
committed
refactor
1 parent 5d370fc commit cdea111

File tree

4 files changed

+204
-70
lines changed

4 files changed

+204
-70
lines changed

‎cli/open.go

+126-61
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,7 @@ func (r *RootCmd) openVSCode() *clibase.Cmd {
4343
cmd := &clibase.Cmd{
4444
Annotations: workspaceCommand,
4545
Use: "vscode <workspace> [<directory in workspace>]",
46-
Short: "Open a workspace in Visual Studio Code",
46+
Short: fmt.Sprintf("Open a workspace in %s", vscodeDesktopName),
4747
Middleware: clibase.Chain(
4848
clibase.RequireRangeArgs(1, 2),
4949
r.InitClient(client),
@@ -73,18 +73,12 @@ func (r *RootCmd) openVSCode() *clibase.Cmd {
7373
insideThisWorkspace := insideAWorkspace && inWorkspaceName == workspaceName
7474

7575
if !insideThisWorkspace {
76-
// We could optionally add a flag to skip wait, like with SSH.
77-
wait := false
78-
for _, script := range workspaceAgent.Scripts {
79-
if script.StartBlocksLogin {
80-
wait = true
81-
break
82-
}
83-
}
76+
// Wait for the agent to connect, we don't care about readiness
77+
// otherwise (e.g. wait).
8478
err = cliui.Agent(ctx, inv.Stderr, workspaceAgent.ID, cliui.AgentOptions{
8579
Fetch: client.WorkspaceAgent,
86-
FetchLogs: client.WorkspaceAgentLogsAfter,
87-
Wait: wait,
80+
FetchLogs: nil,
81+
Wait: false,
8882
})
8983
if err != nil {
9084
if xerrors.Is(err, context.Canceled) {
@@ -93,55 +87,30 @@ func (r *RootCmd) openVSCode() *clibase.Cmd {
9387
return xerrors.Errorf("agent: %w", err)
9488
}
9589

96-
// If the ExpandedDirectory was initially missing, it could mean
97-
// that the agent hadn't reported it in yet. Retry once.
98-
if workspaceAgent.ExpandedDirectory == "" {
99-
autostart = false // Don't retry autostart.
100-
workspace, workspaceAgent, err = getWorkspaceAndAgent(ctx, inv, client, autostart, codersdk.Me, workspaceName)
101-
if err != nil {
102-
return xerrors.Errorf("get workspace and agent retry: %w", err)
103-
}
90+
// The agent will report it's expanded directory before leaving
91+
// the created state, so we need to wait for that to happen.
92+
// However, if no directory is set, the expanded directory will
93+
// not be set either.
94+
if workspaceAgent.Directory != "" {
95+
workspace, workspaceAgent, err = waitForAgentCond(ctx, client, workspace, workspaceAgent, func(a codersdk.WorkspaceAgent) bool {
96+
return workspaceAgent.LifecycleState != codersdk.WorkspaceAgentLifecycleCreated
97+
})
10498
}
10599
}
106100

107-
directory := workspaceAgent.ExpandedDirectory // Empty unless agent directory is set.
101+
var directory string
108102
if len(inv.Args) > 1 {
109-
d := inv.Args[1]
110-
111-
switch {
112-
case insideThisWorkspace:
113-
// TODO(mafredri): Return error if directory doesn't exist?
114-
directory, err = filepath.Abs(d)
115-
if err != nil {
116-
return xerrors.Errorf("expand directory: %w", err)
117-
}
118-
119-
case d == "~" || strings.HasPrefix(d, "~/"):
120-
return xerrors.Errorf("path %q requires expansion and is not supported, use an absolute path instead", d)
121-
122-
case workspaceAgent.OperatingSystem == "windows":
123-
switch {
124-
case directory != "" && !isWindowsAbsPath(d):
125-
directory = windowsJoinPath(directory, d)
126-
case isWindowsAbsPath(d):
127-
directory = d
128-
default:
129-
return xerrors.Errorf("path %q not supported, use an absolute path instead", d)
130-
}
131-
132-
// Note that we use `path` instead of `filepath` since we want Unix behavior.
133-
case directory != "" && !path.IsAbs(d):
134-
directory = path.Join(directory, d)
135-
case path.IsAbs(d):
136-
directory = d
137-
default:
138-
return xerrors.Errorf("path %q not supported, use an absolute path instead", d)
139-
}
103+
directory = inv.Args[1]
140104
}
141-
142-
u, err := url.Parse("vscode://coder.coder-remote/open")
105+
directory, err = resolveAgentAbsPath(workspaceAgent.ExpandedDirectory, directory, workspaceAgent.OperatingSystem, insideThisWorkspace)
143106
if err != nil {
144-
return xerrors.Errorf("parse vscode URI: %w", err)
107+
return xerrors.Errorf("resolve agent path: %w", err)
108+
}
109+
110+
u := &url.URL{
111+
Scheme: "vscode",
112+
Host: "coder.coder-remote",
113+
Path: "/open",
145114
}
146115

147116
qp := url.Values{}
@@ -190,6 +159,16 @@ func (r *RootCmd) openVSCode() *clibase.Cmd {
190159
}
191160
if err != nil {
192161
if !generateToken {
162+
// This is not an important step, so we don't want
163+
// to block the user here.
164+
token := qp.Get("token")
165+
wait := doAsync(func() {
166+
// Best effort, we don't care if this fails.
167+
apiKeyID := strings.SplitN(token, "-", 2)[0]
168+
_ = client.DeleteAPIKey(ctx, codersdk.Me, apiKeyID)
169+
})
170+
defer wait()
171+
193172
qp.Del("token")
194173
u.RawQuery = qp.Encode()
195174
}
@@ -226,19 +205,50 @@ func (r *RootCmd) openVSCode() *clibase.Cmd {
226205
return cmd
227206
}
228207

208+
// waitForAgentCond uses the watch workspace API to update the agent information
209+
// until the condition is met.
210+
func waitForAgentCond(ctx context.Context, client *codersdk.Client, workspace codersdk.Workspace, workspaceAgent codersdk.WorkspaceAgent, cond func(codersdk.WorkspaceAgent) bool) (codersdk.Workspace, codersdk.WorkspaceAgent, error) {
211+
ctx, cancel := context.WithCancel(ctx)
212+
defer cancel()
213+
214+
if cond(workspaceAgent) {
215+
return workspace, workspaceAgent, nil
216+
}
217+
218+
wc, err := client.WatchWorkspace(ctx, workspace.ID)
219+
if err != nil {
220+
return workspace, workspaceAgent, xerrors.Errorf("watch workspace: %w", err)
221+
}
222+
223+
for workspace = range wc {
224+
workspaceAgent, err = getWorkspaceAgent(workspace, workspaceAgent.Name)
225+
if err != nil {
226+
return workspace, workspaceAgent, xerrors.Errorf("get workspace agent: %w", err)
227+
}
228+
if cond(workspaceAgent) {
229+
return workspace, workspaceAgent, nil
230+
}
231+
}
232+
233+
return workspace, workspaceAgent, xerrors.New("watch workspace: unexpected closed channel")
234+
}
235+
229236
// isWindowsAbsPath checks if the path is an absolute path on Windows. On Unix
230237
// systems the check is very simplistic and does not cover edge cases.
231-
//
232-
//nolint:revive // Shadow path variable for readability.
233-
func isWindowsAbsPath(path string) bool {
238+
func isWindowsAbsPath(p string) bool {
234239
if runtime.GOOS == "windows" {
235-
return filepath.IsAbs(path)
240+
return filepath.IsAbs(p)
236241
}
237242

238243
switch {
239-
case len(path) >= 2 && path[1] == ':':
244+
case len(p) < 2:
245+
return false
246+
case p[1] == ':':
240247
// Path starts with a drive letter.
241-
return len(path) == 2 || (len(path) >= 4 && path[2] == '\\' && path[3] == '\\')
248+
return len(p) == 2 || (len(p) >= 3 && p[2] == '\\')
249+
case p[0] == '\\' && p[1] == '\\':
250+
// Path starts with \\.
251+
return true
242252
default:
243253
return false
244254
}
@@ -262,7 +272,62 @@ func windowsJoinPath(elem ...string) string {
262272
s = e
263273
continue
264274
}
265-
s += "\\" + strings.TrimSuffix(s, "\\")
275+
s += "\\" + strings.TrimSuffix(e, "\\")
266276
}
267277
return s
268278
}
279+
280+
// resolveAgentAbsPath resolves the absolute path to a file or directory in the
281+
// workspace. If the path is relative, it will be resolved relative to the
282+
// workspace's expanded directory. If the path is absolute, it will be returned
283+
// as-is. If the path is relative and the workspace directory is not expanded,
284+
// an error will be returned.
285+
//
286+
// If the path is being resolved within the workspace, the path will be resolved
287+
// relative to the current working directory.
288+
func resolveAgentAbsPath(workingDirectory, relOrAbsPath, agentOS string, local bool) (string, error) {
289+
if relOrAbsPath == "" {
290+
return workingDirectory, nil
291+
}
292+
293+
switch {
294+
case relOrAbsPath == "~" || strings.HasPrefix(relOrAbsPath, "~/"):
295+
return "", xerrors.Errorf("path %q requires expansion and is not supported, use an absolute path instead", relOrAbsPath)
296+
297+
case local:
298+
p, err := filepath.Abs(relOrAbsPath)
299+
if err != nil {
300+
return "", xerrors.Errorf("expand path: %w", err)
301+
}
302+
return p, nil
303+
304+
case agentOS == "windows":
305+
switch {
306+
case workingDirectory != "" && !isWindowsAbsPath(relOrAbsPath):
307+
return windowsJoinPath(workingDirectory, relOrAbsPath), nil
308+
case isWindowsAbsPath(relOrAbsPath):
309+
return relOrAbsPath, nil
310+
default:
311+
return "", xerrors.Errorf("path %q not supported, use an absolute path instead", relOrAbsPath)
312+
}
313+
314+
// Note that we use `path` instead of `filepath` since we want Unix behavior.
315+
case workingDirectory != "" && !path.IsAbs(relOrAbsPath):
316+
return path.Join(workingDirectory, relOrAbsPath), nil
317+
case path.IsAbs(relOrAbsPath):
318+
return relOrAbsPath, nil
319+
default:
320+
return "", xerrors.Errorf("path %q not supported, use an absolute path instead", relOrAbsPath)
321+
}
322+
}
323+
324+
func doAsync(f func()) (wait func()) {
325+
done := make(chan struct{})
326+
go func() {
327+
defer close(done)
328+
f()
329+
}()
330+
return func() {
331+
<-done
332+
}
333+
}

‎cli/open_internal_test.go

+54
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
package cli
2+
3+
import "testing"
4+
5+
func Test_resolveAgentAbsPath(t *testing.T) {
6+
t.Parallel()
7+
8+
type args struct {
9+
workingDirectory string
10+
relOrAbsPath string
11+
agentOS string
12+
local bool
13+
}
14+
tests := []struct {
15+
name string
16+
args args
17+
want string
18+
wantErr bool
19+
}{
20+
{"ok no args", args{}, "", false},
21+
{"ok only working directory", args{workingDirectory: "/some/path"}, "/some/path", false},
22+
{"ok with working directory and rel path", args{workingDirectory: "/some/path", relOrAbsPath: "other/path"}, "/some/path/other/path", false},
23+
{"ok with working directory and abs path", args{workingDirectory: "/some/path", relOrAbsPath: "/other/path"}, "/other/path", false},
24+
{"ok with no working directory and abs path", args{relOrAbsPath: "/other/path"}, "/other/path", false},
25+
26+
{"fail tilde", args{relOrAbsPath: "~"}, "", true},
27+
{"fail tilde with working directory", args{workingDirectory: "/some/path", relOrAbsPath: "~"}, "", true},
28+
{"fail tilde path", args{relOrAbsPath: "~/some/path"}, "", true},
29+
{"fail tilde path with working directory", args{workingDirectory: "/some/path", relOrAbsPath: "~/some/path"}, "", true},
30+
{"fail relative dot with no working directory", args{relOrAbsPath: "."}, "", true},
31+
{"fail relative with no working directory", args{relOrAbsPath: "some/path"}, "", true},
32+
33+
{"ok with working directory and rel path on windows", args{workingDirectory: "C:\\some\\path", relOrAbsPath: "other\\path", agentOS: "windows"}, "C:\\some\\path\\othere\\path", false},
34+
{"ok with working directory and abs path on windows", args{workingDirectory: "C:\\some\\path", relOrAbsPath: "C:\\othere\\path", agentOS: "windows"}, "C:\\othere\\path", false},
35+
{"ok with no working directory and abs path on windows", args{relOrAbsPath: "C:\\othere\\path", agentOS: "windows"}, "C:\\othere\\path", false},
36+
37+
{"fail with no working directory and rel path on windows", args{relOrAbsPath: "other\\path", agentOS: "windows"}, "", true},
38+
}
39+
for _, tt := range tests {
40+
tt := tt
41+
t.Run(tt.name, func(t *testing.T) {
42+
t.Parallel()
43+
44+
got, err := resolveAgentAbsPath(tt.args.workingDirectory, tt.args.relOrAbsPath, tt.args.agentOS, tt.args.local)
45+
if (err != nil) != tt.wantErr {
46+
t.Errorf("resolveAgentAbsPath() error = %v, wantErr %v", err, tt.wantErr)
47+
return
48+
}
49+
if got != tt.want {
50+
t.Errorf("resolveAgentAbsPath() = %v, want %v", got, tt.want)
51+
}
52+
})
53+
}
54+
}

‎cli/ssh.go

+19-8
Original file line numberDiff line numberDiff line change
@@ -594,40 +594,51 @@ func getWorkspaceAndAgent(ctx context.Context, inv *clibase.Invocation, client *
594594
return codersdk.Workspace{}, codersdk.WorkspaceAgent{}, xerrors.Errorf("workspace %q is being deleted", workspace.Name)
595595
}
596596

597+
var agentName string
598+
if len(workspaceParts) >= 2 {
599+
agentName = workspaceParts[1]
600+
}
601+
workspaceAgent, err := getWorkspaceAgent(workspace, agentName)
602+
if err != nil {
603+
return codersdk.Workspace{}, codersdk.WorkspaceAgent{}, err
604+
}
605+
606+
return workspace, workspaceAgent, nil
607+
}
608+
609+
func getWorkspaceAgent(workspace codersdk.Workspace, agentName string) (workspaceAgent codersdk.WorkspaceAgent, err error) {
597610
resources := workspace.LatestBuild.Resources
598611

599612
agents := make([]codersdk.WorkspaceAgent, 0)
600613
for _, resource := range resources {
601614
agents = append(agents, resource.Agents...)
602615
}
603616
if len(agents) == 0 {
604-
return codersdk.Workspace{}, codersdk.WorkspaceAgent{}, xerrors.Errorf("workspace %q has no agents", workspace.Name)
617+
return codersdk.WorkspaceAgent{}, xerrors.Errorf("workspace %q has no agents", workspace.Name)
605618
}
606-
var workspaceAgent codersdk.WorkspaceAgent
607-
if len(workspaceParts) >= 2 {
619+
if agentName != "" {
608620
for _, otherAgent := range agents {
609-
if otherAgent.Name != workspaceParts[1] {
621+
if otherAgent.Name != agentName {
610622
continue
611623
}
612624
workspaceAgent = otherAgent
613625
break
614626
}
615627
if workspaceAgent.ID == uuid.Nil {
616-
return codersdk.Workspace{}, codersdk.WorkspaceAgent{}, xerrors.Errorf("agent not found by name %q", workspaceParts[1])
628+
return codersdk.WorkspaceAgent{}, xerrors.Errorf("agent not found by name %q", agentName)
617629
}
618630
}
619631
if workspaceAgent.ID == uuid.Nil {
620632
if len(agents) > 1 {
621633
workspaceAgent, err = cryptorand.Element(agents)
622634
if err != nil {
623-
return codersdk.Workspace{}, codersdk.WorkspaceAgent{}, err
635+
return codersdk.WorkspaceAgent{}, err
624636
}
625637
} else {
626638
workspaceAgent = agents[0]
627639
}
628640
}
629-
630-
return workspace, workspaceAgent, nil
641+
return workspaceAgent, nil
631642
}
632643

633644
// Attempt to poll workspace autostop. We write a per-workspace lockfile to

‎codersdk/workspaces.go

+5-1
Original file line numberDiff line numberDiff line change
@@ -219,7 +219,11 @@ func (c *Client) WatchWorkspace(ctx context.Context, id uuid.UUID) (<-chan Works
219219
if err != nil {
220220
return
221221
}
222-
wc <- ws
222+
select {
223+
case <-ctx.Done():
224+
return
225+
case wc <- ws:
226+
}
223227
}
224228
}
225229
}()

0 commit comments

Comments
 (0)