gRPC improvements and a bug fix for large data uploads (helidon-io#10797)

* Improve buffer management when reading large gRPC messages.
* Fix issue when receiving insufficient data to read gRPC header. Improve buffer management when writing gRPC responses.
* New input stream that implements KnownLength for parsing.
* New test that simulates large uploads using a bidi method that acks each upload.
* Set max limit on gRPC read buffer to prevent attacks. New test showing error if exceeding limit.

Author: spericas

webserver/grpc/src/main/java/io/helidon/webserver/grpc/GrpcConfigBlueprint.java

@@ -78,4 +78,13 @@ interface GrpcConfigBlueprint extends ProtocolConfig {
     default List<GrpcServerService> grpcServices() {
         return List.of();
     }
+
+    /**
+     * Max size of gRPC reading buffer. If receiving an entity larger than this,
+     * processing will be aborted. This can help prevent DoS attacks. Default
+     * set to 2 MB.
+     */
+    @Option.Configured
+    @Option.DefaultInt(2 * 1024 * 1024)
+    int maxReadBufferSize();
 }

webserver/grpc/src/main/java/io/helidon/webserver/grpc/GrpcProtocolHandler.java

@@ -98,6 +98,9 @@ private record MethodMetrics(Counter callStarted,
 
     private static final LazyValue<Map<String, MethodMetrics>> METHOD_METRICS = LazyValue.create(ConcurrentHashMap::new);
 
+    private static final int GRPC_HEADER_SIZE = 5;
+    private static final int INITIAL_BUFFER_SIZE = 16 * 1024;
+
     private final Http2Headers headers;
     private final Http2StreamWriter streamWriter;
     private final int streamId;
@@ -109,6 +112,9 @@ private record MethodMetrics(Counter callStarted,
 
     private volatile ServerCall.Listener<REQ> listener;
     private BufferData entityBytes;
+    private BufferData readBufferData = BufferData.create(INITIAL_BUFFER_SIZE);
+    private BufferData unreadBufferData;
+    private long entityBytesLeft;
     private Compressor compressor;
     private Decompressor decompressor;
     private boolean identityCompressor;
@@ -185,35 +191,61 @@ public void rstStream(Http2RstStream rstStream) {
     public void windowUpdate(Http2WindowUpdate update) {
     }
 
+    /**
+     * Data received from HTTP/2 layer. Data may contain a partial gRPC request
+     * or more than one request, making logic a bit more difficult.
+     *
+     * @param header frame header
+     * @param data   frame data
+     */
     @Override
     public void data(Http2FrameHeader header, BufferData data) {
         try {
             boolean isCompressed = false;
 
-            while (data.available() > 0) {
-                // start of new chunk?
+            // check for any unread data received before
+            BufferData newData;
+            if (unreadBufferData != null) {
+                newData = BufferData.create(unreadBufferData, data);
+                unreadBufferData = null;
+            } else {
+                newData = data;
+            }
+
+            // process 0 or more requests from data
+            while (newData.available() > 0) {
+                // start of new request?
                 if (entityBytes == null) {
-                    isCompressed = (data.read() == 1);
-                    long length = data.readUnsignedInt32();
-                    entityBytes = BufferData.create((int) length);
+                    if (newData.available() >= GRPC_HEADER_SIZE) {
+                        isCompressed = (newData.read() == 1);
+                        entityBytesLeft = newData.readUnsignedInt32();
+                        entityBytes = allocateReadBuffer((int) entityBytesLeft);
+                    } else {
+                        unreadBufferData = newData;
+                        return;     // need more for gRPC header
+                    }
                 }
 
-                // append data to current chunk
-                entityBytes.write(data);
+                // append data to current entity
+                int writableNow = (int) Math.min(entityBytesLeft, newData.available());
+                entityBytes.write(newData, writableNow);
+                entityBytesLeft -= writableNow;
 
-                // is chunk complete?
-                if (entityBytes.capacity() == 0) {
+                // is the entity complete?
+                if (entityBytesLeft == 0) {
                     // fail if compressed and no decompressor
                     if (isCompressed && decompressor == null) {
                         throw new IllegalStateException("Unable to codec for compressed data");
                     }
 
                     // read and possibly decompress data
                     bytesReceived += entityBytes.available();
-                    InputStream is = entityBytes.asInputStream();
+                    InputStream is = new BufferDataInputStream(entityBytes);
                     REQ request = route.method().parseRequest(isCompressed ? decompressor.decompress(is) : is);
                     listenerQueue.add(request);
                     flushQueue();
+
+                    // reset entityBytes
                     entityBytes = null;
                 }
             }
@@ -234,6 +266,18 @@ public void data(Http2FrameHeader header, BufferData data) {
         }
     }
 
+    BufferData allocateReadBuffer(int length) {
+        readBufferData.reset();
+        int capacity = readBufferData.capacity();
+        if (length > capacity) {
+            if (length > grpcConfig.maxReadBufferSize()) {
+                throw new IllegalStateException("gRPC message size exceeds max read buffer size");
+            }
+            readBufferData = BufferData.create(length);
+        }
+        return readBufferData;
+    }
+
     void initCompression(ServerCall<REQ, RES> serverCall, Headers httpHeaders) {
         if (grpcConfig.enableCompression()) {
             // check for encoding and respond using same algorithm
@@ -303,23 +347,22 @@ private void flushQueue() {
     static Http2StreamState nextStreamState(Http2StreamState currentStreamState,
                                             Http2StreamState desiredStreamState) {
         return switch (desiredStreamState) {
-            case HALF_CLOSED_LOCAL ->
-                    currentStreamState == Http2StreamState.HALF_CLOSED_REMOTE
-                            ? Http2StreamState.CLOSED
-                            : Http2StreamState.HALF_CLOSED_LOCAL;
-            case HALF_CLOSED_REMOTE ->
-                    currentStreamState == Http2StreamState.HALF_CLOSED_LOCAL
-                            ? Http2StreamState.CLOSED
-                            : Http2StreamState.HALF_CLOSED_REMOTE;
+            case HALF_CLOSED_LOCAL -> currentStreamState == Http2StreamState.HALF_CLOSED_REMOTE
+                    ? Http2StreamState.CLOSED
+                    : Http2StreamState.HALF_CLOSED_LOCAL;
+            case HALF_CLOSED_REMOTE -> currentStreamState == Http2StreamState.HALF_CLOSED_LOCAL
+                    ? Http2StreamState.CLOSED
+                    : Http2StreamState.HALF_CLOSED_REMOTE;
             default -> desiredStreamState;
         };
     }
 
     private ServerCall<REQ, RES> createServerCall() {
-        return new ServerCall<>() {
+        return new ServerCall<REQ, RES>() {
 
             private long bytesSent;
             private boolean headersSent;
+            private BufferData writeBufferData = BufferData.growing(INITIAL_BUFFER_SIZE);
 
             @Override
             public void request(int numMessages) {
@@ -356,29 +399,24 @@ public void sendMessage(RES message) {
                 try (InputStream inputStream = route.method().streamResponse(message)) {
                     // prepare buffer for writing
                     BufferData bufferData;
-                    if (identityCompressor) {
-                        // avoid buffer copy if length is known
-                        if (inputStream instanceof KnownLength knownLength) {
-                            int bytesLength = knownLength.available();
-                            bufferData = BufferData.create(5 + bytesLength);
-                            bufferData.write(0);        // off for identity compressor
-                            bufferData.writeUnsignedInt32(bytesLength);
-                            bufferData.readFrom(inputStream);
-                        } else {
-                            byte[] bytes = inputStream.readAllBytes();
-                            bufferData = BufferData.create(5 + bytes.length);
-                            bufferData.write(0);        // off for identity compressor
-                            bufferData.writeUnsignedInt32(bytes.length);
-                            bufferData.write(bytes);
-                        }
+                    if (identityCompressor && inputStream instanceof KnownLength knownLength) {
+                        int bytesLength = knownLength.available();
+                        bufferData = allocateWriteBuffer(GRPC_HEADER_SIZE + bytesLength);
+                        bufferData.write(0);        // 0 for identity compressor
+                        bufferData.writeUnsignedInt32(bytesLength);
+                        bufferData.readFrom(inputStream);
                     } else {
                         ByteArrayOutputStream baos = new ByteArrayOutputStream();
-                        try (OutputStream os = compressor.compress(baos)) {
-                            inputStream.transferTo(os);
+                        if (identityCompressor) {
+                            inputStream.transferTo(baos);
+                        } else {
+                            try (OutputStream os = compressor.compress(baos)) {
+                                inputStream.transferTo(os);
+                            }
                         }
                         byte[] bytes = baos.toByteArray();
-                        bufferData = BufferData.create(5 + bytes.length);
-                        bufferData.write(1);
+                        bufferData = allocateWriteBuffer(GRPC_HEADER_SIZE + bytes.length);
+                        bufferData.write(identityCompressor ? 0 : 1);
                         bufferData.writeUnsignedInt32(bytes.length);
                         bufferData.write(bytes);
                     }
@@ -449,6 +487,15 @@ public boolean isCancelled() {
             public MethodDescriptor<REQ, RES> getMethodDescriptor() {
                 return route.method();
             }
+
+            private BufferData allocateWriteBuffer(int length) {
+                writeBufferData.reset();
+                int capacity = writeBufferData.capacity();
+                if (length > capacity) {
+                    writeBufferData = BufferData.create(length);
+                }
+                return writeBufferData;
+            }
         };
     }
 
@@ -490,4 +537,37 @@ private void initMetrics() {
             return new MethodMetrics(callStarted, callDuration, sentMessageSize, recvMessageSize);
         });
     }
+
+    /**
+     * An input stream that can return its length. gRPC parsers can use this extra
+     * knowledge for optimizations. It can also copy a byte array directly on a
+     * single read.
+     */
+    static class BufferDataInputStream extends InputStream implements KnownLength {
+        private final BufferData bufferData;
+
+        BufferDataInputStream(BufferData bufferData) {
+            this.bufferData = bufferData;
+        }
+
+        @Override
+        public int read() {
+            return bufferData.read();
+        }
+
+        @Override
+        public int read(byte[] b) {
+            return bufferData.read(b);
+        }
+
+        @Override
+        public int read(byte[] b, int off, int len) {
+            return bufferData.read(b, off, len);
+        }
+
+        @Override
+        public int available() {
+            return bufferData.available();
+        }
+    }
 }

webserver/tests/grpc/src/main/java/io/helidon/webserver/tests/grpc/UploadService.java

@@ -0,0 +1,59 @@
+/*
+ * Copyright (c) 2025 Oracle and/or its affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.helidon.webserver.tests.grpc;
+
+import io.helidon.webserver.grpc.GrpcService;
+import io.helidon.webserver.grpc.uploads.Uploads;
+import io.helidon.webserver.grpc.uploads.Uploads.Ack;
+import io.helidon.webserver.grpc.uploads.Uploads.Data;
+
+import com.google.protobuf.Descriptors;
+import io.grpc.stub.StreamObserver;
+
+class UploadService implements GrpcService {
+
+    static final Ack TRUE_ACK = Ack.newBuilder().setOk(true).build();
+
+    @Override
+    public Descriptors.FileDescriptor proto() {
+        return Uploads.getDescriptor();
+    }
+
+    @Override
+    public void update(Routing router) {
+        router.bidi("Upload", this::grpcUpload);
+    }
+
+    private StreamObserver<Data> grpcUpload(StreamObserver<Ack> result) {
+        return new StreamObserver<>() {
+            @Override
+            public void onNext(Data data) {
+                result.onNext(TRUE_ACK);
+            }
+
+            @Override
+            public void onError(Throwable t) {
+                result.onError(t);
+            }
+
+            @Override
+            public void onCompleted() {
+                result.onCompleted();
+            }
+        };
+    }
+}

webserver/tests/grpc/src/main/proto/uploads.proto

@@ -0,0 +1,30 @@
+/*
+ * Copyright (c) 2025 Oracle and/or its affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+syntax = "proto3";
+option java_package = "io.helidon.webserver.grpc.uploads";
+
+service UploadService {
+  rpc Upload (stream Data) returns (stream Ack) {}
+}
+
+message Data {
+  bytes payload = 1;
+}
+
+message Ack {
+  bool ok = 1;
+}