go/common/cbor: Use cbor.Decoder in codec's Read

go/common/cbor/cbor_test.go

@@ -1,6 +1,7 @@
 package cbor
 
 import (
+	"bytes"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -21,3 +22,17 @@ func TestOutOfMem2(t *testing.T) {
 	err := Unmarshal([]byte("\x9b\x00\x00\x81112233"), f)
 	require.Error(err, "Invalid CBOR input should fail")
 }
+
+func TestEncoderDecoder(t *testing.T) {
+	require := require.New(t)
+
+	var buf bytes.Buffer
+	enc := NewEncoder(&buf)
+	enc.Encode(42)
+
+	var x int
+	dec := NewDecoder(&buf)
+	err := dec.Decode(&x)
+	require.NoError(err, "Decode")
+	require.EqualValues(42, x, "decoded value should be correct")
+}

go/common/cbor/codec.go

@@ -13,7 +13,8 @@ import (
 const maxMessageSize = 104857600 // 100MB
 
 var (
-	errMessageTooLarge = errors.New("codec: message too large")
+	errMessageTooLarge  = errors.New("codec: message too large")
+	errMessageMalformed = errors.New("codec: message is malformed")
 
 	codecValueSize = prometheus.NewSummaryVec(
 		prometheus.SummaryOpts{
@@ -52,15 +53,14 @@ func (c *MessageReader) Read(msg interface{}) error {
 		return errMessageTooLarge
 	}
 
-	// Read message bytes.
-	rawMessage := make([]byte, length)
-	if _, err := io.ReadFull(c.reader, rawMessage); err != nil {
+	// Decode message bytes.
+	r := io.LimitReader(c.reader, int64(length))
+	dec := NewDecoder(r)
+	if err := dec.Decode(msg); err != nil {
 		return err
 	}
-
-	// Decode CBOR into given message.
-	if err := Unmarshal(rawMessage, msg); err != nil {
-		return err
+	if r.(*io.LimitedReader).N > 0 {
+		return errMessageMalformed
 	}
 
 	return nil

go/common/cbor/codec_test.go

@@ -2,6 +2,7 @@ package cbor
 
 import (
 	"bytes"
+	"encoding/binary"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -34,3 +35,39 @@ func TestCodecRoundTrip(t *testing.T) {
 	require.NoError(t, err, "Read (2nd)")
 	require.EqualValues(t, msg, decodedMsg2, "Decoded message must be equal to source message")
 }
+
+func TestCodecOversized(t *testing.T) {
+	require := require.New(t)
+
+	var buffer bytes.Buffer
+	codec := NewMessageCodec(&buffer)
+
+	err := codec.Write(42)
+	require.NoError(err, "Write")
+
+	// Corrupt the buffer to include a huge length.
+	binary.BigEndian.PutUint32(buffer.Bytes()[:4], maxMessageSize+1)
+
+	var x int
+	err = codec.Read(&x)
+	require.Error(err, "Read should fail with oversized message")
+	require.EqualValues(errMessageTooLarge, err)
+}
+
+func TestCodecMalformed(t *testing.T) {
+	require := require.New(t)
+
+	var buffer bytes.Buffer
+	codec := NewMessageCodec(&buffer)
+
+	err := codec.Write(42)
+	require.NoError(err, "Write")
+
+	// Corrupt the buffer to include an incorrect length (larger than what is really there).
+	binary.BigEndian.PutUint32(buffer.Bytes()[:4], 1024)
+
+	var x int
+	err = codec.Read(&x)
+	require.Error(err, "Read should fail with malformed message")
+	require.EqualValues(errMessageMalformed, err)
+}