Fix remote OOM via delayed deserialization for unary/server-streaming calls

Protobuf unknown-field or repeated field amplification can lead to remote OOM if an attacker
sends a unary request but holds the stream open without half-closing.
This change delays the deserialization of incoming messages for calls where
the client sends at most one message (Unary and Server Streaming) until
the client actually half-closes the stream (sends END_STREAM).

If the call is cancelled before half-close, the buffered raw message is
discarded without being deserialized, preventing the memory explosion.

Author: kannanjgithub

core/src/main/java/io/grpc/internal/ServerCallImpl.java

@@ -34,6 +34,7 @@
 import io.grpc.CompressorRegistry;
 import io.grpc.Context;
 import io.grpc.DecompressorRegistry;
+import io.grpc.Detachable;
 import io.grpc.InternalDecompressorRegistry;
 import io.grpc.InternalStatus;
 import io.grpc.Metadata;
@@ -45,6 +46,9 @@
 import io.perfmark.PerfMark;
 import io.perfmark.Tag;
 import io.perfmark.TaskCloseable;
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
 import java.io.InputStream;
 import java.util.logging.Level;
 import java.util.logging.Logger;
@@ -288,6 +292,7 @@ static final class ServerStreamListenerImpl<ReqT> implements ServerStreamListene
     private final ServerCallImpl<ReqT, ?> call;
     private final ServerCall.Listener<ReqT> listener;
     private final Context.CancellableContext context;
+    private InputStream delayedMessage;
 
     public ServerStreamListenerImpl(
         ServerCallImpl<ReqT, ?> call, ServerCall.Listener<ReqT> listener,
@@ -320,6 +325,20 @@ public void messagesAvailable(MessageProducer producer) {
       }
     }
 
+    private static InputStream bufferMessage(InputStream is) throws IOException {
+      if (is instanceof Detachable) {
+        return ((Detachable) is).detach();
+      }
+      // Fallback: copy to byte array
+      ByteArrayOutputStream baos = new ByteArrayOutputStream();
+      byte[] buffer = new byte[4096];
+      int bytesRead;
+      while ((bytesRead = is.read(buffer)) != -1) {
+        baos.write(buffer, 0, bytesRead);
+      }
+      return new ByteArrayInputStream(baos.toByteArray());
+    }
+
     @SuppressWarnings("Finally") // The code avoids suppressing the exception thrown from try
     private void messagesAvailableInternal(final MessageProducer producer) {
       if (call.cancelled) {
@@ -330,13 +349,32 @@ private void messagesAvailableInternal(final MessageProducer producer) {
       InputStream message;
       try {
         while ((message = producer.next()) != null) {
-          try {
-            listener.onMessage(call.method.parseRequest(message));
-          } catch (Throwable t) {
-            GrpcUtil.closeQuietly(message);
-            throw t;
+          if (call.method.getType().clientSendsOneMessage()) {
+            if (delayedMessage != null) {
+              GrpcUtil.closeQuietly(message);
+              call.close(
+                  Status.INTERNAL.withDescription("Too many requests"),
+                  new Metadata());
+              GrpcUtil.closeQuietly(delayedMessage);
+              delayedMessage = null;
+              return;
+            }
+            try {
+              delayedMessage = bufferMessage(message);
+            } catch (Throwable t) {
+              GrpcUtil.closeQuietly(message);
+              throw t;
+            }
+            message.close();
+          } else {
+            try {
+              listener.onMessage(call.method.parseRequest(message));
+            } catch (Throwable t) {
+              GrpcUtil.closeQuietly(message);
+              throw t;
+            }
+            message.close();
           }
-          message.close();
         }
       } catch (Throwable t) {
         GrpcUtil.closeQuietly(producer);
@@ -353,6 +391,23 @@ public void halfClosed() {
           return;
         }
 
+        if (delayedMessage != null) {
+          InputStream message = delayedMessage;
+          delayedMessage = null;
+          try {
+            listener.onMessage(call.method.parseRequest(message));
+          } catch (Throwable t) {
+            GrpcUtil.closeQuietly(message);
+            Throwables.throwIfUnchecked(t);
+            throw new RuntimeException(t);
+          }
+          try {
+            message.close();
+          } catch (IOException e) {
+            throw new RuntimeException(e);
+          }
+        }
+
         listener.onHalfClose();
       }
     }
@@ -366,6 +421,10 @@ public void closed(Status status) {
     }
 
     private void closedInternal(Status status) {
+      if (delayedMessage != null) {
+        GrpcUtil.closeQuietly(delayedMessage);
+        delayedMessage = null;
+      }
       Throwable cancelCause = null;
       try {
         if (status.isOk()) {

core/src/test/java/io/grpc/internal/ServerCallImplTest.java

@@ -84,7 +84,23 @@ public class ServerCallImplTest {
 
   private static final MethodDescriptor<Long, Long> CLIENT_STREAMING_METHOD =
       MethodDescriptor.<Long, Long>newBuilder()
-          .setType(MethodType.UNARY)
+          .setType(MethodType.CLIENT_STREAMING)
+          .setFullMethodName("service/method")
+          .setRequestMarshaller(new LongMarshaller())
+          .setResponseMarshaller(new LongMarshaller())
+          .build();
+
+  private static final MethodDescriptor<Long, Long> BIDI_STREAMING_METHOD =
+      MethodDescriptor.<Long, Long>newBuilder()
+          .setType(MethodType.BIDI_STREAMING)
+          .setFullMethodName("service/method")
+          .setRequestMarshaller(new LongMarshaller())
+          .setResponseMarshaller(new LongMarshaller())
+          .build();
+
+  private static final MethodDescriptor<Long, Long> SERVER_STREAMING_METHOD =
+      MethodDescriptor.<Long, Long>newBuilder()
+          .setType(MethodType.SERVER_STREAMING)
           .setFullMethodName("service/method")
           .setRequestMarshaller(new LongMarshaller())
           .setResponseMarshaller(new LongMarshaller())
@@ -456,40 +472,97 @@ public void streamListener_onReady_onlyOnce() {
   }
 
   @Test
-  public void streamListener_messageRead() {
+  public void streamListener_messageRead_unary_delayed() {
     ServerStreamListenerImpl<Long> streamListener =
         new ServerCallImpl.ServerStreamListenerImpl<>(call, callListener, context);
     streamListener.messagesAvailable(new SingleMessageProducer(UNARY_METHOD.streamRequest(1234L)));
 
+    // Message should not be delivered yet
+    verify(callListener, never()).onMessage(any(Long.class));
+
+    streamListener.halfClosed();
+
+    // Now it should be delivered
     verify(callListener).onMessage(1234L);
+    verify(callListener).onHalfClose();
   }
 
   @Test
-  public void streamListener_messageRead_onlyOnce() {
+  public void streamListener_messageRead_serverStreaming_delayed() {
+    call = new ServerCallImpl<>(stream, SERVER_STREAMING_METHOD, requestHeaders, context,
+        DecompressorRegistry.getDefaultInstance(), CompressorRegistry.getDefaultInstance(),
+        serverCallTracer, PerfMark.createTag());
+    ServerStreamListenerImpl<Long> streamListener =
+        new ServerCallImpl.ServerStreamListenerImpl<>(call, callListener, context);
+    streamListener.messagesAvailable(new SingleMessageProducer(SERVER_STREAMING_METHOD.streamRequest(1234L)));
+
+    // Message should not be delivered yet
+    verify(callListener, never()).onMessage(any(Long.class));
+
+    streamListener.halfClosed();
+
+    // Now it should be delivered
+    verify(callListener).onMessage(1234L);
+    verify(callListener).onHalfClose();
+  }
+
+  @Test
+  public void streamListener_messageRead_bidi_notDelayed() {
+    call = new ServerCallImpl<>(stream, BIDI_STREAMING_METHOD, requestHeaders, context,
+        DecompressorRegistry.getDefaultInstance(), CompressorRegistry.getDefaultInstance(),
+        serverCallTracer, PerfMark.createTag());
+    ServerStreamListenerImpl<Long> streamListener =
+        new ServerCallImpl.ServerStreamListenerImpl<>(call, callListener, context);
+    streamListener.messagesAvailable(new SingleMessageProducer(BIDI_STREAMING_METHOD.streamRequest(1234L)));
+
+    // Message should be delivered immediately
+    verify(callListener).onMessage(1234L);
+    verify(callListener, never()).onHalfClose();
+  }
+
+  @Test
+  public void streamListener_messageRead_unary_tooManyRequests() {
     ServerStreamListenerImpl<Long> streamListener =
         new ServerCallImpl.ServerStreamListenerImpl<>(call, callListener, context);
     streamListener.messagesAvailable(new SingleMessageProducer(UNARY_METHOD.streamRequest(1234L)));
-    // canceling the call should short circuit future halfClosed() calls.
-    streamListener.closed(Status.CANCELLED);
 
+    // Sending second message should fail
+    streamListener.messagesAvailable(new SingleMessageProducer(UNARY_METHOD.streamRequest(5678L)));
+
+    verify(stream).close(any(Status.class), any(Metadata.class));
+    verify(callListener, never()).onMessage(any(Long.class));
+  }
+
+  @Test
+  public void streamListener_messageRead_onlyOnce_unary() {
+    ServerStreamListenerImpl<Long> streamListener =
+        new ServerCallImpl.ServerStreamListenerImpl<>(call, callListener, context);
     streamListener.messagesAvailable(new SingleMessageProducer(UNARY_METHOD.streamRequest(1234L)));
+    
+    // canceling the call should clean up and prevent delivery
+    streamListener.closed(Status.CANCELLED);
 
-    verify(callListener).onMessage(1234L);
+    streamListener.halfClosed();
+
+    verify(callListener, never()).onMessage(any(Long.class));
   }
 
   @Test
-  public void streamListener_unexpectedRuntimeException() {
+  public void streamListener_unexpectedRuntimeException_unary() {
     ServerStreamListenerImpl<Long> streamListener =
         new ServerCallImpl.ServerStreamListenerImpl<>(call, callListener, context);
     doThrow(new RuntimeException("unexpected exception"))
         .when(callListener)
         .onMessage(any(Long.class));
 
-    InputStream inputStream = UNARY_METHOD.streamRequest(1234L);
+    streamListener.messagesAvailable(new SingleMessageProducer(UNARY_METHOD.streamRequest(1234L)));
+
+    // Exception should not be thrown yet because deserialization/delivery is delayed
+    verify(callListener, never()).onMessage(any(Long.class));
 
-    SingleMessageProducer producer = new SingleMessageProducer(inputStream);
+    // It should be thrown during halfClosed
     RuntimeException e = assertThrows(RuntimeException.class,
-        () -> streamListener.messagesAvailable(producer));
+        () -> streamListener.halfClosed());
     assertThat(e).hasMessageThat().isEqualTo("unexpected exception");
   }