(orc) branch main updated: ORC-2199: Validate union tag against the number of children in `UnionTreeReader`

[do...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

dongjoon-hyun pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/orc.git


The following commit(s) were added to refs/heads/main by this push:
     new f769ad79a ORC-2199: Validate union tag against the number of children in `UnionTreeReader`
f769ad79a is described below

commit f769ad79ad58da1227500164f8848f13ae663b15
Author: Dongjoon Hyun <do...@apache.org>
AuthorDate: Wed Jul 8 21:19:42 2026 -0700

    ORC-2199: Validate union tag against the number of children in `UnionTreeReader`
    
    ### What changes were proposed in this pull request?
    
    This PR aims to validate the union tag against the number of children in `UnionTreeReader`.
    
    Both the `nextVector` and `skipRows` paths now check that each tag read from the `DATA` stream is within `[0, fields.length)` and throw `FileFormatException` with a clear message, consistent with the C++ reader's `getCheckedUnionTag`.
    
    ### Why are the changes needed?
    
    The Java `UnionTreeReader` used the raw tag byte from the `DATA` stream without validation. For a corrupt or malicious file with a tag greater than or equal to the number of union children (or a negative sign-extended byte), `skipRows` fails with `ArrayIndexOutOfBoundsException` and `nextVector` propagates the invalid tag downstream. The C++ reader already rejects this via `getCheckedUnionTag` (`ParseError`), so this closes a C++/Java parity gap.
    
    ### How was this patch tested?
    
    Pass the CIs with a newly added `TestTreeReaderFactory`.
    
    ### Was this patch authored or co-authored using generative AI tooling?
    
    Generated-by: Claude Fable 5
    
    Closes #2682 from dongjoon-hyun/ORC-2199.
    
    Authored-by: Dongjoon Hyun <do...@apache.org>
    Signed-off-by: Dongjoon Hyun <do...@apache.org>
---
 .../org/apache/orc/impl/TreeReaderFactory.java     |  16 ++-
 .../org/apache/orc/impl/TestTreeReaderFactory.java | 129 +++++++++++++++++++++
 2 files changed, 144 insertions(+), 1 deletion(-)

diff --git a/java/core/src/java/org/apache/orc/impl/TreeReaderFactory.java b/java/core/src/java/org/apache/orc/impl/TreeReaderFactory.java
index e413ce6bb..f6a968153 100644
--- a/java/core/src/java/org/apache/orc/impl/TreeReaderFactory.java
+++ b/java/core/src/java/org/apache/orc/impl/TreeReaderFactory.java
@@ -2741,6 +2741,13 @@ public class TreeReaderFactory {
       this.fields = childReaders;
     }
 
+    private void checkUnionTag(int tag) throws FileFormatException {
+      if (tag < 0 || tag >= fields.length) {
+        throw new FileFormatException("Invalid union tag " + tag +
+            " for union with " + fields.length + " children");
+      }
+    }
+
     @Override
     public void seek(PositionProvider[] index, ReadPhase readPhase) throws IOException {
       if (readPhase.contains(this.readerCategory)) {
@@ -2767,6 +2774,11 @@ public class TreeReaderFactory {
           result.isRepeating = false;
           tags.nextVector(result.noNulls ? null : result.isNull, result.tags,
                           batchSize);
+          for (int r = 0; r < batchSize; ++r) {
+            if (result.noNulls || !result.isNull[r]) {
+              checkUnionTag(result.tags[r]);
+            }
+          }
         }
       }
 
@@ -2808,7 +2820,9 @@ public class TreeReaderFactory {
       items = countNonNulls(items);
       long[] counts = new long[fields.length];
       for (int i = 0; i < items; ++i) {
-        counts[tags.next()] += 1;
+        int tag = tags.next();
+        checkUnionTag(tag);
+        counts[tag] += 1;
       }
       for (int i = 0; i < counts.length; ++i) {
         if (TypeReader.shouldProcessChild(fields[i], readPhase)) {
diff --git a/java/core/src/test/org/apache/orc/impl/TestTreeReaderFactory.java b/java/core/src/test/org/apache/orc/impl/TestTreeReaderFactory.java
new file mode 100644
index 000000000..dfc3298cf
--- /dev/null
+++ b/java/core/src/test/org/apache/orc/impl/TestTreeReaderFactory.java
@@ -0,0 +1,129 @@
+/*
+ * 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.orc.impl;
+
+import org.apache.hadoop.hive.ql.exec.vector.LongColumnVector;
+import org.apache.hadoop.hive.ql.exec.vector.UnionColumnVector;
+import org.apache.orc.FileFormatException;
+import org.apache.orc.OrcProto;
+import org.apache.orc.impl.reader.tree.TypeReader;
+import org.apache.orc.impl.writer.StreamOptions;
+import org.junit.jupiter.api.Test;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+public class TestTreeReaderFactory {
+
+  private static final OrcProto.ColumnEncoding DIRECT =
+      OrcProto.ColumnEncoding.newBuilder()
+          .setKind(OrcProto.ColumnEncoding.Kind.DIRECT).build();
+
+  private static InStream toInStream(TestInStream.OutputCollector collect) {
+    ByteBuffer inBuf = ByteBuffer.allocate(collect.buffer.size());
+    collect.buffer.setByteBuffer(inBuf, 0, collect.buffer.size());
+    inBuf.flip();
+    return InStream.create("test", new BufferChunk(inBuf, 0), 0,
+        inBuf.remaining());
+  }
+
+  private static RunLengthByteReader createTagReader(byte... tags)
+      throws IOException {
+    TestInStream.OutputCollector collect = new TestInStream.OutputCollector();
+    RunLengthByteWriter writer = new RunLengthByteWriter(
+        new OutStream("test", new StreamOptions(100), collect));
+    for (byte tag : tags) {
+      writer.write(tag);
+    }
+    writer.flush();
+    return new RunLengthByteReader(toInStream(collect));
+  }
+
+  private static InStream createLongStream(long... values) throws IOException {
+    TestInStream.OutputCollector collect = new TestInStream.OutputCollector();
+    RunLengthIntegerWriter writer = new RunLengthIntegerWriter(
+        new OutStream("test", new StreamOptions(100), collect), true);
+    for (long value : values) {
+      writer.write(value);
+    }
+    writer.flush();
+    return toInStream(collect);
+  }
+
+  @Test
+  public void testUnionWithValidTags() throws Exception {
+    TreeReaderFactory.Context context = new TreeReaderFactory.ReaderContext();
+    TypeReader[] children = new TypeReader[]{
+        new TreeReaderFactory.LongTreeReader(1, null,
+            createLongStream(10, 30), DIRECT, context),
+        new TreeReaderFactory.LongTreeReader(2, null,
+            createLongStream(20), DIRECT, context)};
+    TreeReaderFactory.UnionTreeReader reader =
+        new TreeReaderFactory.UnionTreeReader(0, null, context, null, children);
+    reader.tags = createTagReader((byte) 0, (byte) 1, (byte) 0);
+
+    UnionColumnVector batch = new UnionColumnVector(3,
+        new LongColumnVector(3), new LongColumnVector(3));
+    reader.nextVector(batch, null, 3, null, TypeReader.ReadPhase.ALL);
+
+    assertEquals(0, batch.tags[0]);
+    assertEquals(1, batch.tags[1]);
+    assertEquals(0, batch.tags[2]);
+    assertEquals(10, ((LongColumnVector) batch.fields[0]).vector[0]);
+    assertEquals(20, ((LongColumnVector) batch.fields[1]).vector[1]);
+    assertEquals(30, ((LongColumnVector) batch.fields[0]).vector[2]);
+  }
+
+  @Test
+  public void testUnionWithInvalidTagInNextVector() throws Exception {
+    TreeReaderFactory.Context context = new TreeReaderFactory.ReaderContext();
+    TypeReader[] children = new TypeReader[]{
+        new TreeReaderFactory.LongTreeReader(1, context),
+        new TreeReaderFactory.LongTreeReader(2, context)};
+    TreeReaderFactory.UnionTreeReader reader =
+        new TreeReaderFactory.UnionTreeReader(0, null, context, null, children);
+    reader.tags = createTagReader((byte) 0, (byte) 1, (byte) 5);
+
+    UnionColumnVector batch = new UnionColumnVector(3,
+        new LongColumnVector(3), new LongColumnVector(3));
+    FileFormatException e = assertThrows(FileFormatException.class, () ->
+        reader.nextVector(batch, null, 3, null, TypeReader.ReadPhase.ALL));
+    assertEquals("Invalid union tag 5 for union with 2 children",
+        e.getMessage());
+  }
+
+  @Test
+  public void testUnionWithInvalidTagInSkipRows() throws Exception {
+    TreeReaderFactory.Context context = new TreeReaderFactory.ReaderContext();
+    TypeReader[] children = new TypeReader[]{
+        new TreeReaderFactory.LongTreeReader(1, context),
+        new TreeReaderFactory.LongTreeReader(2, context)};
+    TreeReaderFactory.UnionTreeReader reader =
+        new TreeReaderFactory.UnionTreeReader(0, null, context, null, children);
+    reader.tags = createTagReader((byte) 0, (byte) 1, (byte) 0x80);
+
+    FileFormatException e = assertThrows(FileFormatException.class, () ->
+        reader.skipRows(3, TypeReader.ReadPhase.ALL));
+    assertEquals("Invalid union tag -128 for union with 2 children",
+        e.getMessage());
+  }
+}