crypto: Add buffer management to writer

benchmarks:

    benchmark                         old ns/op     new ns/op     delta
    BenchmarkEncryptWriter            37032949      37754353      +1.95%
    BenchmarkEncrypt                  35989611      36344593      +0.99%
    BenchmarkDecryptReader            38274466      38282979      +0.02%
    BenchmarkEncryptDecryptReader     77506506      77088612      -0.54%
    BenchmarkDecrypt                  36219298      36128875      -0.25%

    benchmark                         old MB/s     new MB/s     speedup
    BenchmarkEncryptWriter            226.52       222.19       0.98x
    BenchmarkEncrypt                  233.08       230.81       0.99x
    BenchmarkDecryptReader            219.17       219.12       1.00x
    BenchmarkEncryptDecryptReader     108.23       108.82       1.01x
    BenchmarkDecrypt                  231.61       232.19       1.00x

    benchmark                         old allocs     new allocs     delta
    BenchmarkEncryptWriter            20             16             -20.00%
    BenchmarkEncrypt                  12             12             +0.00%
    BenchmarkDecryptReader            13             13             +0.00%
    BenchmarkEncryptDecryptReader     32             27             -15.62%
    BenchmarkDecrypt                  10             10             +0.00%

    benchmark                         old bytes     new bytes     delta
    BenchmarkEncryptWriter            1020064       170331        -83.30%
    BenchmarkEncrypt                  1600          1600          +0.00%
    BenchmarkDecryptReader            841070        842094        +0.12%
    BenchmarkEncryptDecryptReader     3027433       845129        -72.08%
    BenchmarkDecrypt                  1573          1573          +0.00%

crypto/buffer_pool.go

@@ -2,7 +2,7 @@ package crypto
 
 import "sync"
 
-const defaultBufSize = 2048
+const defaultBufSize = 32 * 1024 // 32KiB
 
 var bufPool = sync.Pool{
 	New: func() interface{} {

crypto/writer.go

@@ -1,110 +1,84 @@
 package crypto
 
 import (
-	"bytes"
 	"crypto/aes"
 	"crypto/cipher"
+	"errors"
 	"fmt"
 	"io"
-	"sync"
 )
 
 type encryptWriter struct {
-	iv      []byte
-	wroteIV bool
-	data    *bytes.Buffer
-	key     *Key
-	s       cipher.Stream
-	w       io.Writer
-	origWr  io.Writer
-	err     error // remember error writing iv
+	data   []byte
+	key    *Key
+	s      cipher.Stream
+	w      io.Writer
+	closed bool
 }
 
 func (e *encryptWriter) Close() error {
-	// write mac
-	mac := poly1305Sign(e.data.Bytes()[ivSize:], e.data.Bytes()[:ivSize], &e.key.Sign)
-	_, err := e.origWr.Write(mac)
+	if e.closed {
+		return errors.New("Close() called on already closed writer")
+	}
+	e.closed = true
+
+	// encrypt everything
+	iv, c := e.data[:ivSize], e.data[ivSize:]
+	e.s.XORKeyStream(c, c)
+
+	// compute mac
+	mac := poly1305Sign(c, iv, &e.key.Sign)
+	e.data = append(e.data, mac...)
+
+	// write everything
+	n, err := e.w.Write(e.data)
 	if err != nil {
 		return err
 	}
 
-	// return buffer
-	bufPool.Put(e.data.Bytes())
+	if n != len(e.data) {
+		return errors.New("not all bytes written")
+	}
+
+	// return buffer to pool
+	freeBuffer(e.data)
 
 	return nil
 }
 
-const encryptWriterChunkSize = 512 * 1024 // 512 KiB
-var encryptWriterBufPool = sync.Pool{
-	New: func() interface{} {
-		return make([]byte, encryptWriterChunkSize)
-	},
-}
-
 func (e *encryptWriter) Write(p []byte) (int, error) {
-	// write iv first
-	if !e.wroteIV {
-		_, e.err = e.origWr.Write(e.iv[:])
-		e.wroteIV = true
+	// if e.data is too small, return it to the buffer and create new slice
+	if cap(e.data) < len(e.data)+len(p) {
+		b := make([]byte, len(e.data), len(e.data)*2)
+		copy(b, e.data)
+		freeBuffer(e.data)
+		e.data = b
 	}
 
-	if e.err != nil {
-		return 0, e.err
-	}
-
-	buf := encryptWriterBufPool.Get().([]byte)
-	defer encryptWriterBufPool.Put(buf)
-
-	written := 0
-	for len(p) > 0 {
-		max := len(p)
-		if max > encryptWriterChunkSize {
-			max = encryptWriterChunkSize
-		}
-
-		e.s.XORKeyStream(buf, p[:max])
-		n, err := e.w.Write(buf[:max])
-		if n != max {
-			if err == nil { // should never happen
-				err = io.ErrShortWrite
-			}
-		}
-
-		written += n
-		p = p[n:]
-
-		if err != nil {
-			e.err = err
-			return written, err
-		}
-	}
-
-	return written, nil
+	// copy new data to e.data
+	e.data = append(e.data, p...)
+	return len(p), nil
 }
 
 // EncryptTo buffers data written to the returned io.WriteCloser. When Close()
-// is called, the data is encrypted an written to the underlying writer.
+// is called, the data is encrypted and written to the underlying writer.
 func EncryptTo(ks *Key, wr io.Writer) io.WriteCloser {
 	ew := &encryptWriter{
-		iv:     newIV(),
-		data:   bytes.NewBuffer(getBuffer()[:0]),
-		key:    ks,
-		origWr: wr,
+		data: getBuffer(),
+		key:  ks,
 	}
 
 	// buffer iv for mac
-	_, err := ew.data.Write(ew.iv[:])
-	if err != nil {
-		panic(err)
-	}
+	ew.data = ew.data[:ivSize]
+	copy(ew.data, newIV())
 
 	c, err := aes.NewCipher(ks.Encrypt[:])
 	if err != nil {
 		panic(fmt.Sprintf("unable to create cipher: %v", err))
 	}
 
-	ew.s = cipher.NewCTR(c, ew.iv[:])
-	ew.w = io.MultiWriter(ew.data, wr)
+	ew.s = cipher.NewCTR(c, ew.data[:ivSize])
+	ew.w = wr
 
 	return ew
 }