jbonofre commented on code in PR #1422:
URL: https://github.com/apache/arrow-java/pull/1422#discussion_r4436656633
##########
vector/src/main/java/org/apache/arrow/vector/ipc/message/ArrowFieldNode.java:
##########
@@ -31,12 +30,14 @@ public class ArrowFieldNode implements FBSerializable {
* Constructs a new instance.
*
* @param length The number of values written.
- * @param nullCount The number of null values.
+ * @param nullCount The number of null values, or -1 if unknown.
*/
public ArrowFieldNode(long length, long nullCount) {
super();
- this.length = checkedCastToInt(length);
- this.nullCount = checkedCastToInt(nullCount);
+ Preconditions.checkArgument(length >= 0, "Expected length >= 0, got %s",
length);
Review Comment:
This check now throws inside
`MessageSerializer.deserializeRecordBatch(RecordBath, ArrowBuf)`, before that
method releases `body` at the end. A negative node length now leaks the body
buffer. Before this change the same input loaded without error and didn't leak.
I reproduced it with a stream whose `FieldNode.length` is patched to -100.
`ArrowStreamReader` and `ArrowFileReader` both throw and leave 8192 bytes
allocated, and `allocator.close()` fails with "Memory was leaked". Flight goes
through the same call: `ArrowMessage.asRecordBatch()` retains `underlying` and
relies on `deserializeRecordBatch` to release it.
Could we release `body` in a `finally`? `ArrowRecordBatch` retains the
slices it keeps, so the release is correct on both paths:
```java
public static ArrowRecordBatch deserializeRecordBatch(RecordBatch
recordBatchFB, ArrowBuf body)
throws IOException {
try {
// ... existing node/buffer/compression/variadic parsing, unchanged ...
return new ArrowRecordBatch(
checkedCastToInt(recordBatchFB.length()),
nodes,
buffers,
bodyCompression,
variadicBufferCounts,
/*alignBuffers*/ true);
} finally {
body.getReferenceManager().release();
}
}
```
This also fixes the same leak on the existing `IOException` paths in that
method (for example node length > INT_MAX), and covers the dictionary batch and
Flight callers. `ArrowMessage` doesn't need a change. Releasing there too would
release the buffer twice.
##########
vector/src/test/java/org/apache/arrow/vector/ipc/message/TestArrowFieldNode.java:
##########
@@ -0,0 +1,65 @@
+/*
+ * 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.vector.ipc.message;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+import org.junit.jupiter.api.Test;
+
+public class TestArrowFieldNode {
+
+ @Test
+ public void testNegativeLengthRejected() {
Review Comment:
These tests cover the constructor, but not the deserialization path the PR
is about. Could we add a test that deserializes a record batch with a negative
node length and checks that nothing leaks? For example, in
`MessageSerializerTest`:
```java
@Test
public void testNegativeNodeLengthDoesNotLeakBody() {
FlatBufferBuilder builder = new FlatBufferBuilder();
RecordBatch.startNodesVector(builder, 1);
FieldNode.createFieldNode(builder, -100, 0);
int nodes = builder.endVector();
RecordBatch.startBuffersVector(builder, 0);
int buffers = builder.endVector();
RecordBatch.startRecordBatch(builder);
RecordBatch.addLength(builder, 0);
RecordBatch.addNodes(builder, nodes);
RecordBatch.addBuffers(builder, buffers);
builder.finish(RecordBatch.endRecordBatch(builder));
RecordBatch recordBatchFB =
RecordBatch.getRootAsRecordBatch(builder.dataBuffer());
try (BufferAllocator allocator = new RootAllocator()) {
ArrowBuf body = allocator.buffer(8);
assertThrows(
IllegalArgumentException.class,
() -> MessageSerializer.deserializeRecordBatch(recordBatchFB,
body));
assertEquals(0, allocator.getAllocatedMemory());
}
}
```
It needs three imports: `com.google.flatbuffers.FlatBufferBuilder`,
`org.apache.arrow.flatbuf.FieldNode` and `org.apache.arrow.flatbuf.RecordBatch`.
--
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]
