Use pools for obfuscated2 write

This commit is contained in:
9seconds
2020-03-23 10:29:00 +03:00
parent f0d2cd494a
commit d72ed03a43
+30 -2
View File
@@ -1,9 +1,11 @@
package stream package stream
import ( import (
"bytes"
"crypto/cipher" "crypto/cipher"
"fmt" "fmt"
"net" "net"
"sync"
"time" "time"
"go.uber.org/zap" "go.uber.org/zap"
@@ -11,6 +13,26 @@ import (
"github.com/9seconds/mtg/conntypes" "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 { type wrapperObfuscated2 struct {
encryptor cipher.Stream encryptor cipher.Stream
decryptor 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) { 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) copy(buf, p)
w.encryptor.XORKeyStream(buf, buf) 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) { 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) copy(buf, p)
w.encryptor.XORKeyStream(buf, buf) w.encryptor.XORKeyStream(buf, buf)