From 4ebae693e02473976ebf43b07e23d710d417e56a Mon Sep 17 00:00:00 2001 From: 9seconds Date: Fri, 20 Mar 2020 15:54:30 +0300 Subject: [PATCH 1/5] Memory optimizations for direct mode --- proxy/direct.go | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/proxy/direct.go b/proxy/direct.go index 7067cbb..bbdd7bc 100644 --- a/proxy/direct.go +++ b/proxy/direct.go @@ -13,6 +13,15 @@ import ( const directPipeBufferSize = 1024 * 1024 +var ( + directPipePool = sync.Pool{ + New: func() interface{} { + rv := make([]byte, directPipeBufferSize) + return &rv + }, + } +) + func directConnection(request *protocol.TelegramRequest) error { telegramConnRaw, err := obfuscated2.TelegramProtocol(request) if err != nil { @@ -42,8 +51,10 @@ 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 := directPipePool.Get().(*[]byte) + defer directPipePool.Put(buf) + + if _, err := io.CopyBuffer(dst, src, *buf); err != nil { logger.Debugw("Cannot pump sockets", "error", err) } } From b5810c955d266d6111d24b5666c1552b4ab5caee Mon Sep 17 00:00:00 2001 From: 9seconds Date: Mon, 23 Mar 2020 08:42:02 +0300 Subject: [PATCH 2/5] Decrease a size of direct pipe buffer --- proxy/direct.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/proxy/direct.go b/proxy/direct.go index bbdd7bc..6d01c0d 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 var ( directPipePool = sync.Pool{ From f0d2cd494a8d06dfdaa2a31c944e0f7bd7f38177 Mon Sep 17 00:00:00 2001 From: 9seconds Date: Mon, 23 Mar 2020 09:05:18 +0300 Subject: [PATCH 3/5] Use stack allocation instead of heap one --- proxy/direct.go | 14 ++------------ 1 file changed, 2 insertions(+), 12 deletions(-) diff --git a/proxy/direct.go b/proxy/direct.go index 6d01c0d..2d13a13 100644 --- a/proxy/direct.go +++ b/proxy/direct.go @@ -13,15 +13,6 @@ import ( const directPipeBufferSize = 1024 -var ( - directPipePool = sync.Pool{ - New: func() interface{} { - rv := make([]byte, directPipeBufferSize) - return &rv - }, - } -) - func directConnection(request *protocol.TelegramRequest) error { telegramConnRaw, err := obfuscated2.TelegramProtocol(request) if err != nil { @@ -51,10 +42,9 @@ func directPipe(dst io.WriteCloser, src io.ReadCloser, wg *sync.WaitGroup, logge wg.Done() }() - buf := directPipePool.Get().(*[]byte) - defer directPipePool.Put(buf) + buf := [directPipeBufferSize]byte{} - if _, err := io.CopyBuffer(dst, src, *buf); err != nil { + if _, err := io.CopyBuffer(dst, src, buf[:]); err != nil { logger.Debugw("Cannot pump sockets", "error", err) } } From d72ed03a4340b2801f09fe745e186dadc1849b47 Mon Sep 17 00:00:00 2001 From: 9seconds Date: Mon, 23 Mar 2020 10:29:00 +0300 Subject: [PATCH 4/5] Use pools for obfuscated2 write --- wrappers/stream/obfuscated2.go | 32 ++++++++++++++++++++++++++++++-- 1 file changed, 30 insertions(+), 2 deletions(-) diff --git a/wrappers/stream/obfuscated2.go b/wrappers/stream/obfuscated2.go index d94b245..8947d13 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,26 @@ import ( "github.com/9seconds/mtg/conntypes" ) +var ( + poolWrapperObfuscated2WritePool = sync.Pool{ + New: func() interface{} { + return &bytes.Buffer{} + }, + } +) + +func poolWrapperObfuscated2WritePoolAcquire(size int) *bytes.Buffer { + buf := poolWrapperObfuscated2WritePool.Get().(*bytes.Buffer) + buf.Grow(size) + + return buf +} + +func poolWrapperObfuscated2WritePoolRelease(buf *bytes.Buffer) { + buf.Reset() + poolWrapperObfuscated2WritePool.Put(buf) +} + type wrapperObfuscated2 struct { encryptor cipher.Stream decryptor cipher.Stream @@ -40,7 +62,10 @@ func (w *wrapperObfuscated2) Read(p []byte) (int, error) { } func (w *wrapperObfuscated2) WriteTimeout(p []byte, timeout time.Duration) (int, error) { - buf := make([]byte, len(p)) + buffer := poolWrapperObfuscated2WritePoolAcquire(len(p)) + defer poolWrapperObfuscated2WritePoolRelease(buffer) + + buf := buffer.Bytes() copy(buf, p) w.encryptor.XORKeyStream(buf, buf) @@ -48,7 +73,10 @@ func (w *wrapperObfuscated2) WriteTimeout(p []byte, timeout time.Duration) (int, } func (w *wrapperObfuscated2) Write(p []byte) (int, error) { - buf := make([]byte, len(p)) + buffer := poolWrapperObfuscated2WritePoolAcquire(len(p)) + defer poolWrapperObfuscated2WritePoolRelease(buffer) + + buf := buffer.Bytes() copy(buf, p) w.encryptor.XORKeyStream(buf, buf) From 6a9069ca72cdf5118e1a2b30b462184dc790b48f Mon Sep 17 00:00:00 2001 From: 9seconds Date: Mon, 23 Mar 2020 10:53:58 +0300 Subject: [PATCH 5/5] Fixes --- wrappers/stream/obfuscated2.go | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/wrappers/stream/obfuscated2.go b/wrappers/stream/obfuscated2.go index 8947d13..603319f 100644 --- a/wrappers/stream/obfuscated2.go +++ b/wrappers/stream/obfuscated2.go @@ -21,11 +21,8 @@ var ( } ) -func poolWrapperObfuscated2WritePoolAcquire(size int) *bytes.Buffer { - buf := poolWrapperObfuscated2WritePool.Get().(*bytes.Buffer) - buf.Grow(size) - - return buf +func poolWrapperObfuscated2WritePoolAcquire() *bytes.Buffer { + return poolWrapperObfuscated2WritePool.Get().(*bytes.Buffer) } func poolWrapperObfuscated2WritePoolRelease(buf *bytes.Buffer) { @@ -62,22 +59,26 @@ func (w *wrapperObfuscated2) Read(p []byte) (int, error) { } func (w *wrapperObfuscated2) WriteTimeout(p []byte, timeout time.Duration) (int, error) { - buffer := poolWrapperObfuscated2WritePoolAcquire(len(p)) + buffer := poolWrapperObfuscated2WritePoolAcquire() defer poolWrapperObfuscated2WritePoolRelease(buffer) + buffer.Write(p) + buf := buffer.Bytes() - copy(buf, p) + w.encryptor.XORKeyStream(buf, buf) return w.parent.WriteTimeout(buf, timeout) } func (w *wrapperObfuscated2) Write(p []byte) (int, error) { - buffer := poolWrapperObfuscated2WritePoolAcquire(len(p)) + buffer := poolWrapperObfuscated2WritePoolAcquire() defer poolWrapperObfuscated2WritePoolRelease(buffer) + buffer.Write(p) + buf := buffer.Bytes() - copy(buf, p) + w.encryptor.XORKeyStream(buf, buf) return w.parent.Write(buf)