Streaming gRPC client can leak responses if malformed response received (#3354)

Motivation:

When gRPC client receives a response, it first validates that it's
formed according to gRPC specification. If any headers or trailers are
missing, `validateResponseAndGetPayload` will throw an exception. In
this case, we leak undrained response payload body.

Modifications:

1. Try catch logic of `validateResponseAndGetPayload`, subscribe and
cancel response message body in case of unexpected exceptions.
2. Enhance `ProtocolCompatibilityTest` to validate we never leak
responses across all tests.

Result:

Responses are properly drained even when we receive malformed responses.

Author: idelpivnitskiy

servicetalk-grpc-api/src/main/java/io/servicetalk/grpc/api/DefaultGrpcClientCallFactory.java

@@ -19,6 +19,7 @@
 import io.servicetalk.concurrent.api.AsyncContext;
 import io.servicetalk.concurrent.api.Completable;
 import io.servicetalk.concurrent.api.Publisher;
+import io.servicetalk.concurrent.internal.CancelImmediatelySubscriber;
 import io.servicetalk.encoding.api.BufferDecoder;
 import io.servicetalk.encoding.api.BufferDecoderGroup;
 import io.servicetalk.encoding.api.BufferEncoder;
@@ -41,6 +42,7 @@
 import java.util.concurrent.TimeUnit;
 import javax.annotation.Nullable;
 
+import static io.servicetalk.concurrent.api.SourceAdapters.toSource;
 import static io.servicetalk.concurrent.internal.BlockingIterables.singletonBlockingIterable;
 import static io.servicetalk.encoding.api.Identity.identityEncoder;
 import static io.servicetalk.grpc.api.GrpcHeaderValues.GRPC_CONTENT_TYPE_PROTO_SUFFIX;
@@ -154,11 +156,17 @@ public <Req, Resp> StreamingClientCall<Req, Resp> newStreamingCall(
                     streamingHttpClient.executionContext().bufferAllocator()));
             return streamingHttpClient.request(httpRequest)
                     .flatMapPublisher(response -> {
-                        extractResponseContext(response, metadata);
-                        return validateResponseAndGetPayload(response, responseContentType,
-                                streamingHttpClient.executionContext().bufferAllocator(),
-                                readGrpcMessageEncodingRaw(response.headers(), deserializerIdentity, deserializers,
-                                        GrpcStreamingDeserializer::messageEncoding), httpRequest.requestTarget());
+                        try {
+
+                            extractResponseContext(response, metadata);
+                            return validateResponseAndGetPayload(response, responseContentType,
+                                    streamingHttpClient.executionContext().bufferAllocator(),
+                                    readGrpcMessageEncodingRaw(response.headers(), deserializerIdentity, deserializers,
+                                            GrpcStreamingDeserializer::messageEncoding), httpRequest.requestTarget());
+                        } catch (Throwable t) {
+                            toSource(response.messageBody()).subscribe(CancelImmediatelySubscriber.INSTANCE);
+                            return Publisher.failed(GrpcStatusException.fromThrowable(t));
+                        }
                     })
                     .onErrorMap(GrpcStatusException::fromThrowable);
         };
@@ -291,11 +299,16 @@ public <Req, Resp> BlockingStreamingClientCall<Req, Resp> newBlockingStreamingCa
                     streamingHttpClient.executionContext().bufferAllocator()));
             try {
                 final BlockingStreamingHttpResponse response = client.request(httpRequest);
-                extractResponseContext(response, metadata);
-                return validateResponseAndGetPayload(response.toStreamingResponse(), responseContentType,
-                        client.executionContext().bufferAllocator(), readGrpcMessageEncodingRaw(
-                                response.headers(), deserializerIdentity, deserializers,
-                                GrpcStreamingDeserializer::messageEncoding), httpRequest.requestTarget()).toIterable();
+                try {
+                    extractResponseContext(response, metadata);
+                    return validateResponseAndGetPayload(response.toStreamingResponse(), responseContentType,
+                            client.executionContext().bufferAllocator(), readGrpcMessageEncodingRaw(response.headers(),
+                                    deserializerIdentity, deserializers, GrpcStreamingDeserializer::messageEncoding),
+                            httpRequest.requestTarget()).toIterable();
+                } catch (Throwable t) {
+                    response.messageBody().iterator().close();
+                    throw GrpcStatusException.fromThrowable(t);
+                }
             } catch (Throwable cause) {
                 throw GrpcStatusException.fromThrowable(cause);
             }

servicetalk-grpc-netty/src/test/java/io/servicetalk/grpc/netty/ProtocolCompatibilityTest.java

@@ -53,20 +53,27 @@
 import io.servicetalk.grpc.netty.CompatProto.Compat.ServiceFactory;
 import io.servicetalk.grpc.netty.CompatProto.RequestContainer.CompatRequest;
 import io.servicetalk.grpc.netty.CompatProto.ResponseContainer.CompatResponse;
+import io.servicetalk.http.api.FilterableStreamingHttpClient;
+import io.servicetalk.http.api.HttpExecutionStrategies;
+import io.servicetalk.http.api.HttpExecutionStrategy;
 import io.servicetalk.http.api.HttpResponse;
 import io.servicetalk.http.api.HttpResponseStatus;
 import io.servicetalk.http.api.HttpServerBuilder;
 import io.servicetalk.http.api.HttpServerContext;
 import io.servicetalk.http.api.HttpServiceContext;
 import io.servicetalk.http.api.SingleAddressHttpClientBuilder;
 import io.servicetalk.http.api.StreamingHttpClient;
+import io.servicetalk.http.api.StreamingHttpClientFilter;
+import io.servicetalk.http.api.StreamingHttpClientFilterFactory;
 import io.servicetalk.http.api.StreamingHttpRequest;
+import io.servicetalk.http.api.StreamingHttpRequester;
 import io.servicetalk.http.api.StreamingHttpResponse;
 import io.servicetalk.http.api.StreamingHttpResponseFactory;
 import io.servicetalk.http.api.StreamingHttpService;
 import io.servicetalk.http.api.StreamingHttpServiceFilter;
 import io.servicetalk.http.netty.HttpClients;
 import io.servicetalk.http.netty.HttpServers;
+import io.servicetalk.http.utils.BeforeFinallyHttpOperator;
 import io.servicetalk.test.resources.DefaultTestCerts;
 import io.servicetalk.transport.api.ClientSslConfigBuilder;
 import io.servicetalk.transport.api.ServerContext;
@@ -91,6 +98,7 @@
 import io.grpc.stub.StreamObserver;
 import io.netty.handler.ssl.SslContext;
 import io.netty.handler.ssl.util.InsecureTrustManagerFactory;
+import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.function.ThrowingSupplier;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.Arguments;
@@ -200,6 +208,15 @@ public SocketAddress listenAddress() {
     private static final String CUSTOM_ERROR_MESSAGE = "custom error message";
     private static final DeliberateException SERVER_PROCESSED_TOKEN = new DeliberateException();
     private static final Duration DEFAULT_DEADLINE = ofMillis(100);
+    private static final boolean[] TRUE_FALSE = {true, false};
+    private static final String[] COMPRESSION = {"gzip", "identity", null};
+
+    private final ResponseLeakValidator responseLeakValidator = new ResponseLeakValidator();
+
+    @AfterEach
+    void finalChecks() {
+        responseLeakValidator.assertNoPendingRequests();
+    }
 
     private enum ErrorMode {
         NONE,
@@ -211,9 +228,6 @@ private enum ErrorMode {
         STATUS_IN_RESPONSE
     }
 
-    private static final boolean[] TRUE_FALSE = {true, false};
-    private static final String[] COMPRESSION = {"gzip", "identity", null};
-
     private static Collection<Arguments> sslStreamingAndCompressionParams() {
         List<Arguments> args = new ArrayList<>();
         for (boolean ssl : TRUE_FALSE) {
@@ -1217,16 +1231,21 @@ private static CompatResponse computeResponse(final int value) {
                 .build();
     }
 
-    private static CompatClient serviceTalkClient(final SocketAddress serverAddress, final boolean ssl,
-                                                  @Nullable final String compression,
-                                                  @Nullable final Duration timeout) {
+    private CompatClient serviceTalkClient(final SocketAddress serverAddress, final boolean ssl,
+                                           @Nullable final String compression,
+                                           @Nullable final Duration timeout) {
         final GrpcClientBuilder<InetSocketAddress, InetSocketAddress> builder =
                 GrpcClients.forResolvedAddress((InetSocketAddress) serverAddress);
-        if (ssl) {
-            builder.initializeHttp(b -> b.sslConfig(new ClientSslConfigBuilder(
-                    DefaultTestCerts::loadServerCAPem).peerHost(serverPemHostname()).build()));
-        }
-        if (null != timeout) {
+
+        builder.initializeHttp(b -> {
+            if (ssl) {
+                b.sslConfig(new ClientSslConfigBuilder(
+                        DefaultTestCerts::loadServerCAPem).peerHost(serverPemHostname()).build());
+            }
+            b.appendClientFilter(responseLeakValidator);
+        });
+
+        if (timeout != null) {
             builder.defaultTimeout(timeout);
         }
         return builder.build(new Compat.ClientFactory()
@@ -1854,4 +1873,34 @@ private CompatResponse response(final int value) throws Exception {
             return computeResponse(value);
         }
     }
+
+    private static final class ResponseLeakValidator implements StreamingHttpClientFilterFactory {
+
+        private final AtomicInteger pendingRequests = new AtomicInteger();
+
+        @Override
+        public StreamingHttpClientFilter create(FilterableStreamingHttpClient client) {
+            return new StreamingHttpClientFilter(client) {
+                @Override
+                protected Single<StreamingHttpResponse> request(StreamingHttpRequester delegate,
+                                                                StreamingHttpRequest request) {
+                    return Single.defer(() -> {
+                        pendingRequests.incrementAndGet();
+                        return delegate.request(request)
+                                .liftSync(new BeforeFinallyHttpOperator(pendingRequests::decrementAndGet))
+                                .shareContextOnSubscribe();
+                    });
+                }
+            };
+        }
+
+        @Override
+        public HttpExecutionStrategy requiredOffloads() {
+            return HttpExecutionStrategies.offloadNone();
+        }
+
+        void assertNoPendingRequests() {
+            assertThat("Detected pending requests, possible response leak", pendingRequests.get(), is(0));
+        }
+    }
 }