Skip to content

Commit 3c3564a

Browse files
mafredrimafredri
committed
refactor
1 parent 5d370fc commit 3c3564a

File tree

4 files changed

+205
-68
lines changed

4 files changed

+205
-68
lines changed

‎cli/open.go

+127-59
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,33 @@ 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)
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+
})
10198
if err != nil {
102-
return xerrors.Errorf("get workspace and agent retry: %w", err)
99+
return xerrors.Errorf("wait for agent: %w", err)
103100
}
104101
}
105102
}
106103

107-
directory := workspaceAgent.ExpandedDirectory // Empty unless agent directory is set.
104+
var directory string
108105
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-
}
106+
directory = inv.Args[1]
140107
}
141-
142-
u, err := url.Parse("vscode://coder.coder-remote/open")
108+
directory, err = resolveAgentAbsPath(workspaceAgent.ExpandedDirectory, directory, workspaceAgent.OperatingSystem, insideThisWorkspace)
143109
if err != nil {
144-
return xerrors.Errorf("parse vscode URI: %w", err)
110+
return xerrors.Errorf("resolve agent path: %w", err)
111+
}
112+
113+
u := &url.URL{
114+
Scheme: "vscode",
115+
Host: "coder.coder-remote",
116+
Path: "/open",
145117
}
146118

147119
qp := url.Values{}
@@ -190,6 +162,16 @@ func (r *RootCmd) openVSCode() *clibase.Cmd {
190162
}
191163
if err != nil {
192164
if !generateToken {
165+
// This is not an important step, so we don't want
166+
// to block the user here.
167+
token := qp.Get("token")
168+
wait := doAsync(func() {
169+
// Best effort, we don't care if this fails.
170+
apiKeyID := strings.SplitN(token, "-", 2)[0]
171+
_ = client.DeleteAPIKey(ctx, codersdk.Me, apiKeyID)
172+
})
173+
defer wait()
174+
193175
qp.Del("token")
194176
u.RawQuery = qp.Encode()
195177
}
@@ -226,19 +208,50 @@ func (r *RootCmd) openVSCode() *clibase.Cmd {
226208
return cmd
227209
}
228210

211+
// waitForAgentCond uses the watch workspace API to update the agent information
212+
// until the condition is met.
213+
func waitForAgentCond(ctx context.Context, client *codersdk.Client, workspace codersdk.Workspace, workspaceAgent codersdk.WorkspaceAgent, cond func(codersdk.WorkspaceAgent) bool) (codersdk.Workspace, codersdk.WorkspaceAgent, error) {
214+
ctx, cancel := context.WithCancel(ctx)
215+
defer cancel()
216+
217+
if cond(workspaceAgent) {
218+
return workspace, workspaceAgent, nil
219+
}
220+
221+
wc, err := client.WatchWorkspace(ctx, workspace.ID)
222+
if err != nil {
223+
return workspace, workspaceAgent, xerrors.Errorf("watch workspace: %w", err)
224+
}
225+
226+
for workspace = range wc {
227+
workspaceAgent, err = getWorkspaceAgent(workspace, workspaceAgent.Name)
228+
if err != nil {
229+
return workspace, workspaceAgent, xerrors.Errorf("get workspace agent: %w", err)
230+
}
231+
if cond(workspaceAgent) {
232+
return workspace, workspaceAgent, nil
233+
}
234+
}
235+
236+
return workspace, workspaceAgent, xerrors.New("watch workspace: unexpected closed channel")
237+
}
238+
229239
// isWindowsAbsPath checks if the path is an absolute path on Windows. On Unix
230240
// 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 {
241+
func isWindowsAbsPath(p string) bool {
234242
if runtime.GOOS == "windows" {
235-
return filepath.IsAbs(path)
243+
return filepath.IsAbs(p)
236244
}
237245

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

‎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\\other\\path", false},
34+
{"ok with working directory and abs path on windows", args{workingDirectory: "C:\\some\\path", relOrAbsPath: "C:\\other\\path", agentOS: "windows"}, "C:\\other\\path", false},
35+
{"ok with no working directory and abs path on windows", args{relOrAbsPath: "C:\\other\\path", agentOS: "windows"}, "C:\\other\\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)