buffer outgoing writes into a single buffer

rw.go

@@ -2,6 +2,7 @@ package secio
 
 import (
 	"crypto/cipher"
+	"encoding/binary"
 	"errors"
 	"fmt"
 	"io"
@@ -24,17 +25,17 @@ var bufPool = mpool.ByteSlicePool
 
 type etmWriter struct {
 	// params
-	pool mpool.Pool        // for the buffers with encrypted data
-	msg  msgio.WriteCloser // msgio for knowing where boundaries lie
-	str  cipher.Stream     // the stream cipher to encrypt with
-	mac  HMAC              // the mac to authenticate data with
+	pool mpool.Pool    // for the buffers with encrypted data
+	str  cipher.Stream // the stream cipher to encrypt with
+	mac  HMAC          // the mac to authenticate data with
+	w    io.Writer
 
 	sync.Mutex
 }
 
 // NewETMWriter Encrypt-Then-MAC
 func NewETMWriter(w io.Writer, s cipher.Stream, mac HMAC) msgio.WriteCloser {
-	return &etmWriter{msg: msgio.NewWriter(w), str: s, mac: mac, pool: bufPool}
+	return &etmWriter{w: w, str: s, mac: mac, pool: bufPool}
 }
 
 // Write writes passed in buffer as a single message.
@@ -50,9 +51,10 @@ func (w *etmWriter) WriteMsg(b []byte) error {
 	w.Lock()
 	defer w.Unlock()
 
+	bufsize := uint32(4 + len(b) + w.mac.Size())
 	// encrypt.
-	data := w.pool.Get(uint32(len(b))).([]byte)
-	data = data[:len(b)] // the pool's buffer may be larger
+	buf := w.pool.Get(bufsize).([]byte)
+	data := buf[4 : 4+len(b)] // the pool's buffer may be larger
 	w.str.XORKeyStream(data, b)
 
 	// log.Debugf("ENC plaintext (%d): %s %v", len(b), b, b)
@@ -66,15 +68,18 @@ func (w *etmWriter) WriteMsg(b []byte) error {
 	// Sum appends.
 	data = w.mac.Sum(data)
 	w.mac.Reset()
-	// it's sad to append here. our buffers are -- hopefully -- coming from
-	// a shared buffer pool, so the append may not actually cause allocation
-	// one can only hope. i guess we'll see.
+	binary.BigEndian.PutUint32(buf[:4], uint32(len(data)))
 
-	return w.msg.WriteMsg(data)
+	_, err := w.w.Write(buf[:bufsize])
+	w.pool.Put(bufsize, buf)
+	return err
 }
 
 func (w *etmWriter) Close() error {
-	return w.msg.Close()
+	if c, ok := w.w.(io.Closer); ok {
+		return c.Close()
+	}
+	return nil
 }
 
 type etmReader struct {

rw_test.go

@@ -0,0 +1,85 @@
+package secio
+
+import (
+	"bytes"
+	"crypto/aes"
+	"crypto/cipher"
+	"io"
+	"testing"
+)
+
+type keyInfo struct {
+	cipherKey []byte
+	iv        []byte
+	macKey    []byte
+}
+
+func getTestKeyInfo() *keyInfo {
+	return &keyInfo{
+		cipherKey: []byte("this is a test keyaaaaaaaaaaaaaa"),
+		iv:        make([]byte, 16),
+		macKey:    []byte("key for the mac"),
+	}
+}
+
+func getTestingWriter(w io.Writer, ki *keyInfo) (*etmWriter, error) {
+	c, err := aes.NewCipher(ki.cipherKey)
+	if err != nil {
+		return nil, err
+	}
+
+	stream := cipher.NewCFBEncrypter(c, ki.iv)
+
+	mac, err := newMac("SHA256", ki.macKey)
+	if err != nil {
+		return nil, err
+	}
+
+	return NewETMWriter(w, stream, mac).(*etmWriter), nil
+}
+
+func getTestingReader(r io.Reader, ki *keyInfo) (*etmReader, error) {
+	c, err := aes.NewCipher(ki.cipherKey)
+	if err != nil {
+		return nil, err
+	}
+
+	stream := cipher.NewCFBDecrypter(c, ki.iv)
+
+	mac, err := newMac("SHA256", ki.macKey)
+	if err != nil {
+		return nil, err
+	}
+
+	return NewETMReader(r, stream, mac).(*etmReader), nil
+}
+
+func TestBasicETMStream(t *testing.T) {
+	buf := new(bytes.Buffer)
+
+	ki := getTestKeyInfo()
+	w, err := getTestingWriter(buf, ki)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	before := []byte("hello world")
+	err = w.WriteMsg(before)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	r, err := getTestingReader(buf, ki)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	msg, err := r.ReadMsg()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if string(before) != string(msg) {
+		t.Fatal("got wrong message")
+	}
+}