Skip to content

AVRO-4361: [java] Size collection guards from the writer's schema #4047

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
rayokota wants to merge 1 commit into
base: main
Choose a base branch
from
+299 −30
Open

AVRO-4361: [java] Size collection guards from the writer's schema #4047

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
128 changes: 113 additions & 15 deletions lang/java/avro/src/main/java/org/apache/avro/generic/GenericDatumReader.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,10 @@ public class GenericDatumReader<D> implements DatumReader<D> {
private ResolvingDecoder creatorResolver = null;
private final Thread creator;

// Whether the writer's and reader's schemas differ: -1 until computed, then 0
// or 1. Reset whenever either schema changes.
private int resolving = -1;

public GenericDatumReader() {
this(null, null, GenericData.get());
}
Expand Down Expand Up @@ -99,6 +103,7 @@ public void setSchema(Schema writer) {
}
creatorResolver = null;
fastDatumReader = null;
resolving = -1;
}

/** Get the reader's schema. */
Expand All @@ -110,6 +115,7 @@ public Schema getExpected() {
public void setExpected(Schema reader) {
this.expected = reader;
creatorResolver = null;
resolving = -1;
}

private static final ThreadLocal<Map<Schema, Map<Schema, ResolvingDecoder>>> RESOLVER_CACHE = ThreadLocalWithInitial
Expand Down Expand Up @@ -327,15 +333,16 @@ protected Object readArray(Object old, Schema expected, ResolvingDecoder in) thr
long l = in.readArrayStart();
long base = 0;
if (l > 0) {
ensureAvailableCollectionBytes(in, l, expectedType);
int minBytes = minBytesToRead(expectedType);
ensureAvailableCollectionBytes(in, l, minBytes);
// Elements whose minimum encoded size is zero (null, a zero-length fixed, a
// record whose fields are all zero-byte, or a recursive schema where the
// cycle is broken with a 0 minimum) consume no guaranteed input, so
// ensureAvailableCollectionBytes cannot bound their count from the bytes
// remaining. Cap such collections against a heap-aware limit so a tiny
// payload cannot declare a huge block count and drive an unbounded
// backing-array allocation.
boolean zeroByteElements = isZeroByteSchema(expectedType);
boolean zeroByteElements = minBytes == 0;
if (zeroByteElements) {
SystemLimitException.checkMaxCollectionAllocation(l);
}
Expand All @@ -354,7 +361,7 @@ protected Object readArray(Object old, Schema expected, ResolvingDecoder in) thr
}
}
base += l;
l = arrayNext(in, expectedType);
l = arrayNext(in, minBytes);
if (zeroByteElements && l > 0) {
SystemLimitException.checkMaxCollectionAllocation(l);
}
Expand All @@ -369,9 +376,9 @@ protected Object readArray(Object old, Schema expected, ResolvingDecoder in) thr
* Reads the next array block count and validates remaining bytes before the
* caller allocates storage.
*/
private long arrayNext(ResolvingDecoder in, Schema elementType) throws IOException {
private long arrayNext(ResolvingDecoder in, int minBytes) throws IOException {
long l = in.arrayNext();
ensureAvailableCollectionBytes(in, l, elementType);
ensureAvailableCollectionBytes(in, l, minBytes);
return l;
}

Expand Down Expand Up @@ -410,7 +417,8 @@ protected Object readMap(Object old, Schema expected, ResolvingDecoder in) throw
long l = in.readMapStart();
LogicalType logicalType = eValue.getLogicalType();
Conversion<?> conversion = getData().getConversionFor(logicalType);
ensureAvailableMapBytes(in, l, eValue);
int minValueBytes = minBytesToRead(eValue);
ensureAvailableMapBytes(in, l, minValueBytes);
Object map = newMap(old, initialCollectionCapacity(l));
if (l > 0) {
do {
Expand All @@ -424,7 +432,7 @@ protected Object readMap(Object old, Schema expected, ResolvingDecoder in) throw
addToMap(map, readMapKey(null, expected, in), readWithoutConversion(null, eValue, in));
}
}
} while ((l = mapNext(in, eValue)) > 0);
} while ((l = mapNext(in, minValueBytes)) > 0);
}
return map;
}
Expand All @@ -433,23 +441,23 @@ protected Object readMap(Object old, Schema expected, ResolvingDecoder in) throw
* Reads the next map block count and validates remaining bytes before the
* caller allocates storage.
*/
private long mapNext(ResolvingDecoder in, Schema valueType) throws IOException {
private long mapNext(ResolvingDecoder in, int minValueBytes) throws IOException {
long l = in.mapNext();
ensureAvailableMapBytes(in, l, valueType);
ensureAvailableMapBytes(in, l, minValueBytes);
return l;
}

/**
* Validates remaining bytes for a map block. Each map entry has a string key
* (at least 1 byte for the length varint) plus a value, so the minimum bytes
* per entry is {@code 1 + minBytesPerElement(valueSchema)}.
* (at least 1 byte for the length varint) plus a value of at least
* {@code minValueBytes}.
*/
private static void ensureAvailableMapBytes(Decoder decoder, long count, Schema valueSchema) throws EOFException {
private static void ensureAvailableMapBytes(Decoder decoder, long count, int minValueBytes) throws EOFException {
if (count <= 0) {
return;
}
// Map keys are always strings: at least 1 byte for the length varint
long minBytesPerEntry = 1L + minBytesPerElement(valueSchema);
long minBytesPerEntry = 1L + minValueBytes;
int remaining = decoder.remainingBytes();
if (remaining >= 0 && count * minBytesPerEntry > remaining) {
throw new EOFException("Map claims " + count + " entries with at least " + minBytesPerEntry
Expand Down Expand Up @@ -538,6 +546,86 @@ private static int minBytesPerElement(Schema schema, Set<Schema> visited) {
}
}

/**
* Returns the minimum number of bytes required to encode a value of any
* writer's schema that resolves to the given reader's schema. Used when the
* writer's element schema is not at hand, since a value can encode smaller than
* its reader's schema implies: the writer's record may lack a field the reader
* fills from its default, and a {@code float} or {@code double} may be promoted
* from a 1-byte {@code int} or {@code long}.
* <p>
* Returns 0 where a resolved value can encode to no bytes, such as
* {@code null}, a zero-length {@code fixed}, or a record whose fields without
* defaults are all zero-byte.
*/
static int minResolvedBytesPerElement(Schema readerSchema) {
return minResolvedBytesPerElement(readerSchema, Collections.newSetFromMap(new IdentityHashMap<>()));
}

private static int minResolvedBytesPerElement(Schema schema, Set<Schema> visited) {
switch (schema.getType()) {
case NULL:
return 0;
case FIXED:
// A fixed only resolves from a fixed of the same size
return schema.getFixedSize();
case RECORD:
if (!visited.add(schema)) {
return 0; // break recursion for self-referencing schemas
}
long sum = 0;
for (Schema.Field f : schema.getFields()) {
// The writer's record may lack a field that has a default
if (!f.hasDefaultValue()) {
sum += minResolvedBytesPerElement(f.schema(), visited);
if (sum >= Integer.MAX_VALUE) {
sum = Integer.MAX_VALUE;
break;
}
}
}
visited.remove(schema);
return (int) sum;
case UNION:
// A writer's union has a branch index of at least 1 byte, and any other
// writer's schema resolves to one of the branches
int min = 1;
for (Schema branch : schema.getTypes()) {
min = Math.min(min, minResolvedBytesPerElement(branch, visited));
}
return min;
default:
// BOOLEAN, INT, LONG, ENUM, STRING, BYTES, ARRAY, MAP are all >= 1 byte, and
// FLOAT and DOUBLE may be promoted from a 1-byte INT or LONG
return 1;
}
}

/**
* Returns the minimum number of bytes required to encode a value that this
* reader reads as the given reader's schema. When the writer's and reader's
* schemas are the same, this is {@link #minBytesPerElement}; otherwise it is
* {@link #minResolvedBytesPerElement}, since the value was encoded with the
* writer's schema.
*
* @param readerSchema the element (or map value) schema of the reader
* @return the minimum encoded size of a value read as {@code readerSchema}
*/
protected int minBytesToRead(Schema readerSchema) {
return isResolving() ? minResolvedBytesPerElement(readerSchema) : minBytesPerElement(readerSchema);
}

private boolean isResolving() {
int r = resolving;
if (r < 0) {
Schema writer = actual;
Schema reader = expected;
r = writer == reader || (writer != null && writer.equals(reader)) ? 0 : 1;
resolving = r;
}
return r == 1;
}

/**
* Validates that the decoder has enough remaining bytes to hold {@code count}
* elements of the given schema, assuming each element requires at least
Expand All @@ -547,14 +635,24 @@ private static int minBytesPerElement(Schema schema, Set<Schema> visited) {
* This check prevents out-of-memory errors from pre-allocating huge backing
* arrays when the source data is truncated or malicious. It is exposed so the
* fast reader ({@code FastReaderBuilder}) can apply the same guard as this
* classic reader.
* classic reader. The elements were encoded with the writer's schema, so
* {@code elementSchema} must be the writer's element schema.
*/
public static void ensureAvailableCollectionBytes(Decoder decoder, long count, Schema elementSchema)
throws EOFException {
ensureAvailableCollectionBytes(decoder, count, minBytesPerElement(elementSchema));
}

/**
* Validates that the decoder has enough remaining bytes to hold {@code count}
* elements of at least {@code minBytes} bytes each, such as computed by
* {@link #minBytesToRead}. Throws {@link EOFException} if the decoder reports
* fewer remaining bytes than required.
*/
protected static void ensureAvailableCollectionBytes(Decoder decoder, long count, int minBytes) throws EOFException {
if (count <= 0) {
return;
}
int minBytes = minBytesPerElement(elementSchema);
if (minBytes > 0) {
int remaining = decoder.remainingBytes();
if (remaining >= 0 && count * (long) minBytes > remaining) {
Expand Down
4 changes: 3 additions & 1 deletion lang/java/avro/src/main/java/org/apache/avro/io/FastReaderBuilder.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -473,7 +473,9 @@ private Optional<Class<?>> findClass(String clazz) {
@SuppressWarnings("unchecked")
private FieldReader createArrayReader(Schema readerSchema, Container action) throws IOException {
FieldReader elementReader = getReaderFor(action.elementAction, null);
Schema elementType = readerSchema.getElementType();
// The elements were encoded with the writer's schema, which may encode them
// smaller than the reader's (e.g. a field the reader fills from its default)
Schema elementType = action.writer.getElementType();

boolean zeroByteElements = GenericDatumReader.isZeroByteSchema(elementType);

Expand Down
29 changes: 15 additions & 14 deletions lang/java/avro/src/main/java/org/apache/avro/reflect/ReflectDatumReader.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -151,8 +151,9 @@ protected Object readArray(Object old, Schema expected, ResolvingDecoder in) thr
// the bytes-remaining check cannot bound. Without this a small malformed
// record mapped to a Java array field (e.g. long[]) could drive a very large
// eager allocation before any element is read.
ensureAvailableCollectionBytes(in, l, expectedType);
if (isZeroByteSchema(expectedType)) {
int minBytes = minBytesToRead(expectedType);
ensureAvailableCollectionBytes(in, l, minBytes);
if (minBytes == 0) {
SystemLimitException.checkMaxCollectionAllocation(l);
}
Object array = newArray(old, (int) l, expected);
Expand Down Expand Up @@ -199,7 +200,7 @@ private Object readPrimitiveArray(Object array, Class<?> c, long l, ResolvingDec
private Object readObjectArray(Object[] array, Schema expectedType, long l, ResolvingDecoder in) throws IOException {
LogicalType logicalType = expectedType.getLogicalType();
Conversion<?> conversion = getData().getConversionFor(logicalType);
boolean zeroByte = isZeroByteSchema(expectedType);
int minBytes = minBytesToRead(expectedType);
int index = 0;
if (logicalType != null && conversion != null) {
do {
Expand All @@ -209,7 +210,7 @@ private Object readObjectArray(Object[] array, Schema expectedType, long l, Reso
array[index] = element;
index++;
}
} while ((l = nextArrayBlock(in, expectedType, zeroByte)) > 0);
} while ((l = nextArrayBlock(in, minBytes)) > 0);
} else {
do {
int limit = index + (int) l;
Expand All @@ -218,7 +219,7 @@ private Object readObjectArray(Object[] array, Schema expectedType, long l, Reso
array[index] = element;
index++;
}
} while ((l = nextArrayBlock(in, expectedType, zeroByte)) > 0);
} while ((l = nextArrayBlock(in, minBytes)) > 0);
}
return array;
}
Expand All @@ -227,21 +228,21 @@ private Object readCollection(Collection<Object> c, Schema expectedType, long l,
throws IOException {
LogicalType logicalType = expectedType.getLogicalType();
Conversion<?> conversion = getData().getConversionFor(logicalType);
boolean zeroByte = isZeroByteSchema(expectedType);
int minBytes = minBytesToRead(expectedType);
if (logicalType != null && conversion != null) {
do {
for (int i = 0; i < l; i++) {
Object element = readWithConversion(null, expectedType, logicalType, conversion, in);
c.add(element);
}
} while ((l = nextArrayBlock(in, expectedType, zeroByte)) > 0);
} while ((l = nextArrayBlock(in, minBytes)) > 0);
} else {
do {
for (int i = 0; i < l; i++) {
Object element = readWithoutConversion(null, expectedType, in);
c.add(element);
}
} while ((l = nextArrayBlock(in, expectedType, zeroByte)) > 0);
} while ((l = nextArrayBlock(in, minBytes)) > 0);
}
return c;
}
Expand All @@ -257,15 +258,15 @@ private Object readCollection(Collection<Object> c, Schema expectedType, long l,
* logical array split across multiple blocks would otherwise pass only the
* first block's guard.
*
* @param in the decoder
* @param expectedType the array element schema
* @param zeroByte whether the element type's minimum encoded size is zero
* @param in the decoder
* @param minBytes the minimum encoded size of an element, from
* {@link #minBytesToRead}
* @return the validated next block count
*/
private long nextArrayBlock(ResolvingDecoder in, Schema expectedType, boolean zeroByte) throws IOException {
private long nextArrayBlock(ResolvingDecoder in, int minBytes) throws IOException {
long l = in.arrayNext();
ensureAvailableCollectionBytes(in, l, expectedType);
if (zeroByte && l > 0) {
ensureAvailableCollectionBytes(in, l, minBytes);
if (minBytes == 0 && l > 0) {
SystemLimitException.checkMaxCollectionAllocation(l);
}
return l;
Expand Down
Loading
Loading