Merge pull request #90 from Jigsaw-Code/bemasc-pool2

Reduce memory usage in read streams

Author: Benjamin M. Schwartz

shadowsocks/cipher.go

@@ -77,6 +77,16 @@ func newAesGCM(key []byte) (cipher.AEAD, error) {
 	return cipher.NewGCM(blk)
 }
 
+func maxTagSize() int {
+	max := 0
+	for _, spec := range supportedAEADs {
+		if spec.tagSize > max {
+			max = spec.tagSize
+		}
+	}
+	return max
+}
+
 // Cipher encapsulates a Shadowsocks AEAD spec and a secret
 type Cipher struct {
 	aead   aeadSpec

shadowsocks/stream.go

@@ -21,11 +21,17 @@ import (
 	"fmt"
 	"io"
 	"sync"
+
+	"github.com/Jigsaw-Code/outline-ss-server/slicepool"
 )
 
 // payloadSizeMask is the maximum size of payload in bytes.
 const payloadSizeMask = 0x3FFF // 16*1024 - 1
 
+// Buffer pool used for decrypting Shadowsocks streams.
+// The largest buffer we could need is for decrypting a max-length payload.
+var readBufPool = slicepool.MakePool(payloadSizeMask + maxTagSize())
+
 // Writer is an io.Writer that also implements io.ReaderFrom to
 // allow for piping the data without extra allocations and copies.
 // The LazyWrite and Flush methods allow a header to be
@@ -262,7 +268,10 @@ type chunkReader struct {
 	aead cipher.AEAD
 	// Index of the next encrypted chunk to read.
 	counter []byte
-	buf     []byte
+	// Buffer for the uint16 size and its AEAD tag.  Made in init().
+	payloadSizeBuf []byte
+	// Holds a buffer for the payload and its AEAD tag, when needed.
+	payload slicepool.LazySlice
 }
 
 // Reader is an io.Reader that also implements io.WriterTo to
@@ -276,7 +285,11 @@ type Reader interface {
 // the shadowsocks protocol with the given shadowsocks cipher.
 func NewShadowsocksReader(reader io.Reader, ssCipher *Cipher) Reader {
 	return &readConverter{
-		cr: &chunkReader{reader: reader, ssCipher: ssCipher},
+		cr: &chunkReader{
+			reader:   reader,
+			ssCipher: ssCipher,
+			payload:  readBufPool.LazySlice(),
+		},
 	}
 }
 
@@ -296,7 +309,7 @@ func (cr *chunkReader) init() (err error) {
 			return fmt.Errorf("failed to create AEAD: %v", err)
 		}
 		cr.counter = make([]byte, cr.aead.NonceSize())
-		cr.buf = make([]byte, payloadSizeMask+cr.aead.Overhead())
+		cr.payloadSizeBuf = make([]byte, 2+cr.aead.Overhead())
 	}
 	return nil
 }
@@ -318,31 +331,38 @@ func (cr *chunkReader) readMessage(buf []byte) error {
 	return nil
 }
 
+// ReadChunk returns the next chunk from the stream.  Callers must fully
+// consume and discard the previous chunk before calling ReadChunk again.
 func (cr *chunkReader) ReadChunk() ([]byte, error) {
 	if err := cr.init(); err != nil {
 		return nil, err
 	}
+
+	// Release the previous payload buffer.
+	cr.payload.Release()
+
 	// In Shadowsocks-AEAD, each chunk consists of two
 	// encrypted messages.  The first message contains the payload length,
-	// and the second message is the payload.
-	sizeBuf := cr.buf[:2+cr.aead.Overhead()]
-	if err := cr.readMessage(sizeBuf); err != nil {
+	// and the second message is the payload.  Idle read threads will
+	// block here until the next chunk.
+	if err := cr.readMessage(cr.payloadSizeBuf); err != nil {
 		if err != io.EOF && err != io.ErrUnexpectedEOF {
 			err = fmt.Errorf("failed to read payload size: %v", err)
 		}
 		return nil, err
 	}
-	size := int(binary.BigEndian.Uint16(sizeBuf) & payloadSizeMask)
+	size := int(binary.BigEndian.Uint16(cr.payloadSizeBuf) & payloadSizeMask)
 	sizeWithTag := size + cr.aead.Overhead()
-	if cap(cr.buf) < sizeWithTag {
-		// This code is unreachable.
+	payloadBuf := cr.payload.Acquire()
+	if cap(payloadBuf) < sizeWithTag {
+		// This code is unreachable if the constants are set correctly.
 		return nil, io.ErrShortBuffer
 	}
-	payloadBuf := cr.buf[:sizeWithTag]
-	if err := cr.readMessage(payloadBuf); err != nil {
+	if err := cr.readMessage(payloadBuf[:sizeWithTag]); err != nil {
 		if err == io.EOF { // EOF is not expected mid-chunk.
 			err = io.ErrUnexpectedEOF
 		}
+		cr.payload.Release()
 		return nil, err
 	}
 	return payloadBuf[:size], nil
@@ -388,6 +408,7 @@ func (c *readConverter) ensureLeftover() error {
 	if len(c.leftover) > 0 {
 		return nil
 	}
+	c.leftover = nil
 	payload, err := c.cr.ReadChunk()
 	if err != nil {
 		return err