From b8a07e96482b380121c8b4a7d0a2b6a059e1bf07 Mon Sep 17 00:00:00 2001
From: Russ Magee <rmagee@gmail.com>
Date: Tue, 1 May 2018 02:39:45 -0700
Subject: [PATCH] Chaff packets w/sync.Mutex to allow chaff & main goroutine to
 both input to server. TODO: smart chaff, mutating or mimicking recent input.

---
 hkexnet.go         | 15 ++++++++---
 hkexsh/hkexsh.go   | 64 +++++++++++++++++++++++++++++++++++++++++++++-
 hkexshd/hkexshd.go |  2 --
 3 files changed, 74 insertions(+), 7 deletions(-)

diff --git a/hkexnet.go b/hkexnet.go
index 81f8bc8..f714ab7 100644
--- a/hkexnet.go
+++ b/hkexnet.go
@@ -391,7 +391,7 @@ func (c Conn) Read(b []byte) (n int, err error) {
 
 		// Throw away pkt if it's chaff (ie., caller to Read() won't see this data)
 		if ctrlStatOp == CSOChaff {
-			log.Printf("[Chaff pkt]\n")
+			log.Printf("[Chaff pkt, discarded]\n")
 		} else if ctrlStatOp == CSOTermSize {
 			fmt.Sscanf(string(payloadBytes), "%d %d", &c.Rows, &c.Cols)
 			log.Printf("[TermSize pkt: rows %v cols %v]\n", c.Rows, c.Cols)
@@ -428,7 +428,8 @@ func (c Conn) Read(b []byte) (n int, err error) {
 //
 // See go doc io.Writer
 func (c Conn) Write(b []byte) (n int, err error) {
-	return c.WritePacket(b, CSONone)
+	n, err = c.WritePacket(b, CSONone)
+	return n, err
 }
 
 // Write a byte slice with specified ctrlStatusOp byte
@@ -441,6 +442,11 @@ func (c Conn) WritePacket(b []byte, op byte) (n int, err error) {
 
 	payloadLen = uint32(len(b))
 
+	// Testing: '`1' will trigger a chaff packet
+	//if payloadLen == 2 && string(b) == "`1" {
+	//	op = CSOChaff
+	//}
+
 	// Calculate hmac on payload
 	c.wm.Write(b)
 	hmacOut = c.wm.Sum(nil)[0:4]
@@ -458,13 +464,14 @@ func (c Conn) WritePacket(b []byte, op byte) (n int, err error) {
 	log.Printf("  ->ctext:\r\n%s\r\n", hex.Dump(wb.Bytes()))
 
 	ctrlStatOp := op
-	_ = binary.Write(c.c, binary.BigEndian, &ctrlStatOp)
 
+	//{
+	_ = binary.Write(c.c, binary.BigEndian, &ctrlStatOp)
 	// Write hmac LSB, payloadLen followed by payload
 	_ = binary.Write(c.c, binary.BigEndian, hmacOut)
 	_ = binary.Write(c.c, binary.BigEndian, payloadLen)
-
 	n, err = c.c.Write(wb.Bytes())
+	//}
 	if err != nil {
 		//panic(err)
 		log.Println(err)
diff --git a/hkexsh/hkexsh.go b/hkexsh/hkexsh.go
index 0de4ce2..df0ff77 100644
--- a/hkexsh/hkexsh.go
+++ b/hkexsh/hkexsh.go
@@ -20,6 +20,7 @@ import (
 	"strings"
 	"sync"
 	"syscall"
+	"time"
 
 	hkexsh "blitter.com/go/hkexsh"
 	isatty "github.com/mattn/go-isatty"
@@ -192,6 +193,8 @@ func main() {
 		}
 	}()
 
+	m := &sync.Mutex{}
+
 	if isInteractive {
 		// Handle pty resizes (notify server side)
 		ch := make(chan os.Signal, 1)
@@ -209,11 +212,26 @@ func main() {
 					panic(err)
 				}
 				termSzPacket := fmt.Sprintf("%d %d", rows, cols)
+				m.Lock()
 				conn.WritePacket([]byte(termSzPacket), hkexsh.CSOTermSize)
+				m.Unlock()
 			}
 		}()
 		ch <- syscall.SIGWINCH // Initial resize.
 
+		// client chaffing goroutine
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+
+			for {
+				m.Lock()
+				conn.WritePacket([]byte("CHAFF"), hkexsh.CSOChaff)
+				m.Unlock()
+				time.Sleep(10 * time.Millisecond)
+			}
+		}()
+
 		// client writer (to server) goroutine
 		wg.Add(1)
 		go func() {
@@ -221,7 +239,11 @@ func main() {
 
 			// io.Copy() expects EOF so this will
 			// exit with outerr == nil
-			_, outerr := io.Copy(conn, os.Stdin)
+			//!_, outerr := io.Copy(conn, os.Stdin)
+			_, outerr := func(m *sync.Mutex, conn *hkexsh.Conn, r io.Reader) (w int64, e error) {
+				return safeCopy(m, conn, r)
+			}(m, conn, os.Stdin)
+
 			if outerr != nil {
 				log.Println(outerr)
 				if outerr.Error() != "EOF" {
@@ -238,3 +260,43 @@ func main() {
 	// Wait until both stdin and stdout goroutines finish
 	wg.Wait()
 }
+
+func safeCopy(m *sync.Mutex, dst io.Writer, src io.Reader) (written int64, err error) {
+	//	// If the reader has a WriteTo method, use it to do the copy.
+	//	// Avoids an allocation and a copy.
+	//	if wt, ok := src.(io.WriterTo); ok {
+	//		return wt.WriteTo(dst)
+	//	}
+	//	// Similarly, if the writer has a ReadFrom method, use it to do the copy.
+	//	if rt, ok := dst.(io.ReaderFrom); ok {
+	//		return rt.ReadFrom(src)
+	//	}
+
+	buf := make([]byte, 32*1024)
+	for {
+		nr, er := src.Read(buf)
+		if nr > 0 {
+			m.Lock()
+			nw, ew := dst.Write(buf[0:nr])
+			m.Unlock()
+			if nw > 0 {
+				written += int64(nw)
+			}
+			if ew != nil {
+				err = ew
+				break
+			}
+			if nr != nw {
+				err = io.ErrShortWrite
+				break
+			}
+		}
+		if er != nil {
+			if er != io.EOF {
+				err = er
+			}
+			break
+		}
+	}
+	return written, err
+}
diff --git a/hkexshd/hkexshd.go b/hkexshd/hkexshd.go
index eea06aa..bad499b 100644
--- a/hkexshd/hkexshd.go
+++ b/hkexshd/hkexshd.go
@@ -28,8 +28,6 @@ type cmdSpec struct {
 	who        []byte
 	cmd        []byte
 	authCookie []byte
-	termRows   []byte
-	termCols   []byte
 	status     int
 }