From f6d1cd81d8f8605cff9956b6341e02b4a2ea2d7e Mon Sep 17 00:00:00 2001
From: Jernej Kos <jernej@kos.mx>
Date: Thu, 9 Apr 2020 14:14:32 +0200
Subject: [PATCH] go/common/cbor: Use cbor.Decoder in codec's Read

---
 go/common/cbor/cbor_test.go  | 15 +++++++++++++++
 go/common/cbor/codec.go      | 16 ++++++++--------
 go/common/cbor/codec_test.go | 37 ++++++++++++++++++++++++++++++++++++
 3 files changed, 60 insertions(+), 8 deletions(-)

diff --git a/go/common/cbor/cbor_test.go b/go/common/cbor/cbor_test.go
index 23034a10159..bc1a375ab41 100644
--- a/go/common/cbor/cbor_test.go
+++ b/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")
+}
diff --git a/go/common/cbor/codec.go b/go/common/cbor/codec.go
index c34b81445e6..bfcb91d8b91 100644
--- a/go/common/cbor/codec.go
+++ b/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
diff --git a/go/common/cbor/codec_test.go b/go/common/cbor/codec_test.go
index 9340de13a35..708b4e6737f 100644
--- a/go/common/cbor/codec_test.go
+++ b/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)
+}