Make encryption buffer size adjustable

shadowsocks/stream.go

@@ -51,6 +51,8 @@ type Writer struct {
  byteWrapper bytes.Reader
  // Number of plaintext bytes that are currently buffered.
  pending int
+ // Maximum number of plaintext bytes to write in each chunk.
+ maxChunkSize int
  // These are populated by init():
  buf []byte
  aead cipher.AEAD
@@ -61,14 +63,29 @@ type Writer struct {
 // NewShadowsocksWriter creates a Writer that encrypts the given Writer using
 // the shadowsocks protocol with the given shadowsocks cipher.
 func NewShadowsocksWriter(writer io.Writer, ssCipher *Cipher) *Writer {
- return &Writer{writer: writer, ssCipher: ssCipher, saltGenerator: RandomSaltGenerator}
+ return &Writer{
+ writer: writer,
+ ssCipher: ssCipher,
+ saltGenerator: RandomSaltGenerator,
+ maxChunkSize: payloadSizeMask,
+ }
 }
 
 // SetSaltGenerator sets the salt generator to be used. Must be called before the first write.
 func (sw *Writer) SetSaltGenerator(saltGenerator SaltGenerator) {
  sw.saltGenerator = saltGenerator
 }
 
+// SetMaxChunkSize sets the maximum number of bytes to encrypt as a chunk.
+// Defaults to 2^16-1. Smaller values save memory but increase framing overhead.
+// Must be called before the first write.
+func (sw *Writer) SetMaxChunkSize(maxChunkSize int) {
+ if maxChunkSize > payloadSizeMask || maxChunkSize <= 0 {
+ maxChunkSize = payloadSizeMask
+ }
+ sw.maxChunkSize = maxChunkSize
+}
+
 // init generates a random salt, sets up the AEAD object and writes
 // the salt to the inner Writer.
 func (sw *Writer) init() (err error) {
@@ -86,7 +103,7 @@ func (sw *Writer) init() (err error) {
  // The maximum length message is the salt (first message only), length, length tag,
  // payload, and payload tag.
  sizeBufSize := 2 + sw.aead.Overhead()
- maxPayloadBufSize := payloadSizeMask + sw.aead.Overhead()
+ maxPayloadBufSize := sw.maxChunkSize + sw.aead.Overhead()
  sw.buf = make([]byte, len(salt)+sizeBufSize+maxPayloadBufSize)
  // Store the salt at the start of sw.buf.
  copy(sw.buf, salt)
@@ -165,7 +182,7 @@ func (sw *Writer) buffers() (sizeBuf, payloadBuf []byte) {
  // followed by a variable-length payload block.
  sizeBuf = sw.buf[saltSize : saltSize+2]
  payloadStart := saltSize + 2 + sw.aead.Overhead()
- payloadBuf = sw.buf[payloadStart : payloadStart+payloadSizeMask]
+ payloadBuf = sw.buf[payloadStart : payloadStart+sw.maxChunkSize]
  return
 }
 

shadowsocks/stream_test.go

@@ -411,3 +411,57 @@ func TestLazyWriteConcurrentFlush(t *testing.T) {
  t.Errorf("Wrong final content: %v", decrypted)
  }
 }
+
+func TestChunkSizeIntegrity(t *testing.T) {
+ t.Parallel()
+ cipher := newTestCipher(t)
+
+ // Test extreme and reasonable values.
+ testChunkSizes := []int{
+ 1, 2, 3, 4, 250, 256, 257, 1000, 4000, 8192, 16383,
+ }
+
+ const numWrites = 5
+
+ input := make([]byte, numWrites*16383)
+ for i := range input {
+ input[i] = byte(i) // Arbitrary test contents
+ }
+
+ for _, maxChunkSize := range testChunkSizes {
+ maxChunkSize := maxChunkSize
+ t.Run(fmt.Sprintf("maxChunkSize=%d", maxChunkSize), func(t *testing.T) {
+ t.Parallel()
+ connReader, connWriter := io.Pipe()
+ writer := NewShadowsocksWriter(connWriter, cipher)
+ writer.SetMaxChunkSize(maxChunkSize)
+ reader := NewShadowsocksReader(connReader, cipher)
+ go func() {
+ defer connWriter.Close()
+ if _, err := writer.Write(input[:numWrites*maxChunkSize]); err != nil {
+ t.Errorf("Failed Write: %v", err)
+ }
+ }()
+
+ // Check that all writes have the expected size and contents.
+ buf := make([]byte, 2*maxChunkSize)
+ for i := 0; i < numWrites; i++ {
+ n, err := reader.Read(buf)
+ if err != nil {
+ t.Errorf("Read failed at chunk %d: %v", i, err)
+ } else if n != maxChunkSize {
+ t.Errorf("Chunk %d has wrong size: %d", i, n)
+ } else if !bytes.Equal(input[i*maxChunkSize:][:maxChunkSize], buf[:n]) {
+ t.Errorf("Data mismatch at chunk %d", i)
+ }
+ }
+ // Check that the stream is closed cleanly.
+ n, err := reader.Read(buf)
+ if n != 0 {
+ t.Errorf("Got %d extra bytes", n)
+ } else if err != io.EOF {
+ t.Errorf("Wanted EOF, got %v", err)
+ }
+ })
+ }
+}