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));