Remove manual process management in favor of CommandContext()

Golang provides the function exec.CommandContext(). Using this, spawned
resources can be linked to the session context and are cleaned up
automatically at session close.

Author: rumpelsepp

ssh_session.go

@@ -36,9 +36,10 @@ func createSSHSessionHandler(shell string) ssh.Handler {
 			createPty(s, shell)
 
 		case len(s.Command()) > 0:
-			log.Printf("No PTY requested, executing command: '%s'", s.RawCommand())
+			log.Printf("Command execution requested: '%s'", s.RawCommand())
+
+			cmd := exec.CommandContext(s.Context(), s.Command()[0], s.Command()[1:]...)
 
-			cmd := exec.Command(s.Command()[0], s.Command()[1:]...)
 			// We use StdinPipe to avoid blocking on missing input
 			if stdin, err := cmd.StdinPipe(); err != nil {
 				log.Println("Could not initialize stdinPipe", err)
@@ -49,11 +50,10 @@ func createSSHSessionHandler(shell string) ssh.Handler {
 					if _, err := io.Copy(stdin, s); err != nil {
 						log.Printf("Error while copying input from %s to stdin: %s", s.RemoteAddr().String(), err)
 					}
-					if err := stdin.Close(); err != nil {
-						log.Println("Error while closing stdinPipe:", err)
-					}
+					s.Close()
 				}()
 			}
+
 			cmd.Stdout = s
 			cmd.Stderr = s
 
@@ -64,27 +64,24 @@ func createSSHSessionHandler(shell string) ssh.Handler {
 			case err := <-done:
 				if err != nil {
 					log.Println("Command execution failed:", err)
-					io.WriteString(s, "Command execution failed: "+err.Error())
+					io.WriteString(s, "Command execution failed: "+err.Error()+"\n")
 				} else {
 					log.Println("Command execution successful")
 				}
 				s.Exit(cmd.ProcessState.ExitCode())
 
 			case <-s.Context().Done():
-				log.Println("Session closed by remote, killing dangling process")
-				if cmd.Process != nil && cmd.ProcessState == nil {
-					if err := cmd.Process.Kill(); err != nil {
-						log.Println("Failed to kill process:", err)
-					}
-				}
+				log.Printf("Session terminated: %s", s.Context().Err())
+				return
 			}
 
 		default:
 			log.Println("No PTY requested, no command supplied")
 
+			// Keep this open until the session exits, could e.g. be port forwarding
 			select {
 			case <-s.Context().Done():
-				log.Println("Session closed")
+				log.Printf("Session terminated: %s", s.Context().Err())
 			}
 		}
 	}

ssh_session_unix.go

@@ -30,9 +30,11 @@ import (
 )
 
 func createPty(s ssh.Session, shell string) {
-	ptyReq, winCh, _ := s.Pty()
+	var (
+		ptyReq, winCh, _ = s.Pty()
+		cmd              = exec.CommandContext(s.Context(), shell)
+	)
 
-	cmd := exec.Command(shell)
 	cmd.Env = append(cmd.Env, fmt.Sprintf("TERM=%s", ptyReq.Term))
 	f, err := pty.Start(cmd)
 	if err != nil {
@@ -61,11 +63,7 @@ func createPty(s ssh.Session, shell string) {
 		s.Exit(cmd.ProcessState.ExitCode())
 
 	case <-s.Context().Done():
-		log.Println("Session closed by remote, killing dangling process")
-		if cmd.Process != nil && cmd.ProcessState == nil {
-			if err := cmd.Process.Kill(); err != nil {
-				log.Println("Failed to kill process:", err)
-			}
-		}
+		log.Printf("Session terminated: %s", s.Context().Err())
+		return
 	}
 }

ssh_session_windows.go

@@ -51,7 +51,7 @@ func createPty(s ssh.Session, shell string) {
 		}
 		log.Println("Launching shell with ssh-shellhost.exe")
 
-		cmd := exec.Command(shell)
+		cmd := exec.CommandContext(s.Context(), shell)
 		cmd.SysProcAttr = &syscall.SysProcAttr{
 			HideWindow:    true,
 			CmdLine:       " " + "---pty cmd", // Must leave a space to the beginning
@@ -67,9 +67,7 @@ func createPty(s ssh.Session, shell string) {
 				if _, err := io.Copy(stdin, s); err != nil {
 					log.Printf("Error while copying input from %s to stdin: %s", s.RemoteAddr().String(), err)
 				}
-				if err := stdin.Close(); err != nil {
-					log.Println("Error while closing stdinPipe:", err)
-				}
+				s.Close()
 			}()
 		}
 		cmd.Stdout = s
@@ -88,12 +86,8 @@ func createPty(s ssh.Session, shell string) {
 			s.Exit(cmd.ProcessState.ExitCode())
 
 		case <-s.Context().Done():
-			log.Println("Session closed by remote, killing dangling process")
-			if cmd.Process != nil && cmd.ProcessState == nil {
-				if err := cmd.Process.Kill(); err != nil {
-					log.Println("Failed to kill process:", err)
-				}
-			}
+			log.Printf("Session terminated: %s", s.Context().Err())
+			return
 		}
 
 	} else {
@@ -128,6 +122,7 @@ func createPty(s ssh.Session, shell string) {
 		if err != nil {
 			log.Fatalf("Failed to find process: %v", err)
 		}
+		defer process.Kill()
 
 		// Link data streams of ssh session and conpty
 		go io.Copy(s, cpty.OutPipe())
@@ -156,10 +151,8 @@ func createPty(s ssh.Session, shell string) {
 			s.Exit(result.ProcessState.ExitCode())
 
 		case <-s.Context().Done():
-			log.Println("Session closed by remote, killing process")
-			if err := process.Kill(); err != nil {
-				log.Println("Failed to kill process:", err)
-			}
+			log.Printf("Session terminated: %s", s.Context().Err())
+			return
 		}
 	}
 }