From edf30fc0aa6fcb92ad74aec24237cfe03242009f Mon Sep 17 00:00:00 2001
From: Jeff Lindsay <progrium@gmail.com>
Date: Thu, 16 Feb 2017 15:59:07 -0600
Subject: [PATCH] Fix short write (#28)

* session: pty normalization hack needs to return expected bytes written or chaos ensues, such as short write errors
* _example: renaming to avoid editors that auto-install from making docker binaries in your path that aren't docker
* session: keep extra calculations limited to pty case

Signed-off-by: Jeff Lindsay <progrium@gmail.com>
---
 _example/{docker => ssh-docker}/docker.go            | 2 +-
 _example/{pty => ssh-pty}/pty.go                     | 0
 _example/{public_key => ssh-publickey}/public_key.go | 0
 _example/{simple => ssh-simple}/simple.go            | 0
 session.go                                           | 9 ++++++++-
 5 files changed, 9 insertions(+), 2 deletions(-)
 rename _example/{docker => ssh-docker}/docker.go (98%)
 rename _example/{pty => ssh-pty}/pty.go (100%)
 rename _example/{public_key => ssh-publickey}/public_key.go (100%)
 rename _example/{simple => ssh-simple}/simple.go (100%)

diff --git a/_example/docker/docker.go b/_example/ssh-docker/docker.go
similarity index 98%
rename from _example/docker/docker.go
rename to _example/ssh-docker/docker.go
index 1306965..68df6e3 100644
--- a/_example/docker/docker.go
+++ b/_example/ssh-docker/docker.go
@@ -74,7 +74,7 @@ func dockerRun(cfg *container.Config, sess ssh.Session) (err error, status int64
 	go func() {
 		var err error
 		if cfg.Tty {
-			_, err = io.Copy(sess, stream.Conn)
+			_, err = io.Copy(sess, stream.Reader)
 		} else {
 			_, err = stdcopy.StdCopy(sess, sess.Stderr(), stream.Reader)
 		}
diff --git a/_example/pty/pty.go b/_example/ssh-pty/pty.go
similarity index 100%
rename from _example/pty/pty.go
rename to _example/ssh-pty/pty.go
diff --git a/_example/public_key/public_key.go b/_example/ssh-publickey/public_key.go
similarity index 100%
rename from _example/public_key/public_key.go
rename to _example/ssh-publickey/public_key.go
diff --git a/_example/simple/simple.go b/_example/ssh-simple/simple.go
similarity index 100%
rename from _example/simple/simple.go
rename to _example/ssh-simple/simple.go
diff --git a/session.go b/session.go
index 65c8b6a..8228ce1 100644
--- a/session.go
+++ b/session.go
@@ -65,9 +65,16 @@ type session struct {
 
 func (sess *session) Write(p []byte) (n int, err error) {
 	if sess.pty != nil {
-		// normalize \n to \r\n when pty is accepted
+		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 sess.Channel.Write(p)
 }