GH-1337: Skip unknown fields in FlightData parsing (ArrowMessage.frame)

flight/flight-core/src/main/java/org/apache/arrow/flight/ArrowMessage.java

@@ -55,17 +55,26 @@
 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;
   // If true, serialize Arrow data by giving gRPC a reference to the underlying Arrow buffer
   // 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<FlightData> 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