diff --git a/proxy/direct.go b/proxy/direct.go index 7067cbb..2d13a13 100644 --- a/proxy/direct.go +++ b/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) } } diff --git a/wrappers/stream/obfuscated2.go b/wrappers/stream/obfuscated2.go index d94b245..603319f 100644 --- a/wrappers/stream/obfuscated2.go +++ b/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)