From 67ddbe40371a486c21666b7f2564d34f7c29aab7 Mon Sep 17 00:00:00 2001
From: Chen Yufei <cyfdecyf@gmail.com>
Date: Sun, 26 Apr 2015 22:11:54 +0800
Subject: [PATCH] Reuse buffer when possible for salsa20 cipher.

---
 shadowsocks/encrypt.go  | 35 +++++++++++++++++++++++++----------
 shadowsocks/leakybuf.go |  4 ++--
 2 files changed, 27 insertions(+), 12 deletions(-)

diff --git a/shadowsocks/encrypt.go b/shadowsocks/encrypt.go
index 6ddf131..e0bdc14 100644
--- a/shadowsocks/encrypt.go
+++ b/shadowsocks/encrypt.go
@@ -2,9 +2,6 @@ package shadowsocks
 
 import (
 	"bytes"
-	"golang.org/x/crypto/blowfish"
-	"golang.org/x/crypto/cast5"
-	"golang.org/x/crypto/salsa20/salsa"
 	"crypto/aes"
 	"crypto/cipher"
 	"crypto/des"
@@ -13,6 +10,9 @@ import (
 	"crypto/rc4"
 	"encoding/binary"
 	"errors"
+	"golang.org/x/crypto/blowfish"
+	"golang.org/x/crypto/cast5"
+	"golang.org/x/crypto/salsa20/salsa"
 	"io"
 )
 
@@ -139,20 +139,35 @@ func newRC4MD5Stream(key, iv []byte, _ DecOrEnc) (cipher.Stream, error) {
 }
 
 type salsaStreamCipher struct {
-	nonce [8]byte
-	key [32]byte
+	nonce   [8]byte
+	key     [32]byte
 	counter int
 }
 
 func (c *salsaStreamCipher) XORKeyStream(dst, src []byte) {
+	var buf []byte
+	padLen := c.counter % 64
+	dataSize := len(src) + padLen
+	if cap(dst) >= dataSize {
+		buf = dst[:dataSize]
+	} else if leakyBufSize >= dataSize {
+		buf = leakyBuf.Get()
+		defer leakyBuf.Put(buf)
+		buf = buf[:dataSize]
+	} else {
+		buf = make([]byte, dataSize)
+	}
+
 	var subNonce [16]byte
-	padlen := c.counter % 64
-	buf := make([]byte, padlen+len(src))
-	copy(buf[padlen:], src[:])
 	copy(subNonce[:], c.nonce[:])
-	binary.LittleEndian.PutUint64(subNonce[len(c.nonce):], uint64(c.counter / 64))
+	binary.LittleEndian.PutUint64(subNonce[len(c.nonce):], uint64(c.counter/64))
+
+	// It's difficult to avoid data copy here. src or dst maybe slice from
+	// Conn.Read/Write, which can't have padding.
+	copy(buf[padLen:], src[:])
 	salsa.XORKeyStream(buf, buf, &subNonce, &c.key)
-	copy(dst, buf[padlen:])
+	copy(dst, buf[padLen:])
+
 	c.counter += len(src)
 }
 
diff --git a/shadowsocks/leakybuf.go b/shadowsocks/leakybuf.go
index f673cc3..3b55832 100644
--- a/shadowsocks/leakybuf.go
+++ b/shadowsocks/leakybuf.go
@@ -6,10 +6,10 @@ type LeakyBuf struct {
 	freeList chan []byte
 }
 
-const bufSize = 4096
+const leakyBufSize = 4096
 const maxNBuf = 2048
 
-var leakyBuf = NewLeakyBuf(maxNBuf, bufSize)
+var leakyBuf = NewLeakyBuf(maxNBuf, leakyBufSize)
 
 // NewLeakyBuf creates a leaky buffer which can hold at most n buffer, each
 // with bufSize bytes.