feat: safe stderr

Author: caarlos0

pty.go

@@ -0,0 +1,31 @@
+package ssh
+
+import (
+	"bytes"
+	"io"
+)
+
+// NewPtyWriter creates a writer that handles when the session has a active
+// PTY, replacing the \n with \r\n.
+func NewPtyWriter(w io.Writer) io.Writer {
+	return ptyWriter{
+		w: w,
+	}
+}
+
+type ptyWriter struct {
+	w io.Writer
+}
+
+func (w ptyWriter) Write(p []byte) (int, error) {
+	m := len(p)
+	// normalize \n to \r\n when pty is accepted.
+	// this is a hardcoded shortcut since we don't support terminal modes.
+	p = bytes.Replace(p, []byte{'\n'}, []byte{'\r', '\n'}, -1)
+	p = bytes.Replace(p, []byte{'\r', '\r', '\n'}, []byte{'\r', '\n'}, -1)
+	n, err := w.w.Write(p)
+	if n > m {
+		n = m
+	}
+	return n, err
+}

pty_test.go

@@ -0,0 +1,24 @@
+package ssh_test
+
+import (
+	"bytes"
+	"testing"
+
+	"github.com/gliderlabs/ssh"
+)
+
+func TestNewPtyWriter(t *testing.T) {
+	in := "\nfoo\r\nbar\nmore text\rmore\r\r\r\nfoo\n\n"
+	out := "\r\nfoo\r\nbar\r\nmore text\rmore\r\r\r\nfoo\r\n\r\n"
+	var b bytes.Buffer
+	n, err := ssh.NewPtyWriter(&b).Write([]byte(in))
+	if err != nil {
+		t.Error("did not expect an error", err)
+	}
+	if out != b.String() {
+		t.Errorf("outputs do not match, expected %q got %q", out, b.String())
+	}
+	if n != len(in) {
+		t.Errorf("expected to write %d bytes, wrote %d", len(in), n)
+	}
+}

session.go

@@ -1,9 +1,9 @@
 package ssh
 
 import (
-	"bytes"
 	"errors"
 	"fmt"
+	"io"
 	"net"
 	"sync"
 
@@ -82,6 +82,10 @@ type Session interface {
 	// the request handling loop. Registering nil will unregister the channel.
 	// During the time that no channel is registered, breaks are ignored.
 	Break(c chan<- bool)
+
+	// SafeStderr returns the Stderr io.Writer that handles replacing \n with
+	// \r\n when there's an active Pty.
+	SafeStderr() io.Writer
 }
 
 // maxSigBufSize is how many signals will be buffered
@@ -127,18 +131,16 @@ type session struct {
 	breakCh           chan<- bool
 }
 
-func (sess *session) Write(p []byte) (n int, err error) {
+func (sess *session) SafeStderr() io.Writer {
 	if sess.pty != nil {
-		m := len(p)
-		// normalize \n to \r\n when pty is accepted.
-		// this is a hardcoded shortcut since we don't support terminal modes.
-		p = bytes.Replace(p, []byte{'\n'}, []byte{'\r', '\n'}, -1)
-		p = bytes.Replace(p, []byte{'\r', '\r', '\n'}, []byte{'\r', '\n'}, -1)
-		n, err = sess.Channel.Write(p)
-		if n > m {
-			n = m
-		}
-		return
+		return NewPtyWriter(sess.Stderr())
+	}
+	return sess.Stderr()
+}
+
+func (sess *session) Write(p []byte) (int, error) {
+	if sess.pty != nil {
+		return NewPtyWriter(sess.Channel).Write(p)
 	}
 	return sess.Channel.Write(p)
 }
@@ -242,7 +244,7 @@ func (sess *session) handleRequests(reqs <-chan *gossh.Request) {
 				continue
 			}
 
-			var payload = struct{ Value string }{}
+			payload := struct{ Value string }{}
 			gossh.Unmarshal(req.Payload, &payload)
 			sess.rawCmd = payload.Value
 
@@ -267,7 +269,7 @@ func (sess *session) handleRequests(reqs <-chan *gossh.Request) {
 				continue
 			}
 
-			var payload = struct{ Value string }{}
+			payload := struct{ Value string }{}
 			gossh.Unmarshal(req.Payload, &payload)
 			sess.subsystem = payload.Value
 

session_test.go

@@ -228,6 +228,33 @@ func TestPty(t *testing.T) {
 	<-done
 }
 
+func TestPtyWriter(t *testing.T) {
+	t.Parallel()
+	term := "xterm"
+	winWidth := 40
+	winHeight := 80
+	session, _, cleanup := newTestSession(t, &Server{
+		Handler: func(s Session) {
+			_, _ = fmt.Fprintln(s, "foo\nbar")
+			_, _ = fmt.Fprintln(s.SafeStderr(), "many\nerrors")
+			_ = s.Exit(0)
+		},
+	}, nil)
+	defer cleanup()
+	if err := session.RequestPty(term, winHeight, winWidth, gossh.TerminalModes{}); err != nil {
+		t.Fatalf("expected nil but got %v", err)
+	}
+	bts, err := session.CombinedOutput("")
+	if err != nil {
+		t.Fatalf("expected nil but got %v", err)
+	}
+
+	expected := "foo\r\nbar\r\nmany\r\nerrors\r\n"
+	if expected != string(bts) {
+		t.Fatalf("expected output to be %q, got %q", expected, string(bts))
+	}
+}
+
 func TestPtyResize(t *testing.T) {
 	t.Parallel()
 	winch0 := Window{40, 80}