Adding maxMessageSize config option

Fixes #832

Author: nmittler

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

@@ -32,19 +32,15 @@
 package io.grpc.internal;
 
 import static com.google.common.base.Preconditions.checkArgument;
-import static io.grpc.Status.Code.CANCELLED;
-import static io.grpc.Status.Code.DEADLINE_EXCEEDED;
+import static io.grpc.internal.GrpcUtil.CANCEL_REASONS;
 
 import com.google.common.base.Objects;
 import com.google.common.base.Preconditions;
 
 import io.grpc.Metadata;
 import io.grpc.Status;
-import io.grpc.Status.Code;
 
 import java.io.InputStream;
-import java.util.EnumSet;
-import java.util.Set;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 
@@ -55,8 +51,6 @@ public abstract class AbstractClientStream<IdT> extends AbstractStream<IdT>
     implements ClientStream {
 
   private static final Logger log = Logger.getLogger(AbstractClientStream.class.getName());
-  private static final Set<Code> CANCEL_REASONS =
-      EnumSet.of(CANCELLED, DEADLINE_EXCEEDED, Code.INTERNAL, Code.UNKNOWN);
 
   private final ClientStreamListener listener;
   private boolean listenerClosed;
@@ -67,15 +61,10 @@ public abstract class AbstractClientStream<IdT> extends AbstractStream<IdT>
   private Metadata trailers;
   private Runnable closeListenerTask;
 
-
-  /**
-   * Constructor used by subclasses.
-   *
-   * @param listener the listener to receive notifications
-   */
   protected AbstractClientStream(WritableBufferAllocator bufferAllocator,
-                                 ClientStreamListener listener) {
-    super(bufferAllocator);
+                                 ClientStreamListener listener,
+                                 int maxMessageSize) {
+    super(bufferAllocator, maxMessageSize);
     this.listener = Preconditions.checkNotNull(listener);
   }
 

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

@@ -63,8 +63,9 @@ public abstract class AbstractServerStream<IdT> extends AbstractStream<IdT>
   /** Saved trailers from close() that need to be sent once the framer has sent all messages. */
   private Metadata stashedTrailers;
 
-  protected AbstractServerStream(WritableBufferAllocator bufferAllocator) {
-    super(bufferAllocator);
+  protected AbstractServerStream(WritableBufferAllocator bufferAllocator,
+                                 int maxMessageSize) {
+    super(bufferAllocator, maxMessageSize);
   }
 
   /**

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

@@ -100,7 +100,7 @@ protected enum Phase {
 
   private final Object onReadyLock = new Object();
 
-  AbstractStream(WritableBufferAllocator bufferAllocator) {
+  AbstractStream(WritableBufferAllocator bufferAllocator, int maxMessageSize) {
     MessageDeframer.Listener inboundMessageHandler = new MessageDeframer.Listener() {
       @Override
       public void bytesRead(int numBytes) {
@@ -130,7 +130,7 @@ public void deliverFrame(WritableBuffer frame, boolean endOfStream, boolean flus
     };
 
     framer = new MessageFramer(outboundFrameHandler, bufferAllocator);
-    deframer = new MessageDeframer(inboundMessageHandler);
+    deframer = new MessageDeframer(inboundMessageHandler, MessageEncoding.NONE, maxMessageSize);
   }
 
   @Override

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

@@ -31,13 +31,18 @@
 
 package io.grpc.internal;
 
+import static io.grpc.Status.Code.CANCELLED;
+import static io.grpc.Status.Code.DEADLINE_EXCEEDED;
+
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
 
 import io.grpc.Metadata;
 import io.grpc.Status;
 
 import java.net.HttpURLConnection;
+import java.util.EnumSet;
+import java.util.Set;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.ThreadFactory;
@@ -109,6 +114,17 @@ public final class GrpcUtil {
    */
   public static final String MESSAGE_ENCODING = "grpc-encoding";
 
+  /**
+   * The default maximum uncompressed size (in bytes) for inbound messages. Defaults to 100 MiB.
+   */
+  public static final int DEFAULT_MAX_MESSAGE_SIZE = 100 * 1024 * 1024;
+
+  /**
+   * The set of valid status codes for client cancellation.
+   */
+  public static final Set<Status.Code> CANCEL_REASONS =
+          EnumSet.of(CANCELLED, DEADLINE_EXCEEDED, Status.Code.INTERNAL, Status.Code.UNKNOWN);
+
   /**
    * Maps HTTP error response status codes to transport codes.
    */

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

@@ -71,8 +71,9 @@ public Integer parseAsciiString(String serialized) {
   private boolean contentTypeChecked;
 
   protected Http2ClientStream(WritableBufferAllocator bufferAllocator,
-                              ClientStreamListener listener) {
-    super(bufferAllocator, listener);
+                              ClientStreamListener listener,
+                              int maxMessageSize) {
+    super(bufferAllocator, listener, maxMessageSize);
   }
 
   /**

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

@@ -39,6 +39,7 @@
 import io.grpc.Status;
 
 import java.io.Closeable;
+import java.io.FilterInputStream;
 import java.io.IOException;
 import java.io.InputStream;
 
@@ -94,6 +95,7 @@ private enum State {
   }
 
   private final Listener listener;
+  private final int maxMessageSize;
   private MessageEncoding.Decompressor decompressor;
   private State state = State.HEADER;
   private int requiredLength = HEADER_LENGTH;
@@ -105,25 +107,19 @@ private enum State {
   private boolean deliveryStalled = true;
   private boolean inDelivery = false;
 
-  /**
-   * Creates a deframer. Compression will not be supported.
-   *
-   * @param listener listener for deframer events.
-   */
-  public MessageDeframer(Listener listener) {
-    this(listener, MessageEncoding.NONE);
-  }
-
   /**
    * Create a deframer.
    *
    * @param listener listener for deframer events.
    * @param decompressor the compression used if a compressed frame is encountered, with
    *  {@code NONE} meaning unsupported
+   * @param maxMessageSize the maximum allowed size for received messages.
    */
-  public MessageDeframer(Listener listener, MessageEncoding.Decompressor decompressor) {
-    this.listener = Preconditions.checkNotNull(listener, "listener");
+  public MessageDeframer(Listener listener, MessageEncoding.Decompressor decompressor,
+                         int maxMessageSize) {
+    this.listener = Preconditions.checkNotNull(listener, "sink");
     this.decompressor = Preconditions.checkNotNull(decompressor, "decompressor");
+    this.maxMessageSize = maxMessageSize;
   }
 
   /**
@@ -162,8 +158,7 @@ public void request(int numMessages) {
    *        the remote endpoint.  End of stream should not be used in the event of a transport
    *        error, such as a stream reset.
    * @throws IllegalStateException if {@link #close()} has been called previously or if
-   *         {@link #deframe(ReadableBuffer, boolean)} has previously been called with
-   *         {@code endOfStream=true}.
+   *         this method has previously been called with {@code endOfStream=true}.
    */
   public void deframe(ReadableBuffer data, boolean endOfStream) {
     Preconditions.checkNotNull(data, "data");
@@ -291,10 +286,6 @@ private void deliver() {
     }
   }
 
-  private boolean isDataAvailable() {
-    return unprocessed.readableBytes() > 0 || (nextFrame != null && nextFrame.readableBytes() > 0);
-  }
-
   /**
    * Attempts to read the required bytes into nextFrame.
    *
@@ -340,6 +331,10 @@ private void processHeader() {
 
     // Update the required length to include the length of the frame.
     requiredLength = nextFrame.readInt();
+    if (requiredLength < 0 || requiredLength > maxMessageSize) {
+      throw Status.INTERNAL.withDescription(String.format("Frame size %d exceeds maximum: %d, ",
+              requiredLength, maxMessageSize)).asRuntimeException();
+    }
 
     // Continue reading the frame body.
     state = State.BODY;
@@ -370,9 +365,79 @@ private InputStream getCompressedBody() {
     }
 
     try {
-      return decompressor.decompress(ReadableBuffers.openStream(nextFrame, true));
+      // Enforce the maxMessageSize limit on the returned stream.
+      return new SizeEnforcingInputStream(decompressor.decompress(
+              ReadableBuffers.openStream(nextFrame, true)));
     } catch (IOException e) {
       throw new RuntimeException(e);
     }
   }
+
+  /**
+   * An {@link InputStream} that enforces the {@link #maxMessageSize} limit for compressed frames.
+   */
+  private final class SizeEnforcingInputStream extends FilterInputStream {
+    private long count;
+    private long mark = -1;
+
+    public SizeEnforcingInputStream(InputStream in) {
+      super(in);
+    }
+
+    @Override
+    public int read() throws IOException {
+      int result = in.read();
+      if (result != -1) {
+        count++;
+      }
+      verifySize();
+      return result;
+    }
+
+    @Override
+    public int read(byte[] b, int off, int len) throws IOException {
+      int result = in.read(b, off, len);
+      if (result != -1) {
+        count += result;
+      }
+      verifySize();
+      return result;
+    }
+
+    @Override
+    public long skip(long n) throws IOException {
+      long result = in.skip(n);
+      count += result;
+      verifySize();
+      return result;
+    }
+
+    @Override
+    public synchronized void mark(int readlimit) {
+      in.mark(readlimit);
+      mark = count;
+      // it's okay to mark even if mark isn't supported, as reset won't work
+    }
+
+    @Override
+    public synchronized void reset() throws IOException {
+      if (!in.markSupported()) {
+        throw new IOException("Mark not supported");
+      }
+      if (mark == -1) {
+        throw new IOException("Mark not set");
+      }
+
+      in.reset();
+      count = mark;
+    }
+
+    private void verifySize() {
+      if (count > maxMessageSize) {
+        throw Status.INTERNAL.withDescription(String.format(
+                "Compressed frame exceeds maximum frame size: %d. Bytes read: %d",
+                maxMessageSize, count)).asRuntimeException();
+      }
+    }
+  }
 }

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

@@ -31,6 +31,7 @@
 
 package io.grpc.internal;
 
+import static io.grpc.internal.GrpcUtil.DEFAULT_MAX_MESSAGE_SIZE;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.fail;
 import static org.mockito.Matchers.isA;
@@ -242,7 +243,7 @@ public void inboundHeadersReceived_notifiesListenerOnBadEncoding() {
   private static class BaseAbstractClientStream<T> extends AbstractClientStream<T> {
     protected BaseAbstractClientStream(
         WritableBufferAllocator allocator, ClientStreamListener listener) {
-      super(allocator, listener);
+      super(allocator, listener, DEFAULT_MAX_MESSAGE_SIZE);
     }
 
     @Override

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

@@ -32,6 +32,7 @@
 
 package io.grpc.internal;
 
+import static io.grpc.internal.GrpcUtil.DEFAULT_MAX_MESSAGE_SIZE;
 import static org.junit.Assert.fail;
 import static org.mockito.Mockito.verify;
 
@@ -102,7 +103,7 @@ public void validPhaseTransitions() {
    */
   private class AbstractStreamBase<IdT> extends AbstractStream<IdT> {
     private AbstractStreamBase(WritableBufferAllocator bufferAllocator) {
-      super(bufferAllocator);
+      super(bufferAllocator, DEFAULT_MAX_MESSAGE_SIZE);
     }
 
     @Override

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

@@ -31,6 +31,7 @@
 
 package io.grpc.internal;
 
+import static io.grpc.internal.GrpcUtil.DEFAULT_MAX_MESSAGE_SIZE;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
 import static org.mockito.Matchers.anyInt;
@@ -67,7 +68,8 @@
 @RunWith(JUnit4.class)
 public class MessageDeframerTest {
   private Listener listener = mock(Listener.class);
-  private MessageDeframer deframer = new MessageDeframer(listener);
+  private MessageDeframer deframer = new MessageDeframer(listener, MessageEncoding.NONE,
+          DEFAULT_MAX_MESSAGE_SIZE);
   private ArgumentCaptor<InputStream> messages = ArgumentCaptor.forClass(InputStream.class);
 
   @Test
@@ -176,7 +178,7 @@ public void endOfStreamCallbackShouldWaitForMessageDelivery() {
 
   @Test
   public void compressed() {
-    deframer = new MessageDeframer(listener, new MessageEncoding.Gzip());
+    deframer = new MessageDeframer(listener, new MessageEncoding.Gzip(), DEFAULT_MAX_MESSAGE_SIZE);
     deframer.request(1);
 
     byte[] payload = compress(new byte[1000]);

netty/src/main/java/io/grpc/netty/CancelClientStreamCommand.java

@@ -31,15 +31,12 @@
 
 package io.grpc.netty;
 
-import static io.grpc.Status.Code.CANCELLED;
-import static io.grpc.Status.Code.DEADLINE_EXCEEDED;
+import static io.grpc.internal.GrpcUtil.CANCEL_REASONS;
 
 import com.google.common.base.Preconditions;
 
 import io.grpc.Status;
 
-import java.util.EnumSet;
-
 /**
  * Command sent from a Netty client stream to the handler to cancel the stream.
  */
@@ -50,8 +47,8 @@ class CancelClientStreamCommand {
   CancelClientStreamCommand(NettyClientStream stream, Status reason) {
     this.stream = Preconditions.checkNotNull(stream, "stream");
     Preconditions.checkNotNull(reason);
-    Preconditions.checkArgument(EnumSet.of(CANCELLED, DEADLINE_EXCEEDED).contains(reason.getCode()),
-        "Invalid cancellation reason");
+    Preconditions.checkArgument(CANCEL_REASONS.contains(reason.getCode()),
+            "Invalid cancellation reason");
     this.reason = reason;
   }