From 02f4b08929d4083e01a8bd8a7ddb24db3de979c5 Mon Sep 17 00:00:00 2001
From: Zack Scholl <zack.scholl@gmail.com>
Date: Mon, 24 Sep 2018 07:51:24 -0700
Subject: [PATCH] not working dunno why

---
 src/comm/comm.go      | 41 +++++++++++++++++++++++++++++++----------
 src/croc/croc_test.go |  2 +-
 src/sender/sender.go  |  7 +++----
 src/tcp/tcp.go        |  2 +-
 4 files changed, 36 insertions(+), 16 deletions(-)

diff --git a/src/comm/comm.go b/src/comm/comm.go
index e7bbf1b..3ff1640 100644
--- a/src/comm/comm.go
+++ b/src/comm/comm.go
@@ -3,6 +3,7 @@ package comm
 import (
 	"bytes"
 	"fmt"
+	"log"
 	"net"
 	"strconv"
 	"strings"
@@ -39,20 +40,33 @@ func (c Comm) Write(b []byte) (int, error) {
 
 func (c Comm) Read() (buf []byte, numBytes int, bs []byte, err error) {
 	// read until we get 5 bytes
-	bs = make([]byte, 5)
-	_, err = c.connection.Read(bs)
+	tmp := make([]byte, 5)
+	n, err := c.connection.Read(tmp)
 	if err != nil {
 		return
 	}
-	tmp := make([]byte, 1)
+	tmpCopy := make([]byte, n)
+	// Copy the buffer so it doesn't get changed while read by the recipient.
+	copy(tmpCopy, tmp[:n])
+	bs = tmpCopy
+
+	tmp = make([]byte, 1)
 	for {
+		// see if we have enough bytes
 		bs = bytes.Trim(bs, "\x00")
 		if len(bs) == 5 {
 			break
 		}
-		c.connection.Read(tmp)
-		bs = append(bs, tmp...)
+		n, err := c.connection.Read(tmp)
+		if err != nil {
+			return nil, 0, nil, err
+		}
+		tmpCopy = make([]byte, n)
+		// Copy the buffer so it doesn't get changed while read by the recipient.
+		copy(tmpCopy, tmp[:n])
+		bs = append(bs, tmpCopy...)
 	}
+
 	numBytes, err = strconv.Atoi(strings.TrimLeft(string(bs), "0"))
 	if err != nil {
 		return nil, 0, nil, err
@@ -61,20 +75,27 @@ func (c Comm) Read() (buf []byte, numBytes int, bs []byte, err error) {
 	tmp = make([]byte, numBytes)
 	bufStart := 0
 	for {
-		_, err = c.connection.Read(tmp)
+		n, err := c.connection.Read(tmp)
 		if err != nil {
-			return nil, numBytes, bs, err
+			return nil, 0, nil, err
 		}
-		tmp = bytes.TrimRight(tmp, "\x00")
-		copy(buf[bufStart:bufStart+len(tmp)], tmp[:])
+		tmpCopy := make([]byte, n)
+		// Copy the buffer so it doesn't get changed while read by the recipient.
+		copy(tmpCopy, tmp[:n])
+
+		tmpCopy = bytes.TrimSpace(tmpCopy)
+		tmpCopy = bytes.Replace(tmpCopy, []byte(" "), []byte{}, -1)
+		tmpCopy = bytes.Trim(tmpCopy, "\x00")
+		copy(buf[bufStart:bufStart+len(tmpCopy)], tmpCopy[:])
 		bufStart = len(buf)
 		if bufStart < numBytes {
+			// shrink the amount we need to read
 			tmp = tmp[:numBytes-bufStart]
 		} else {
 			break
 		}
 	}
-	// log.Printf("wanted %d and got %d", numBytes, len(buf))
+	log.Printf("wanted %d and got %d", numBytes, len(buf))
 	return
 }
 
diff --git a/src/croc/croc_test.go b/src/croc/croc_test.go
index 7fc92da..f9a3ea6 100644
--- a/src/croc/croc_test.go
+++ b/src/croc/croc_test.go
@@ -47,7 +47,7 @@ func TestSendReceiveWebsockets(t *testing.T) {
 	sendAndReceive(t, 1)
 }
 func TestSendReceiveTCP(t *testing.T) {
-	sendAndReceive(t, 1)
+	sendAndReceive(t, 2)
 }
 
 func generateRandomFile(megabytes int) (fname string) {
diff --git a/src/sender/sender.go b/src/sender/sender.go
index a7a01b1..19f5d6f 100644
--- a/src/sender/sender.go
+++ b/src/sender/sender.go
@@ -207,6 +207,7 @@ func send(forceSend int, serverAddress, serverTCP string, isLocal bool, c *webso
 				return errors.New("recipient refused file")
 			}
 
+			buffer := make([]byte, models.WEBSOCKET_BUFFER_SIZE/8)
 			if !useWebsockets {
 				// connection to TCP
 				tcpConnection, err = connectToTCPServer(utils.SHA256(fmt.Sprintf("%x", sessionKey)), serverAddress+":"+serverTCP)
@@ -214,15 +215,13 @@ func send(forceSend int, serverAddress, serverTCP string, isLocal bool, c *webso
 					log.Error(err)
 					return
 				}
+				buffer = make([]byte, models.TCP_BUFFER_SIZE/2)
 			}
 
 			fmt.Fprintf(os.Stderr, "\rSending (->%s)...\n", otherIP)
 			// send file, compure hash simultaneously
 			startTransfer = time.Now()
-			buffer := make([]byte, models.WEBSOCKET_BUFFER_SIZE/8)
-			if !useWebsockets {
-				buffer = make([]byte, models.TCP_BUFFER_SIZE/2)
-			}
+
 			bar := progressbar.NewOptions(
 				int(fstats.Size),
 				progressbar.OptionSetRenderBlankState(true),
diff --git a/src/tcp/tcp.go b/src/tcp/tcp.go
index 393e037..4177f77 100644
--- a/src/tcp/tcp.go
+++ b/src/tcp/tcp.go
@@ -126,7 +126,7 @@ func chanFromConn(conn net.Conn) chan []byte {
 	c := make(chan []byte)
 
 	go func() {
-		b := make([]byte, models.WEBSOCKET_BUFFER_SIZE)
+		b := make([]byte, models.TCP_BUFFER_SIZE)
 
 		for {
 			n, err := conn.Read(b)