Skip to content

Commit 4b2d53e

Browse files
authored
GH-1205: Fix byte-array element leak in FromSchemaByteArray (#1249)
## What's Changed `FromSchemaByteArray()` in `dataset/src/main/cpp/jni_util.cc` acquires the Java byte-array elements via `GetByteArrayElements()`, but the matching `ReleaseByteArrayElements()` was only reached on the success path. When `arrow::ipc::ReadSchema()` failed, `ARROW_ASSIGN_OR_RAISE` returned early and skipped the release, leaking the pinned/copied array buffer on every call with malformed or incompatible serialized schema bytes (e.g. via `createDataset()`). This releases the elements through an RAII guard so they are always freed on every exit path - success or error, and stays correct if additional early returns are added later. Closes #1205.
1 parent ae00889 commit 4b2d53e

2 files changed

Lines changed: 87 additions & 4 deletions

File tree

‎dataset/src/main/cpp/jni_util.cc‎

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -352,13 +352,17 @@ arrow::Result<std::shared_ptr<arrow::Schema>> FromSchemaByteArray(
352352
arrow::ipc::DictionaryMemo in_memo;
353353
int schemaBytes_len = env->GetArrayLength(schemaBytes);
354354
jbyte* schemaBytes_data = env->GetByteArrayElements(schemaBytes, nullptr);
355+
// Ensure the pinned/copied Java array elements are always released, even when
356+
// schema parsing fails and an error Result is returned early below.
357+
auto release_elements = [&](jbyte* data) {
358+
env->ReleaseByteArrayElements(schemaBytes, data, JNI_ABORT);
359+
};
360+
std::unique_ptr<jbyte, decltype(release_elements)> elements_guard(schemaBytes_data,
361+
release_elements);
355362
auto serialized_schema = std::make_shared<arrow::Buffer>(
356363
reinterpret_cast<uint8_t*>(schemaBytes_data), schemaBytes_len);
357364
arrow::io::BufferReader buf_reader(serialized_schema);
358-
ARROW_ASSIGN_OR_RAISE(std::shared_ptr<arrow::Schema> schema,
359-
arrow::ipc::ReadSchema(&buf_reader, &in_memo))
360-
env->ReleaseByteArrayElements(schemaBytes, schemaBytes_data, JNI_ABORT);
361-
return schema;
365+
return arrow::ipc::ReadSchema(&buf_reader, &in_memo);
362366
}
363367
arrow::Status ExportRecordBatch(JNIEnv* env, const std::shared_ptr<RecordBatch>& batch,
364368
jlong struct_array) {
Lines changed: 79 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,79 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one or more
3+
* contributor license agreements. See the NOTICE file distributed with
4+
* this work for additional information regarding copyright ownership.
5+
* The ASF licenses this file to You under the Apache License, Version 2.0
6+
* (the "License"); you may not use this file except in compliance with
7+
* the License. You may obtain a copy of the License at
8+
*
9+
* http://www.apache.org/licenses/LICENSE-2.0
10+
*
11+
* Unless required by applicable law or agreed to in writing, software
12+
* distributed under the License is distributed on an "AS IS" BASIS,
13+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
* See the License for the specific language governing permissions and
15+
* limitations under the License.
16+
*/
17+
package org.apache.arrow.dataset.jni;
18+
19+
import static org.junit.jupiter.api.Assertions.assertThrows;
20+
21+
import java.io.File;
22+
import java.lang.reflect.Field;
23+
import org.apache.arrow.dataset.ParquetWriteSupport;
24+
import org.apache.arrow.dataset.file.FileFormat;
25+
import org.apache.arrow.dataset.file.FileSystemDatasetFactory;
26+
import org.junit.jupiter.api.Test;
27+
import org.junit.jupiter.api.io.TempDir;
28+
29+
/**
30+
* Regression test for the native {@code FromSchemaByteArray} helper (GH-1205). Passing malformed
31+
* serialized schema bytes to {@code JniWrapper#createDataset} exercises the former leak branch in
32+
* {@code FromSchemaByteArray} (where {@code ReleaseByteArrayElements} was skipped on error) and
33+
* asserts it fails gracefully with a Java exception rather than crashing.
34+
*
35+
* <p>This guards the error path; it does not directly assert that the native byte-array elements
36+
* were released. The leaked bytes are a JVM-internal copy made by {@code GetByteArrayElements},
37+
* which is not tracked by {@code NativeMemoryPool} nor observable through any portable Java API, so
38+
* a deterministic leak assertion isn't feasible here. See the PR #1249 discussion for details.
39+
*/
40+
public class TestFromSchemaByteArray extends TestNativeDataset {
41+
42+
@TempDir public File TMP;
43+
44+
public static final String AVRO_SCHEMA_USER = "user.avsc";
45+
46+
private static long factoryId(NativeDatasetFactory factory) throws Exception {
47+
Field field = NativeDatasetFactory.class.getDeclaredField("datasetFactoryId");
48+
field.setAccessible(true);
49+
return field.getLong(factory);
50+
}
51+
52+
@Test
53+
public void testCreateDatasetWithMalformedSchemaBytes() throws Exception {
54+
ParquetWriteSupport writeSupport =
55+
ParquetWriteSupport.writeTempFile(AVRO_SCHEMA_USER, TMP, 1, "a");
56+
FileSystemDatasetFactory factory =
57+
new FileSystemDatasetFactory(
58+
rootAllocator(),
59+
NativeMemoryPool.getDefault(),
60+
FileFormat.PARQUET,
61+
writeSupport.getOutputURI());
62+
try {
63+
final long datasetFactoryId = factoryId(factory);
64+
// Bytes that are not a valid serialized Arrow schema, so native ReadSchema fails and the
65+
// error path of FromSchemaByteArray is taken.
66+
final byte[] malformedSchemaBytes = new byte[] {0, 1, 2, 3, 4, 5, 6, 7};
67+
68+
// Repeat many times to keep hitting the error path that previously skipped the element
69+
// release; each iteration must fail gracefully rather than crash.
70+
for (int i = 0; i < 1000; i++) {
71+
assertThrows(
72+
RuntimeException.class,
73+
() -> JniWrapper.get().createDataset(datasetFactoryId, malformedSchemaBytes));
74+
}
75+
} finally {
76+
factory.close();
77+
}
78+
}
79+
}

0 commit comments

Comments
 (0)