Arawoof06 commented on code in PR #1318:
URL: https://github.com/apache/arrow-java/pull/1318#discussion_r4190063208


##########
flight/flight-core/src/main/java/org/apache/arrow/flight/ArrowMessage.java:
##########
@@ -377,6 +381,24 @@ private static int readRawVarint32(int firstByte, 
InputStream is) throws IOExcep
     return CodedInputStream.readRawVarint32(firstByte, is);
   }
 
+  /**
+   * Reject a field whose declared length is negative or larger than the bytes 
left in the message.
+   *
+   * <p>The length prefix is read straight off the wire, and a field can never 
be longer than the
+   * bytes still buffered for the message. Without this check an oversized 
value drives an unbounded
+   * allocation before any content is read; the {@code new byte[size]} paths 
above do so on the JVM
+   * heap, bypassing the {@link BufferAllocator} limit entirely.
+   */
+  private static void checkFieldLength(int size, InputStream stream) throws 
IOException {
+    final int remaining = stream.available();
+    if (size < 0 || size > remaining) {

Review Comment:
   Done, split into "negative field length N" and "field length N exceeds M 
bytes remaining in the message".



##########
flight/flight-core/src/test/java/org/apache/arrow/flight/TestArrowMessage.java:
##########
@@ -0,0 +1,101 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.arrow.flight;
+
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import com.google.protobuf.WireFormat;
+import io.grpc.MethodDescriptor;
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import org.apache.arrow.flight.impl.Flight.FlightData;
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.memory.RootAllocator;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+public class TestArrowMessage {
+
+  private static final int HEADER_TAG =
+      (FlightData.DATA_HEADER_FIELD_NUMBER << 3) | 
WireFormat.WIRETYPE_LENGTH_DELIMITED;
+  private static final int APP_METADATA_TAG =
+      (FlightData.APP_METADATA_FIELD_NUMBER << 3) | 
WireFormat.WIRETYPE_LENGTH_DELIMITED;
+
+  private BufferAllocator allocator;
+
+  @BeforeEach
+  public void setUp() {
+    allocator = new RootAllocator(Long.MAX_VALUE);
+  }
+
+  @AfterEach
+  public void tearDown() {
+    allocator.close();
+  }
+
+  /**
+   * A field whose declared length is far larger than the bytes actually 
present in the frame must
+   * be rejected before anything is allocated for it, rather than driving an 
allocation sized by the
+   * attacker-controlled length prefix.
+   */
+  @Test
+  public void frameRejectsOversizedFieldLength() {
+    final ByteArrayOutputStream frame = new ByteArrayOutputStream();
+    writeRawVarint32(frame, HEADER_TAG);
+    // Claim a much larger length than the (zero) bytes that follow.
+    writeRawVarint32(frame, 1 << 20);
+
+    final MethodDescriptor.Marshaller<ArrowMessage> marshaller =
+        ArrowMessage.createMarshaller(allocator);
+    final Exception e =
+        assertThrows(
+            Exception.class, () -> marshaller.parse(new 
ByteArrayInputStream(frame.toByteArray())));
+    assertTrue(
+        e.getMessage() != null && e.getMessage().contains("exceeds"),
+        "unexpected failure: " + e.getMessage());

Review Comment:
   Both added: `frameRejectsNegativeFieldLength` covers all four fields with a 
5-byte varint, and `frameReleasesBuffersWhenLaterFieldIsRejected` parses a 
valid app_metadata and body followed by an oversized field and asserts 
`getAllocatedMemory()` is 0 before `tearDown` closes the allocator.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]

Reply via email to