jbonofre commented on code in PR #1318: URL: https://github.com/apache/arrow-java/pull/1318#discussion_r4178110580
########## 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: Could we also add tests to verify: 1. negative field length prefixes (`size < 0`) are rejected property. 2. leak prevention: when a frame with a valid `APP_METADATA_TAG` followed by an oversized `BODY_TAG` is parsed and fails, `allocator.close()` in `@AfterEach tearDown()` passes without throwing `IllegalStateException: Memory was leaked by query`. ########## flight/flight-core/src/main/java/org/apache/arrow/flight/ArrowMessage.java: ########## @@ -312,6 +314,7 @@ private static ArrowMessage frame(BufferAllocator allocator, final InputStream s case APP_METADATA_TAG: { int size = readRawVarint32(stream); + checkFieldLength(size, stream); Review Comment: Since `appMetadata = allocator.buffer(size);` allocates off-heap right below here, if a subsequent field (such as `BODY_TAG`) fails `checkFieldLength` or throws an IO error, `frame()` aborts without returning an `ArrowMessage`. Because `ArrowMessage.close()` is never called, `appMetadata` is leaked in `allocator`. Also, if a frame contains multiple `APP_METADATA_TAG` fields, any previous buffer is overwritten without being closed (unlike `BODY_TAG`, which explicitly releases any prior reference). Could we make `frame()` exception-safe by: 1. Releasing any existing `appMetadata` before allocating a new one: ```java if (appMetadata != null) { appMetadata.close(); appMetadata = null; } ``` 2. In the `catch (Exception ioe)` block the `frame()`, closing any allocated buffers before rethrowing: ```java } catch (Exception ioe) { try { AutoCloseables.close(appMetadata, body); } catch (Exception e) { ioe.addSuppressed(e); } throw new RuntimeException(ioe); } ``` ########## 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: If `size < 0`, reporting `"field length -1 exceeds %d bytes remaining in the message"` is a bit confusing since the length didn't exceed the bytes remaining, but is negative. Could we split this or review the message? For example: ```java if (size < 0) { throw new IOException("Malformed FlightData frame: field length is negative: " + size); } if (size > remaining) { throw new IOException( String.format( "Malformed FlightData frame: field length %d exceeds %d bytes remaining in the message", size, remaining)); } ``` -- 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]
