Refactored all the things!

wrappers/blockcipher.go

@@ -1,6 +1,7 @@
 package wrappers
 
 import (
+	"bytes"
 	"crypto/aes"
 	"crypto/cipher"
 	"net"
@@ -9,77 +10,89 @@ import (
 	"github.com/juju/errors"
 )
 
-type WrapBlockCipher struct {
-	BufferedReader
+type BlockCipher struct {
+	buf *bytes.Buffer
 
 	conn      WrapStreamReadWriteCloser
 	encryptor cipher.BlockMode
 	decryptor cipher.BlockMode
 }
 
-func (w *WrapBlockCipher) Read(p []byte) (int, error) {
-	return w.BufferedRead(p, func() error {
-		var buf []byte
+func (b *BlockCipher) Read(p []byte) (int, error) {
+	if b.buf.Len() > 0 {
+		return b.flush(p)
+	}
 
-		for len(buf) == 0 || len(buf)%aes.BlockSize != 0 {
-			rv, err := utils.ReadCurrentData(w.conn)
-			if err != nil {
-				return errors.Annotate(err, "Cannot read from socket")
-			}
-			buf = append(buf, rv...)
+	buf := []byte{}
+	for len(buf) == 0 || len(buf)%aes.BlockSize != 0 {
+		rv, err := utils.ReadCurrentData(b.conn)
+		if err != nil {
+			return 0, errors.Annotate(err, "Cannot read from socket")
 		}
+		buf = append(buf, rv...)
+	}
 
-		w.decryptor.CryptBlocks(buf, buf)
-		w.Buffer.Write(buf)
+	b.decryptor.CryptBlocks(buf, buf)
+	b.buf.Write(buf)
 
-		return nil
-	})
+	return b.flush(p)
 }
 
-func (w *WrapBlockCipher) Write(p []byte) (int, error) {
+func (b *BlockCipher) flush(p []byte) (int, error) {
+	if b.buf.Len() <= len(p) {
+		sizeToReturn := b.buf.Len()
+		copy(p, b.buf.Bytes())
+		b.buf.Reset()
+		return sizeToReturn, nil
+	}
+
+	return b.buf.Read(p)
+}
+
+func (b *BlockCipher) Write(p []byte) (int, error) {
 	if len(p)%aes.BlockSize > 0 {
 		return 0, errors.Errorf("Incorrect block size %d", len(p))
 	}
 
 	encrypted := make([]byte, len(p))
-	w.encryptor.CryptBlocks(encrypted, p)
+	b.encryptor.CryptBlocks(encrypted, p)
 
-	return w.conn.Write(encrypted)
+	return b.conn.Write(encrypted)
 }
 
-func (w *WrapBlockCipher) LogDebug(msg string, data ...interface{}) {
-	w.conn.LogDebug(msg, data...)
+func (b *BlockCipher) LogDebug(msg string, data ...interface{}) {
+	b.conn.LogDebug(msg, data...)
 }
 
-func (w *WrapBlockCipher) LogInfo(msg string, data ...interface{}) {
-	w.conn.LogInfo(msg, data...)
+func (b *BlockCipher) LogInfo(msg string, data ...interface{}) {
+	b.conn.LogInfo(msg, data...)
 }
 
-func (w *WrapBlockCipher) LogWarn(msg string, data ...interface{}) {
-	w.conn.LogWarn(msg, data...)
+func (b *BlockCipher) LogWarn(msg string, data ...interface{}) {
+	b.conn.LogWarn(msg, data...)
 }
 
-func (w *WrapBlockCipher) LogError(msg string, data ...interface{}) {
-	w.conn.LogError(msg, data...)
+func (b *BlockCipher) LogError(msg string, data ...interface{}) {
+	b.conn.LogError(msg, data...)
 }
 
-func (w *WrapBlockCipher) LocalAddr() *net.TCPAddr {
-	return w.conn.LocalAddr()
+func (b *BlockCipher) LocalAddr() *net.TCPAddr {
+	return b.conn.LocalAddr()
 }
 
-func (w *WrapBlockCipher) RemoteAddr() *net.TCPAddr {
-	return w.conn.RemoteAddr()
+func (b *BlockCipher) RemoteAddr() *net.TCPAddr {
+	return b.conn.RemoteAddr()
 }
 
-func (w *WrapBlockCipher) Close() error {
-	return w.conn.Close()
+func (b *BlockCipher) Close() error {
+	return b.conn.Close()
 }
 
-func NewWrapBlockCipher(conn WrapStreamReadWriteCloser, encryptor, decryptor cipher.BlockMode) WrapStreamReadWriteCloser {
-	return &WrapBlockCipher{
-		BufferedReader: NewBufferedReader(),
-		conn:           conn,
-		encryptor:      encryptor,
-		decryptor:      decryptor,
+func NewBlockCipher(conn WrapStreamReadWriteCloser, encryptor, decryptor cipher.BlockMode) WrapStreamReadWriteCloser {
+	return &BlockCipher{
+		buf:       &bytes.Buffer{},
+		conn:      conn,
+		encryptor: encryptor,
+		decryptor: decryptor,
 	}
 }