From 2a1eadb9a9db28dbed27e37904e23e73d2209d3e Mon Sep 17 00:00:00 2001
From: Chris Nokleberg <chrisn@google.com>
Date: Mon, 23 Mar 2020 09:56:28 -0700
Subject: [PATCH] core: preserve KnownLength when wrapping InputStream

useInputStreamMessages ensures that the InputStream supports marking by wrapping the stream in a BufferedInputStream if markSupported() returns false. This change uses a new subclass of BufferedInputStream that also implements KnownLength, when the original stream also implements KnownLength.
---
 .../main/java/io/grpc/ServerInterceptors.java | 36 ++++++++++++-------
 1 file changed, 23 insertions(+), 13 deletions(-)

diff --git a/api/src/main/java/io/grpc/ServerInterceptors.java b/api/src/main/java/io/grpc/ServerInterceptors.java
index dc1b0609d9b..430b90b9d0b 100644
--- a/api/src/main/java/io/grpc/ServerInterceptors.java
+++ b/api/src/main/java/io/grpc/ServerInterceptors.java
@@ -139,24 +139,34 @@ public static ServerServiceDefinition useInputStreamMessages(
       final ServerServiceDefinition serviceDef) {
     final MethodDescriptor.Marshaller<InputStream> marshaller =
         new MethodDescriptor.Marshaller<InputStream>() {
-      @Override
-      public InputStream stream(final InputStream value) {
-        return value;
-      }
+          @Override
+          public InputStream stream(final InputStream value) {
+            return value;
+          }
 
-      @Override
-      public InputStream parse(final InputStream stream) {
-        if (stream.markSupported()) {
-          return stream;
-        } else {
-          return new BufferedInputStream(stream);
-        }
-      }
-    };
+          @Override
+          public InputStream parse(final InputStream stream) {
+            if (stream.markSupported()) {
+              return stream;
+            } else if (stream instanceof KnownLength) {
+              return new KnownLengthBufferedInputStream(stream);
+            } else {
+              return new BufferedInputStream(stream);
+            }
+          }
+        };
 
     return useMarshalledMessages(serviceDef, marshaller);
   }
 
+  /** {@link BufferedInputStream} that also implements {@link KnownLength}. */
+  private static final class KnownLengthBufferedInputStream extends BufferedInputStream
+      implements KnownLength {
+    KnownLengthBufferedInputStream(InputStream in) {
+      super(in);
+    }
+  }
+
   /**
    * Create a new {@code ServerServiceDefinition} whose {@link MethodDescriptor} serializes to
    * and from T for all methods.  The {@code ServerCallHandler} created will automatically