Skip to content

Commit 198d705

Browse files
theFongtheFong
committed
refactor: deduplicate SSH/polling utils and fix bounds check
Address PR review comments: - Extract waitForSSHToBeAvailable, pollUntil, startWorkspaceIfStopped into shared pkg/cmd/util/ssh.go to eliminate duplication between exec.go and shell.go - Add bounds check on strings.Split before accessing index 1 in waitForSSHToBeAvailable (prevents potential panic) - Add 10-minute timeout to pollUntil to prevent infinite hangs - Clean up openTerminalWithTmux to use _ in function signature
1 parent 4b6003e commit 198d705

4 files changed

Lines changed: 125 additions & 157 deletions

File tree

‎pkg/cmd/exec/exec.go‎

Lines changed: 6 additions & 77 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,6 @@ package exec
22

33
import (
44
"bufio"
5-
"errors"
65
"fmt"
76
"os"
87
"os/exec"
@@ -18,7 +17,6 @@ import (
1817
"github.com/brevdev/brev-cli/pkg/store"
1918
"github.com/brevdev/brev-cli/pkg/terminal"
2019
"github.com/brevdev/brev-cli/pkg/writeconnectionevent"
21-
"github.com/briandowns/spinner"
2220
"github.com/hashicorp/go-multierror"
2321

2422
"github.com/spf13/cobra"
@@ -51,13 +49,10 @@ var (
5149
)
5250

5351
type ExecStore interface {
54-
util.GetWorkspaceByNameOrIDErrStore
52+
util.WorkspaceStartStore
5553
refresh.RefreshStore
5654
GetOrganizations(options *store.GetOrganizationsOptions) ([]entity.Organization, error)
5755
GetWorkspaces(organizationID string, options *store.GetWorkspacesOptions) ([]entity.Workspace, error)
58-
StartWorkspace(workspaceID string) (*entity.Workspace, error)
59-
GetWorkspace(workspaceID string) (*entity.Workspace, error)
60-
GetCurrentUserKeys() (*entity.UserKeys, error)
6156
}
6257

6358
func NewCmdExec(t *terminal.Terminal, store ExecStore, noLoginStartStore ExecStore) *cobra.Command {
@@ -178,6 +173,8 @@ func parseCommand(command string) (string, error) {
178173
return command, nil
179174
}
180175

176+
const pollTimeout = 10 * time.Minute
177+
181178
func runExecCommand(t *terminal.Terminal, sstore ExecStore, workspaceNameOrID string, host bool, command string) error {
182179
s := t.NewSpinner()
183180
workspace, err := util.GetUserWorkspaceByNameOrIDErr(sstore, workspaceNameOrID)
@@ -186,12 +183,12 @@ func runExecCommand(t *terminal.Terminal, sstore ExecStore, workspaceNameOrID st
186183
}
187184

188185
if workspace.Status == "STOPPED" { // we start the env for the user
189-
err = startWorkspaceIfStopped(t, s, sstore, workspaceNameOrID, workspace)
186+
err = util.StartWorkspaceIfStopped(t, s, sstore, workspaceNameOrID, workspace, pollTimeout)
190187
if err != nil {
191188
return breverrors.WrapAndTrace(err)
192189
}
193190
}
194-
err = pollUntil(s, workspace.ID, "RUNNING", sstore, " waiting for instance to be ready...")
191+
err = util.PollUntil(s, workspace.ID, "RUNNING", sstore, " waiting for instance to be ready...", pollTimeout)
195192
if err != nil {
196193
return breverrors.WrapAndTrace(err)
197194
}
@@ -216,7 +213,7 @@ func runExecCommand(t *terminal.Terminal, sstore ExecStore, workspaceNameOrID st
216213
if err != nil {
217214
return breverrors.WrapAndTrace(err)
218215
}
219-
err = waitForSSHToBeAvailable(sshName, s)
216+
err = util.WaitForSSHToBeAvailable(sshName, s)
220217
if err != nil {
221218
return breverrors.WrapAndTrace(err)
222219
}
@@ -248,30 +245,6 @@ func runExecCommand(t *terminal.Terminal, sstore ExecStore, workspaceNameOrID st
248245
return nil
249246
}
250247

251-
func waitForSSHToBeAvailable(sshAlias string, s *spinner.Spinner) error {
252-
counter := 0
253-
s.Suffix = " waiting for SSH connection to be available"
254-
s.Start()
255-
for {
256-
cmd := exec.Command("ssh", "-o", "ConnectTimeout=10", sshAlias, "echo", " ")
257-
out, err := cmd.CombinedOutput()
258-
if err == nil {
259-
s.Stop()
260-
return nil
261-
}
262-
263-
outputStr := string(out)
264-
stdErr := strings.Split(outputStr, "\n")[1]
265-
266-
if counter == 40 || !store.SatisfactorySSHErrMessage(stdErr) {
267-
return breverrors.WrapAndTrace(errors.New("\n" + stdErr))
268-
}
269-
270-
counter++
271-
time.Sleep(1 * time.Second)
272-
}
273-
}
274-
275248
func runSSH(sshAlias string, command string) error {
276249
sshAgentEval := "eval $(ssh-agent -s)"
277250

@@ -292,47 +265,3 @@ func runSSH(sshAlias string, command string) error {
292265
}
293266
return nil
294267
}
295-
296-
func startWorkspaceIfStopped(t *terminal.Terminal, s *spinner.Spinner, tstore ExecStore, wsIDOrName string, workspace *entity.Workspace) error {
297-
activeOrg, err := tstore.GetActiveOrganizationOrDefault()
298-
if err != nil {
299-
return breverrors.WrapAndTrace(err)
300-
}
301-
workspaces, err := tstore.GetWorkspaceByNameOrID(activeOrg.ID, wsIDOrName)
302-
if err != nil {
303-
return breverrors.WrapAndTrace(err)
304-
}
305-
startedWorkspace, err := tstore.StartWorkspace(workspaces[0].ID)
306-
if err != nil {
307-
return breverrors.WrapAndTrace(err)
308-
}
309-
t.Vprintf("%s", t.Yellow("Instance %s is starting. \n\n", startedWorkspace.Name))
310-
err = pollUntil(s, workspace.ID, entity.Running, tstore, " hang tight")
311-
if err != nil {
312-
return breverrors.WrapAndTrace(err)
313-
}
314-
workspace, err = util.GetUserWorkspaceByNameOrIDErr(tstore, wsIDOrName)
315-
if err != nil {
316-
return breverrors.WrapAndTrace(err)
317-
}
318-
return nil
319-
}
320-
321-
func pollUntil(s *spinner.Spinner, wsid string, state string, execStore ExecStore, waitMsg string) error {
322-
isReady := false
323-
s.Suffix = waitMsg
324-
s.Start()
325-
for !isReady {
326-
time.Sleep(5 * time.Second)
327-
ws, err := execStore.GetWorkspace(wsid)
328-
if err != nil {
329-
return breverrors.WrapAndTrace(err)
330-
}
331-
s.Suffix = waitMsg
332-
if ws.Status == state {
333-
isReady = true
334-
}
335-
}
336-
s.Stop()
337-
return nil
338-
}

‎pkg/cmd/open/open.go‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -596,7 +596,11 @@ func waitForSSHToBeAvailable(t *terminal.Terminal, s *spinner.Spinner, sshAlias
596596
}
597597

598598
outputStr := string(out)
599-
stdErr := strings.Split(outputStr, "\n")[1]
599+
lines := strings.Split(outputStr, "\n")
600+
stdErr := outputStr
601+
if len(lines) > 1 {
602+
stdErr = lines[1]
603+
}
600604

601605
if counter == 160 || !store.SatisfactorySSHErrMessage(stdErr) {
602606
return breverrors.WrapAndTrace(errors.New("\n" + stdErr))
@@ -732,8 +736,7 @@ func openTerminal(sshAlias string, _ string, _ OpenStore) error {
732736
return nil
733737
}
734738

735-
func openTerminalWithTmux(sshAlias string, path string, store OpenStore) error {
736-
_ = store // unused parameter required by interface
739+
func openTerminalWithTmux(sshAlias string, path string, _ OpenStore) error {
737740

738741
err := ensureTmuxInstalled(sshAlias)
739742
if err != nil {

‎pkg/cmd/shell/shell.go‎

Lines changed: 6 additions & 77 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,9 @@
11
package shell
22

33
import (
4-
"errors"
54
"fmt"
65
"os"
76
"os/exec"
8-
"strings"
97
"time"
108

119
"github.com/brevdev/brev-cli/pkg/analytics"
@@ -18,7 +16,6 @@ import (
1816
"github.com/brevdev/brev-cli/pkg/store"
1917
"github.com/brevdev/brev-cli/pkg/terminal"
2018
"github.com/brevdev/brev-cli/pkg/writeconnectionevent"
21-
"github.com/briandowns/spinner"
2219

2320
"github.com/spf13/cobra"
2421
)
@@ -43,13 +40,10 @@ var (
4340
)
4441

4542
type ShellStore interface {
46-
util.GetWorkspaceByNameOrIDErrStore
43+
util.WorkspaceStartStore
4744
refresh.RefreshStore
4845
GetOrganizations(options *store.GetOrganizationsOptions) ([]entity.Organization, error)
4946
GetWorkspaces(organizationID string, options *store.GetWorkspacesOptions) ([]entity.Workspace, error)
50-
StartWorkspace(workspaceID string) (*entity.Workspace, error)
51-
GetWorkspace(workspaceID string) (*entity.Workspace, error)
52-
GetCurrentUserKeys() (*entity.UserKeys, error)
5347
}
5448

5549
func NewCmdShell(t *terminal.Terminal, store ShellStore, noLoginStartStore ShellStore) *cobra.Command {
@@ -78,6 +72,8 @@ func NewCmdShell(t *terminal.Terminal, store ShellStore, noLoginStartStore Shell
7872
return cmd
7973
}
8074

75+
const pollTimeout = 10 * time.Minute
76+
8177
func runShellCommand(t *terminal.Terminal, sstore ShellStore, workspaceNameOrID string, host bool) error {
8278
s := t.NewSpinner()
8379
workspace, err := util.GetUserWorkspaceByNameOrIDErr(sstore, workspaceNameOrID)
@@ -86,12 +82,12 @@ func runShellCommand(t *terminal.Terminal, sstore ShellStore, workspaceNameOrID
8682
}
8783

8884
if workspace.Status == "STOPPED" { // we start the env for the user
89-
err = startWorkspaceIfStopped(t, s, sstore, workspaceNameOrID, workspace)
85+
err = util.StartWorkspaceIfStopped(t, s, sstore, workspaceNameOrID, workspace, pollTimeout)
9086
if err != nil {
9187
return breverrors.WrapAndTrace(err)
9288
}
9389
}
94-
err = pollUntil(s, workspace.ID, "RUNNING", sstore, " waiting for instance to be ready...")
90+
err = util.PollUntil(s, workspace.ID, "RUNNING", sstore, " waiting for instance to be ready...", pollTimeout)
9591
if err != nil {
9692
return breverrors.WrapAndTrace(err)
9793
}
@@ -116,7 +112,7 @@ func runShellCommand(t *terminal.Terminal, sstore ShellStore, workspaceNameOrID
116112
if err != nil {
117113
return breverrors.WrapAndTrace(err)
118114
}
119-
err = waitForSSHToBeAvailable(sshName, s)
115+
err = util.WaitForSSHToBeAvailable(sshName, s)
120116
if err != nil {
121117
return breverrors.WrapAndTrace(err)
122118
}
@@ -148,30 +144,6 @@ func runShellCommand(t *terminal.Terminal, sstore ShellStore, workspaceNameOrID
148144
return nil
149145
}
150146

151-
func waitForSSHToBeAvailable(sshAlias string, s *spinner.Spinner) error {
152-
counter := 0
153-
s.Suffix = " waiting for SSH connection to be available"
154-
s.Start()
155-
for {
156-
cmd := exec.Command("ssh", "-o", "ConnectTimeout=10", sshAlias, "echo", " ")
157-
out, err := cmd.CombinedOutput()
158-
if err == nil {
159-
s.Stop()
160-
return nil
161-
}
162-
163-
outputStr := string(out)
164-
stdErr := strings.Split(outputStr, "\n")[1]
165-
166-
if counter == 40 || !store.SatisfactorySSHErrMessage(stdErr) {
167-
return breverrors.WrapAndTrace(errors.New("\n" + stdErr))
168-
}
169-
170-
counter++
171-
time.Sleep(1 * time.Second)
172-
}
173-
}
174-
175147
func runSSH(sshAlias string) error {
176148
sshAgentEval := "eval $(ssh-agent -s)"
177149
cmd := fmt.Sprintf("%s && ssh %s", sshAgentEval, sshAlias)
@@ -192,46 +164,3 @@ func runSSH(sshAlias string) error {
192164
}
193165
return nil
194166
}
195-
196-
func startWorkspaceIfStopped(t *terminal.Terminal, s *spinner.Spinner, tstore ShellStore, wsIDOrName string, workspace *entity.Workspace) error {
197-
activeOrg, err := tstore.GetActiveOrganizationOrDefault()
198-
if err != nil {
199-
return breverrors.WrapAndTrace(err)
200-
}
201-
workspaces, err := tstore.GetWorkspaceByNameOrID(activeOrg.ID, wsIDOrName)
202-
if err != nil {
203-
return breverrors.WrapAndTrace(err)
204-
}
205-
startedWorkspace, err := tstore.StartWorkspace(workspaces[0].ID)
206-
if err != nil {
207-
return breverrors.WrapAndTrace(err)
208-
}
209-
t.Vprintf("%s", t.Yellow("Instance %s is starting. \n\n", startedWorkspace.Name))
210-
err = pollUntil(s, workspace.ID, entity.Running, tstore, " hang tight 🤙")
211-
if err != nil {
212-
return breverrors.WrapAndTrace(err)
213-
}
214-
workspace, err = util.GetUserWorkspaceByNameOrIDErr(tstore, wsIDOrName)
215-
if err != nil {
216-
return breverrors.WrapAndTrace(err)
217-
}
218-
return nil
219-
}
220-
221-
func pollUntil(s *spinner.Spinner, wsid string, state string, shellStore ShellStore, waitMsg string) error {
222-
isReady := false
223-
s.Suffix = waitMsg
224-
s.Start()
225-
for !isReady {
226-
time.Sleep(5 * time.Second)
227-
ws, err := shellStore.GetWorkspace(wsid)
228-
if err != nil {
229-
return breverrors.WrapAndTrace(err)
230-
}
231-
s.Suffix = waitMsg
232-
if ws.Status == state {
233-
isReady = true
234-
}
235-
}
236-
return nil
237-
}

0 commit comments

Comments
 (0)