Fix decoder panics on malformed frame headers and lengths (v4 backport)

message.go

@@ -22,6 +22,7 @@ import (
 	"encoding/binary"
 	"fmt"
 	"io"
+	"math"
 	"time"
 
 	"github.com/michaelquigley/pfxlog"
@@ -53,6 +54,11 @@ const (
 	IsFirstGroupConnection          = 11
 	UnderlayTypeHeader              = 12
 
+	// replyForHeaderLen is the wire width of ReplyForHeader. The message layer decodes this
+	// header itself, so unmarshalHeaders rejects any other length rather than leaving a value
+	// the getters cannot interpret.
+	replyForHeaderLen = 4
+
 	// Headers in the range 128-255 inclusive will be reflected when creating replies
 	ReflectedHeaderBitMask = 1 << 7
 	MaxReflectedHeader     = (1 << 8) - 1
@@ -82,22 +88,22 @@ func (header *MessageHeader) Sequence() int32 {
 	return header.sequence
 }
 
+// cacheReplyFor populates the cached replyFor value from the ReplyFor header, defaulting to -1
+// (not a reply) when the header is absent or not the expected 4 bytes. Every path assigns
+// header.replyFor, so ReplyFor, IsReply and IsReplyingTo can dereference it unconditionally.
+// Messages read off the wire are length-checked by unmarshalHeaders; the fallback here covers
+// headers set programmatically.
 func (header *MessageHeader) cacheReplyFor() {
 	if header.replyFor == nil {
-		replyFor, found := header.Headers[ReplyForHeader]
-		if found {
-			if len(replyFor) != 4 {
-				pfxlog.Logger().Warnf("incorrect replyFor encoding. length should be 4 not %v", len(replyFor))
+		val := int32(-1)
+		if replyFor, found := header.Headers[ReplyForHeader]; found {
+			if len(replyFor) == replyForHeaderLen {
+				val = int32(binary.LittleEndian.Uint32(replyFor))
 			} else {
-				val := int32(binary.LittleEndian.Uint32(replyFor))
-				header.replyFor = &val
+				pfxlog.Logger().Warnf("incorrect replyFor encoding. length should be 4 not %v", len(replyFor))
 			}
 		}
-
-		if replyFor == nil {
-			val := int32(-1)
-			header.replyFor = &val
-		}
+		header.replyFor = &val
 	}
 }
 
@@ -510,13 +516,22 @@ func ReadV2(peer io.Reader) (*Message, error) {
 
 // unmarshalV2 converts a block of V2 wire format data into a *Message.
 func unmarshalV2(peer io.Reader, messageSectionData []byte, headersLength, bodyLength uint32) (*Message, error) {
-	dataSectionData := make([]byte, headersLength+bodyLength)
+	// summed as uint64: in uint32 the sum wraps, and a wrapped total passes every check
+	// that follows before the header slice panics. The bound keeps the total exact as an
+	// int on 32-bit platforms; it is a representability limit, not a policy limit on
+	// message size.
+	dataSectionLen := uint64(headersLength) + uint64(bodyLength)
+	if dataSectionLen > math.MaxInt32 {
+		return nil, fmt.Errorf("declared data section of %d bytes exceeds the maximum of %d", dataSectionLen, math.MaxInt32)
+	}
+
+	dataSectionData := make([]byte, dataSectionLen)
 	read, err := io.ReadFull(peer, dataSectionData)
 	if err != nil {
 		return nil, err
 	}
 
-	if read != int(headersLength+bodyLength) {
+	if uint64(read) != dataSectionLen {
 		return nil, errors.New("short read")
 	}
 
@@ -569,8 +584,14 @@ func unmarshalHeaders(headerData []byte) (map[int32][]byte, error) {
 
 		key := readInt32(headerData[i : i+4])
 		length := readUint32(headerData[i+4 : i+8])
-		if (i + 8 + int(length)) > len(headerData) {
-			return nil, fmt.Errorf("short header data (%d >= %d)", i+8+int(length), len(headerData))
+		// compared as uint64: where int is 32 bits, int(length) is negative above MaxInt32,
+		// which passes this check and then panics on the slice below
+		end := uint64(i) + 8 + uint64(length)
+		if end > uint64(len(headerData)) {
+			return nil, fmt.Errorf("short header data (%d >= %d)", end, len(headerData))
+		}
+		if key == ReplyForHeader && length != replyForHeaderLen {
+			return nil, fmt.Errorf("invalid replyFor header length (%d), must be %d", length, replyForHeaderLen)
 		}
 		data := headerData[i+8 : i+8+int(length)]
 		out[key] = data

message_test.go

@@ -17,9 +17,14 @@
 package channel
 
 import (
+	"bytes"
+	"encoding/binary"
 	"errors"
-	"github.com/stretchr/testify/assert"
+	"fmt"
 	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func Test_getRetryVersionFor(t *testing.T) {
@@ -164,3 +169,127 @@ func Test_StringToStringMapEncodeDecode(t *testing.T) {
 	req.NoError(err)
 	req.Equal((map[string]string)(nil), decoded)
 }
+
+// Test_ReplyForMalformed covers a ReplyFor header that is present but not 4 bytes wide.
+// Every getter must report the not-a-reply default for it.
+func Test_ReplyForMalformed(t *testing.T) {
+	// backing array so a zero-length header value is a non-nil empty slice, as it is
+	// when sliced out of wire data
+	backing := make([]byte, 16)
+
+	for _, length := range []int{0, 1, 3, 5, 8} {
+		t.Run(fmt.Sprintf("len-%d", length), func(t *testing.T) {
+			m := NewMessage(1, nil)
+			m.Headers[ReplyForHeader] = backing[:length]
+
+			require.NotPanics(t, func() {
+				assert.False(t, m.IsReply())
+				assert.Equal(t, int32(-1), m.ReplyFor())
+				assert.False(t, m.IsReplyingTo(1))
+			})
+		})
+	}
+}
+
+func Test_ReplyForWellFormed(t *testing.T) {
+	m := NewMessage(1, nil)
+	m.PutUint32Header(ReplyForHeader, 7)
+
+	assert.True(t, m.IsReply())
+	assert.Equal(t, int32(7), m.ReplyFor())
+	assert.True(t, m.IsReplyingTo(7))
+	assert.False(t, m.IsReplyingTo(8))
+}
+
+func Test_ReplyForAbsent(t *testing.T) {
+	m := NewMessage(1, nil)
+
+	assert.False(t, m.IsReply())
+	assert.Equal(t, int32(-1), m.ReplyFor())
+}
+
+// Test_unmarshalHeadersRejectsBadReplyFor asserts a malformed ReplyFor is rejected at
+// unmarshal, so the frame is dropped rather than silently treated as a non-reply.
+func Test_unmarshalHeadersRejectsBadReplyFor(t *testing.T) {
+	buildHeader := func(key int32, val []byte) []byte {
+		buf := make([]byte, 8+len(val))
+		binary.LittleEndian.PutUint32(buf[0:4], uint32(key))
+		binary.LittleEndian.PutUint32(buf[4:8], uint32(len(val)))
+		copy(buf[8:], val)
+		return buf
+	}
+
+	for _, length := range []int{0, 1, 3, 5, 8} {
+		t.Run(fmt.Sprintf("len-%d", length), func(t *testing.T) {
+			_, err := unmarshalHeaders(buildHeader(ReplyForHeader, make([]byte, length)))
+			require.Error(t, err)
+			assert.Contains(t, err.Error(), "invalid replyFor header length")
+		})
+	}
+
+	t.Run("valid", func(t *testing.T) {
+		headers, err := unmarshalHeaders(buildHeader(ReplyForHeader, []byte{1, 0, 0, 0}))
+		require.NoError(t, err)
+		assert.Equal(t, []byte{1, 0, 0, 0}, headers[ReplyForHeader])
+	})
+
+	t.Run("other header any length", func(t *testing.T) {
+		headers, err := unmarshalHeaders(buildHeader(TypeHeader, []byte{1, 2, 3}))
+		require.NoError(t, err)
+		assert.Equal(t, []byte{1, 2, 3}, headers[TypeHeader])
+	})
+}
+
+// buildHeaderBlock lays out one header in the on-wire format: key, declared length, value.
+// declaredLen is passed separately from the value so a test can declare a length the value
+// does not have.
+func buildHeaderBlock(key int32, declaredLen uint32, val []byte) []byte {
+	buf := make([]byte, 8+len(val))
+	binary.LittleEndian.PutUint32(buf[0:4], uint32(key))
+	binary.LittleEndian.PutUint32(buf[4:8], declaredLen)
+	copy(buf[8:], val)
+	return buf
+}
+
+// Test_unmarshalHeadersLengthOverflow covers a header whose declared length exceeds what an
+// int holds exactly on a 32-bit platform. The bounds check must reject it at the declared
+// width rather than converting first.
+//
+// NOTE: this only exercises the defect it guards under GOARCH=386. Where int is 64 bits the
+// conversion is exact and the case is reached by the ordinary short-header path.
+func Test_unmarshalHeadersLengthOverflow(t *testing.T) {
+	for _, declared := range []uint32{0x80000000, 0xFFFFFFF0, 0xFFFFFFFF} {
+		t.Run(fmt.Sprintf("declared-%#x", declared), func(t *testing.T) {
+			require.NotPanics(t, func() {
+				_, err := unmarshalHeaders(buildHeaderBlock(7, declared, nil))
+				require.Error(t, err)
+				assert.Contains(t, err.Error(), "short header data")
+			})
+		})
+	}
+}
+
+// Test_unmarshalV2LengthWrap covers declared lengths whose sum wraps at uint32 width. A
+// wrapped total is self-consistent with the checks that follow it, so it must be rejected
+// at the point the two lengths are combined.
+func Test_unmarshalV2LengthWrap(t *testing.T) {
+	messageSection := make([]byte, dataSectionV2)
+	copy(messageSection[0:magicLength], magicV2)
+
+	for _, tc := range []struct {
+		name          string
+		headersLength uint32
+		bodyLength    uint32
+	}{
+		{name: "sum wraps to zero", headersLength: 0xFFFFFFFF, bodyLength: 1},
+		{name: "sum wraps to small", headersLength: 0xFFFFFFF0, bodyLength: 0x20},
+		{name: "sum exceeds MaxInt32 without wrapping", headersLength: 0x7FFFFFFF, bodyLength: 0x7FFFFFFF},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			require.NotPanics(t, func() {
+				_, err := unmarshalV2(bytes.NewReader(nil), messageSection, tc.headersLength, tc.bodyLength)
+				require.Error(t, err)
+			})
+		})
+	}
+}