From a6969c046a1983a32c83d3282606cc93c6971201 Mon Sep 17 00:00:00 2001 From: Laurent Goujon Date: Sun, 11 Oct 2026 10:56:03 +0100 Subject: [PATCH] GH-1337: Skip unknown fields in FlightData parsing (ArrowMessage.frame) ArrowMessage.frame() ignored unknown fields by reading their tag but not their payload, so the payload was parsed as the next tag. Skip the value according to its wire type instead, as protobuf requires: varints, fixed32, fixed64, length-delimited fields and groups (recursively, up to a configurable nesting depth). Co-Authored-By: Claude Sonnet 5.5 --- .../org/apache/arrow/flight/ArrowMessage.java | 149 ++++++++++- .../apache/arrow/flight/TestArrowMessage.java | 231 ++++++++++++++++++ 2 files changed, 378 insertions(+), 2 deletions(-) diff --git a/flight/flight-core/src/main/java/org/apache/arrow/flight/ArrowMessage.java b/flight/flight-core/src/main/java/org/apache/arrow/flight/ArrowMessage.java index d1c70031ea..29188ca1d6 100644 --- a/flight/flight-core/src/main/java/org/apache/arrow/flight/ArrowMessage.java +++ b/flight/flight-core/src/main/java/org/apache/arrow/flight/ArrowMessage.java @@ -55,10 +55,15 @@ import org.apache.arrow.vector.ipc.message.MessageSerializer; import org.apache.arrow.vector.types.MetadataVersion; import org.apache.arrow.vector.types.pojo.Schema; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; /** The in-memory representation of FlightData used to manage a stream of Arrow messages. */ class ArrowMessage implements AutoCloseable { + // Must be declared before the static initializer below, which logs invalid settings. + private static final Logger logger = LoggerFactory.getLogger(ArrowMessage.class); + // If true, deserialize Arrow data by giving Arrow a reference to the underlying gRPC buffer // instead of copying the data. Defaults to true. public static final boolean ENABLE_ZERO_COPY_READ; @@ -66,6 +71,10 @@ class ArrowMessage implements AutoCloseable { // instead of copying the data. Defaults to false. public static final boolean ENABLE_ZERO_COPY_WRITE; + // Maximum nesting depth of groups in unknown fields that are skipped. Defaults to 100, like the + // recursion limit of protobuf-java. + static final int MAX_PROTOBUF_SKIP_GROUP_DEPTH; + static { String zeroCopyReadFlag = System.getProperty("arrow.flight.enable_zero_copy_read"); if (zeroCopyReadFlag == null) { @@ -77,6 +86,33 @@ class ArrowMessage implements AutoCloseable { } ENABLE_ZERO_COPY_READ = !"false".equalsIgnoreCase(zeroCopyReadFlag); ENABLE_ZERO_COPY_WRITE = "true".equalsIgnoreCase(zeroCopyWriteFlag); + + String maxGroupDepth = System.getProperty("arrow.flight.max_protobuf_skip_group_depth"); + if (maxGroupDepth == null) { + maxGroupDepth = System.getenv("ARROW_FLIGHT_MAX_PROTOBUF_SKIP_GROUP_DEPTH"); + } + MAX_PROTOBUF_SKIP_GROUP_DEPTH = parseMaxGroupDepth(maxGroupDepth); + } + + private static final int DEFAULT_MAX_PROTOBUF_SKIP_GROUP_DEPTH = 100; + + private static int parseMaxGroupDepth(String value) { + if (value == null) { + return DEFAULT_MAX_PROTOBUF_SKIP_GROUP_DEPTH; + } + try { + final int depth = Integer.parseInt(value.trim()); + if (depth >= 0) { + return depth; + } + } catch (NumberFormatException e) { + // fall through to the default + } + logger.warn( + "Invalid maximum protobuf group depth '{}', using the default of {}", + value, + DEFAULT_MAX_PROTOBUF_SKIP_GROUP_DEPTH); + return DEFAULT_MAX_PROTOBUF_SKIP_GROUP_DEPTH; } private static final int DESCRIPTOR_TAG = @@ -88,6 +124,18 @@ class ArrowMessage implements AutoCloseable { private static final int APP_METADATA_TAG = (FlightData.APP_METADATA_FIELD_NUMBER << 3) | WireFormat.WIRETYPE_LENGTH_DELIMITED; + // Wire sizes of protobuf values. protobuf-java declares these in WireFormat with the same names, + // but they are package-private. + + /** Size in bytes of a fixed32 value. */ + private static final int FIXED32_SIZE = 4; + + /** Size in bytes of a fixed64 value. */ + private static final int FIXED64_SIZE = 8; + + /** Maximum size in bytes of an encoded varint, i.e. of a 64-bit value. */ + private static final int MAX_VARINT_SIZE = 10; + private static final Marshaller NO_BODY_MARSHALLER = ProtoUtils.marshaller(FlightData.getDefaultInstance()); @@ -325,7 +373,10 @@ private static ArrowMessage frame(BufferAllocator allocator, final InputStream s break; default: - // ignore unknown fields. + // Anything else is an unknown field, including a known field number sent with another + // wire type than expected. As in protobuf, it is skipped, not rejected. + skipField(stream, tag); + break; } } // Protobuf implementations can omit empty fields, such as body; for some message types, like @@ -432,7 +483,101 @@ private static ArrowBuf readFieldBuffer(BufferAllocator allocator, InputStream s return buf; } - private static IOException fieldTooLong(int size, int remaining) { + /** Skip the value of a field whose tag has just been read, without buffering it. */ + private static void skipField(InputStream stream, int tag) throws IOException { + skipField(stream, tag, 0); + } + + private static void skipField(InputStream stream, int tag, int depth) throws IOException { + switch (WireFormat.getTagWireType(tag)) { + case WireFormat.WIRETYPE_VARINT: + skipVarint(stream); + break; + case WireFormat.WIRETYPE_FIXED64: + skipFully(stream, FIXED64_SIZE); + break; + case WireFormat.WIRETYPE_LENGTH_DELIMITED: + skipFully(stream, readFieldLength(stream)); + break; + case WireFormat.WIRETYPE_START_GROUP: + if (depth >= MAX_PROTOBUF_SKIP_GROUP_DEPTH) { + throw new IOException( + "Malformed FlightData frame: groups nested deeper than " + + MAX_PROTOBUF_SKIP_GROUP_DEPTH + + " levels"); + } + skipGroup(stream, WireFormat.getTagFieldNumber(tag), depth + 1); + break; + case WireFormat.WIRETYPE_FIXED32: + skipFully(stream, FIXED32_SIZE); + break; + default: + // An end group without a matching start group, or an invalid wire type. + throw new IOException( + "Malformed FlightData frame: unexpected wire type " + + WireFormat.getTagWireType(tag) + + " for field " + + WireFormat.getTagFieldNumber(tag)); + } + } + + /** Skip the fields of a group up to and including its matching end group tag. */ + private static void skipGroup(InputStream stream, int fieldNumber, int depth) throws IOException { + while (true) { + final int firstByte = stream.read(); + if (firstByte == -1) { + throw new IOException("Malformed FlightData frame: unterminated group " + fieldNumber); + } + final int tag = readRawVarint32(firstByte, stream); + if (tag == 0) { + throw new IOException("Malformed FlightData frame: invalid tag 0 in group " + fieldNumber); + } + if (WireFormat.getTagWireType(tag) == WireFormat.WIRETYPE_END_GROUP) { + if (WireFormat.getTagFieldNumber(tag) != fieldNumber) { + throw new IOException( + "Malformed FlightData frame: group " + + fieldNumber + + " ended by group " + + WireFormat.getTagFieldNumber(tag)); + } + return; + } + skipField(stream, tag, depth); + } + } + + /** Skip a varint of up to {@link #MAX_VARINT_SIZE} bytes. */ + private static void skipVarint(InputStream stream) throws IOException { + for (int i = 0; i < MAX_VARINT_SIZE; i++) { + final int b = stream.read(); + if (b == -1) { + throw new IOException("Malformed FlightData frame: truncated varint"); + } + if ((b & 0x80) == 0) { + return; + } + } + throw new IOException( + "Malformed FlightData frame: varint longer than " + MAX_VARINT_SIZE + " bytes"); + } + + /** Skip exactly {@code size} bytes, failing if the stream ends first. */ + private static void skipFully(InputStream stream, long size) throws IOException { + long remaining = size; + while (remaining > 0) { + final long skipped = stream.skip(remaining); + if (skipped > 0) { + remaining -= skipped; + } else if (stream.read() == -1) { + // skip() may return 0 without being at EOF, so probe with a single-byte read. + throw fieldTooLong(size, size - remaining); + } else { + remaining--; + } + } + } + + private static IOException fieldTooLong(long size, long remaining) { return new IOException( "Malformed FlightData frame: field length " + size diff --git a/flight/flight-core/src/test/java/org/apache/arrow/flight/TestArrowMessage.java b/flight/flight-core/src/test/java/org/apache/arrow/flight/TestArrowMessage.java index 41e39ad2d6..ab1fd9c68d 100644 --- a/flight/flight-core/src/test/java/org/apache/arrow/flight/TestArrowMessage.java +++ b/flight/flight-core/src/test/java/org/apache/arrow/flight/TestArrowMessage.java @@ -19,6 +19,7 @@ import static org.junit.jupiter.api.Assertions.assertArrayEquals; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertInstanceOf; +import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertThrows; import com.google.common.collect.Iterables; @@ -31,7 +32,11 @@ import java.io.ByteArrayOutputStream; import java.io.IOException; import java.io.InputStream; +import java.util.ArrayList; +import java.util.Arrays; import java.util.Collections; +import java.util.List; +import java.util.stream.Stream; import java.util.zip.GZIPInputStream; import java.util.zip.GZIPOutputStream; import org.apache.arrow.flight.impl.Flight.FlightData; @@ -45,8 +50,11 @@ import org.apache.arrow.vector.types.pojo.Schema; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Named; import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; import org.junit.jupiter.params.provider.EnumSource; +import org.junit.jupiter.params.provider.MethodSource; /** Tests for deframing FlightData messages in {@link ArrowMessage}. */ public class TestArrowMessage { @@ -184,6 +192,229 @@ public void frameReleasesBuffersWhenLaterFieldIsRejected(StreamType streamType) } } + /** Writes one field into a frame. */ + @FunctionalInterface + interface FieldWriter { + void write(CodedOutputStream out) throws IOException; + } + + private static final Schema SCHEMA = + new Schema(Collections.singletonList(Field.nullable("foo", new ArrowType.Int(32, true)))); + + /** + * Payload whose bytes look like the descriptor, header and app_metadata tags (0x0A, 0x12, 0x1A). + * A parser that doesn't skip the field reads these as tags, and the following byte as a length. + */ + private static byte[] tagLikePayload(int length) { + final byte[] tags = {0x0A, 0x12, 0x1A}; + final byte[] payload = new byte[length]; + for (int i = 0; i < length; i++) { + payload[i] = tags[i % tags.length]; + } + return payload; + } + + /** + * Fields the parser doesn't know, or knows with another wire type. Every value is chosen so that + * its bytes can be mistaken for a tag if the field is not skipped properly. + */ + static Stream unknownFields() { + final List> fields = new ArrayList<>(); + // Varints equal to a descriptor, header and app_metadata tag; and values that need 10 bytes. + for (long value : new long[] {10, 18, 26, -1L, Long.MIN_VALUE}) { + fields.add(Named.of("varint " + value, out -> out.writeUInt64(4, value))); + } + for (int length : new int[] {0, 10, 18, 26}) { + fields.add( + Named.of( + "length-delimited of " + length + " bytes", + out -> out.writeByteArray(5, tagLikePayload(length)))); + } + fields.add(Named.of("fixed64", out -> out.writeFixed64(6, 0x0A121A0A121A0A12L))); + fields.add(Named.of("fixed32", out -> out.writeFixed32(7, 0x0A121A0A))); + fields.add( + Named.of( + "group", + out -> { + out.writeTag(8, WireFormat.WIRETYPE_START_GROUP); + out.writeUInt64(1, 10); + out.writeByteArray(2, tagLikePayload(18)); + out.writeTag(8, WireFormat.WIRETYPE_END_GROUP); + })); + fields.add( + Named.of( + "group in group", + out -> { + out.writeTag(8, WireFormat.WIRETYPE_START_GROUP); + out.writeTag(9, WireFormat.WIRETYPE_START_GROUP); + out.writeFixed32(1, 0x0A121A0A); + out.writeTag(9, WireFormat.WIRETYPE_END_GROUP); + out.writeUInt64(1, 26); + out.writeTag(8, WireFormat.WIRETYPE_END_GROUP); + })); + // Known field numbers with a wire type other than the expected one. As in protobuf, these are + // treated as unknown fields: they are skipped, not rejected. + fields.add( + Named.of( + "descriptor field as varint", + out -> out.writeUInt64(FlightData.FLIGHT_DESCRIPTOR_FIELD_NUMBER, 10))); + fields.add( + Named.of( + "header field as fixed32", + out -> out.writeFixed32(FlightData.DATA_HEADER_FIELD_NUMBER, 18))); + fields.add( + Named.of( + "app_metadata field as fixed64", + out -> out.writeFixed64(FlightData.APP_METADATA_FIELD_NUMBER, 26))); + fields.add( + Named.of( + "body field as varint", out -> out.writeUInt64(FlightData.DATA_BODY_FIELD_NUMBER, 26))); + + // Create a test case for each stream type. + return Arrays.stream(StreamType.values()) + .flatMap(streamType -> fields.stream().map(field -> Arguments.of(streamType, field))); + } + + /** A frame of the known fields, with {@code unknown} written before the known field at index. */ + private static byte[] frameWithUnknownField(FieldWriter unknown, int index) throws IOException { + final FieldWriter[] known = { + out -> + out.writeMessage( + FlightData.FLIGHT_DESCRIPTOR_FIELD_NUMBER, + FlightDescriptor.command(FIRST).toProtocol()), + out -> + out.writeBytes( + FlightData.DATA_HEADER_FIELD_NUMBER, + ByteString.copyFrom(MessageSerializer.serializeMetadata(SCHEMA, IpcOption.DEFAULT))), + out -> out.writeByteArray(FlightData.APP_METADATA_FIELD_NUMBER, LAST), + }; + final ByteArrayOutputStream frame = new ByteArrayOutputStream(); + final CodedOutputStream out = CodedOutputStream.newInstance(frame); + for (int i = 0; i <= known.length; i++) { + if (i == index) { + unknown.write(out); + } + if (i < known.length) { + known[i].write(out); + } + } + out.flush(); + return frame.toByteArray(); + } + + private void assertParsesAsKnownFields(StreamType streamType, byte[] frame) throws Exception { + try (ArrowMessage message = marshaller.parse(streamType.open(frame))) { + assertEquals(FlightDescriptor.command(FIRST).toProtocol(), message.getDescriptor()); + assertEquals(SCHEMA, message.asSchema()); + assertArrayEquals(LAST, toByteArray(message.getApplicationMetadata())); + } + assertEquals(0, allocator.getAllocatedMemory()); + } + + /** An unknown field after the known ones is skipped. */ + @ParameterizedTest + @MethodSource("unknownFields") + public void frameSkipsUnknownField(StreamType streamType, FieldWriter unknown) throws Exception { + assertParsesAsKnownFields(streamType, frameWithUnknownField(unknown, 3)); + } + + /** An unknown field is skipped wherever it sits among the known ones. */ + @ParameterizedTest + @EnumSource + public void frameSkipsUnknownFieldAtAnyPosition(StreamType streamType) throws Exception { + // The payload looks like a descriptor, a header and an app_metadata tag in a row. + final FieldWriter unknown = out -> out.writeByteArray(5, tagLikePayload(26)); + for (int index = 0; index <= 3; index++) { + assertParsesAsKnownFields(streamType, frameWithUnknownField(unknown, index)); + } + } + + /** Groups nested as deeply as allowed are skipped. */ + @ParameterizedTest + @EnumSource + public void frameSkipsGroupsNestedUpToLimit(StreamType streamType) throws Exception { + final int limit = ArrowMessage.MAX_PROTOBUF_SKIP_GROUP_DEPTH; + try (ArrowMessage message = marshaller.parse(streamType.open(nestedGroups(limit, limit)))) { + assertNull(message.getApplicationMetadata()); + } + } + + /** Groups nested more deeply than allowed are rejected. */ + @ParameterizedTest + @EnumSource + public void frameRejectsGroupsNestedBeyondLimit(StreamType streamType) throws Exception { + final int limit = ArrowMessage.MAX_PROTOBUF_SKIP_GROUP_DEPTH; + assertRejected(streamType, nestedGroups(limit + 1, limit + 1)); + } + + /** A group must be closed by an end group with the same field number. */ + @ParameterizedTest + @EnumSource + public void frameRejectsGroupEndedByOtherGroup(StreamType streamType) throws Exception { + final ByteArrayOutputStream frame = new ByteArrayOutputStream(); + final CodedOutputStream out = CodedOutputStream.newInstance(frame); + out.writeTag(8, WireFormat.WIRETYPE_START_GROUP); + out.writeTag(9, WireFormat.WIRETYPE_END_GROUP); + out.flush(); + assertRejected(streamType, frame.toByteArray()); + } + + /** A group that is still open at the end of the frame is rejected. */ + @ParameterizedTest + @EnumSource + public void frameRejectsUnterminatedGroup(StreamType streamType) throws Exception { + assertRejected(streamType, nestedGroups(2, 0)); + } + + /** An end group without a start group is rejected. */ + @ParameterizedTest + @EnumSource + public void frameRejectsEndGroupWithoutStartGroup(StreamType streamType) throws Exception { + assertRejected(streamType, nestedGroups(0, 1)); + } + + /** {@code opened} start group tags (field numbers 4..) followed by {@code closed} end tags. */ + private static byte[] nestedGroups(int opened, int closed) throws IOException { + final ByteArrayOutputStream frame = new ByteArrayOutputStream(); + final CodedOutputStream out = CodedOutputStream.newInstance(frame); + for (int i = 1; i <= opened; i++) { + out.writeTag(i + 3, WireFormat.WIRETYPE_START_GROUP); + } + for (int i = closed; i >= 1; i--) { + out.writeTag(i + 3, WireFormat.WIRETYPE_END_GROUP); + } + out.flush(); + return frame.toByteArray(); + } + + /** Invalid wire types cannot be skipped and are rejected. */ + @ParameterizedTest + @EnumSource + public void frameRejectsUnskippableWireTypes(StreamType streamType) throws Exception { + for (int wireType : new int[] {WireFormat.WIRETYPE_END_GROUP, 6, 7}) { + final ByteArrayOutputStream frame = new ByteArrayOutputStream(); + final CodedOutputStream out = CodedOutputStream.newInstance(frame); + out.writeTag(5, wireType); + out.flush(); + assertRejected(streamType, frame.toByteArray()); + } + } + + /** A skipped field that runs past the end of the frame is rejected. */ + @ParameterizedTest + @EnumSource + public void frameRejectsTruncatedUnknownField(StreamType streamType) throws Exception { + final byte[][] frames = { + fieldPrefix(5, 10), + {(byte) ((6 << 3) | WireFormat.WIRETYPE_FIXED64), 1, 2, 3}, + {(byte) ((7 << 3) | WireFormat.WIRETYPE_FIXED32), 1}, + {(byte) ((4 << 3) | WireFormat.WIRETYPE_VARINT), (byte) 0x80}, + }; + for (byte[] frame : frames) { + assertRejected(streamType, frame); + } + } + private void assertRejected(StreamType streamType, byte[] frame) throws IOException { final InputStream stream = streamType.open(frame); final RuntimeException e = assertThrows(RuntimeException.class, () -> marshaller.parse(stream));