Refactor maxInMemory limit handling for async XML parsing

The limit was previously enforced in XmlEventDecoder, because it is
what parses incoming buffers. However, the actual caching is in
Jaxb2Decoder, which holds on to XML events, but has no good way to
estimate their size.

After this commit XmlEventDecoder no longer enforces memory limits
for async parsing. It releases each buffer immediately anyway.

Instead XmlEventDecoder is only responsible to update the number
of bytes received via a new ReceivedByteTracker type while
Jaxb2XmlDecoder uses the same to perform limit and reset the
count depending on when it is aggregating XML events.

Closes gh-37031

spring-web/src/main/java/org/springframework/http/codec/xml/Jaxb2Helper.java

@@ -34,6 +34,7 @@ import org.jspecify.annotations.Nullable;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.SynchronousSink;
 
+import org.springframework.core.io.buffer.DataBufferLimitException;
 import org.springframework.util.Assert;
 import org.springframework.util.ClassUtils;
 
@@ -141,8 +142,10 @@ abstract class Jaxb2Helper {
 	 * </li>
 	 * </ol>
 	 */
-	public static Flux<List<XMLEvent>> split(Flux<XMLEvent> xmlEventFlux, Set<QName> names) {
-		return xmlEventFlux.handle(new SplitHandler(names));
+	public static Flux<List<XMLEvent>> split(
+			Flux<XMLEvent> xmlEventFlux, Set<QName> names, XmlEventDecoder.@Nullable ReceivedByteTracker byteTracker) {
+
+		return xmlEventFlux.handle(new SplitHandler(names, byteTracker));
 	}
 
 
@@ -150,14 +153,17 @@ abstract class Jaxb2Helper {
 
 		private final Set<QName> names;
 
+		private final XmlEventDecoder.ReceivedByteTracker byteTracker;
+
 		private @Nullable List<XMLEvent> events;
 
 		private int elementDepth = 0;
 
 		private int barrier = Integer.MAX_VALUE;
 
-		public SplitHandler(Set<QName> names) {
+		public SplitHandler(Set<QName> names, XmlEventDecoder.@Nullable ReceivedByteTracker byteTracker) {
 			this.names = names;
+			this.byteTracker = (byteTracker != null ? byteTracker : XmlEventDecoder.ReceivedByteTracker.NO_OP);
 		}
 
 		@Override
@@ -179,11 +185,19 @@ abstract class Jaxb2Helper {
 			if (event.isEndElement()) {
 				this.elementDepth--;
 				if (this.elementDepth == this.barrier) {
-					this.barrier = Integer.MAX_VALUE;
 					Assert.state(this.events != null, "No XMLEvent List");
 					sink.next(this.events);
+					this.barrier = Integer.MAX_VALUE;
+					this.events = null;
 				}
 			}
+			if (this.events == null) {
+				this.byteTracker.reset();
+			}
+			else if (this.byteTracker.isMaxInMemorySizeExceeded()) {
+				throw new DataBufferLimitException(
+						"Exceeded limit on max bytes per XML node: " + this.byteTracker.getMaxInMemorySize());
+			}
 		}
 	}
 

spring-web/src/main/java/org/springframework/http/codec/xml/Jaxb2XmlDecoder.java

@@ -145,12 +145,16 @@ public class Jaxb2XmlDecoder extends AbstractDecoder<Object> {
 	public Flux<Object> decode(Publisher<DataBuffer> inputStream, ResolvableType elementType,
 			@Nullable MimeType mimeType, @Nullable Map<String, Object> hints) {
 
+		XmlEventDecoder.ReceivedByteTracker byteTracker =
+				new XmlEventDecoder.ReceivedByteTracker(this.maxInMemorySize);
+
 		Flux<XMLEvent> xmlEventFlux = this.xmlEventDecoder.decode(
-				inputStream, ResolvableType.forClass(XMLEvent.class), mimeType, hints);
+				inputStream, ResolvableType.forClass(XMLEvent.class), mimeType,
+				Hints.merge(hints, XmlEventDecoder.BYTE_TRACKER_HINT, byteTracker));
 
 		Class<?> outputClass = elementType.toClass();
 		Set<QName> typeNames = Jaxb2Helper.toQNames(outputClass);
-		Flux<List<XMLEvent>> splitEvents = Jaxb2Helper.split(xmlEventFlux, typeNames);
+		Flux<List<XMLEvent>> splitEvents = Jaxb2Helper.split(xmlEventFlux, typeNames, byteTracker);
 
 		return splitEvents.map(events -> {
 			Object value = unmarshal(events, outputClass);

spring-web/src/main/java/org/springframework/http/codec/xml/XmlEventDecoder.java

@@ -84,6 +84,12 @@ import org.springframework.util.xml.StaxUtils;
  */
 public class XmlEventDecoder extends AbstractDecoder<XMLEvent> {
 
+	/**
+	 * Hint key for a {@link ReceivedByteTracker} instance that callers can use
+	 * to track the number of received bytes.
+	 */
+	public static final String BYTE_TRACKER_HINT = XmlEventDecoder.class.getName() + ".byteTracker";
+
 	private static final XMLInputFactory inputFactory = StaxUtils.createDefensiveInputFactory();
 
 	private static final boolean AALTO_PRESENT = ClassUtils.isPresent(
@@ -100,10 +106,13 @@ public class XmlEventDecoder extends AbstractDecoder<XMLEvent> {
 
 
 	/**
-	 * Set the max number of bytes that can be buffered by this decoder. This
-	 * is either the size the entire input when decoding as a whole, or when
-	 * using async parsing via Aalto XML, it is size one top-level XML tree.
-	 * When the limit is exceeded, {@link DataBufferLimitException} is raised.
+	 * Set the max number of bytes this decoder should buffer in memory resulting
+	 * in a {@link DataBufferLimitException} when the limit is exceeded.
+	 * <p>When joining all buffers and decoding as a whole, the limit is applied
+	 * to the entire input.
+	 * <p>>When using Aalto XML async parsing, the limit does not apply at the
+	 * level of this decoder because the XML events parsed from each buffer are
+	 * emitted immediately and the buffer is released.
 	 * <p>By default this is set to 256K.
 	 * @param byteCount the max number of bytes to buffer, or -1 for unlimited
 	 * @since 5.1.11
@@ -126,7 +135,7 @@ public class XmlEventDecoder extends AbstractDecoder<XMLEvent> {
 			@Nullable MimeType mimeType, @Nullable Map<String, Object> hints) {
 
 		if (this.useAalto) {
-			AaltoDataBufferToXmlEvent mapper = new AaltoDataBufferToXmlEvent(this.maxInMemorySize);
+			AaltoDataBufferToXmlEvent mapper = new AaltoDataBufferToXmlEvent(hints);
 			return Flux.from(input)
 					.flatMapIterable(mapper)
 					.doFinally(signalType -> mapper.endOfInput());
@@ -155,7 +164,7 @@ public class XmlEventDecoder extends AbstractDecoder<XMLEvent> {
 	/*
 	 * Separate static class to isolate Aalto dependency.
 	 */
-	private static class AaltoDataBufferToXmlEvent implements Function<DataBuffer, List<? extends XMLEvent>> {
+	private static final class AaltoDataBufferToXmlEvent implements Function<DataBuffer, List<? extends XMLEvent>> {
 
 		private static final AsyncXMLInputFactory inputFactory =
 				StaxUtils.createDefensiveInputFactory(InputFactoryImpl::new);
@@ -165,22 +174,19 @@ public class XmlEventDecoder extends AbstractDecoder<XMLEvent> {
 
 		private final XMLEventAllocator eventAllocator = EventAllocatorImpl.getDefaultInstance();
 
-		private final int maxInMemorySize;
+		@Nullable
+		private final ReceivedByteTracker byteTracker;
 
-		private int byteCount;
-
-		private int elementDepth;
-
-
-		public AaltoDataBufferToXmlEvent(int maxInMemorySize) {
-			this.maxInMemorySize = maxInMemorySize;
+		private AaltoDataBufferToXmlEvent(@Nullable Map<String, Object> hints) {
+			this.byteTracker = (hints != null ? (ReceivedByteTracker) hints.get(BYTE_TRACKER_HINT) : null);
 		}
 
-
 		@Override
 		public List<? extends XMLEvent> apply(DataBuffer dataBuffer) {
 			try {
-				increaseByteCount(dataBuffer);
+				if (this.byteTracker != null) {
+					this.byteTracker.incrementByteCount(dataBuffer);
+				}
 				AsyncByteBufferFeeder inputFeeder = this.streamReader.getInputFeeder();
 				try (DataBuffer.ByteBufferIterator iterator = dataBuffer.readableByteBuffers()) {
 					while (iterator.hasNext()) {
@@ -199,12 +205,8 @@ public class XmlEventDecoder extends AbstractDecoder<XMLEvent> {
 						if (event.isEndDocument()) {
 							break;
 						}
-						checkDepthAndResetByteCount(event);
 					}
 				}
-				if (this.maxInMemorySize > 0 && this.byteCount > this.maxInMemorySize) {
-					raiseLimitException();
-				}
 				return events;
 			}
 			catch (XMLStreamException ex) {
@@ -215,40 +217,55 @@ public class XmlEventDecoder extends AbstractDecoder<XMLEvent> {
 			}
 		}
 
-		private void increaseByteCount(DataBuffer dataBuffer) {
-			if (this.maxInMemorySize > 0) {
-				if (dataBuffer.readableByteCount() > Integer.MAX_VALUE - this.byteCount) {
-					raiseLimitException();
-				}
-				else {
-					this.byteCount += dataBuffer.readableByteCount();
-				}
-			}
-		}
-
-		private void checkDepthAndResetByteCount(XMLEvent event) {
-			if (this.maxInMemorySize > 0) {
-				if (event.isStartElement()) {
-					this.byteCount = this.elementDepth == 1 ? 0 : this.byteCount;
-					this.elementDepth++;
-				}
-				else if (event.isEndElement()) {
-					this.elementDepth--;
-					this.byteCount = this.elementDepth == 1 ? 0 : this.byteCount;
-				}
-			}
-		}
-
-		private void raiseLimitException() {
-			throw new DataBufferLimitException(
-					"Exceeded limit on max bytes per XML top-level node: " + this.maxInMemorySize);
-		}
-
 		public void endOfInput() {
 			this.streamReader.getInputFeeder().endOfInput();
 		}
 	}
 
 
+	/**
+	 * Callers of {@link XmlEventDecoder} that buffer emitted XML events at a
+	 * higher level, can pass an instance of this tracker as an
+	 * {@link XmlEventDecoder#BYTE_TRACKER_HINT} to monitor the total number of
+	 * bytes received, and to reset periodically.
+	 * <p>For use with Aalto XML async parsing only, in which case this decoder
+	 * parses releases each buffer immediately.
+	 */
+	public static class ReceivedByteTracker {
+
+		/** An instance to use when there is no limit. */
+		public static final ReceivedByteTracker NO_OP = new ReceivedByteTracker(-1);
+
+		private final int maxInMemorySize;
+
+		private int byteCount;
+
+		public ReceivedByteTracker(int maxInMemorySize) {
+			this.maxInMemorySize = maxInMemorySize;
+		}
+
+		public int getMaxInMemorySize() {
+			return this.maxInMemorySize;
+		}
+
+		public boolean isMaxInMemorySizeExceeded() {
+			return (this.maxInMemorySize != -1 && this.byteCount > this.maxInMemorySize);
+		}
+
+		public void reset() {
+			this.byteCount = 0;
+		}
+
+		private void incrementByteCount(DataBuffer buffer) {
+			if (this.maxInMemorySize != -1) {
+				this.byteCount += buffer.readableByteCount();
+			}
+		}
+
+		@Override
+		public String toString() {
+			return this.byteCount + " bytes";
+		}
+	}
 
 }

spring-web/src/test/java/org/springframework/http/codec/xml/Jaxb2XmlDecoderTests.java

@@ -35,7 +35,9 @@ import reactor.test.StepVerifier;
 
 import org.springframework.core.ResolvableType;
 import org.springframework.core.codec.DecodingException;
+import org.springframework.core.codec.Hints;
 import org.springframework.core.io.buffer.DataBuffer;
+import org.springframework.core.io.buffer.DataBufferLimitException;
 import org.springframework.core.testfixture.io.buffer.AbstractLeakCheckingTests;
 import org.springframework.http.MediaType;
 import org.springframework.util.MimeType;
@@ -92,7 +94,7 @@ class Jaxb2XmlDecoderTests extends AbstractLeakCheckingTests {
 	@Test
 	void splitOneBranches() {
 		Flux<XMLEvent> xmlEvents = this.xmlEventDecoder.decode(toDataBufferMono(POJO_ROOT), null, null, HINTS);
-		Flux<List<XMLEvent>> result = Jaxb2Helper.split(xmlEvents, Set.of(new QName("pojo")));
+		Flux<List<XMLEvent>> result = Jaxb2Helper.split(xmlEvents, Set.of(new QName("pojo")), null);
 
 		StepVerifier.create(result)
 				.consumeNextWith(events -> {
@@ -113,8 +115,7 @@ class Jaxb2XmlDecoderTests extends AbstractLeakCheckingTests {
 	@Test
 	void splitMultipleBranches() {
 		Flux<XMLEvent> xmlEvents = this.xmlEventDecoder.decode(toDataBufferMono(POJO_CHILD), null, null, HINTS);
-		Flux<List<XMLEvent>> result = Jaxb2Helper.split(xmlEvents, Set.of(new QName("pojo")));
-
+		Flux<List<XMLEvent>> result = Jaxb2Helper.split(xmlEvents, Set.of(new QName("pojo")), null);
 
 		StepVerifier.create(result)
 				.consumeNextWith(events -> {
@@ -143,6 +144,34 @@ class Jaxb2XmlDecoderTests extends AbstractLeakCheckingTests {
 				.verify();
 	}
 
+	@Test
+	void splitMultipleBranchesLimitExceeded() {
+
+		Flux<String> source = Flux.just(
+				"<?xml version=\"1.0\" encoding=\"UTF-8\"?>",
+				"<root><pojo><foo>foo</foo></pojo>",
+				"<pojo><foo>foofoo</foo>",
+				"<bar>barbar</bar></pojo>",
+				"<root/>");
+
+		XmlEventDecoder.ReceivedByteTracker byteTracker = new XmlEventDecoder.ReceivedByteTracker(30);
+		Map<String, Object> hints = Hints.from(XmlEventDecoder.BYTE_TRACKER_HINT, byteTracker);
+		Flux<XMLEvent> xmlEvents = this.xmlEventDecoder.decode(source.map(this::toToDataBuffer), null, null, hints);
+		Flux<List<XMLEvent>> result = Jaxb2Helper.split(xmlEvents, Set.of(new QName("pojo")), byteTracker);
+
+		StepVerifier.create(result)
+				.consumeNextWith(events -> {
+					assertThat(events).hasSize(5);
+					assertStartElement(events.get(0), "pojo");
+					assertStartElement(events.get(1), "foo");
+					assertCharacters(events.get(2), "foo");
+					assertEndElement(events.get(3), "foo");
+					assertEndElement(events.get(4), "pojo");
+				})
+				.expectError(DataBufferLimitException.class)
+				.verify();
+	}
+
 	private static void assertStartElement(XMLEvent event, String expectedLocalName) {
 		assertThat(event.isStartElement()).isTrue();
 		assertThat(event.asStartElement().getName().getLocalPart()).isEqualTo(expectedLocalName);
@@ -263,13 +292,18 @@ class Jaxb2XmlDecoderTests extends AbstractLeakCheckingTests {
 
 	private Mono<DataBuffer> toDataBufferMono(String value) {
 		return Mono.defer(() -> {
-			byte[] bytes = value.getBytes(StandardCharsets.UTF_8);
-			DataBuffer buffer = this.bufferFactory.allocateBuffer(bytes.length);
-			buffer.write(bytes);
+			DataBuffer buffer = toToDataBuffer(value);
 			return Mono.just(buffer);
 		});
 	}
 
+	private DataBuffer toToDataBuffer(String value) {
+		byte[] bytes = value.getBytes(StandardCharsets.UTF_8);
+		DataBuffer buffer = this.bufferFactory.allocateBuffer(bytes.length);
+		buffer.write(bytes);
+		return buffer;
+	}
+
 	@jakarta.xml.bind.annotation.XmlType
 	@XmlSeeAlso(Child.class)
 	public abstract static class Parent {

spring-web/src/test/java/org/springframework/http/codec/xml/XmlEventDecoderTests.java

@@ -27,7 +27,6 @@ import reactor.core.publisher.Mono;
 import reactor.test.StepVerifier;
 
 import org.springframework.core.io.buffer.DataBuffer;
-import org.springframework.core.io.buffer.DataBufferLimitException;
 import org.springframework.core.testfixture.io.buffer.AbstractLeakCheckingTests;
 
 import static org.assertj.core.api.Assertions.assertThat;
@@ -88,28 +87,6 @@ class XmlEventDecoderTests extends AbstractLeakCheckingTests {
 				.verify();
 	}
 
-	@Test
-	void toXMLEventsWithLimit() {
-
-		this.decoder.setMaxInMemorySize(6);
-
-		Flux<String> source = Flux.just(
-				"<pojo>", "<foo>", "foofoo", "</foo>", "<bar>", "barbarbar", "</bar>", "</pojo>");
-
-		Flux<XMLEvent> events = this.decoder.decode(
-				source.map(this::stringBuffer), null, null, Collections.emptyMap());
-
-		StepVerifier.create(events)
-				.consumeNextWith(e -> assertThat(e.isStartDocument()).isTrue())
-				.consumeNextWith(e -> assertStartElement(e, "pojo"))
-				.consumeNextWith(e -> assertStartElement(e, "foo"))
-				.consumeNextWith(e -> assertCharacters(e, "foofoo"))
-				.consumeNextWith(e -> assertEndElement(e, "foo"))
-				.consumeNextWith(e -> assertStartElement(e, "bar"))
-				.expectError(DataBufferLimitException.class)
-				.verify();
-	}
-
 	@Test
 	void decodeErrorAalto() {
 		Flux<DataBuffer> source = Flux.concat(