Support StreamReadConstraints.maxDocumentLength and maxTokenCount in Avro parser

avro/src/main/java/tools/jackson/dataformat/avro/AvroFactory.java

@@ -260,6 +260,8 @@ protected AvroParser _createParser(ObjectReadContext readCtxt, IOContext ioCtxt,
             byte[] data, int offset, int len)
         throws JacksonException
     {
+        // [core#1548] Validate doc length up front for fixed buffers
+        _streamReadConstraints.validateDocumentLength(len);
         if (_useApacheLibDecoder) {
             return new ApacheAvroParserImpl(readCtxt, ioCtxt,
                     readCtxt.getStreamReadFeatures(_streamReadFeatures),

avro/src/main/java/tools/jackson/dataformat/avro/apacheimpl/ApacheAvroParserImpl.java

@@ -96,6 +96,11 @@ public ApacheAvroParserImpl(ObjectReadContext readCtxt, IOContext ioCtxt,
 
         final boolean buffering = AvroReadFeature.AVRO_BUFFERING.enabledIn(avroFeatures);
         BinaryDecoder decoderToReuse = apacheCodecRecycler.acquireDecoder();
+        // Apache decoder does its own buffering, so document length constraint
+        // has to be applied by counting bytes it pulls from the stream
+        if (_streamReadConstraints.hasMaxDocumentLength()) {
+            in = new LengthCheckingInputStream(in, _streamReadConstraints);
+        }
         _decoder = buffering
                 ? DECODER_FACTORY.binaryDecoder(in, decoderToReuse)
                 : DECODER_FACTORY.directBinaryDecoder(in, decoderToReuse);
@@ -418,4 +423,56 @@ protected JsonToken setString(String str) {
         _textValue = str;
         return JsonToken.VALUE_STRING;
     }
+
+    /*
+    /**********************************************************************
+    /* Helper classes
+    /**********************************************************************
+     */
+
+    /**
+     * {@link InputStream} wrapper that applies {@link StreamReadConstraints#validateDocumentLength}
+     * to the number of bytes read so far, for use with Apache {@link BinaryDecoder} which
+     * reads from the stream directly.
+     */
+    private final static class LengthCheckingInputStream extends FilterInputStream
+    {
+        private final StreamReadConstraints _constraints;
+
+        private long _bytesRead;
+
+        LengthCheckingInputStream(InputStream in, StreamReadConstraints constraints) {
+            super(in);
+            _constraints = constraints;
+        }
+
+        @Override
+        public int read() throws IOException {
+            int b = in.read();
+            if (b >= 0) {
+                _constraints.validateDocumentLength(++_bytesRead);
+            }
+            return b;
+        }
+
+        @Override
+        public int read(byte[] b, int off, int len) throws IOException {
+            int count = in.read(b, off, len);
+            if (count > 0) {
+                _bytesRead += count;
+                _constraints.validateDocumentLength(_bytesRead);
+            }
+            return count;
+        }
+
+        @Override
+        public long skip(long n) throws IOException {
+            long count = in.skip(n);
+            if (count > 0) {
+                _bytesRead += count;
+                _constraints.validateDocumentLength(_bytesRead);
+            }
+            return count;
+        }
+    }
 }

avro/src/main/java/tools/jackson/dataformat/avro/deser/AvroParserImpl.java

@@ -75,8 +75,7 @@ public JsonToken nextToken() throws JacksonException
         } catch (IOException e) {
             throw _wrapIOFailure(e);
         }
-        _currToken = t;
-        return t;
+        return _nullSafeUpdateToken(t);
     }
 
     /**
@@ -162,7 +161,7 @@ public int nextNameMatch(PropertyNameMatcher matcher) throws JacksonException
             throw _wrapIOFailure(e);
         }
         // 20-Dec-2017, tatu: not sure check would be any faster
-        _currToken = _avroContext.currentToken();
+        _nullSafeUpdateToken(_avroContext.currentToken());
 /*
         if (match < 0) { // END_OBJECT, mismatching PROPERTY_NAME or something else:
             _currToken = _avroContext.currentToken();

avro/src/main/java/tools/jackson/dataformat/avro/deser/JacksonAvroParserImpl.java

@@ -1116,6 +1116,7 @@ protected final boolean _loadMore() throws IOException
                 throw _wrapIOFailure(e); 
             }
             _currInputProcessed += _inputEnd;
+            _streamReadConstraints.validateDocumentLength(_currInputProcessed);
             _inputPtr = 0;
             if (count > 0) {
                 _inputEnd = count;
@@ -1142,6 +1143,7 @@ protected final void _loadToHaveAtLeast(int minAvailable) throws IOException
         // Need to move remaining data in front?
         int amount = _inputEnd - _inputPtr;
         _currInputProcessed += _inputPtr;
+        _streamReadConstraints.validateDocumentLength(_currInputProcessed);
         if (_inputPtr > 0) {
             if (amount > 0) {
                 //_currInputRowStart -= _inputPtr;

avro/src/test/java/tools/jackson/dataformat/avro/constraints/LongDocumentAvroReadTest.java

@@ -0,0 +1,104 @@
+package tools.jackson.dataformat.avro.constraints;
+
+import java.io.ByteArrayInputStream;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.UUID;
+
+import org.junit.jupiter.api.Test;
+
+import tools.jackson.core.JsonParser;
+import tools.jackson.core.StreamReadConstraints;
+import tools.jackson.core.exc.StreamConstraintsException;
+
+import tools.jackson.dataformat.avro.*;
+
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.junit.jupiter.api.Assertions.fail;
+
+// [dataformats-binary#785]: `StreamReadConstraints.maxDocumentLength` for Avro
+public class LongDocumentAvroReadTest extends AvroTestBase
+{
+    public static class Item {
+        public String id;
+        public int size;
+        public long stuff;
+    }
+
+    public static class Items {
+        public List<Item> items = new ArrayList<>();
+    }
+
+    private final static int MAX_DOC_LEN = 50_000;
+
+    private final AvroMapper MAPPER_VANILLA = newMapper();
+
+    private final AvroSchema ITEMS_SCHEMA = MAPPER_VANILLA.schemaFor(Items.class);
+
+    @Test
+    public void testLongDocumentConstraint() throws Exception
+    {
+        // Need a bit longer than minimum since checking is approximate, not exact
+        byte[] doc = createBigDoc(60_000);
+        for (boolean apache : new boolean[] { false, true }) {
+            AvroMapper mapper = constrainedMapper(apache);
+            _testLongDocumentConstraint(mapper, doc, true);
+            _testLongDocumentConstraint(mapper, doc, false);
+        }
+    }
+
+    @Test
+    public void testLongDocumentNoConstraint() throws Exception
+    {
+        byte[] doc = createBigDoc(60_000);
+        for (AvroMapper mapper : new AvroMapper[] { MAPPER_VANILLA, newApacheMapper() }) {
+            try (JsonParser p = mapper.reader().with(ITEMS_SCHEMA).createParser(new ByteArrayInputStream(doc))) {
+                while (p.nextToken() != null) { }
+            }
+            try (JsonParser p = mapper.reader().with(ITEMS_SCHEMA).createParser(doc)) {
+                while (p.nextToken() != null) { }
+            }
+        }
+    }
+
+    private void _testLongDocumentConstraint(AvroMapper mapper, byte[] doc, boolean stream)
+        throws Exception
+    {
+        // note: fixed-buffer case fails already on `createParser()`
+        try (JsonParser p = stream
+                ? mapper.reader().with(ITEMS_SCHEMA).createParser(new ByteArrayInputStream(doc))
+                : mapper.reader().with(ITEMS_SCHEMA).createParser(doc)) {
+            while (p.nextToken() != null) { }
+            fail("expected StreamConstraintsException");
+        } catch (StreamConstraintsException e) {
+            final String msg = e.getMessage();
+            assertTrue(msg.contains("Document length ("), "unexpected message: "+msg);
+            assertTrue(msg.contains("exceeds the maximum allowed ("+MAX_DOC_LEN), "unexpected message: "+msg);
+        }
+    }
+
+    private AvroMapper constrainedMapper(boolean apacheDecoder) {
+        AvroFactoryBuilder b = apacheDecoder
+                ? AvroFactory.builderWithApacheDecoder() : AvroFactory.builder();
+        return new AvroMapper(b
+                .streamReadConstraints(StreamReadConstraints.builder()
+                    .maxDocumentLength(MAX_DOC_LEN).build())
+                .build());
+    }
+
+    private byte[] createBigDoc(final int size) throws Exception
+    {
+        Items items = new Items();
+        // Each Item is ~50 bytes encoded; over-estimate count to be safe
+        for (int i = 0, len = size / 40; i < len; ++i) {
+            Item item = new Item();
+            item.id = UUID.randomUUID().toString();
+            item.size = i;
+            item.stuff = Long.MAX_VALUE;
+            items.items.add(item);
+        }
+        byte[] doc = MAPPER_VANILLA.writer(ITEMS_SCHEMA).writeValueAsBytes(items);
+        assertTrue(doc.length > size, "doc.length="+doc.length);
+        return doc;
+    }
+}

avro/src/test/java/tools/jackson/dataformat/avro/constraints/TokenCountAvroReadTest.java

@@ -0,0 +1,86 @@
+package tools.jackson.dataformat.avro.constraints;
+
+import java.io.ByteArrayInputStream;
+import java.util.ArrayList;
+import java.util.List;
+
+import org.junit.jupiter.api.Test;
+
+import tools.jackson.core.JsonParser;
+import tools.jackson.core.StreamReadConstraints;
+import tools.jackson.core.exc.StreamConstraintsException;
+
+import tools.jackson.dataformat.avro.*;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.fail;
+
+// [dataformats-binary#785]: `StreamReadConstraints.maxTokenCount` for Avro
+public class TokenCountAvroReadTest extends AvroTestBase
+{
+    public static class Ints {
+        public List<Integer> values = new ArrayList<>();
+    }
+
+    private final AvroMapper MAPPER = newMapper();
+
+    private final AvroSchema SCHEMA = MAPPER.schemaFor(Ints.class);
+
+    // Verify token count is tracked accurately
+    @Test
+    public void testTokenCountIsTracked() throws Exception
+    {
+        // {"values":[1,2,3]}: START_OBJECT, PROPERTY_NAME, START_ARRAY,
+        // VALUE_NUMBER_INT x3, END_ARRAY, END_OBJECT = 8 tokens
+        byte[] doc = createDoc(3);
+        for (AvroMapper mapper : new AvroMapper[] {
+                mapperWithMaxTokenCount(false, Long.MAX_VALUE),
+                mapperWithMaxTokenCount(true, Long.MAX_VALUE) }) {
+            try (JsonParser p = mapper.reader().with(SCHEMA).createParser(doc)) {
+                assertEquals(0L, p.currentTokenCount());
+                while (p.nextToken() != null) { }
+                assertEquals(8L, p.currentTokenCount());
+            }
+        }
+    }
+
+    @Test
+    public void testTokenCountLimit() throws Exception
+    {
+        // createDoc(100) produces 100 + 5 tokens
+        byte[] doc = createDoc(100);
+        for (boolean apache : new boolean[] { false, true }) {
+            AvroMapper mapper = mapperWithMaxTokenCount(apache, 10);
+            _testTokenCountLimit(mapper.reader().with(SCHEMA).createParser(doc));
+            _testTokenCountLimit(mapper.reader().with(SCHEMA).createParser(new ByteArrayInputStream(doc)));
+        }
+    }
+
+    private void _testTokenCountLimit(JsonParser p) throws Exception
+    {
+        try (p) {
+            while (p.nextToken() != null) { }
+            fail("expected StreamConstraintsException");
+        } catch (StreamConstraintsException e) {
+            verifyException(e, "Token count");
+            verifyException(e, "exceeds the maximum allowed (10,");
+        }
+    }
+
+    private AvroMapper mapperWithMaxTokenCount(boolean apacheDecoder, long maxTokenCount) {
+        AvroFactoryBuilder b = apacheDecoder
+                ? AvroFactory.builderWithApacheDecoder() : AvroFactory.builder();
+        return new AvroMapper(b
+                .streamReadConstraints(StreamReadConstraints.builder()
+                    .maxTokenCount(maxTokenCount).build())
+                .build());
+    }
+
+    private byte[] createDoc(int numValues) throws Exception {
+        Ints ints = new Ints();
+        for (int i = 0; i < numValues; i++) {
+            ints.values.add(i);
+        }
+        return MAPPER.writer(SCHEMA).writeValueAsBytes(ints);
+    }
+}

release-notes/CREDITS

@@ -76,3 +76,6 @@ PJ Fanning (@pjfanning)
 * Contributed #763: (protobuf) Use `VarHandle` for multi-byte primitive reads
   and writes in `ProtobufParser` / `ProtobufGenerator`
  (3.3.0)
+* Contributed #785: (avro) Support `StreamReadConstraints.maxDocumentLength` and
+  `maxTokenCount` in Avro parser
+ (3.3.0)

release-notes/VERSION

@@ -38,6 +38,9 @@ implementations)
 #767: (smile) Use more efficient `String` construction wrt "Compact Strings"
   for "short" ASCII text values of async parser
  (fix by @cowtowncoder, w/ Claude code)
+#785: (avro) Support `StreamReadConstraints.maxDocumentLength` and `maxTokenCount`
+  in Avro parser
+ (contributed by @pjfanning)
 
 3.2.3 (not yet released)