Add pool support everywhere

tlstypes/client_hello.go

@@ -25,7 +25,7 @@ func (c ClientHello) Digest() []byte {
 	}
 
 	mac := hmac.New(sha256.New, config.C.Secret)
-	mac.Write(rec.Bytes()) // nolint: errcheck
+	rec.WriteBytes(mac)
 	computedDigest := mac.Sum(nil)
 
 	for i := range computedDigest {

tlstypes/consts.go

@@ -1,5 +1,7 @@
 package tlstypes
 
+import "io"
+
 type RecordType uint8
 
 const (
@@ -69,11 +71,16 @@ var (
 )
 
 type Byter interface {
-	Bytes() []byte
+	WriteBytes(io.Writer)
+	Len() int
 }
 
 type RawBytes []byte
 
-func (r RawBytes) Bytes() []byte {
-	return []byte(r)
+func (r RawBytes) WriteBytes(writer io.Writer) {
+	writer.Write(r) // nolint: errcheck
+}
+
+func (r RawBytes) Len() int {
+	return len(r)
 }

tlstypes/handshake.go

@@ -1,7 +1,7 @@
 package tlstypes
 
 import (
-	"bytes"
+	"io"
 
 	"github.com/9seconds/mtg/utils"
 )
@@ -14,24 +14,31 @@ type Handshake struct {
 	Tail      Byter
 }
 
-func (h *Handshake) Bytes() []byte {
-	buf := bytes.Buffer{}
-	packetBuf := bytes.Buffer{}
+func (h *Handshake) WriteBytes(writer io.Writer) {
+	packetBuf := acquireBytesBuffer()
+	defer releaseBytesBuffer(packetBuf)
 
-	buf.WriteByte(byte(h.Type))
+	writer.Write([]byte{byte(h.Type)}) // nolint: errcheck
 
 	packetBuf.Write(h.Version.Bytes())
 	packetBuf.Write(h.Random[:])
 	packetBuf.WriteByte(byte(len(h.SessionID)))
 	packetBuf.Write(h.SessionID)
-	packetBuf.Write(h.Tail.Bytes())
+	h.Tail.WriteBytes(packetBuf)
 
 	sizeUint24 := utils.ToUint24(uint32(packetBuf.Len()))
 	sizeUint24Bytes := sizeUint24[:]
 	sizeUint24Bytes[0], sizeUint24Bytes[2] = sizeUint24Bytes[2], sizeUint24Bytes[0]
 
-	buf.Write(sizeUint24Bytes)
-	packetBuf.WriteTo(&buf) // nolint: errcheck
+	writer.Write(sizeUint24Bytes) // nolint: errcheck
+	packetBuf.WriteTo(writer)     // nolint: errcheck
+}
 
-	return buf.Bytes()
+func (h *Handshake) Len() int {
+	buf := acquireBytesBuffer()
+	defer releaseBytesBuffer(buf)
+
+	h.WriteBytes(buf)
+
+	return buf.Len()
 }

tlstypes/pools.go

@@ -0,0 +1,23 @@
+package tlstypes
+
+import (
+	"bytes"
+	"sync"
+)
+
+var (
+	poolBytesBuffer = sync.Pool{
+		New: func() interface{} {
+			return &bytes.Buffer{}
+		},
+	}
+)
+
+func acquireBytesBuffer() *bytes.Buffer {
+	return poolBytesBuffer.Get().(*bytes.Buffer)
+}
+
+func releaseBytesBuffer(buf *bytes.Buffer) {
+	buf.Reset()
+	poolBytesBuffer.Put(buf)
+}

tlstypes/record.go

@@ -15,16 +15,15 @@ type Record struct {
 	Data    Byter
 }
 
-func (r Record) Bytes() []byte {
-	buf := bytes.Buffer{}
-	data := r.Data.Bytes()
+func (r Record) WriteBytes(writer io.Writer) {
+	writer.Write([]byte{byte(r.Type)})                           // nolint: errcheck
+	writer.Write(r.Version.Bytes())                              // nolint: errcheck
+	binary.Write(writer, binary.BigEndian, uint16(r.Data.Len())) // nolint: errcheck
+	r.Data.WriteBytes(writer)
+}
 
-	buf.WriteByte(byte(r.Type))
-	buf.Write(r.Version.Bytes())
-	binary.Write(&buf, binary.BigEndian, uint16(len(data))) // nolint: errcheck
-	buf.Write(data)
-
-	return buf.Bytes()
+func (r Record) Len() int {
+	return 1 + 2 + 2 + r.Data.Len()
 }
 
 func ReadRecord(reader io.Reader) (Record, error) {

tlstypes/server_hello.go

@@ -20,20 +20,22 @@ type ServerHello struct {
 }
 
 func (s ServerHello) WelcomePacket() []byte {
+	buf := &bytes.Buffer{}
+
 	s.Random = [32]byte{}
 	rec := Record{
 		Type:    RecordTypeHandshake,
 		Version: Version12,
 		Data:    &s,
 	}
-	buf := bytes.NewBuffer(rec.Bytes())
+	rec.WriteBytes(buf)
 
 	recChangeCipher := Record{
 		Type:    RecordTypeChangeCipherSpec,
 		Version: Version12,
 		Data:    RawBytes([]byte{0x01}),
 	}
-	buf.Write(recChangeCipher.Bytes())
+	recChangeCipher.WriteBytes(buf)
 
 	hostCert := make([]byte, 1024+mrand.Intn(3092))
 	rand.Read(hostCert) // nolint: errcheck
@@ -43,7 +45,8 @@ func (s ServerHello) WelcomePacket() []byte {
 		Version: Version12,
 		Data:    RawBytes(hostCert),
 	}
-	buf.Write(recData.Bytes())
+	recData.WriteBytes(buf)
+
 	packet := buf.Bytes()
 
 	mac := hmac.New(sha256.New, config.C.Secret)