Code review updates

framed.go

@@ -1,6 +1,6 @@
 /*
-Package framed provides an implementation of io.ReadWriteCloser that reads and
-writes whole frames only.
+Package framed provides an implementations of io.Writer and io.Reader that write
+and read whole frames only.
 
 Frames are length-prefixed.  The first two bytes are an unsigned 16 bit int
 stored in little-endian byte order indicating the length of the content.  The
@@ -22,12 +22,14 @@ const (
 	// FrameHeaderBits is the size of the frame header in bits
 	FrameHeaderBits = 16
 
-	// FrameHeaderSize is the size of the frame header in bytes
-	FrameHeaderSize = 16 / 2
+	// FrameHeaderLength is the size of the frame header in bytes
+	FrameHeaderLength = FrameHeaderBits / 8
 
-	// MaxFrameSize is the maximum possible size of a frame (not including the
+	// MaxFrameLength is the maximum possible size of a frame (not including the
 	// length prefix)
-	MaxFrameSize = 1<<FrameHeaderBits - 1
+	MaxFrameLength = 1<<FrameHeaderBits - 1
+
+	tooLongError = "Attempted to write frame of length %s which is longer than maximum allowed length of %s"
 )
 
 var endianness = binary.LittleEndian
@@ -109,10 +111,7 @@ func (framed *Reader) ReadFrame() (frame []byte, err error) {
 	frame = make([]byte, l)
 
 	// Read into buffer
-	n, err := io.ReadFull(framed.Stream, frame)
-	if n != l {
-		return nil, fmt.Errorf("Read wrong number of bytes. Expected frame of length %d, got %d", n, l)
-	}
+	_, err = io.ReadFull(framed.Stream, frame)
 	return
 }
 
@@ -125,6 +124,9 @@ func (framed *Writer) Write(frame []byte) (n int, err error) {
 	defer framed.mutex.Unlock()
 
 	n = len(frame)
+	if n > MaxFrameLength {
+		return 0, fmt.Errorf(tooLongError, n, MaxFrameLength)
+	}
 
 	// Write the length header
 	if err = binary.Write(framed.Stream, endianness, uint16(n)); err != nil {
@@ -150,6 +152,10 @@ func (framed *Writer) WritePieces(pieces ...[]byte) (n int, err error) {
 		n = n + len(piece)
 	}
 
+	if n > MaxFrameLength {
+		return 0, fmt.Errorf(tooLongError, n, MaxFrameLength)
+	}
+
 	// Write the length header
 	if err = binary.Write(framed.Stream, endianness, uint16(n)); err != nil {
 		return

framed_test.go

@@ -2,9 +2,12 @@ package framed
 
 import (
 	"bytes"
+	"io/ioutil"
 	"sync"
 	"testing"
 	"time"
+
+	"github.com/getlantern/testify/assert"
 )
 
 const (
@@ -59,8 +62,8 @@ func TestFraming(t *testing.T) {
 			}
 			if err != nil {
 				t.Errorf("Unable to write: %s", err)
-			} else if n != len(testMessage) {
-				t.Errorf("%d bytes written did not match length of test message %d", n, len(testMessage))
+			} else {
+				assert.Equal(t, len(testMessage), n, "Bytes written should match length of test message")
 			}
 			wg.Done()
 		}()
@@ -81,19 +84,39 @@ func TestFraming(t *testing.T) {
 				if n, err = reader.Read(buffer); err != nil {
 					t.Errorf("Unable to read: %s", err)
 					return
-				} else if n != len(testMessage) {
-					t.Errorf("%d bytes read did not match length of test message %d", n, len(testMessage))
-					return
+				} else {
+					assert.Equal(t, len(testMessage), n, "Bytes read should match length of test message")
 				}
 				frame = buffer[:n]
 			}
 
-			if !bytes.Equal(frame, testMessage) {
-				t.Errorf("Received did not match expected.  Expected: '%q', Received: '%q'", testMessage, frame)
-			}
+			assert.Equal(t, testMessage, frame, "Received should match sent")
 			wg.Done()
 		}()
 	}
 
 	wg.Wait()
 }
+
+func TestWriteTooLong(t *testing.T) {
+	w := NewWriter(ioutil.Discard)
+	b := make([]byte, MaxFrameLength+1)
+	n, err := w.Write(b)
+	assert.Error(t, err, "Writing too long message should result in error")
+	assert.Equal(t, 0, n, "Writing too long message should result in 0 bytes written")
+	n, err = w.Write(b[:len(b)-1])
+	assert.NoError(t, err, "Writing message of MaxFrameLength should be allowed")
+	assert.Equal(t, MaxFrameLength, n, "Writing message of MaxFrameLength should have written MaxFrameLength bytes")
+}
+
+func TestWritePiecesTooLong(t *testing.T) {
+	w := NewWriter(ioutil.Discard)
+	b1 := make([]byte, MaxFrameLength)
+	b2 := make([]byte, 1)
+	n, err := w.WritePieces(b1, b2)
+	assert.Error(t, err, "Writing too long message should result in error")
+	assert.Equal(t, 0, n, "Writing too long message should result in 0 bytes written")
+	n, err = w.WritePieces(b1[:len(b1)-1], b2)
+	assert.NoError(t, err, "Writing message of MaxFrameLength should be allowed")
+	assert.Equal(t, MaxFrameLength, n, "Writing message of MaxFrameLength should have written MaxFrameLength bytes")
+}