Merge pull request cli#154 from github/jg/fix-ctx-error

josebalius · web-flow · commit 10ad85486f96 · 2021-09-15T15:45:35.000-04:00
ghcs create/ssh: fix ctx cancellation errors &amp; fix todo for X11 forwarding
diff --git a/cmd/ghcs/create.go b/cmd/ghcs/create.go
@@ -105,12 +105,16 @@ func create(opts *createOptions) error {
 	return nil
 }
 
+// showStatus polls the codespace for a list of post create states and their status. It will keep polling
+// until all states have finished. Once all states have finished, we poll once more to check if any new
+// states have been introduced and stop polling otherwise.
 func showStatus(ctx context.Context, log *output.Logger, apiClient *api.API, user *api.User, codespace *api.Codespace) error {
 	var lastState codespaces.PostCreateState
 	var breakNextState bool
 
 	finishedStates := make(map[string]bool)
 	ctx, stopPolling := context.WithCancel(ctx)
+	defer stopPolling()
 
 	poller := func(states []codespaces.PostCreateState) {
 		var inProgress bool
@@ -153,7 +157,12 @@ func showStatus(ctx context.Context, log *output.Logger, apiClient *api.API, use
 		}
 	}
 
-	if err := codespaces.PollPostCreateStates(ctx, log, apiClient, user, codespace, poller); err != nil {
+	err := codespaces.PollPostCreateStates(ctx, log, apiClient, user, codespace, poller)
+	if err != nil {
+		if errors.Is(err, context.Canceled) && breakNextState {
+			return nil // we cancelled the context to stop polling, we can ignore the error
+		}
+
 		return fmt.Errorf("failed to poll state changes from codespace: %v", err)
 	}
 
diff --git a/internal/codespaces/ssh.go b/internal/codespaces/ssh.go
@@ -32,9 +32,14 @@ func NewRemoteCommand(ctx context.Context, tunnelPort int, destination, command
 // an interactive shell) over ssh.
 func newSSHCommand(ctx context.Context, port int, dst, command string) (*exec.Cmd, []string) {
 	connArgs := []string{"-p", strconv.Itoa(port), "-o", "NoHostAuthenticationForLocalhost=yes"}
-	// TODO(adonovan): eliminate X11 and X11Trust flags where unneeded.
-	cmdArgs := append([]string{dst, "-X", "-Y", "-C"}, connArgs...) // X11, X11Trust, Compression
 
+	cmdArgs := []string{dst, "-C"} // Always use Compression
+	if command == "" {
+		// if we are in a shell send X11 and X11Trust
+		cmdArgs = append(cmdArgs, "-X", "-Y")
+	}
+
+	cmdArgs = append(cmdArgs, connArgs...)
 	if command != "" {
 		cmdArgs = append(cmdArgs, command)
 	}
diff --git a/internal/codespaces/states.go b/internal/codespaces/states.go
@@ -39,12 +39,12 @@ type PostCreateState struct {
 func PollPostCreateStates(ctx context.Context, log logger, apiClient *api.API, user *api.User, codespace *api.Codespace, poller func([]PostCreateState)) error {
 	token, err := apiClient.GetCodespaceToken(ctx, user.Login, codespace.Name)
 	if err != nil {
-		return fmt.Errorf("getting codespace token: %v", err)
+		return fmt.Errorf("getting codespace token: %w", err)
 	}
 
 	session, err := ConnectToLiveshare(ctx, log, apiClient, user.Login, token, codespace)
 	if err != nil {
-		return fmt.Errorf("connect to Live Share: %v", err)
+		return fmt.Errorf("connect to Live Share: %w", err)
 	}
 
 	// Ensure local port is listening before client (getPostCreateOutput) connects.
@@ -57,7 +57,7 @@ func PollPostCreateStates(ctx context.Context, log logger, apiClient *api.API, u
 	log.Println("Fetching SSH Details...")
 	remoteSSHServerPort, sshUser, err := session.StartSSHServer(ctx)
 	if err != nil {
-		return fmt.Errorf("error getting ssh server details: %v", err)
+		return fmt.Errorf("error getting ssh server details: %w", err)
 	}
 
 	tunnelClosed := make(chan error, 1) // buffered to avoid sender stuckness
@@ -75,12 +75,12 @@ func PollPostCreateStates(ctx context.Context, log logger, apiClient *api.API, u
 			return ctx.Err()
 
 		case err := <-tunnelClosed:
-			return fmt.Errorf("connection failed: %v", err)
+			return fmt.Errorf("connection failed: %w", err)
 
 		case <-t.C:
 			states, err := getPostCreateOutput(ctx, localPort, codespace, sshUser)
 			if err != nil {
-				return fmt.Errorf("get post create output: %v", err)
+				return fmt.Errorf("get post create output: %w", err)
 			}
 
 			poller(states)
@@ -96,13 +96,13 @@ func getPostCreateOutput(ctx context.Context, tunnelPort int, codespace *api.Cod
 	stdout := new(bytes.Buffer)
 	cmd.Stdout = stdout
 	if err := cmd.Run(); err != nil {
-		return nil, fmt.Errorf("run command: %v", err)
+		return nil, fmt.Errorf("run command: %w", err)
 	}
 	var output struct {
 		Steps []PostCreateState `json:"steps"`
 	}
 	if err := json.Unmarshal(stdout.Bytes(), &output); err != nil {
-		return nil, fmt.Errorf("unmarshal output: %v", err)
+		return nil, fmt.Errorf("unmarshal output: %w", err)
 	}
 
 	return output.Steps, nil
