Add pool support everywhere

wrappers/stream/faketls.go

@@ -1,6 +1,7 @@
 package stream
 
 import (
+	"bytes"
 	"errors"
 	"fmt"
 	"net"
@@ -39,13 +40,19 @@ func (w *wrapperFakeTLS) WriteTimeout(p []byte, timeout time.Duration) (int, err
 func (w *wrapperFakeTLS) write(p []byte, writeFunc func([]byte) (int, error)) (int, error) {
 	sum := 0
 
+	buf := acquireBytesBuffer()
+	defer releaseBytesBuffer(buf)
+
 	for _, v := range tlstypes.MakeRecords(p) {
-		_, err := writeFunc(v.Bytes())
+		buf.Reset()
+		v.WriteBytes(buf)
+
+		_, err := writeFunc(buf.Bytes())
 		if err != nil {
 			return sum, err
 		}
 
-		sum += len(v.Data.Bytes())
+		sum += v.Data.Len()
 	}
 
 	return sum, nil
@@ -86,7 +93,10 @@ func NewFakeTLS(socket conntypes.StreamReadWriteCloser) conntypes.StreamReadWrit
 			switch rec.Type {
 			case tlstypes.RecordTypeChangeCipherSpec:
 			case tlstypes.RecordTypeApplicationData:
-				return rec.Data.Bytes(), nil
+				buf := &bytes.Buffer{}
+				rec.Data.WriteBytes(buf)
+
+				return buf.Bytes(), nil
 			default:
 				return nil, fmt.Errorf("unsupported record type %v", rec.Type)
 			}

wrappers/stream/mtproto_cipher.go

@@ -1,7 +1,6 @@
 package stream
 
 import (
-	"bytes"
 	"crypto/aes"
 	"crypto/cipher"
 	"crypto/md5"  // nolint: gosec
@@ -54,7 +53,9 @@ func mtprotoDeriveKeys(purpose mtprotoCipherPurpose,
 	resp *rpc.NonceResponse,
 	client, remote *net.TCPAddr,
 	secret []byte) ([]byte, []byte) {
-	message := bytes.Buffer{}
+	message := acquireBytesBuffer()
+	defer releaseBytesBuffer(message)
+
 	message.Write(resp.Nonce)   // nolint: gosec
 	message.Write(req.Nonce)    // nolint: gosec
 	message.Write(req.CryptoTS) // nolint: gosec

wrappers/stream/obfuscated2.go

@@ -1,11 +1,9 @@
 package stream
 
 import (
-	"bytes"
 	"crypto/cipher"
 	"fmt"
 	"net"
-	"sync"
 	"time"
 
 	"go.uber.org/zap"
@@ -13,23 +11,6 @@ import (
 	"github.com/9seconds/mtg/conntypes"
 )
 
-var (
-	poolWrapperObfuscated2WritePool = sync.Pool{
-		New: func() interface{} {
-			return &bytes.Buffer{}
-		},
-	}
-)
-
-func poolWrapperObfuscated2WritePoolAcquire() *bytes.Buffer {
-	return poolWrapperObfuscated2WritePool.Get().(*bytes.Buffer)
-}
-
-func poolWrapperObfuscated2WritePoolRelease(buf *bytes.Buffer) {
-	buf.Reset()
-	poolWrapperObfuscated2WritePool.Put(buf)
-}
-
 type wrapperObfuscated2 struct {
 	encryptor cipher.Stream
 	decryptor cipher.Stream
@@ -59,8 +40,8 @@ func (w *wrapperObfuscated2) Read(p []byte) (int, error) {
 }
 
 func (w *wrapperObfuscated2) WriteTimeout(p []byte, timeout time.Duration) (int, error) {
-	buffer := poolWrapperObfuscated2WritePoolAcquire()
-	defer poolWrapperObfuscated2WritePoolRelease(buffer)
+	buffer := acquireBytesBuffer()
+	defer releaseBytesBuffer(buffer)
 
 	buffer.Write(p)
 
@@ -72,8 +53,8 @@ func (w *wrapperObfuscated2) WriteTimeout(p []byte, timeout time.Duration) (int,
 }
 
 func (w *wrapperObfuscated2) Write(p []byte) (int, error) {
-	buffer := poolWrapperObfuscated2WritePoolAcquire()
-	defer poolWrapperObfuscated2WritePoolRelease(buffer)
+	buffer := acquireBytesBuffer()
+	defer releaseBytesBuffer(buffer)
 
 	buffer.Write(p)
 

wrappers/stream/pools.go

@@ -0,0 +1,23 @@
+package stream
+
+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)
+}