GODRIVER-150 Raise error for read message length too large (#545)

Author: benjirewis

x/mongo/driver/topology/connection.go

@@ -30,6 +30,11 @@ import (
 
 var globalConnectionID uint64 = 1
 
+var (
+	defaultMaxMessageSize uint32 = 48000000
+	errResponseTooLarge   error  = errors.New("length of read message too large")
+)
+
 func nextConnectionID() uint64 { return atomic.AddUint64(&globalConnectionID, 1) }
 
 type connection struct {
@@ -406,6 +411,16 @@ func (c *connection) read(ctx context.Context, dst []byte) (bytesRead []byte, er
 	// read the length as an int32
 	size := (int32(sizeBuf[0])) | (int32(sizeBuf[1]) << 8) | (int32(sizeBuf[2]) << 16) | (int32(sizeBuf[3]) << 24)
 
+	// In the case of an isMaster response where MaxMessageSize has not yet been set, use the hard-coded
+	// defaultMaxMessageSize instead.
+	maxMessageSize := c.desc.MaxMessageSize
+	if maxMessageSize == 0 {
+		maxMessageSize = defaultMaxMessageSize
+	}
+	if uint32(size) > maxMessageSize {
+		return nil, errResponseTooLarge.Error(), errResponseTooLarge
+	}
+
 	if int(size) > cap(dst) {
 		// Since we can't grow this slice without allocating, just allocate an entirely new slice.
 		dst = make([]byte, 0, size)

x/mongo/driver/topology/connection_test.go

@@ -485,6 +485,41 @@ func TestConnection(t *testing.T) {
 					}
 					listener.assertMethodsCalled(t, 1, 1)
 				})
+				t.Run("message too large errors", func(t *testing.T) {
+					testCases := []struct {
+						name   string
+						buffer []byte
+						desc   description.Server
+					}{
+						{
+							"message too large errors with small max message size",
+							[]byte{0x0A, 0x00, 0x00, 0x00}, // defines a message size of 10 in hex with the first four bytes.
+							description.Server{MaxMessageSize: 9},
+						},
+						{
+							"message too large errors with default max message size",
+							[]byte{0x01, 0x6C, 0xDC, 0x02}, // defines a message size of 48000001 in hex with the first four bytes.
+							description.Server{},
+						},
+					}
+					for _, tc := range testCases {
+						t.Run(tc.name, func(t *testing.T) {
+							err := errors.New("length of read message too large")
+							tnc := &testNetConn{buf: make([]byte, len(tc.buffer))}
+							copy(tnc.buf, tc.buffer)
+							conn := &connection{id: "foobar", nc: tnc, connected: connected, desc: tc.desc}
+							listener := newTestCancellationListener(false)
+							conn.cancellationListener = listener
+
+							want := ConnectionError{ConnectionID: "foobar", Wrapped: err, message: err.Error()}
+							_, got := conn.readWireMessage(context.Background(), nil)
+							if !cmp.Equal(got, want, cmp.Comparer(compareErrors)) {
+								t.Errorf("errors do not match. got %v; want %v", got, want)
+							}
+							listener.assertMethodsCalled(t, 1, 1)
+						})
+					}
+				})
 				t.Run("success", func(t *testing.T) {
 					want := []byte{0x0A, 0x00, 0x00, 0x00, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0A}
 					tnc := &testNetConn{buf: make([]byte, len(want))}