From 0d0bdd9fde62977be3a1578ff55f5ff850299c66 Mon Sep 17 00:00:00 2001
From: 9seconds <nineseconds@yandex.ru>
Date: Fri, 20 Jul 2018 14:49:22 +0300
Subject: [PATCH] Reuse buffers for stream cipher

---
 wrappers/streamcipher.go      | 14 +++++++++++---
 wrappers/streamcipher_pool.go | 16 ++++++++++++++++
 2 files changed, 27 insertions(+), 3 deletions(-)
 create mode 100644 wrappers/streamcipher_pool.go

diff --git a/wrappers/streamcipher.go b/wrappers/streamcipher.go
index b50a351..eaef556 100644
--- a/wrappers/streamcipher.go
+++ b/wrappers/streamcipher.go
@@ -1,6 +1,7 @@
 package wrappers
 
 import (
+	"bytes"
 	"crypto/cipher"
 	"net"
 
@@ -28,10 +29,17 @@ func (s *StreamCipher) Read(p []byte) (int, error) {
 }
 
 func (s *StreamCipher) Write(p []byte) (int, error) {
-	encrypted := make([]byte, len(p))
-	s.encryptor.XORKeyStream(encrypted, p)
+	buf := streamCipherBufferPool.Get().(*bytes.Buffer)
+	defer streamCipherBufferPool.Put(buf)
 
-	return s.conn.Write(encrypted)
+	buf.Reset()
+	buf.Grow(len(p))
+	buf.Write(p)
+
+	data := buf.Bytes()
+	s.encryptor.XORKeyStream(data, data)
+
+	return s.conn.Write(data)
 }
 
 // Logger returns an instance of the logger for this wrapper.
diff --git a/wrappers/streamcipher_pool.go b/wrappers/streamcipher_pool.go
new file mode 100644
index 0000000..03e4791
--- /dev/null
+++ b/wrappers/streamcipher_pool.go
@@ -0,0 +1,16 @@
+package wrappers
+
+import (
+	"bytes"
+	"sync"
+)
+
+var streamCipherBufferPool sync.Pool
+
+func init() {
+	streamCipherBufferPool = sync.Pool{
+		New: func() interface{} {
+			return &bytes.Buffer{}
+		},
+	}
+}