From 86e293560771eab2cf19bebe01ed50a7b9aecbcd Mon Sep 17 00:00:00 2001
From: Andy Wang <cbeuw.andy@gmail.com>
Date: Sat, 17 Oct 2020 13:46:22 +0100
Subject: [PATCH] Use a pre-made buffer for TLSConn.Write

---
 internal/client/TLS.go            |  2 +-
 internal/common/tls.go            | 23 ++++++++++++++++++++---
 internal/multiplex/mux_test.go    |  4 ++--
 internal/multiplex/stream_test.go |  4 ++--
 internal/server/TLS.go            |  2 +-
 5 files changed, 26 insertions(+), 9 deletions(-)

diff --git a/internal/client/TLS.go b/internal/client/TLS.go
index 9693a6b..39e8bb3 100644
--- a/internal/client/TLS.go
+++ b/internal/client/TLS.go
@@ -71,7 +71,7 @@ func (tls *DirectTLS) Handshake(rawConn net.Conn, authInfo AuthInfo) (sessionKey
 		return
 	}
 	log.Trace("client hello sent successfully")
-	tls.TLSConn = &common.TLSConn{Conn: rawConn}
+	tls.TLSConn = common.NewTLSConn(rawConn)
 
 	buf := make([]byte, 1024)
 	log.Trace("waiting for ServerHello")
diff --git a/internal/common/tls.go b/internal/common/tls.go
index 2f95905..9af2c3c 100644
--- a/internal/common/tls.go
+++ b/internal/common/tls.go
@@ -4,6 +4,7 @@ import (
 	"encoding/binary"
 	"io"
 	"net"
+	"sync"
 	"time"
 )
 
@@ -37,6 +38,15 @@ func AddRecordLayer(input []byte, typ byte, ver uint16) []byte {
 
 type TLSConn struct {
 	net.Conn
+	writeM   sync.Mutex
+	writeBuf []byte
+}
+
+func NewTLSConn(conn net.Conn) *TLSConn {
+	return &TLSConn{
+		Conn:     conn,
+		writeBuf: make([]byte, 15000),
+	}
 }
 
 func (tls *TLSConn) LocalAddr() net.Addr {
@@ -79,9 +89,16 @@ func (tls *TLSConn) Read(buffer []byte) (n int, err error) {
 }
 
 func (tls *TLSConn) Write(in []byte) (n int, err error) {
-	// TODO: write record layer directly first?
-	toWrite := AddRecordLayer(in, ApplicationData, VersionTLS13)
-	n, err = tls.Conn.Write(toWrite)
+	msgLen := len(in)
+	tls.writeM.Lock()
+	tls.writeBuf = append(tls.writeBuf[:5], in...)
+	tls.writeBuf[0] = ApplicationData
+	tls.writeBuf[1] = byte(VersionTLS13 >> 8)
+	tls.writeBuf[2] = byte(VersionTLS13 & 0xFF)
+	tls.writeBuf[3] = byte(msgLen >> 8)
+	tls.writeBuf[4] = byte(msgLen & 0xFF)
+	n, err = tls.Conn.Write(tls.writeBuf[:recordLayerLength+msgLen])
+	tls.writeM.Unlock()
 	return n - recordLayerLength, err
 }
 
diff --git a/internal/multiplex/mux_test.go b/internal/multiplex/mux_test.go
index 9d992b1..5930583 100644
--- a/internal/multiplex/mux_test.go
+++ b/internal/multiplex/mux_test.go
@@ -51,8 +51,8 @@ func makeSessionPair(numConn int) (*Session, *Session, []*connPair) {
 	paris := make([]*connPair, numConn)
 	for i := 0; i < numConn; i++ {
 		c, s := connutil.AsyncPipe()
-		clientConn := &common.TLSConn{Conn: c}
-		serverConn := &common.TLSConn{Conn: s}
+		clientConn := common.NewTLSConn(c)
+		serverConn := common.NewTLSConn(s)
 		paris[i] = &connPair{
 			clientConn: clientConn,
 			serverConn: serverConn,
diff --git a/internal/multiplex/stream_test.go b/internal/multiplex/stream_test.go
index c8cdf31..6c89515 100644
--- a/internal/multiplex/stream_test.go
+++ b/internal/multiplex/stream_test.go
@@ -132,8 +132,8 @@ func TestStream_WriteSync(t *testing.T) {
 	clientSesh := setupSesh(false, sessionKey, E_METHOD_PLAIN)
 	serverSesh := setupSesh(false, sessionKey, E_METHOD_PLAIN)
 	w, r := connutil.AsyncPipe()
-	clientSesh.AddConnection(&common.TLSConn{Conn: w})
-	serverSesh.AddConnection(&common.TLSConn{Conn: r})
+	clientSesh.AddConnection(common.NewTLSConn(w))
+	serverSesh.AddConnection(common.NewTLSConn(r))
 	testData := make([]byte, payloadLen)
 	rand.Read(testData)
 
diff --git a/internal/server/TLS.go b/internal/server/TLS.go
index 49463fa..8a0ea6a 100644
--- a/internal/server/TLS.go
+++ b/internal/server/TLS.go
@@ -65,7 +65,7 @@ func (TLS) makeResponder(clientHelloSessionId []byte, sharedSecret [32]byte) Res
 			originalConn.Close()
 			return
 		}
-		preparedConn = &common.TLSConn{Conn: originalConn}
+		preparedConn = common.NewTLSConn(originalConn)
 		return
 	}
 	return respond