client: close client.Std{err,out} if they are not os.Std{err,out}

cpu clients are implemented in the client package. The cpu
command uses this package. As in os/exec, code can call
client.Command(); followed by Client.Start() and client.Wait(),
or use client.Run(), which combines Start and Wait.

The client struct has io.WriteCloser for its Stdout and Stderr.
The question of how these gets closed is important.

Normally, when a cpu process runs, kernel process exit handling performs
a close on stdout and stderr. User code need not call close.

When the client package is used as  part of a program that does not exit, as
in a test, code might set the client struct
Stderr and Stdout to a pipe, to avoid littering test output with cpu command
output. In this case, the pipe reader will hang until the Stderr and Stdout
are closed, as no process exits.

It is not safe to just blindly close these files; users might unintentionally
close Stdout, leading to confusing results.

The actual IO is handled by two io.Copy goroutines started in cpu.Start, one for
Stdout, one for Stderr. These two goroutines
exit when the remote process closes the ssh Session Std{err,out}.
At that point, code could decide whether to close the Std{err,out}.
An ideal time to do this check is when the io.Copy exits.

This change tests the type of client.Std{err,out}; IF they are an *os.File,
AND if their Fd() matches the Fd() of os.Std{err,out}, respectively, then they are
left alone; otherwise, they are closed.

This seems a safe test. It has been tested both in a test and interactively
and the behavior is sensible.

Signed-off-by: Ronald G Minnich <rminnich@gmail.com>

Author: rminnich

client/client.go

@@ -61,8 +61,8 @@ type Cmd struct {
 	SessionOut        io.Reader
 	SessionErr        io.Reader
 	Stdin             io.Reader
-	Stdout            io.Writer
-	Stderr            io.Writer
+	Stdout            io.WriteCloser
+	Stderr            io.WriteCloser
 	Row               int
 	Col               int
 	hasTTY            bool // Set if we have a TTY
@@ -102,6 +102,13 @@ func (c *Cmd) Listen(n, addr string) (net.Listener, error) {
 	return c.client.Listen(n, addr)
 }
 
+func sameFD(w io.WriteCloser, std *os.File) bool {
+	if file, ok := w.(*os.File); ok {
+		return file.Fd() == std.Fd()
+	}
+	return false
+}
+
 // Command implements exec.Command. The required parameter is a host.
 // The args arg args to $SHELL. If there are no args, then starting $SHELL
 // is assumed.
@@ -512,6 +519,7 @@ func (c *Cmd) Start() error {
 		}
 		go c.TTYIn(c.session, c.SessionIn, c.Stdin)
 	} else {
+		verbose("Setup batch input")
 		go func() {
 			if _, err := io.Copy(c.SessionIn, c.Stdin); err != nil && !errors.Is(err, io.EOF) {
 				log.Printf("copying stdin: %v", err)
@@ -522,14 +530,26 @@ func (c *Cmd) Start() error {
 		}()
 	}
 	go func() {
+		verbose("set up copying to c.Stdout %s", c.Stdout)
 		if _, err := io.Copy(c.Stdout, c.SessionOut); err != nil && !errors.Is(err, io.EOF) {
 			log.Printf("copying stdout: %v", err)
 		}
+
+		// If the file is NOT stdout, close it.
+		// This is needed when programmers have
+		// set c.Stdout to be some other WriteCloser, e.g. a pipe.
+		if !sameFD(c.Stdout, os.Stdout) {
+			c.Stdout.Close()
+		}
 	}()
 	go func() {
+		verbose("set up copying to c.Stderr %s", c.Stderr)
 		if _, err := io.Copy(c.Stderr, c.SessionErr); err != nil && !errors.Is(err, io.EOF) {
 			log.Printf("copying stderr: %v", err)
 		}
+		if !sameFD(c.Stdout, os.Stderr) {
+			c.Stderr.Close()
+		}
 	}()
 
 	return nil