Close SSH client as well due to use of multiple sessions during forwarding

mafredri · mafredri · commit 03db373a8f0b · 2023-04-19T12:32:31.000Z
diff --git a/cli/ssh.go b/cli/ssh.go
@@ -114,10 +114,10 @@ func (r *RootCmd) ssh() *clibase.Cmd {
 			if err != nil {
 				return err
 			}
-			watchAndClose := func(c io.Closer) {
+			watchAndClose := func(closer func() error) {
 				// Ensure session is ended on both context cancellation
 				// and workspace stop.
-				defer c.Close()
+				defer closer()
 
 				for {
 					select {
@@ -145,7 +145,7 @@ func (r *RootCmd) ssh() *clibase.Cmd {
 					return err
 				}
 				defer rawSSH.Close()
-				go watchAndClose(rawSSH)
+				go watchAndClose(rawSSH.Close)
 
 				go func() {
 					// Ensure stdout copy closes incase stdin is closed
@@ -170,7 +170,11 @@ func (r *RootCmd) ssh() *clibase.Cmd {
 				return err
 			}
 			defer sshSession.Close()
-			go watchAndClose(sshSession)
+			go watchAndClose(func() error {
+				_ = sshSession.Close()
+				_ = sshClient.Close()
+				return nil
+			})
 
 			if identityAgent == "" {
 				identityAgent = os.Getenv("SSH_AUTH_SOCK")
