Fix decoding of message headers for messages with lengths >= 2^24.
recvMsg was interpreting buf[1] as the payload format instead of buf[0];
since compressionNone (the only thing supported) == 0, recvMsg got lucky
for message lengths under 2^24, which has buf[1] == 0.
Fix the error, and ditch the constants in recvMsg. I think they were the
cause of the bug.
Also make encode fail more clearly if someone tries to transmit a
message that exceeds 2^32 bytes.
Fixes #399.
diff --git a/rpc_util.go b/rpc_util.go
index 78c5f2e..f7d0ea5 100644
--- a/rpc_util.go
+++ b/rpc_util.go
@@ -38,6 +38,7 @@
"encoding/binary"
"fmt"
"io"
+ "math"
"math/rand"
"os"
"time"
@@ -139,42 +140,36 @@
// EOF is returned with nil msg and 0 pf if the entire stream is done. Other
// non-nil error is returned if something is wrong on reading.
func (p *parser) recvMsg() (pf payloadFormat, msg []byte, err error) {
- const (
- headerSize = 5
- formatIndex = 1
- )
-
- var hdr msgFixedHeader
- var buf [headerSize]byte
+ var buf [5]byte // see msgFixedHeader
if _, err := io.ReadFull(p.s, buf[:]); err != nil {
return 0, nil, err
}
- hdr.T = payloadFormat(buf[formatIndex])
- hdr.Length = binary.BigEndian.Uint32(buf[formatIndex:])
+ pf = payloadFormat(buf[0])
+ length := binary.BigEndian.Uint32(buf[1:])
- if hdr.Length == 0 {
- return hdr.T, nil, nil
+ if length == 0 {
+ return pf, nil, nil
}
- msg = make([]byte, int(hdr.Length))
+ msg = make([]byte, int(length))
if _, err := io.ReadFull(p.s, msg); err != nil {
if err == io.EOF {
err = io.ErrUnexpectedEOF
}
return 0, nil, err
}
- return hdr.T, msg, nil
+ return pf, msg, nil
}
// encode serializes msg and prepends the message header. If msg is nil, it
// generates the message header of 0 message length.
func encode(c Codec, msg interface{}, pf payloadFormat) ([]byte, error) {
var buf bytes.Buffer
- // Write message fixed header.
+ // Write message into the fixed header.
buf.WriteByte(uint8(pf))
var b []byte
- var length uint32
+ var length int
if msg != nil {
var err error
// TODO(zhaoq): optimize to reduce memory alloc and copying.
@@ -182,10 +177,13 @@
if err != nil {
return nil, err
}
- length = uint32(len(b))
+ length = len(b)
+ }
+ if length > math.MaxUint32 {
+ return nil, Errorf(codes.InvalidArgument, "grpc: message too large (%d bytes)", length)
}
var szHdr [4]byte
- binary.BigEndian.PutUint32(szHdr[:], length)
+ binary.BigEndian.PutUint32(szHdr[:], uint32(length))
buf.Write(szHdr[:])
buf.Write(b)
return buf.Bytes(), nil
diff --git a/rpc_util_test.go b/rpc_util_test.go
index 1a29684..2673cd0 100644
--- a/rpc_util_test.go
+++ b/rpc_util_test.go
@@ -47,6 +47,7 @@
)
func TestSimpleParsing(t *testing.T) {
+ bigMsg := bytes.Repeat([]byte{'x'}, 1<<24)
for _, test := range []struct {
// input
p []byte
@@ -60,6 +61,8 @@
{[]byte{0, 0, 0, 0, 1, 'a'}, nil, []byte{'a'}, compressionNone},
{[]byte{1, 0}, io.ErrUnexpectedEOF, nil, compressionNone},
{[]byte{0, 0, 0, 0, 10, 'a'}, io.ErrUnexpectedEOF, nil, compressionNone},
+ // Check that messages with length >= 2^24 are parsed.
+ {append([]byte{0, 1, 0, 0, 0}, bigMsg...), nil, bigMsg, compressionNone},
} {
buf := bytes.NewReader(test.p)
parser := &parser{buf}
