From f933c7b453c61b0c8ea46fd2aaa2f1593457134f Mon Sep 17 00:00:00 2001
From: Qian Wang <cbeuw.andy@gmail.com>
Date: Sat, 3 Aug 2019 14:58:48 +0100
Subject: [PATCH] Move pipe to util

---
 cmd/ck-client/ck-client.go | 25 ++-----------------------
 cmd/ck-server/ck-server.go | 29 ++++-------------------------
 internal/util/util.go      | 21 +++++++++++++++++++++
 3 files changed, 27 insertions(+), 48 deletions(-)

diff --git a/cmd/ck-client/ck-client.go b/cmd/ck-client/ck-client.go
index e177924..72cf478 100644
--- a/cmd/ck-client/ck-client.go
+++ b/cmd/ck-client/ck-client.go
@@ -23,27 +23,6 @@ import (
 
 var version string
 
-func pipe(dst io.ReadWriteCloser, src io.ReadWriteCloser) {
-	// The maximum size of TLS message will be 16380+12+16. 12 because of the stream header and 16
-	// because of the salt/mac
-	// 16408 is the max TLS message size on Firefox
-	buf := make([]byte, 16380)
-	for {
-		i, err := io.ReadAtLeast(src, buf, 1)
-		if err != nil {
-			go dst.Close()
-			go src.Close()
-			return
-		}
-		i, err = dst.Write(buf[:i])
-		if err != nil {
-			go dst.Close()
-			go src.Close()
-			return
-		}
-	}
-}
-
 // This establishes a connection with ckserver and performs a handshake
 func makeRemoteConn(sta *client.State) (net.Conn, []byte, error) {
 
@@ -253,8 +232,8 @@ func main() {
 				stream.Close()
 				return
 			}
-			go pipe(localConn, stream)
-			pipe(stream, localConn)
+			go util.Pipe(localConn, stream)
+			util.Pipe(stream, localConn)
 		}()
 	}
 
diff --git a/cmd/ck-server/ck-server.go b/cmd/ck-server/ck-server.go
index 8de6533..1a059f0 100644
--- a/cmd/ck-server/ck-server.go
+++ b/cmd/ck-server/ck-server.go
@@ -24,27 +24,6 @@ import (
 var b64 = base64.StdEncoding.EncodeToString
 var version string
 
-func pipe(dst io.ReadWriteCloser, src io.ReadWriteCloser) {
-	// The maximum size of TLS message will be 16380+12+16. 12 because of the stream header and 16
-	// because of the salt/mac
-	// 16408 is the max TLS message size on Firefox
-	buf := make([]byte, 16380)
-	for {
-		i, err := io.ReadAtLeast(src, buf, 1)
-		if err != nil {
-			go dst.Close()
-			go src.Close()
-			return
-		}
-		i, err = dst.Write(buf[:i])
-		if err != nil {
-			go dst.Close()
-			go src.Close()
-			return
-		}
-	}
-}
-
 func dispatchConnection(conn net.Conn, sta *server.State) {
 	remoteAddr := conn.RemoteAddr()
 	var err error
@@ -73,8 +52,8 @@ func dispatchConnection(conn net.Conn, sta *server.State) {
 		if err != nil {
 			log.Error("Failed to send first packet to redirection server", err)
 		}
-		go pipe(webConn, conn)
-		go pipe(conn, webConn)
+		go util.Pipe(webConn, conn)
+		go util.Pipe(conn, webConn)
 	}
 
 	ch, err := server.ParseClientHello(data)
@@ -197,8 +176,8 @@ func dispatchConnection(conn net.Conn, sta *server.State) {
 			sesh.Close()
 			continue
 		}
-		go pipe(localConn, newStream)
-		go pipe(newStream, localConn)
+		go util.Pipe(localConn, newStream)
+		go util.Pipe(newStream, localConn)
 	}
 
 }
diff --git a/internal/util/util.go b/internal/util/util.go
index ad6a1c1..4e82d66 100644
--- a/internal/util/util.go
+++ b/internal/util/util.go
@@ -124,3 +124,24 @@ func AddRecordLayer(input []byte, typ []byte, ver []byte) []byte {
 	copy(ret[5:], input)
 	return ret
 }
+
+func Pipe(dst io.ReadWriteCloser, src io.ReadWriteCloser) {
+	// The maximum size of TLS message will be 16380+12+16. 12 because of the stream header and 16
+	// because of the salt/mac
+	// 16408 is the max TLS message size on Firefox
+	buf := make([]byte, 16380)
+	for {
+		i, err := io.ReadAtLeast(src, buf, 1)
+		if err != nil {
+			go dst.Close()
+			go src.Close()
+			return
+		}
+		i, err = dst.Write(buf[:i])
+		if err != nil {
+			go dst.Close()
+			go src.Close()
+			return
+		}
+	}
+}