From 09d8c90ef4af9866def10f0f0f020005519bdcf1 Mon Sep 17 00:00:00 2001 From: Kannan J Date: Tue, 25 Aug 2026 11:09:45 +0000 Subject: [PATCH 1/4] 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. --- .../java/io/grpc/internal/ServerCallImpl.java | 71 +++++- .../io/grpc/internal/ServerCallImplTest.java | 235 +++++++++++++++++- 2 files changed, 290 insertions(+), 16 deletions(-) diff --git a/core/src/main/java/io/grpc/internal/ServerCallImpl.java b/core/src/main/java/io/grpc/internal/ServerCallImpl.java index e224384ce8f..bb3b6c5402b 100644 --- a/core/src/main/java/io/grpc/internal/ServerCallImpl.java +++ b/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 implements ServerStreamListene private final ServerCallImpl call; private final ServerCall.Listener listener; private final Context.CancellableContext context; + private InputStream delayedMessage; public ServerStreamListenerImpl( ServerCallImpl call, ServerCall.Listener 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()) { diff --git a/core/src/test/java/io/grpc/internal/ServerCallImplTest.java b/core/src/test/java/io/grpc/internal/ServerCallImplTest.java index 7394c83eab2..b3a91377994 100644 --- a/core/src/test/java/io/grpc/internal/ServerCallImplTest.java +++ b/core/src/test/java/io/grpc/internal/ServerCallImplTest.java @@ -40,6 +40,7 @@ import io.grpc.CompressorRegistry; import io.grpc.Context; import io.grpc.DecompressorRegistry; +import io.grpc.Detachable; import io.grpc.InternalChannelz.ServerStats; import io.grpc.Metadata; import io.grpc.MethodDescriptor; @@ -51,6 +52,7 @@ import io.grpc.internal.ServerCallImpl.ServerStreamListenerImpl; import io.perfmark.PerfMark; import java.io.ByteArrayInputStream; +import java.io.IOException; import java.io.InputStream; import java.io.InputStreamReader; import org.junit.Before; @@ -84,7 +86,23 @@ public class ServerCallImplTest { private static final MethodDescriptor CLIENT_STREAMING_METHOD = MethodDescriptor.newBuilder() - .setType(MethodType.UNARY) + .setType(MethodType.CLIENT_STREAMING) + .setFullMethodName("service/method") + .setRequestMarshaller(new LongMarshaller()) + .setResponseMarshaller(new LongMarshaller()) + .build(); + + private static final MethodDescriptor BIDI_STREAMING_METHOD = + MethodDescriptor.newBuilder() + .setType(MethodType.BIDI_STREAMING) + .setFullMethodName("service/method") + .setRequestMarshaller(new LongMarshaller()) + .setResponseMarshaller(new LongMarshaller()) + .build(); + + private static final MethodDescriptor SERVER_STREAMING_METHOD = + MethodDescriptor.newBuilder() + .setType(MethodType.SERVER_STREAMING) .setFullMethodName("service/method") .setRequestMarshaller(new LongMarshaller()) .setResponseMarshaller(new LongMarshaller()) @@ -456,41 +474,169 @@ public void streamListener_onReady_onlyOnce() { } @Test - public void streamListener_messageRead() { + public void streamListener_messageRead_unary_delayed() { ServerStreamListenerImpl 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_unary_detachable() { + ServerStreamListenerImpl streamListener = + new ServerCallImpl.ServerStreamListenerImpl<>(call, callListener, context); + + InputStream delegate = UNARY_METHOD.streamRequest(1234L); + FakeDetachableInputStream detachableStream = new FakeDetachableInputStream(delegate); + + streamListener.messagesAvailable(new SingleMessageProducer(detachableStream)); + + // It should have been detached immediately + assertTrue(detachableStream.detached); + 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_unary_bufferMessageException() { + ServerStreamListenerImpl streamListener = + new ServerCallImpl.ServerStreamListenerImpl<>(call, callListener, context); + + InputStream delegate = UNARY_METHOD.streamRequest(1234L); + FakeDetachableInputStream detachableStream = new FakeDetachableInputStream(delegate, true); + + RuntimeException e = assertThrows(RuntimeException.class, + () -> streamListener.messagesAvailable(new SingleMessageProducer(detachableStream))); + assertThat(e).hasMessageThat().isEqualTo("detach failed"); + + // The stream should have been closed in the catch block + assertTrue(detachableStream.closed); + verify(callListener, never()).onMessage(any(Long.class)); + } + + @Test + public void streamListener_messageRead_serverStreaming_delayed() { + call = new ServerCallImpl<>(stream, SERVER_STREAMING_METHOD, requestHeaders, context, + DecompressorRegistry.getDefaultInstance(), CompressorRegistry.getDefaultInstance(), + serverCallTracer, PerfMark.createTag()); + ServerStreamListenerImpl 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 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 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 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 streamListener = new ServerCallImpl.ServerStreamListenerImpl<>(call, callListener, context); doThrow(new RuntimeException("unexpected exception")) .when(callListener) .onMessage(any(Long.class)); - InputStream inputStream = UNARY_METHOD.streamRequest(1234L); + InputStream delegate = UNARY_METHOD.streamRequest(1234L); + FakeDetachableInputStream detachableStream = new FakeDetachableInputStream(delegate); + + streamListener.messagesAvailable(new SingleMessageProducer(detachableStream)); - SingleMessageProducer producer = new SingleMessageProducer(inputStream); + // Exception should not be thrown yet because deserialization/delivery is delayed + verify(callListener, never()).onMessage(any(Long.class)); + + // It should be thrown during halfClosed RuntimeException e = assertThrows(RuntimeException.class, - () -> streamListener.messagesAvailable(producer)); + () -> streamListener.halfClosed()); assertThat(e).hasMessageThat().isEqualTo("unexpected exception"); + + // The detached stream should have been closed in the catch block + assertTrue(detachableStream.detachedStream.closed); + } + + @Test + public void streamListener_halfClosed_closeException() { + ServerStreamListenerImpl streamListener = + new ServerCallImpl.ServerStreamListenerImpl<>(call, callListener, context); + + InputStream delegate = UNARY_METHOD.streamRequest(1234L); + FakeDetachableInputStream detachableStream = + new FakeDetachableInputStream(delegate, false, true); + + streamListener.messagesAvailable(new SingleMessageProducer(detachableStream)); + + // Message should not be delivered yet + verify(callListener, never()).onMessage(any(Long.class)); + + // halfClosed should throw RuntimeException wrapping IOException + RuntimeException e = assertThrows(RuntimeException.class, + () -> streamListener.halfClosed()); + assertThat(e).hasCauseThat().isInstanceOf(IOException.class); + assertThat(e.getCause()).hasMessageThat().isEqualTo("close failed"); + + // The message was delivered before close failed + verify(callListener).onMessage(1234L); + assertTrue(detachableStream.detachedStream.closed); } private static class LongMarshaller implements Marshaller { @@ -508,4 +654,73 @@ public Long parse(InputStream stream) { } } } + + private static class FakeCloseTrackerInputStream extends InputStream { + boolean closed = false; + private final InputStream delegate; + private final boolean throwOnClose; + + FakeCloseTrackerInputStream(InputStream delegate, boolean throwOnClose) { + this.delegate = delegate; + this.throwOnClose = throwOnClose; + } + + @Override + public int read() throws IOException { + return delegate.read(); + } + + @Override + public void close() throws IOException { + closed = true; + if (throwOnClose) { + throw new IOException("close failed"); + } + delegate.close(); + } + } + + private static class FakeDetachableInputStream extends InputStream implements Detachable { + boolean detached = false; + boolean closed = false; + private final InputStream delegate; + private final boolean throwOnDetach; + private final boolean throwOnClose; + FakeCloseTrackerInputStream detachedStream; + + FakeDetachableInputStream(InputStream delegate) { + this(delegate, false, false); + } + + FakeDetachableInputStream(InputStream delegate, boolean throwOnDetach) { + this(delegate, throwOnDetach, false); + } + + FakeDetachableInputStream(InputStream delegate, boolean throwOnDetach, boolean throwOnClose) { + this.delegate = delegate; + this.throwOnDetach = throwOnDetach; + this.throwOnClose = throwOnClose; + } + + @Override + public InputStream detach() { + if (throwOnDetach) { + throw new RuntimeException("detach failed"); + } + detached = true; + detachedStream = new FakeCloseTrackerInputStream(delegate, throwOnClose); + return detachedStream; + } + + @Override + public int read() throws IOException { + return delegate.read(); + } + + @Override + public void close() throws IOException { + closed = true; + delegate.close(); + } + } } From 553b33a65ac5abb529f02542ce8a3ff79006b262 Mon Sep 17 00:00:00 2001 From: Kannan J Date: Wed, 26 Aug 2026 06:40:17 +0000 Subject: [PATCH 2/4] Address review comments. - Use GrpcUtil.closeQuietly to close the buffered message in halfClosed() to prevent unnecessary exception propagation if close fails after successful message delivery. - Use stream.cancel instead of call.close when detecting too many requests for unary calls. This ensures the transport is notified to abort the stream (sending RST_STREAM) and immediately releases resources, preventing leaks from clients that withhold END_STREAM. --- .../main/java/io/grpc/internal/ServerCallImpl.java | 10 ++-------- .../java/io/grpc/internal/ServerCallImplTest.java | 12 +++++------- 2 files changed, 7 insertions(+), 15 deletions(-) diff --git a/core/src/main/java/io/grpc/internal/ServerCallImpl.java b/core/src/main/java/io/grpc/internal/ServerCallImpl.java index bb3b6c5402b..fb4c966f439 100644 --- a/core/src/main/java/io/grpc/internal/ServerCallImpl.java +++ b/core/src/main/java/io/grpc/internal/ServerCallImpl.java @@ -352,9 +352,7 @@ private void messagesAvailableInternal(final MessageProducer producer) { if (call.method.getType().clientSendsOneMessage()) { if (delayedMessage != null) { GrpcUtil.closeQuietly(message); - call.close( - Status.INTERNAL.withDescription("Too many requests"), - new Metadata()); + call.stream.cancel(Status.INTERNAL.withDescription("Too many requests")); GrpcUtil.closeQuietly(delayedMessage); delayedMessage = null; return; @@ -401,11 +399,7 @@ public void halfClosed() { Throwables.throwIfUnchecked(t); throw new RuntimeException(t); } - try { - message.close(); - } catch (IOException e) { - throw new RuntimeException(e); - } + GrpcUtil.closeQuietly(message); } listener.onHalfClose(); diff --git a/core/src/test/java/io/grpc/internal/ServerCallImplTest.java b/core/src/test/java/io/grpc/internal/ServerCallImplTest.java index b3a91377994..edc11ac804a 100644 --- a/core/src/test/java/io/grpc/internal/ServerCallImplTest.java +++ b/core/src/test/java/io/grpc/internal/ServerCallImplTest.java @@ -571,7 +571,7 @@ public void streamListener_messageRead_unary_tooManyRequests() { // Sending second message should fail streamListener.messagesAvailable(new SingleMessageProducer(UNARY_METHOD.streamRequest(5678L))); - verify(stream).close(any(Status.class), any(Metadata.class)); + verify(stream).cancel(any(Status.class)); verify(callListener, never()).onMessage(any(Long.class)); } @@ -628,14 +628,12 @@ public void streamListener_halfClosed_closeException() { // Message should not be delivered yet verify(callListener, never()).onMessage(any(Long.class)); - // halfClosed should throw RuntimeException wrapping IOException - RuntimeException e = assertThrows(RuntimeException.class, - () -> streamListener.halfClosed()); - assertThat(e).hasCauseThat().isInstanceOf(IOException.class); - assertThat(e.getCause()).hasMessageThat().isEqualTo("close failed"); + // halfClosed should not throw because we use closeQuietly + streamListener.halfClosed(); - // The message was delivered before close failed + // The message was delivered and halfClosed completed verify(callListener).onMessage(1234L); + verify(callListener).onHalfClose(); assertTrue(detachableStream.detachedStream.closed); } From 50197f1877a7a7cd4845e0f2055e8d1f2aa74f2d Mon Sep 17 00:00:00 2001 From: Kannan J Date: Thu, 27 Aug 2026 08:07:16 +0000 Subject: [PATCH 3/4] Synchronously mark the call as cancelled by calling `closedInternal` to avoid races with more messages or halfClose after the error scenario was observed. --- core/src/main/java/io/grpc/internal/ServerCallImpl.java | 1 + core/src/test/java/io/grpc/internal/ServerCallImplTest.java | 2 ++ 2 files changed, 3 insertions(+) diff --git a/core/src/main/java/io/grpc/internal/ServerCallImpl.java b/core/src/main/java/io/grpc/internal/ServerCallImpl.java index fb4c966f439..fafd74308c9 100644 --- a/core/src/main/java/io/grpc/internal/ServerCallImpl.java +++ b/core/src/main/java/io/grpc/internal/ServerCallImpl.java @@ -355,6 +355,7 @@ private void messagesAvailableInternal(final MessageProducer producer) { call.stream.cancel(Status.INTERNAL.withDescription("Too many requests")); GrpcUtil.closeQuietly(delayedMessage); delayedMessage = null; + closedInternal(Status.INTERNAL.withDescription("Too many requests")); return; } try { diff --git a/core/src/test/java/io/grpc/internal/ServerCallImplTest.java b/core/src/test/java/io/grpc/internal/ServerCallImplTest.java index edc11ac804a..53ccf5bfe6a 100644 --- a/core/src/test/java/io/grpc/internal/ServerCallImplTest.java +++ b/core/src/test/java/io/grpc/internal/ServerCallImplTest.java @@ -573,6 +573,8 @@ public void streamListener_messageRead_unary_tooManyRequests() { verify(stream).cancel(any(Status.class)); verify(callListener, never()).onMessage(any(Long.class)); + verify(callListener).onCancel(); + assertTrue(context.isCancelled()); } @Test From f237d2fc5372fca0d19447f006f0d071870acc66 Mon Sep 17 00:00:00 2001 From: Kannan J Date: Mon, 7 Sep 2026 05:22:10 +0000 Subject: [PATCH 4/4] Address review comments. --- .../io/grpc/internal/MessageDeframer.java | 104 ++++++++++++- .../java/io/grpc/internal/ServerCallImpl.java | 30 +--- .../io/grpc/internal/MessageDeframerTest.java | 147 ++++++++++++++++++ .../io/grpc/internal/ServerCallImplTest.java | 143 +++++------------ 4 files changed, 288 insertions(+), 136 deletions(-) diff --git a/core/src/main/java/io/grpc/internal/MessageDeframer.java b/core/src/main/java/io/grpc/internal/MessageDeframer.java index f388c006e97..dbb0d88b87a 100644 --- a/core/src/main/java/io/grpc/internal/MessageDeframer.java +++ b/core/src/main/java/io/grpc/internal/MessageDeframer.java @@ -437,14 +437,102 @@ private InputStream getCompressedBody() { .asRuntimeException(); } - try { - // Enforce the maxMessageSize limit on the returned stream. - InputStream unlimitedStream = - decompressor.decompress(ReadableBuffers.openStream(nextFrame, true)); - return new SizeEnforcingInputStream( - unlimitedStream, maxInboundMessageSize, statsTraceCtx); - } catch (IOException e) { - throw new RuntimeException(e); + return new LazyDecompressingInputStream( + ReadableBuffers.openStream(nextFrame, true), + maxInboundMessageSize, + statsTraceCtx, + decompressor); + } + + /** + * An {@link InputStream} that delays decompressing a compressed frame until data is first read. + */ + @VisibleForTesting + static final class LazyDecompressingInputStream extends FilterInputStream { + private final Decompressor decompressor; + private final int maxMessageSize; + private final StatsTraceContext statsTraceCtx; + private boolean initialized; + private boolean closed; + + LazyDecompressingInputStream( + InputStream rawStream, + int maxMessageSize, + StatsTraceContext statsTraceCtx, + Decompressor decompressor) { + super(rawStream); + this.decompressor = decompressor; + this.maxMessageSize = maxMessageSize; + this.statsTraceCtx = statsTraceCtx; + } + + private synchronized void ensureInitialized() throws IOException { + if (closed) { + throw new IOException("Stream closed"); + } + if (!initialized) { + InputStream decompressed = decompressor.decompress(in); + in = new SizeEnforcingInputStream(decompressed, maxMessageSize, statsTraceCtx); + initialized = true; + } + } + + @Override + public int read() throws IOException { + ensureInitialized(); + return super.read(); + } + + @Override + public int read(byte[] b, int off, int len) throws IOException { + ensureInitialized(); + return super.read(b, off, len); + } + + @Override + public long skip(long n) throws IOException { + ensureInitialized(); + return super.skip(n); + } + + @Override + public int available() throws IOException { + ensureInitialized(); + return super.available(); + } + + @Override + public synchronized void close() throws IOException { + if (!closed) { + closed = true; + super.close(); + } + } + + @Override + public synchronized void mark(int readlimit) { + try { + ensureInitialized(); + super.mark(readlimit); + } catch (IOException e) { + throw new RuntimeException(e); + } + } + + @Override + public synchronized void reset() throws IOException { + ensureInitialized(); + super.reset(); + } + + @Override + public boolean markSupported() { + try { + ensureInitialized(); + return super.markSupported(); + } catch (IOException e) { + throw new RuntimeException(e); + } } } diff --git a/core/src/main/java/io/grpc/internal/ServerCallImpl.java b/core/src/main/java/io/grpc/internal/ServerCallImpl.java index fafd74308c9..823f528be98 100644 --- a/core/src/main/java/io/grpc/internal/ServerCallImpl.java +++ b/core/src/main/java/io/grpc/internal/ServerCallImpl.java @@ -34,7 +34,6 @@ 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; @@ -46,9 +45,6 @@ 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; @@ -325,20 +321,6 @@ 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) { @@ -349,22 +331,18 @@ private void messagesAvailableInternal(final MessageProducer producer) { InputStream message; try { while ((message = producer.next()) != null) { + // TODO: Consider forcing this check to be done in the transport (MessageDeframer) + // https://github.com/grpc/grpc-java/pull/13004/changes#r3939373996 if (call.method.getType().clientSendsOneMessage()) { if (delayedMessage != null) { GrpcUtil.closeQuietly(message); call.stream.cancel(Status.INTERNAL.withDescription("Too many requests")); GrpcUtil.closeQuietly(delayedMessage); delayedMessage = null; - closedInternal(Status.INTERNAL.withDescription("Too many requests")); + call.cancelled = true; return; } - try { - delayedMessage = bufferMessage(message); - } catch (Throwable t) { - GrpcUtil.closeQuietly(message); - throw t; - } - message.close(); + delayedMessage = message; } else { try { listener.onMessage(call.method.parseRequest(message)); diff --git a/core/src/test/java/io/grpc/internal/MessageDeframerTest.java b/core/src/test/java/io/grpc/internal/MessageDeframerTest.java index 54758bc096f..ea6b83bb6e6 100644 --- a/core/src/test/java/io/grpc/internal/MessageDeframerTest.java +++ b/core/src/test/java/io/grpc/internal/MessageDeframerTest.java @@ -19,6 +19,8 @@ import static com.google.common.truth.Truth.assertThat; import static io.grpc.internal.GrpcUtil.DEFAULT_MAX_MESSAGE_SIZE; import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotNull; import static org.junit.Assert.assertNull; import static org.junit.Assert.assertThrows; import static org.junit.Assert.assertTrue; @@ -35,9 +37,11 @@ import com.google.common.io.ByteStreams; import com.google.common.primitives.Bytes; import io.grpc.Codec; +import io.grpc.Decompressor; import io.grpc.InternalChannelz.TransportStats; import io.grpc.StatusRuntimeException; import io.grpc.StreamTracer; +import io.grpc.internal.MessageDeframer.LazyDecompressingInputStream; import io.grpc.internal.MessageDeframer.Listener; import io.grpc.internal.MessageDeframer.SizeEnforcingInputStream; import io.grpc.internal.testing.TestStreamTracer.TestBaseStreamTracer; @@ -52,6 +56,7 @@ import java.util.List; import java.util.Locale; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.zip.GZIPOutputStream; import org.junit.Before; import org.junit.Test; @@ -313,6 +318,75 @@ public void compressed() { verifyNoMoreInteractions(listener); } + @Test + public void compressed_lazyDecompression() throws IOException { + final AtomicBoolean decompressCalled = new AtomicBoolean(false); + Decompressor countingDecompressor = new Decompressor() { + @Override + public String getMessageEncoding() { + return "gzip"; + } + + @Override + public InputStream decompress(InputStream is) throws IOException { + decompressCalled.set(true); + return new Codec.Gzip().decompress(is); + } + }; + + deframer = new MessageDeframer(listener, countingDecompressor, DEFAULT_MAX_MESSAGE_SIZE, + statsTraceCtx, transportTracer); + deframer.request(1); + + byte[] payload = compress(new byte[1000]); + byte[] header = new byte[]{1, 0, 0, 0, (byte) payload.length}; + deframer.deframe(buffer(Bytes.concat(header, payload))); + + verify(listener).messagesAvailable(producer.capture()); + InputStream stream = producer.getValue().next(); + assertNotNull(stream); + + // Decompressor should not be invoked before bytes are read + assertFalse(decompressCalled.get()); + + // Reading a byte triggers decompression + assertEquals(0, stream.read()); + assertTrue(decompressCalled.get()); + } + + @Test + public void compressed_closeWithoutReading_noDecompression() throws IOException { + final AtomicBoolean decompressCalled = new AtomicBoolean(false); + Decompressor countingDecompressor = new Decompressor() { + @Override + public String getMessageEncoding() { + return "gzip"; + } + + @Override + public InputStream decompress(InputStream is) throws IOException { + decompressCalled.set(true); + return new Codec.Gzip().decompress(is); + } + }; + + deframer = new MessageDeframer(listener, countingDecompressor, DEFAULT_MAX_MESSAGE_SIZE, + statsTraceCtx, transportTracer); + deframer.request(1); + + byte[] payload = compress(new byte[1000]); + byte[] header = new byte[]{1, 0, 0, 0, (byte) payload.length}; + deframer.deframe(buffer(Bytes.concat(header, payload))); + + verify(listener).messagesAvailable(producer.capture()); + InputStream stream = producer.getValue().next(); + assertNotNull(stream); + + // Closing without reading should not decompress + stream.close(); + assertFalse(decompressCalled.get()); + } + @Test public void deliverIsReentrantSafe() { doAnswer( @@ -493,6 +567,79 @@ public void sizeEnforcingInputStream_markReset() throws IOException { } } + @RunWith(JUnit4.class) + public static class LazyDecompressingInputStreamTests { + private TestBaseStreamTracer tracer = new TestBaseStreamTracer(); + private StatsTraceContext statsTraceCtx = new StatsTraceContext(new StreamTracer[]{tracer}); + + @Test + public void lazyDecompressingInputStream_doesNotInitializeUntilRead() throws IOException { + final AtomicBoolean decompressCalled = new AtomicBoolean(false); + Decompressor countingDecompressor = new Decompressor() { + @Override + public String getMessageEncoding() { + return "gzip"; + } + + @Override + public InputStream decompress(InputStream is) throws IOException { + decompressCalled.set(true); + return new Codec.Gzip().decompress(is); + } + }; + + ByteArrayInputStream in = + new ByteArrayInputStream(compress("hello".getBytes(StandardCharsets.UTF_8))); + LazyDecompressingInputStream stream = new LazyDecompressingInputStream( + in, 100, statsTraceCtx, countingDecompressor); + + assertFalse(decompressCalled.get()); + byte[] buf = new byte[5]; + int read = stream.read(buf); + assertEquals(5, read); + assertEquals("hello", new String(buf, StandardCharsets.UTF_8)); + assertTrue(decompressCalled.get()); + stream.close(); + } + + @Test + public void lazyDecompressingInputStream_closeWithoutRead() throws IOException { + final AtomicBoolean decompressCalled = new AtomicBoolean(false); + final AtomicBoolean inClosed = new AtomicBoolean(false); + Decompressor countingDecompressor = new Decompressor() { + @Override + public String getMessageEncoding() { + return "gzip"; + } + + @Override + public InputStream decompress(InputStream is) throws IOException { + decompressCalled.set(true); + return new Codec.Gzip().decompress(is); + } + }; + + ByteArrayInputStream in = + new ByteArrayInputStream(compress("hello".getBytes(StandardCharsets.UTF_8))) { + @Override + public void close() throws IOException { + inClosed.set(true); + super.close(); + } + }; + LazyDecompressingInputStream stream = new LazyDecompressingInputStream( + in, 100, statsTraceCtx, countingDecompressor); + + assertFalse(decompressCalled.get()); + stream.close(); + assertTrue(inClosed.get()); + assertFalse(decompressCalled.get()); + + // Reading after close should throw IOException + assertThrows(IOException.class, () -> stream.read()); + } + } + /** * Verify stats were published through the tracer. * diff --git a/core/src/test/java/io/grpc/internal/ServerCallImplTest.java b/core/src/test/java/io/grpc/internal/ServerCallImplTest.java index 53ccf5bfe6a..9eb3e5b959b 100644 --- a/core/src/test/java/io/grpc/internal/ServerCallImplTest.java +++ b/core/src/test/java/io/grpc/internal/ServerCallImplTest.java @@ -40,7 +40,6 @@ import io.grpc.CompressorRegistry; import io.grpc.Context; import io.grpc.DecompressorRegistry; -import io.grpc.Detachable; import io.grpc.InternalChannelz.ServerStats; import io.grpc.Metadata; import io.grpc.MethodDescriptor; @@ -477,54 +476,20 @@ public void streamListener_onReady_onlyOnce() { public void streamListener_messageRead_unary_delayed() { ServerStreamListenerImpl streamListener = new ServerCallImpl.ServerStreamListenerImpl<>(call, callListener, context); - streamListener.messagesAvailable(new SingleMessageProducer(UNARY_METHOD.streamRequest(1234L))); + FakeCloseTrackerInputStream messageStream = + new FakeCloseTrackerInputStream(UNARY_METHOD.streamRequest(1234L), false); + streamListener.messagesAvailable(new SingleMessageProducer(messageStream)); - // Message should not be delivered yet + // Message should not be delivered or closed yet verify(callListener, never()).onMessage(any(Long.class)); + assertFalse(messageStream.closed); streamListener.halfClosed(); - // Now it should be delivered + // Now it should be delivered and closed verify(callListener).onMessage(1234L); verify(callListener).onHalfClose(); - } - - @Test - public void streamListener_messageRead_unary_detachable() { - ServerStreamListenerImpl streamListener = - new ServerCallImpl.ServerStreamListenerImpl<>(call, callListener, context); - - InputStream delegate = UNARY_METHOD.streamRequest(1234L); - FakeDetachableInputStream detachableStream = new FakeDetachableInputStream(delegate); - - streamListener.messagesAvailable(new SingleMessageProducer(detachableStream)); - - // It should have been detached immediately - assertTrue(detachableStream.detached); - 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_unary_bufferMessageException() { - ServerStreamListenerImpl streamListener = - new ServerCallImpl.ServerStreamListenerImpl<>(call, callListener, context); - - InputStream delegate = UNARY_METHOD.streamRequest(1234L); - FakeDetachableInputStream detachableStream = new FakeDetachableInputStream(delegate, true); - - RuntimeException e = assertThrows(RuntimeException.class, - () -> streamListener.messagesAvailable(new SingleMessageProducer(detachableStream))); - assertThat(e).hasMessageThat().isEqualTo("detach failed"); - - // The stream should have been closed in the catch block - assertTrue(detachableStream.closed); - verify(callListener, never()).onMessage(any(Long.class)); + assertTrue(messageStream.closed); } @Test @@ -534,17 +499,21 @@ public void streamListener_messageRead_serverStreaming_delayed() { serverCallTracer, PerfMark.createTag()); ServerStreamListenerImpl streamListener = new ServerCallImpl.ServerStreamListenerImpl<>(call, callListener, context); - streamListener.messagesAvailable( - new SingleMessageProducer(SERVER_STREAMING_METHOD.streamRequest(1234L))); + FakeCloseTrackerInputStream messageStream = + new FakeCloseTrackerInputStream( + SERVER_STREAMING_METHOD.streamRequest(1234L), false); + streamListener.messagesAvailable(new SingleMessageProducer(messageStream)); - // Message should not be delivered yet + // Message should not be delivered or closed yet verify(callListener, never()).onMessage(any(Long.class)); + assertFalse(messageStream.closed); streamListener.halfClosed(); - // Now it should be delivered + // Now it should be delivered and closed verify(callListener).onMessage(1234L); verify(callListener).onHalfClose(); + assertTrue(messageStream.closed); } @Test @@ -573,6 +542,14 @@ public void streamListener_messageRead_unary_tooManyRequests() { verify(stream).cancel(any(Status.class)); verify(callListener, never()).onMessage(any(Long.class)); + assertTrue(call.isCancelled()); + + // Subsequent halfClosed should be ignored because call is cancelled + streamListener.halfClosed(); + verify(callListener, never()).onHalfClose(); + + // When transport closes, onCancel is called once + streamListener.closed(Status.CANCELLED); verify(callListener).onCancel(); assertTrue(context.isCancelled()); } @@ -581,14 +558,19 @@ public void streamListener_messageRead_unary_tooManyRequests() { public void streamListener_messageRead_onlyOnce_unary() { ServerStreamListenerImpl streamListener = new ServerCallImpl.ServerStreamListenerImpl<>(call, callListener, context); - streamListener.messagesAvailable(new SingleMessageProducer(UNARY_METHOD.streamRequest(1234L))); - + FakeCloseTrackerInputStream messageStream = + new FakeCloseTrackerInputStream(UNARY_METHOD.streamRequest(1234L), false); + streamListener.messagesAvailable(new SingleMessageProducer(messageStream)); + + assertFalse(messageStream.closed); + // canceling the call should clean up and prevent delivery streamListener.closed(Status.CANCELLED); streamListener.halfClosed(); verify(callListener, never()).onMessage(any(Long.class)); + assertTrue(messageStream.closed); } @Test @@ -599,21 +581,22 @@ public void streamListener_unexpectedRuntimeException_unary() { .when(callListener) .onMessage(any(Long.class)); - InputStream delegate = UNARY_METHOD.streamRequest(1234L); - FakeDetachableInputStream detachableStream = new FakeDetachableInputStream(delegate); + FakeCloseTrackerInputStream messageStream = + new FakeCloseTrackerInputStream(UNARY_METHOD.streamRequest(1234L), false); - streamListener.messagesAvailable(new SingleMessageProducer(detachableStream)); + streamListener.messagesAvailable(new SingleMessageProducer(messageStream)); // Exception should not be thrown yet because deserialization/delivery is delayed verify(callListener, never()).onMessage(any(Long.class)); + assertFalse(messageStream.closed); // It should be thrown during halfClosed RuntimeException e = assertThrows(RuntimeException.class, () -> streamListener.halfClosed()); assertThat(e).hasMessageThat().isEqualTo("unexpected exception"); - // The detached stream should have been closed in the catch block - assertTrue(detachableStream.detachedStream.closed); + // The stream should have been closed in the catch block + assertTrue(messageStream.closed); } @Test @@ -621,14 +604,14 @@ public void streamListener_halfClosed_closeException() { ServerStreamListenerImpl streamListener = new ServerCallImpl.ServerStreamListenerImpl<>(call, callListener, context); - InputStream delegate = UNARY_METHOD.streamRequest(1234L); - FakeDetachableInputStream detachableStream = - new FakeDetachableInputStream(delegate, false, true); + FakeCloseTrackerInputStream messageStream = + new FakeCloseTrackerInputStream(UNARY_METHOD.streamRequest(1234L), true); - streamListener.messagesAvailable(new SingleMessageProducer(detachableStream)); + streamListener.messagesAvailable(new SingleMessageProducer(messageStream)); // Message should not be delivered yet verify(callListener, never()).onMessage(any(Long.class)); + assertFalse(messageStream.closed); // halfClosed should not throw because we use closeQuietly streamListener.halfClosed(); @@ -636,7 +619,7 @@ public void streamListener_halfClosed_closeException() { // The message was delivered and halfClosed completed verify(callListener).onMessage(1234L); verify(callListener).onHalfClose(); - assertTrue(detachableStream.detachedStream.closed); + assertTrue(messageStream.closed); } private static class LongMarshaller implements Marshaller { @@ -679,48 +662,4 @@ public void close() throws IOException { delegate.close(); } } - - private static class FakeDetachableInputStream extends InputStream implements Detachable { - boolean detached = false; - boolean closed = false; - private final InputStream delegate; - private final boolean throwOnDetach; - private final boolean throwOnClose; - FakeCloseTrackerInputStream detachedStream; - - FakeDetachableInputStream(InputStream delegate) { - this(delegate, false, false); - } - - FakeDetachableInputStream(InputStream delegate, boolean throwOnDetach) { - this(delegate, throwOnDetach, false); - } - - FakeDetachableInputStream(InputStream delegate, boolean throwOnDetach, boolean throwOnClose) { - this.delegate = delegate; - this.throwOnDetach = throwOnDetach; - this.throwOnClose = throwOnClose; - } - - @Override - public InputStream detach() { - if (throwOnDetach) { - throw new RuntimeException("detach failed"); - } - detached = true; - detachedStream = new FakeCloseTrackerInputStream(delegate, throwOnClose); - return detachedStream; - } - - @Override - public int read() throws IOException { - return delegate.read(); - } - - @Override - public void close() throws IOException { - closed = true; - delegate.close(); - } - } }