Merge pull request #139 from 9seconds/memory-direct

Memory optimizations for direct mode

proxy/direct.go

@@ -11,7 +11,7 @@ import (
 	"github.com/9seconds/mtg/protocol"
 )
 
-const directPipeBufferSize = 1024 * 1024
+const directPipeBufferSize = 1024
 
 func directConnection(request *protocol.TelegramRequest) error {
 	telegramConnRaw, err := obfuscated2.TelegramProtocol(request)
@@ -42,8 +42,9 @@ func directPipe(dst io.WriteCloser, src io.ReadCloser, wg *sync.WaitGroup, logge
 		wg.Done()
 	}()
 
-	buf := make([]byte, directPipeBufferSize)
-	if _, err := io.CopyBuffer(dst, src, buf); err != nil {
+	buf := [directPipeBufferSize]byte{}
+
+	if _, err := io.CopyBuffer(dst, src, buf[:]); err != nil {
 		logger.Debugw("Cannot pump sockets", "error", err)
 	}
 }

wrappers/stream/obfuscated2.go

@@ -1,9 +1,11 @@
 package stream
 
 import (
+	"bytes"
 	"crypto/cipher"
 	"fmt"
 	"net"
+	"sync"
 	"time"
 
 	"go.uber.org/zap"
@@ -11,6 +13,23 @@ 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
@@ -40,16 +59,26 @@ func (w *wrapperObfuscated2) Read(p []byte) (int, error) {
 }
 
 func (w *wrapperObfuscated2) WriteTimeout(p []byte, timeout time.Duration) (int, error) {
-	buf := make([]byte, len(p))
-	copy(buf, p)
+	buffer := poolWrapperObfuscated2WritePoolAcquire()
+	defer poolWrapperObfuscated2WritePoolRelease(buffer)
+
+	buffer.Write(p)
+
+	buf := buffer.Bytes()
+
 	w.encryptor.XORKeyStream(buf, buf)
 
 	return w.parent.WriteTimeout(buf, timeout)
 }
 
 func (w *wrapperObfuscated2) Write(p []byte) (int, error) {
-	buf := make([]byte, len(p))
-	copy(buf, p)
+	buffer := poolWrapperObfuscated2WritePoolAcquire()
+	defer poolWrapperObfuscated2WritePoolRelease(buffer)
+
+	buffer.Write(p)
+
+	buf := buffer.Bytes()
+
 	w.encryptor.XORKeyStream(buf, buf)
 
 	return w.parent.Write(buf)