Add buffered reader

wrappers/blockcipherrwc.go

@@ -1,7 +1,6 @@
 package wrappers
 
 import (
-	"bytes"
 	"crypto/aes"
 	"crypto/cipher"
 	"net"
@@ -10,7 +9,7 @@ import (
 )
 
 type BlockCipherReadWriteCloserWithAddr struct {
-	buf *bytes.Buffer
+	BufferedReader
 
 	conn      ReadWriteCloserWithAddr
 	encryptor cipher.BlockMode
@@ -18,20 +17,19 @@ type BlockCipherReadWriteCloserWithAddr struct {
 }
 
 func (c *BlockCipherReadWriteCloserWithAddr) Read(p []byte) (int, error) {
-	if c.buf.Len() > 0 {
-		return c.flush(p)
-	}
-
-	for c.buf.Len() == 0 || c.buf.Len()%aes.BlockSize != 0 {
-		n, err := c.conn.Read(p)
-		if err != nil {
-			return 0, errors.Annotate(err, "Cannot read from socket")
+	return c.BufferedRead(p, func() error {
+		bufferLength := c.Buffer.Len()
+		for bufferLength%aes.BlockSize != 0 || bufferLength == 0 {
+			n, err := c.conn.Read(p)
+			if err != nil {
+				return errors.Annotate(err, "Cannot read from socket")
+			}
+			c.Buffer.Write(p[:n])
 		}
-		c.buf.Write(p[:n])
-	}
-	c.decryptor.CryptBlocks(c.buf.Bytes(), c.buf.Bytes())
+		c.decryptor.CryptBlocks(c.Buffer.Bytes(), c.Buffer.Bytes())
 
-	return c.flush(p)
+		return nil
+	})
 }
 
 func (c *BlockCipherReadWriteCloserWithAddr) Write(p []byte) (int, error) {
@@ -39,19 +37,13 @@ func (c *BlockCipherReadWriteCloserWithAddr) Write(p []byte) (int, error) {
 		return 0, errors.Errorf("Incorrect block size %d", len(p))
 	}
 
-	buf := getBuffer()
-	defer putBuffer(buf)
-	buf.Grow(len(p))
-	buf.Write(p)
-
-	encrypted := buf.Bytes()
+	encrypted := make([]byte, len(p))
 	c.encryptor.CryptBlocks(encrypted, p)
 
 	return c.conn.Write(encrypted)
 }
 
 func (c *BlockCipherReadWriteCloserWithAddr) Close() error {
-	defer putBuffer(c.buf)
 	return c.conn.Close()
 }
 
@@ -59,32 +51,11 @@ func (c *BlockCipherReadWriteCloserWithAddr) Addr() *net.TCPAddr {
 	return c.conn.Addr()
 }
 
-func (c *BlockCipherReadWriteCloserWithAddr) flush(p []byte) (int, error) {
-	sizeToRead := len(p)
-	if c.buf.Len() < sizeToRead {
-		sizeToRead = c.buf.Len()
-	}
-
-	data := c.buf.Bytes()
-	copy(p, data[:sizeToRead])
-	if sizeToRead == c.buf.Len() {
-		c.buf.Reset()
-	} else {
-		newBuf := getBuffer()
-		newBuf.Write(data[sizeToRead:])
-
-		putBuffer(c.buf)
-		c.buf = newBuf
-	}
-
-	return sizeToRead, nil
-}
-
 func NewBlockCipherRWC(conn ReadWriteCloserWithAddr, encryptor, decryptor cipher.BlockMode) ReadWriteCloserWithAddr {
 	return &BlockCipherReadWriteCloserWithAddr{
-		buf:       getBuffer(),
-		conn:      conn,
-		encryptor: encryptor,
-		decryptor: decryptor,
+		BufferedReader: NewBufferedReader(),
+		conn:           conn,
+		encryptor:      encryptor,
+		decryptor:      decryptor,
 	}
 }