Skip to content

Commit b64073d

Browse files
jbonofrejbonofre
committed
fix(vector): release JSON reader buffers when parsing a buffer fails
A buffer allocated while reading a JSON buffer leaked if parsing failed before the buffer was returned, for example on an invalid value or an extra value before the closing bracket. Release it on failure, including for view and variadic buffers. Also reject a null count with IllegalArgumentException instead of a NullPointerException.
1 parent 11adc4a commit b64073d

2 files changed

Lines changed: 102 additions & 39 deletions

File tree

‎vector/src/main/java/org/apache/arrow/vector/ipc/JsonFileReader.java‎

Lines changed: 62 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -272,13 +272,29 @@ public int skip(int numBatches) throws IOException {
272272
}
273273

274274
private abstract class BufferReader {
275+
// Buffer allocated by the current read, released if the read fails before returning it
276+
private ArrowBuf allocated;
277+
275278
protected abstract ArrowBuf read(BufferAllocator allocator, int count) throws IOException;
276279

280+
protected ArrowBuf allocate(BufferAllocator allocator, long size) {
281+
allocated = allocator.buffer(size);
282+
return allocated;
283+
}
284+
277285
ArrowBuf readBuffer(BufferAllocator allocator, int count) throws IOException {
278286
readToken(START_ARRAY);
279-
ArrowBuf buf = read(allocator, count);
280-
readToken(END_ARRAY);
281-
return buf;
287+
ArrowBuf buf = null;
288+
try {
289+
buf = read(allocator, count);
290+
readToken(END_ARRAY);
291+
return buf;
292+
} catch (Throwable e) {
293+
AutoCloseables.close(e, buf != null ? buf : allocated);
294+
throw e;
295+
} finally {
296+
allocated = null;
297+
}
282298
}
283299
}
284300

@@ -294,22 +310,27 @@ List<ArrowBuf> readVariadicBuffers(BufferAllocator allocator, int variadicBuffer
294310
throws IOException {
295311
readToken(START_ARRAY);
296312
ArrayList<ArrowBuf> dataBuffers = new ArrayList<>(variadicBuffersCount);
297-
for (int i = 0; i < variadicBuffersCount; i++) {
298-
parser.nextToken();
299-
final byte[] value;
313+
try {
314+
for (int i = 0; i < variadicBuffersCount; i++) {
315+
parser.nextToken();
316+
final byte[] value;
300317

301-
String variadicStr = parser.readValueAs(String.class);
302-
if (variadicStr == null) {
303-
value = new byte[0];
304-
} else {
305-
value = decodeHexSafe(variadicStr);
306-
}
318+
String variadicStr = parser.readValueAs(String.class);
319+
if (variadicStr == null) {
320+
value = new byte[0];
321+
} else {
322+
value = decodeHexSafe(variadicStr);
323+
}
307324

308-
ArrowBuf buf = allocator.buffer(value.length);
309-
buf.writeBytes(value);
310-
dataBuffers.add(buf);
325+
ArrowBuf buf = allocator.buffer(value.length);
326+
dataBuffers.add(buf);
327+
buf.writeBytes(value);
328+
}
329+
readToken(END_ARRAY);
330+
} catch (Throwable e) {
331+
AutoCloseables.close(e, dataBuffers);
332+
throw e;
311333
}
312-
readToken(END_ARRAY);
313334
return dataBuffers;
314335
}
315336

@@ -372,11 +393,15 @@ private ArrowBuf readViewBuffers(
372393
}
373394

374395
ArrowBuf buf = allocator.buffer(bufferSize);
375-
376-
for (byte[] value : values) {
377-
buf.writeBytes(value);
396+
try {
397+
for (byte[] value : values) {
398+
buf.writeBytes(value);
399+
}
400+
readToken(END_ARRAY);
401+
} catch (Throwable e) {
402+
AutoCloseables.close(e, buf);
403+
throw e;
378404
}
379-
readToken(END_ARRAY);
380405
return buf;
381406
}
382407

@@ -387,7 +412,7 @@ private class BufferHelper {
387412
@Override
388413
protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException {
389414
final int bufferSize = getValidityBufferSizeFromCount(count);
390-
ArrowBuf buf = allocator.buffer(bufferSize);
415+
ArrowBuf buf = allocate(allocator, bufferSize);
391416

392417
// C++ integration test fails without this.
393418
buf.setZero(0, bufferSize);
@@ -407,7 +432,7 @@ protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException
407432
@Override
408433
protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException {
409434
final long size = (long) count * IntervalDayVector.TYPE_WIDTH;
410-
ArrowBuf buf = allocator.buffer(size);
435+
ArrowBuf buf = allocate(allocator, size);
411436

412437
for (int i = 0; i < count; i++) {
413438
readToken(START_OBJECT);
@@ -425,7 +450,7 @@ protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException
425450
@Override
426451
protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException {
427452
final long size = (long) count * IntervalMonthDayNanoVector.TYPE_WIDTH;
428-
ArrowBuf buf = allocator.buffer(size);
453+
ArrowBuf buf = allocate(allocator, size);
429454

430455
for (int i = 0; i < count; i++) {
431456
readToken(START_OBJECT);
@@ -444,7 +469,7 @@ protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException
444469
@Override
445470
protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException {
446471
final long size = (long) count * TinyIntVector.TYPE_WIDTH;
447-
ArrowBuf buf = allocator.buffer(size);
472+
ArrowBuf buf = allocate(allocator, size);
448473

449474
for (int i = 0; i < count; i++) {
450475
parser.nextToken();
@@ -460,7 +485,7 @@ protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException
460485
@Override
461486
protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException {
462487
final long size = (long) count * SmallIntVector.TYPE_WIDTH;
463-
ArrowBuf buf = allocator.buffer(size);
488+
ArrowBuf buf = allocate(allocator, size);
464489

465490
for (int i = 0; i < count; i++) {
466491
parser.nextToken();
@@ -476,7 +501,7 @@ protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException
476501
@Override
477502
protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException {
478503
final long size = (long) count * IntVector.TYPE_WIDTH;
479-
ArrowBuf buf = allocator.buffer(size);
504+
ArrowBuf buf = allocate(allocator, size);
480505

481506
for (int i = 0; i < count; i++) {
482507
parser.nextToken();
@@ -492,7 +517,7 @@ protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException
492517
@Override
493518
protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException {
494519
final long size = (long) count * BigIntVector.TYPE_WIDTH;
495-
ArrowBuf buf = allocator.buffer(size);
520+
ArrowBuf buf = allocate(allocator, size);
496521

497522
for (int i = 0; i < count; i++) {
498523
parser.nextToken();
@@ -509,7 +534,7 @@ protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException
509534
@Override
510535
protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException {
511536
final long size = (long) count * TinyIntVector.TYPE_WIDTH;
512-
ArrowBuf buf = allocator.buffer(size);
537+
ArrowBuf buf = allocate(allocator, size);
513538

514539
for (int i = 0; i < count; i++) {
515540
parser.nextToken();
@@ -525,7 +550,7 @@ protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException
525550
@Override
526551
protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException {
527552
final long size = (long) count * SmallIntVector.TYPE_WIDTH;
528-
ArrowBuf buf = allocator.buffer(size);
553+
ArrowBuf buf = allocate(allocator, size);
529554

530555
for (int i = 0; i < count; i++) {
531556
parser.nextToken();
@@ -541,7 +566,7 @@ protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException
541566
@Override
542567
protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException {
543568
final long size = (long) count * IntVector.TYPE_WIDTH;
544-
ArrowBuf buf = allocator.buffer(size);
569+
ArrowBuf buf = allocate(allocator, size);
545570

546571
for (int i = 0; i < count; i++) {
547572
parser.nextToken();
@@ -557,7 +582,7 @@ protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException
557582
@Override
558583
protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException {
559584
final long size = (long) count * BigIntVector.TYPE_WIDTH;
560-
ArrowBuf buf = allocator.buffer(size);
585+
ArrowBuf buf = allocate(allocator, size);
561586

562587
for (int i = 0; i < count; i++) {
563588
parser.nextToken();
@@ -574,7 +599,7 @@ protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException
574599
@Override
575600
protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException {
576601
final long size = (long) count * Float4Vector.TYPE_WIDTH;
577-
ArrowBuf buf = allocator.buffer(size);
602+
ArrowBuf buf = allocate(allocator, size);
578603

579604
for (int i = 0; i < count; i++) {
580605
parser.nextToken();
@@ -590,7 +615,7 @@ protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException
590615
@Override
591616
protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException {
592617
final long size = (long) count * Float8Vector.TYPE_WIDTH;
593-
ArrowBuf buf = allocator.buffer(size);
618+
ArrowBuf buf = allocate(allocator, size);
594619

595620
for (int i = 0; i < count; i++) {
596621
parser.nextToken();
@@ -606,7 +631,7 @@ protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException
606631
@Override
607632
protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException {
608633
final long size = (long) count * DecimalVector.TYPE_WIDTH;
609-
ArrowBuf buf = allocator.buffer(size);
634+
ArrowBuf buf = allocate(allocator, size);
610635

611636
for (int i = 0; i < count; i++) {
612637
parser.nextToken();
@@ -625,7 +650,7 @@ protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException
625650
@Override
626651
protected ArrowBuf read(BufferAllocator allocator, int count) throws IOException {
627652
final long size = (long) count * Decimal256Vector.TYPE_WIDTH;
628-
ArrowBuf buf = allocator.buffer(size);
653+
ArrowBuf buf = allocate(allocator, size);
629654

630655
for (int i = 0; i < count; i++) {
631656
parser.nextToken();
@@ -980,8 +1005,8 @@ private <T> T readNextField(String expectedFieldName, Class<T> c)
9801005
}
9811006

9821007
private int readCount() throws IOException {
983-
int count = readNextField("count", Integer.class);
984-
if (count < 0) {
1008+
Integer count = readNextField("count", Integer.class);
1009+
if (count == null || count < 0) {
9851010
throw new IllegalArgumentException("Expected count >= 0, got " + count);
9861011
}
9871012
return count;

‎vector/src/test/java/org/apache/arrow/vector/ipc/TestJSONFile.java‎

Lines changed: 40 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -597,18 +597,56 @@ public void testNegativeChildCountDoesNotLeak() throws IOException {
597597
assertReadRejectedWithoutLeak(file);
598598
}
599599

600+
@Test
601+
public void testNullColumnCountRejected() throws IOException {
602+
File file =
603+
writeJson(
604+
"{\"schema\": {\"fields\": ["
605+
+ INT_FIELD_JSON.replace("NAME", "a")
606+
+ "]}, \"batches\": [{\"count\": 1, \"columns\": ["
607+
+ "{\"name\": \"a\", \"count\": null, \"VALIDITY\": [1], \"DATA\": [5]}]}]}");
608+
assertReadRejectedWithoutLeak(file);
609+
}
610+
611+
@Test
612+
public void testExtraBufferValueDoesNotLeak() throws IOException {
613+
File file =
614+
writeJson(
615+
"{\"schema\": {\"fields\": ["
616+
+ INT_FIELD_JSON.replace("NAME", "a")
617+
+ "]}, \"batches\": [{\"count\": 1, \"columns\": ["
618+
+ "{\"name\": \"a\", \"count\": 1, \"VALIDITY\": [1, 1], \"DATA\": [5]}]}]}");
619+
assertReadRejectedWithoutLeak(file, IllegalStateException.class);
620+
}
621+
622+
@Test
623+
public void testInvalidBufferValueDoesNotLeak() throws IOException {
624+
File file =
625+
writeJson(
626+
"{\"schema\": {\"fields\": ["
627+
+ INT_FIELD_JSON.replace("NAME", "a")
628+
+ "]}, \"batches\": [{\"count\": 1, \"columns\": ["
629+
+ "{\"name\": \"a\", \"count\": 1, \"VALIDITY\": [1], \"DATA\": [\"x\"]}]}]}");
630+
assertReadRejectedWithoutLeak(file, IOException.class);
631+
}
632+
600633
/** Checks that both read methods reject the file and release everything they allocated. */
601634
private void assertReadRejectedWithoutLeak(File file) throws IOException {
635+
assertReadRejectedWithoutLeak(file, IllegalArgumentException.class);
636+
}
637+
638+
private void assertReadRejectedWithoutLeak(File file, Class<? extends Throwable> expectedType)
639+
throws IOException {
602640
try (BufferAllocator readerAllocator = allocator.newChildAllocator("reader", 0, 1 << 20);
603641
JsonFileReader reader = new JsonFileReader(file, readerAllocator)) {
604642
reader.start();
605-
assertThrows(IllegalArgumentException.class, () -> reader.read());
643+
assertThrows(expectedType, () -> reader.read());
606644
}
607645
try (BufferAllocator readerAllocator = allocator.newChildAllocator("reader", 0, 1 << 20);
608646
JsonFileReader reader = new JsonFileReader(file, readerAllocator)) {
609647
Schema schema = reader.start();
610648
try (VectorSchemaRoot root = VectorSchemaRoot.create(schema, readerAllocator)) {
611-
assertThrows(IllegalArgumentException.class, () -> reader.read(root));
649+
assertThrows(expectedType, () -> reader.read(root));
612650
}
613651
}
614652
}

0 commit comments

Comments
 (0)