Enforce receive message size limits during decompression

src/Grpc.AspNetCore.Server/GrpcServiceOptions.cs

@@ -69,6 +69,7 @@ public bool MaxSendMessageSizeSpecified
     /// <summary>
     /// Gets or sets the maximum message size in bytes that can be received by the server.
     /// If the server receives a message that exceeds this limit, it throws an exception.
+    /// For compressed messages, the limit applies to both the compressed and decompressed message sizes.
     /// <para>
     /// A <c>null</c> value removes the maximum message size limit. Defaults to 4,194,304 (4 MB).
     /// </para>

src/Grpc.AspNetCore.Server/Internal/PipeExtensions.cs

@@ -23,6 +23,7 @@
 using System.Runtime.CompilerServices;
 using Grpc.Core;
 using Grpc.Net.Compression;
+using Grpc.Shared;
 using Microsoft.Extensions.Logging;
 
 namespace Grpc.AspNetCore.Server.Internal;
@@ -406,7 +407,7 @@ private static bool TryReadMessage(ref ReadOnlySequence<byte> buffer, HttpContex
             }
 
             // Performance improvement would be to decompress without converting to an intermediary byte array
-            if (!TryDecompressMessage(context.Logger, encoding, context.Options.CompressionProviders, messageBuffer, out var decompressedMessage))
+            if (!TryDecompressMessage(context.Logger, encoding, context.Options.CompressionProviders, messageBuffer, context.Options.MaxReceiveMessageSize, out var decompressedMessage))
             {
                 // https://github.com/grpc/grpc/blob/master/doc/compression.md#test-cases
                 // A message compressed by a client in a way not supported by its server MUST fail with status UNIMPLEMENTED,
@@ -439,19 +440,20 @@ private static bool TryReadMessage(ref ReadOnlySequence<byte> buffer, HttpContex
         return true;
     }
 
-    private static bool TryDecompressMessage(ILogger logger, string compressionEncoding, IReadOnlyDictionary<string, ICompressionProvider> compressionProviders, in ReadOnlySequence<byte> messageData, out ReadOnlySequence<byte> result)
+    private static bool TryDecompressMessage(ILogger logger, string compressionEncoding, IReadOnlyDictionary<string, ICompressionProvider> compressionProviders, in ReadOnlySequence<byte> messageData, int? maxReceiveMessageSize, out ReadOnlySequence<byte> result)
     {
         if (compressionProviders.TryGetValue(compressionEncoding, out var compressionProvider))
         {
             GrpcServerLog.DecompressingMessage(logger, compressionProvider.EncodingName);
 
-            var output = new MemoryStream();
             using (var compressionStream = compressionProvider.CreateDecompressionStream(new ReadOnlySequenceStream(messageData)))
             {
-                compressionStream.CopyTo(output);
+                if (!CommonGrpcProtocolHelpers.TryReadCompressedMessage(compressionStream, maxReceiveMessageSize, out result))
+                {
+                    throw new RpcException(ReceivedMessageExceedsLimitStatus);
+                }
             }
 
-            result = new ReadOnlySequence<byte>(output.GetBuffer(), 0, (int)output.Length);
             return true;
         }
 

src/Grpc.Net.Client/GrpcChannelOptions.cs

@@ -62,6 +62,7 @@ public sealed class GrpcChannelOptions
     /// <summary>
     /// Gets or sets the maximum message size in bytes that can be received by the client. If the client receives a
     /// message that exceeds this limit, it throws an exception.
+    /// For compressed messages, the limit applies to both the compressed and decompressed message sizes.
     /// <para>
     /// A <c>null</c> value removes the maximum message size limit. Defaults to 4,194,304 (4 MB).
     /// </para>

src/Grpc.Net.Client/Internal/StreamExtensions.cs

@@ -22,9 +22,7 @@
 using System.Diagnostics.CodeAnalysis;
 using System.Runtime.InteropServices;
 using Grpc.Core;
-using Grpc.Net.Compression;
 using Grpc.Shared;
-using Microsoft.Extensions.Logging;
 
 namespace Grpc.Net.Client.Internal;
 
@@ -131,7 +129,7 @@ private static Status CreateUnknownMessageEncodingMessageStatus(string unsupport
                 }
 
                 // Performance improvement would be to decompress without converting to an intermediary byte array
-                if (!TryDecompressMessage(call.Logger, grpcEncoding, call.Channel.CompressionProviders, buffer, length, out var decompressedMessage))
+                if (!TryDecompressMessage(call, grpcEncoding, buffer, length, out var decompressedMessage))
                 {
                     var supportedEncodings = new List<string>();
                     supportedEncodings.Add(GrpcProtocolConstants.IdentityGrpcEncoding);
@@ -251,19 +249,20 @@ private static async Task ReadMessageContentAsync(Stream responseStream, Memory<
         }
     }
 
-    private static bool TryDecompressMessage(ILogger logger, string compressionEncoding, Dictionary<string, ICompressionProvider> compressionProviders, byte[] messageData, int length, out ReadOnlySequence<byte> result)
+    private static bool TryDecompressMessage(GrpcCall call, string compressionEncoding, byte[] messageData, int length, out ReadOnlySequence<byte> result)
     {
-        if (compressionProviders.TryGetValue(compressionEncoding, out var compressionProvider))
+        if (call.Channel.CompressionProviders.TryGetValue(compressionEncoding, out var compressionProvider))
         {
-            GrpcCallLog.DecompressingMessage(logger, compressionProvider.EncodingName);
+            GrpcCallLog.DecompressingMessage(call.Logger, compressionProvider.EncodingName);
 
-            var output = new MemoryStream();
             using (var compressionStream = compressionProvider.CreateDecompressionStream(new MemoryStream(messageData, 0, length, writable: true, publiclyVisible: true)))
             {
-                compressionStream.CopyTo(output);
+                if (!CommonGrpcProtocolHelpers.TryReadCompressedMessage(compressionStream, call.Channel.ReceiveMaxMessageSize, out result))
+                {
+                    throw call.CreateRpcException(ReceivedMessageExceedsLimitStatus);
+                }
             }
 
-            result = new ReadOnlySequence<byte>(output.GetBuffer(), 0, (int)output.Length);
             return true;
         }
 

src/Shared/CommonGrpcProtocolHelpers.cs

@@ -16,6 +16,7 @@
 
 #endregion
 
+using System.Buffers;
 using System.Text;
 
 namespace Grpc.Shared;
@@ -25,6 +26,41 @@ internal static class CommonGrpcProtocolHelpers
     public static readonly Task<bool> TrueTask = Task.FromResult(true);
     public static readonly Task<bool> FalseTask = Task.FromResult(false);
 
+    public static bool TryReadCompressedMessage(Stream compressionStream, int? maxReceiveMessageSize, out ReadOnlySequence<byte> message)
+    {
+        using var output = new MemoryStream();
+        if (maxReceiveMessageSize == null)
+        {
+            compressionStream.CopyTo(output);
+        }
+        else
+        {
+            var buffer = ArrayPool<byte>.Shared.Rent(8192);
+            try
+            {
+                int read;
+                // Read at most one byte beyond the limit to detect oversized messages without fully decompressing them.
+                while ((read = compressionStream.Read(buffer, 0, (int)Math.Min(buffer.Length, maxReceiveMessageSize.Value - output.Length + 1))) > 0)
+                {
+                    if (output.Length + read > maxReceiveMessageSize.Value)
+                    {
+                        message = default;
+                        return false;
+                    }
+
+                    output.Write(buffer, 0, read);
+                }
+            }
+            finally
+            {
+                ArrayPool<byte>.Shared.Return(buffer);
+            }
+        }
+
+        message = new ReadOnlySequence<byte>(output.GetBuffer(), 0, (int)output.Length);
+        return true;
+    }
+
     // Timer and DateTime.UtcNow have a 14ms precision. Add a small delay when scheduling deadline
     // timer that tests if exceeded or not. This avoids rescheduling the deadline callback multiple
     // times when timer is triggered before DateTime.UtcNow reports the deadline has been exceeded.

test/FunctionalTests/Server/MaxMessageSizeTests.cs

@@ -18,6 +18,7 @@
 
 using Greet;
 using Grpc.AspNetCore.FunctionalTests.Infrastructure;
+using Grpc.AspNetCore.Server.Internal;
 using Grpc.Core;
 using Grpc.Tests.Shared;
 using NUnit.Framework;
@@ -27,6 +28,34 @@ namespace Grpc.AspNetCore.FunctionalTests.Server;
 [TestFixture]
 public class MaxMessageSizeTests : FunctionalTestBase
 {
+    [TestCase("gzip")]
+    [TestCase("deflate")]
+    public async Task ReceivedCompressedMessageExceedsSize_ThrowError(string encoding)
+    {
+        SetExpectedErrorsFilter(writeContext =>
+            writeContext.EventId.Name == "ErrorReadingMessage" &&
+            GetRpcExceptionDetail(writeContext.Exception) == "Received message exceeds the maximum configured message size.");
+
+        var invoked = false;
+        var method = Fixture.DynamicGrpc.AddUnaryMethod<HelloRequest, HelloReply>((request, context) =>
+        {
+            invoked = true;
+            return Task.FromResult(new HelloReply());
+        });
+        using var ms = new MemoryStream();
+        MessageHelpers.WriteMessage(ms, new HelloRequest { Name = new string('!', 4 * 1024 * 1024 + 1) }, encoding);
+        Assert.Less(ms.Length, 64 * 1024);
+        using var request = GrpcHttpHelper.Create(method.FullName);
+        request.Headers.Add(GrpcProtocolConstants.MessageEncodingHeader, encoding);
+        request.Content = new GrpcStreamContent(ms);
+
+        using var response = await Fixture.Client.SendAsync(request).DefaultTimeout();
+
+        response.AssertTrailerStatus(StatusCode.ResourceExhausted, "Received message exceeds the maximum configured message size.");
+        Assert.IsFalse(invoked);
+        AssertHasLogRpcConnectionError(StatusCode.ResourceExhausted, "Received message exceeds the maximum configured message size.");
+    }
+
     [Test]
     public async Task ReceivedMessageExceedsSize_ThrowError()
     {

test/Grpc.AspNetCore.Server.Tests/Grpc.AspNetCore.Server.Tests.csproj

@@ -22,6 +22,8 @@
     <Compile Include="..\Shared\SyncPoint.cs" Link="Infrastructure\SyncPoint.cs" />
     <Compile Include="..\Shared\SyncPointMemoryStream.cs" Link="Infrastructure\SyncPointMemoryStream.cs" />
     <Compile Include="..\Shared\TaskExtensions.cs" Link="Infrastructure\TaskExtensions.cs" />
+    <Compile Include="..\Shared\TestCompressionProvider.cs" Link="Infrastructure\TestCompressionProvider.cs" />
+    <Compile Include="..\Shared\TrackingDecompressionStream.cs" Link="Infrastructure\TrackingDecompressionStream.cs" />
     <Compile Include="..\Shared\TestHelpers.cs" Link="Infrastructure\TestHelpers.cs" />
     <Compile Include="..\Shared\TestHttpContextAccessor.cs" Link="Infrastructure\TestHttpContextAccessor.cs" />
     <Compile Include="..\Shared\TestRequestBodyPipeFeature.cs" Link="Infrastructure\TestRequestBodyPipeFeature.cs" />

test/Grpc.AspNetCore.Server.Tests/PipeExtensionsTestsBase.cs

@@ -17,6 +17,7 @@
 #endregion
 
 using System.Buffers;
+using System.Buffers.Binary;
 using System.IO.Pipelines;
 using System.Text;
 using Grpc.AspNetCore.Server.Internal;
@@ -133,6 +134,165 @@ public async Task ReadSingleMessageAsync_ExceedReceiveSize_ReturnData()
         Assert.AreEqual(StatusCode.ResourceExhausted, ex.StatusCode);
     }
 
+    [TestCase("gzip", false)]
+    [TestCase("gzip", true)]
+    [TestCase("deflate", false)]
+    [TestCase("deflate", true)]
+    public async Task ReadMessageAsync_CompressedMessage_EnforceReceiveSize(string encoding, bool streaming)
+    {
+        var provider = encoding == "gzip"
+            ? (ICompressionProvider)new GzipCompressionProvider(System.IO.Compression.CompressionLevel.Fastest)
+            : new DeflateCompressionProvider(System.IO.Compression.CompressionLevel.Fastest);
+        const int limit = 16 * 1024;
+
+        foreach (var size in new[] { limit - 1, limit, limit + 1 })
+        {
+            using var ms = CreateCompressedMessage(provider, new byte[size]);
+            Assert.Less(ms.Length - 5, limit);
+            var httpContext = new DefaultHttpContext();
+            httpContext.Request.Headers[GrpcProtocolConstants.MessageEncodingHeader] = encoding;
+            var context = HttpContextServerCallContextHelper.CreateServerCallContext(
+                httpContext,
+                compressionProviders: new List<ICompressionProvider> { provider },
+                maxReceiveMessageSize: limit);
+            var pipeReader = PipeReader.Create(ms);
+            var deserialized = false;
+            TestData Deserialize(DeserializationContext c)
+            {
+                deserialized = true;
+                return TestDataMarshaller.ContextualDeserializer(c);
+            }
+
+            if (size > limit)
+            {
+                var ex = await ExceptionAssert.ThrowsAsync<RpcException>(async () =>
+                {
+                    if (streaming)
+                    {
+                        await pipeReader.ReadStreamMessageAsync(context, Deserialize);
+                    }
+                    else
+                    {
+                        await pipeReader.ReadSingleMessageAsync(context, Deserialize);
+                    }
+                }).DefaultTimeout();
+
+                Assert.AreEqual(StatusCode.ResourceExhausted, ex.StatusCode);
+                Assert.AreEqual("Received message exceeds the maximum configured message size.", ex.Status.Detail);
+                Assert.IsFalse(deserialized);
+            }
+            else
+            {
+                var message = streaming
+                    ? await pipeReader.ReadStreamMessageAsync(context, Deserialize)
+                    : await pipeReader.ReadSingleMessageAsync(context, Deserialize);
+
+                Assert.IsTrue(deserialized);
+                Assert.AreEqual(size, message!.Span.Length);
+            }
+        }
+    }
+
+    [TestCase(false, null)]
+    [TestCase(true, null)]
+    [TestCase(false, int.MaxValue)]
+    [TestCase(true, int.MaxValue)]
+    public async Task ReadMessageAsync_CompressedMessage_LargeOrUnlimitedReceiveSize_ReturnData(bool streaming, int? limit)
+    {
+        var provider = new GzipCompressionProvider(System.IO.Compression.CompressionLevel.Fastest);
+        var content = new byte[32 * 1024];
+        using var ms = CreateCompressedMessage(provider, content);
+        var httpContext = new DefaultHttpContext();
+        httpContext.Request.Headers[GrpcProtocolConstants.MessageEncodingHeader] = "gzip";
+        var context = HttpContextServerCallContextHelper.CreateServerCallContext(
+            httpContext,
+            compressionProviders: new List<ICompressionProvider> { provider },
+            maxReceiveMessageSize: limit);
+        var pipeReader = PipeReader.Create(ms);
+
+        var message = streaming
+            ? await pipeReader.ReadStreamMessageAsync(context, TestDataMarshaller.ContextualDeserializer)
+            : await pipeReader.ReadSingleMessageAsync(context, TestDataMarshaller.ContextualDeserializer);
+
+        CollectionAssert.AreEqual(content, message!.Span.ToArray());
+    }
+
+    [TestCase(false, 0)]
+    [TestCase(true, 0)]
+    [TestCase(false, 16 * 1024)]
+    [TestCase(true, 16 * 1024)]
+    public async Task ReadMessageAsync_CompressedMessage_StopsDecompressingAtReceiveSize(bool streaming, int limit)
+    {
+        using var decompressionStream = new TrackingDecompressionStream(new byte[128 * 1024]);
+        var provider = new TestCompressionProvider(decompressionStream);
+        var httpContext = new DefaultHttpContext();
+        httpContext.Request.Headers[GrpcProtocolConstants.MessageEncodingHeader] = "test";
+        var context = HttpContextServerCallContextHelper.CreateServerCallContext(
+            httpContext,
+            compressionProviders: new List<ICompressionProvider> { provider },
+            maxReceiveMessageSize: limit);
+        using var ms = new MemoryStream(new byte[] { 1, 0, 0, 0, 0 });
+        var pipeReader = PipeReader.Create(ms);
+        TestData Deserialize(DeserializationContext c) => throw new InvalidOperationException("Deserializer must not be called.");
+
+        var ex = await ExceptionAssert.ThrowsAsync<RpcException>(async () =>
+        {
+            if (streaming)
+            {
+                await pipeReader.ReadStreamMessageAsync(context, Deserialize);
+            }
+            else
+            {
+                await pipeReader.ReadSingleMessageAsync(context, Deserialize);
+            }
+        }).DefaultTimeout();
+
+        Assert.AreEqual(StatusCode.ResourceExhausted, ex.StatusCode);
+        Assert.AreEqual(limit + 1, decompressionStream.BytesRead);
+        Assert.IsTrue(decompressionStream.Disposed);
+    }
+
+    [TestCase(false)]
+    [TestCase(true)]
+    public async Task ReadMessageAsync_CompressedFrameExceedsReceiveSize_ThrowBeforeDecompressing(bool streaming)
+    {
+        var provider = new TestCompressionProvider(null);
+        var context = HttpContextServerCallContextHelper.CreateServerCallContext(
+            compressionProviders: new List<ICompressionProvider> { provider },
+            maxReceiveMessageSize: 1);
+        using var ms = new MemoryStream(new byte[] { 1, 0, 0, 0, 2, 0, 0 });
+        var pipeReader = PipeReader.Create(ms);
+
+        var ex = await ExceptionAssert.ThrowsAsync<RpcException>(async () =>
+        {
+            if (streaming)
+            {
+                await pipeReader.ReadStreamMessageAsync(context, TestDataMarshaller.ContextualDeserializer);
+            }
+            else
+            {
+                await pipeReader.ReadSingleMessageAsync(context, TestDataMarshaller.ContextualDeserializer);
+            }
+        }).DefaultTimeout();
+
+        Assert.AreEqual(StatusCode.ResourceExhausted, ex.StatusCode);
+        Assert.IsFalse(provider.DecompressionStreamCreated);
+    }
+
+    private static MemoryStream CreateCompressedMessage(ICompressionProvider provider, byte[] content)
+    {
+        var ms = new MemoryStream();
+        ms.Write(new byte[] { 1, 0, 0, 0, 0 });
+        using (var compressionStream = provider.CreateCompressionStream(ms, null))
+        {
+            compressionStream.Write(content);
+        }
+
+        BinaryPrimitives.WriteInt32BigEndian(ms.GetBuffer().AsSpan(1), (int)ms.Length - 5);
+        ms.Position = 0;
+        return ms;
+    }
+
     [Test]
     public async Task ReadSingleMessageAsync_LongMessage_ReturnData()
     {

test/Grpc.Net.Client.Tests/Grpc.Net.Client.Tests.csproj

@@ -26,6 +26,8 @@
     <Compile Include="..\Shared\ResponseUtils.cs" Link="Infrastructure\ResponseUtils.cs" />
     <Compile Include="..\Shared\SyncPoint.cs" Link="Infrastructure\SyncPoint.cs" />
     <Compile Include="..\Shared\SyncPointMemoryStream.cs" Link="Infrastructure\SyncPointMemoryStream.cs" />
+    <Compile Include="..\Shared\TestCompressionProvider.cs" Link="Infrastructure\TestCompressionProvider.cs" />
+    <Compile Include="..\Shared\TrackingDecompressionStream.cs" Link="Infrastructure\TrackingDecompressionStream.cs" />
     <Compile Include="..\Shared\TestHttpMessageHandler.cs" Link="Infrastructure\TestHttpMessageHandler.cs" />
     <Compile Include="..\Shared\TestHelpers.cs" Link="Infrastructure\TestHelpers.cs" />
     <Compile Include="..\Shared\NUnitLogger.cs" Link="Infrastructure\NUnitLogger.cs" />