chore: PR comments

deansheather · deansheather · commit a90d3d634cb5 · 2023-01-05T05:52:22.000Z
diff --git a/agent/ssh.go b/agent/ssh.go
@@ -44,7 +44,7 @@ func (h *forwardedUnixHandler) HandleSSHRequest(ctx ssh.Context, _ *ssh.Server,
 	conn, ok := ctx.Value(ssh.ContextKeyConn).(*gossh.ServerConn)
 	if !ok {
 		h.log.Warn(ctx, "SSH unix forward request from client with no gossh connection")
-		return false, []byte{}
+		return false, nil
 	}
 
 	switch req.Type {
@@ -53,7 +53,7 @@ func (h *forwardedUnixHandler) HandleSSHRequest(ctx ssh.Context, _ *ssh.Server,
 		err := gossh.Unmarshal(req.Payload, &reqPayload)
 		if err != nil {
 			h.log.Warn(ctx, "parse streamlocal-forward@openssh.com request payload from client", slog.Error(err))
-			return false, []byte{}
+			return false, nil
 		}
 
 		addr := reqPayload.SocketPath
@@ -64,7 +64,7 @@ func (h *forwardedUnixHandler) HandleSSHRequest(ctx ssh.Context, _ *ssh.Server,
 			h.log.Warn(ctx, "SSH unix forward request for socket path that is already being forwarded (maybe to another client?)",
 				slog.F("socket_path", addr),
 			)
-			return false, []byte{}
+			return false, nil
 		}
 
 		// Create socket parent dir if not exists.
@@ -76,7 +76,7 @@ func (h *forwardedUnixHandler) HandleSSHRequest(ctx ssh.Context, _ *ssh.Server,
 				slog.F("socket_path", addr),
 				slog.Error(err),
 			)
-			return false, []byte{}
+			return false, nil
 		}
 
 		ln, err := net.Listen("unix", addr)
@@ -85,19 +85,20 @@ func (h *forwardedUnixHandler) HandleSSHRequest(ctx ssh.Context, _ *ssh.Server,
 				slog.F("socket_path", addr),
 				slog.Error(err),
 			)
-			return false, []byte{}
+			return false, nil
 		}
+
+		// The listener needs to successfully start before it can be added to
+		// the map, so we don't have to worry about checking for an existing
+		// listener.
+		//
+		// This is also what the upstream TCP version of this code does.
 		h.Lock()
 		h.forwards[addr] = ln
 		h.Unlock()
 		go func() {
 			<-ctx.Done()
-			h.Lock()
-			ln, ok := h.forwards[addr]
-			h.Unlock()
-			if ok {
-				_ = ln.Close()
-			}
+			_ = ln.Close()
 		}()
 		go func() {
 			for {
@@ -109,6 +110,7 @@ func (h *forwardedUnixHandler) HandleSSHRequest(ctx ssh.Context, _ *ssh.Server,
 							slog.Error(err),
 						)
 					}
+					// closed below
 					break
 				}
 				payload := gossh.Marshal(&forwardedStreamLocalPayload{
@@ -129,9 +131,14 @@ func (h *forwardedUnixHandler) HandleSSHRequest(ctx ssh.Context, _ *ssh.Server,
 					Bicopy(ctx, ch, c)
 				}()
 			}
+
 			h.Lock()
-			delete(h.forwards, addr)
+			ln2, ok := h.forwards[addr]
+			if ok && ln2 == ln {
+				delete(h.forwards, addr)
+			}
 			h.Unlock()
+			_ = ln.Close()
 		}()
 
 		return true, nil
diff --git a/cli/ssh.go b/cli/ssh.go
@@ -221,30 +221,18 @@ func ssh() *cobra.Command {
 				}
 			}
 
-			// Wait for the context to be canceled, or the SSH session to end.
-			sshErr := make(chan error)
-			go func() {
-				defer close(sshErr)
-
-				err = sshSession.Wait()
-				if err != nil {
-					// If the connection drops unexpectedly, we get an ExitMissingError but no other
-					// error details, so try to at least give the user a better message
-					if errors.Is(err, &gossh.ExitMissingError{}) {
-						sshErr <- xerrors.New("SSH connection ended unexpectedly")
-						return
-					}
-					sshErr <- err
+			err = sshSession.Wait()
+			if err != nil {
+				// If the connection drops unexpectedly, we get an
+				// ExitMissingError but no other error details, so try to at
+				// least give the user a better message
+				if errors.Is(err, &gossh.ExitMissingError{}) {
+					return xerrors.New("SSH connection ended unexpectedly")
 				}
-			}()
-
-			select {
-			case <-ctx.Done():
-				_ = sshSession.Close()
-				return ctx.Err()
-			case err := <-sshErr:
 				return err
 			}
+
+			return nil
 		},
 	}
 	cliflag.BoolVarP(cmd.Flags(), &stdio, "stdio", "", "CODER_SSH_STDIO", false, "Specifies whether to emit SSH output over stdin/stdout.")
@@ -456,7 +444,12 @@ func uploadGPGKeys(ctx context.Context, sshClient *gossh.Client) error {
 	// Check if the agent is running in the workspace already.
 	// Note: we don't support windows in the workspace for GPG forwarding so
 	//       using shell commands is fine.
-	agentSocketBytes, err := runRemoteSSH(sshClient, nil, "set -eux; agent_socket=$(gpgconf --list-dir agent-socket); echo $agent_socket; test ! -S $agent_socket")
+	agentSocketBytes, err := runRemoteSSH(sshClient, nil, `
+set -eux
+agent_socket=$(gpgconf --list-dir agent-socket)
+echo "$agent_socket"
+test ! -S "$agent_socket"
+`)
 	agentSocket := strings.TrimSpace(string(agentSocketBytes))
 	if err != nil {
 		return xerrors.Errorf("check if agent socket is running (check if %q exists): %w", agentSocket, err)
@@ -540,24 +533,30 @@ func sshForwardRemote(ctx context.Context, stderr io.Writer, sshClient *gossh.Cl
 				return
 			}
 
-			localConn, err := net.Dial(localAddr.Network(), localAddr.String())
-			if err != nil {
-				_, _ = fmt.Fprintf(stderr, "Dial local address %s: %+v\n", localAddr.String(), err)
-				_ = remoteConn.Close()
-				continue
-			}
+			go func() {
+				defer func() {
+					_ = remoteConn.Close()
+				}()
 
-			if c, ok := localAddr.(cookieAddr); ok {
-				_, err = localConn.Write(c.cookie)
+				localConn, err := net.Dial(localAddr.Network(), localAddr.String())
 				if err != nil {
-					_, _ = fmt.Fprintf(stderr, "Write cookie to local connection: %+v\n", err)
+					_, _ = fmt.Fprintf(stderr, "Dial local address %s: %+v\n", localAddr.String(), err)
+					return
+				}
+				defer func() {
 					_ = localConn.Close()
-					_ = remoteConn.Close()
-					continue
+				}()
+
+				if c, ok := localAddr.(cookieAddr); ok {
+					_, err = localConn.Write(c.cookie)
+					if err != nil {
+						_, _ = fmt.Fprintf(stderr, "Write cookie to local connection: %+v\n", err)
+						return
+					}
 				}
-			}
 
-			go agent.Bicopy(ctx, localConn, remoteConn)
+				agent.Bicopy(ctx, localConn, remoteConn)
+			}()
 		}
 	}()
 
