CSHARP-1501: Drivers must raise an error if response messageLength > ismaster.maxMessageSizeBytes

Author: ksadoff

src/MongoDB.Driver.Core/Core/Connections/BinaryConnection.cs

@@ -220,6 +220,16 @@ private void Dispose(bool disposing)
             }
         }
 
+        private void EnsureMessageSizeIsValid(int messageSize)
+        {
+            var maxMessageSize = _description?.MaxMessageSize ?? 48000000;
+
+            if (messageSize < 0 || messageSize > maxMessageSize)
+            {
+                throw new FormatException("The size of the message is invalid.");
+            }
+        }
+
         public void Open(CancellationToken cancellationToken)
         {
             ThrowIfDisposed();
@@ -321,6 +331,7 @@ private IByteBuffer ReceiveBuffer()
                 var messageSizeBytes = new byte[4];
                 _stream.ReadBytes(messageSizeBytes, 0, 4, _backgroundTaskCancellationToken);
                 var messageSize = BitConverter.ToInt32(messageSizeBytes, 0);
+                EnsureMessageSizeIsValid(messageSize);
                 var inputBufferChunkSource = new InputBufferChunkSource(BsonChunkPool.Default);
                 var buffer = ByteBufferFactory.Create(inputBufferChunkSource, messageSize);
                 buffer.Length = messageSize;
@@ -382,6 +393,7 @@ private async Task<IByteBuffer> ReceiveBufferAsync()
                 var messageSizeBytes = new byte[4];
                 await _stream.ReadBytesAsync(messageSizeBytes, 0, 4, _backgroundTaskCancellationToken).ConfigureAwait(false);
                 var messageSize = BitConverter.ToInt32(messageSizeBytes, 0);
+                EnsureMessageSizeIsValid(messageSize);
                 var inputBufferChunkSource = new InputBufferChunkSource(BsonChunkPool.Default);
                 var buffer = ByteBufferFactory.Create(inputBufferChunkSource, messageSize);
                 buffer.Length = messageSize;

tests/MongoDB.Driver.Core.Tests/Core/Connections/BinaryConnectionTests.cs

@@ -53,18 +53,18 @@ public BinaryConnectionTests()
 
             _endPoint = new DnsEndPoint("localhost", 27017);
             var serverId = new ServerId(new ClusterId(), _endPoint);
+            var connectionId = new ConnectionId(serverId);
+            var isMasterResult = new IsMasterResult(new BsonDocument { { "ok", 1 }, { "maxMessageSizeBytes", 48000000 } });
+            var buildInfoResult = new BuildInfoResult(new BsonDocument { { "ok", 1 }, { "version", "2.6.3" } });
+            var connectionDescription = new ConnectionDescription(connectionId, isMasterResult, buildInfoResult);
 
             _mockConnectionInitializer = new Mock<IConnectionInitializer>();
-            _mockConnectionInitializer.Setup(i => i.InitializeConnection(It.IsAny<IConnection>(), CancellationToken.None))
-                .Returns(() => new ConnectionDescription(
-                    new ConnectionId(serverId),
-                    new IsMasterResult(new BsonDocument()),
-                    new BuildInfoResult(new BsonDocument("version", "2.6.3"))));
-            _mockConnectionInitializer.Setup(i => i.InitializeConnectionAsync(It.IsAny<IConnection>(), CancellationToken.None))
-                .Returns(() => Task.FromResult(new ConnectionDescription(
-                    new ConnectionId(serverId),
-                    new IsMasterResult(new BsonDocument()),
-                    new BuildInfoResult(new BsonDocument("version", "2.6.3")))));
+            _mockConnectionInitializer
+                .Setup(i => i.InitializeConnection(It.IsAny<IConnection>(), CancellationToken.None))
+                .Returns(connectionDescription);
+            _mockConnectionInitializer
+                .Setup(i => i.InitializeConnectionAsync(It.IsAny<IConnection>(), CancellationToken.None))
+                .ReturnsAsync(connectionDescription);
 
             _subject = new BinaryConnection(
                 serverId: serverId,
@@ -212,6 +212,47 @@ public void Open_should_not_complete_the_second_call_until_the_first_is_complete
             _capturedEvents.Any().Should().BeFalse();
         }
 
+        [Theory]
+        [ParameterAttributeData]
+        public void ReceiveMessage_should_throw_a_FormatException_when_message_is_an_invalid_size(
+            [Values(-1, 48000001)]
+            int length,
+            [Values(false, true)]
+            bool async)
+        {
+            using (var stream = new BlockingMemoryStream())
+            {
+                var bytes = BitConverter.GetBytes(length);
+                stream.Write(bytes, 0, bytes.Length);
+                stream.Seek(0, SeekOrigin.Begin);
+                var encoderSelector = new ReplyMessageEncoderSelector<BsonDocument>(BsonDocumentSerializer.Instance);
+
+                Exception exception;
+                if (async)
+                {
+                    _mockStreamFactory
+                        .Setup(f => f.CreateStreamAsync(_endPoint, CancellationToken.None))
+                        .ReturnsAsync(stream);
+                    _subject.OpenAsync(CancellationToken.None).GetAwaiter().GetResult();
+                    exception = Record
+                        .Exception(() => _subject.ReceiveMessageAsync(10, encoderSelector, _messageEncoderSettings, CancellationToken.None)
+                        .GetAwaiter()
+                        .GetResult());
+                }
+                else
+                {
+                    _mockStreamFactory.Setup(f => f.CreateStream(_endPoint, CancellationToken.None))
+                        .Returns(stream);
+                    _subject.Open(CancellationToken.None);
+                    exception = Record.Exception( () =>  _subject.ReceiveMessage(10, encoderSelector, _messageEncoderSettings, CancellationToken.None));
+                }
+
+                exception.Should().BeOfType<MongoConnectionException>();
+                var e = exception.InnerException.Should().BeOfType<FormatException>().Subject;
+                e.Message.Should().Be("The size of the message is invalid.");
+            }
+        }
+
         [Theory]
         [ParameterAttributeData]
         public void ReceiveMessage_should_throw_an_ArgumentNullException_when_the_encoderSelector_is_null(