diff --git a/spring-web/src/main/java/org/springframework/http/codec/xml/Jaxb2Helper.java b/spring-web/src/main/java/org/springframework/http/codec/xml/Jaxb2Helper.java index 23e5ff34124..996acd7cf2f 100644 --- a/spring-web/src/main/java/org/springframework/http/codec/xml/Jaxb2Helper.java +++ b/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 { * * */ - public static Flux> split(Flux xmlEventFlux, Set names) { - return xmlEventFlux.handle(new SplitHandler(names)); + public static Flux> split( + Flux xmlEventFlux, Set names, XmlEventDecoder.@Nullable ReceivedByteTracker byteTracker) { + + return xmlEventFlux.handle(new SplitHandler(names, byteTracker)); } @@ -150,14 +153,17 @@ abstract class Jaxb2Helper { private final Set names; + private final XmlEventDecoder.ReceivedByteTracker byteTracker; + private @Nullable List events; private int elementDepth = 0; private int barrier = Integer.MAX_VALUE; - public SplitHandler(Set names) { + public SplitHandler(Set 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()); + } } } diff --git a/spring-web/src/main/java/org/springframework/http/codec/xml/Jaxb2XmlDecoder.java b/spring-web/src/main/java/org/springframework/http/codec/xml/Jaxb2XmlDecoder.java index c7721fd9b81..1dac0277c02 100644 --- a/spring-web/src/main/java/org/springframework/http/codec/xml/Jaxb2XmlDecoder.java +++ b/spring-web/src/main/java/org/springframework/http/codec/xml/Jaxb2XmlDecoder.java @@ -145,12 +145,16 @@ public class Jaxb2XmlDecoder extends AbstractDecoder { public Flux decode(Publisher inputStream, ResolvableType elementType, @Nullable MimeType mimeType, @Nullable Map hints) { + XmlEventDecoder.ReceivedByteTracker byteTracker = + new XmlEventDecoder.ReceivedByteTracker(this.maxInMemorySize); + Flux 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 typeNames = Jaxb2Helper.toQNames(outputClass); - Flux> splitEvents = Jaxb2Helper.split(xmlEventFlux, typeNames); + Flux> splitEvents = Jaxb2Helper.split(xmlEventFlux, typeNames, byteTracker); return splitEvents.map(events -> { Object value = unmarshal(events, outputClass); diff --git a/spring-web/src/main/java/org/springframework/http/codec/xml/XmlEventDecoder.java b/spring-web/src/main/java/org/springframework/http/codec/xml/XmlEventDecoder.java index d766b3ee4e4..8ded46a9c3c 100644 --- a/spring-web/src/main/java/org/springframework/http/codec/xml/XmlEventDecoder.java +++ b/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 { + /** + * 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 { /** - * 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. + *

When joining all buffers and decoding as a whole, the limit is applied + * to the entire input. + *

>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. *

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 { @Nullable MimeType mimeType, @Nullable Map 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 { /* * Separate static class to isolate Aalto dependency. */ - private static class AaltoDataBufferToXmlEvent implements Function> { + private static final class AaltoDataBufferToXmlEvent implements Function> { private static final AsyncXMLInputFactory inputFactory = StaxUtils.createDefensiveInputFactory(InputFactoryImpl::new); @@ -165,22 +174,19 @@ public class XmlEventDecoder extends AbstractDecoder { 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 hints) { + this.byteTracker = (hints != null ? (ReceivedByteTracker) hints.get(BYTE_TRACKER_HINT) : null); } - @Override public List 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 { 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 { } } - 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. + *

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"; + } + } } diff --git a/spring-web/src/test/java/org/springframework/http/codec/xml/Jaxb2XmlDecoderTests.java b/spring-web/src/test/java/org/springframework/http/codec/xml/Jaxb2XmlDecoderTests.java index 8df28da8b96..108c2e113fd 100644 --- a/spring-web/src/test/java/org/springframework/http/codec/xml/Jaxb2XmlDecoderTests.java +++ b/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 xmlEvents = this.xmlEventDecoder.decode(toDataBufferMono(POJO_ROOT), null, null, HINTS); - Flux> result = Jaxb2Helper.split(xmlEvents, Set.of(new QName("pojo"))); + Flux> 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 xmlEvents = this.xmlEventDecoder.decode(toDataBufferMono(POJO_CHILD), null, null, HINTS); - Flux> result = Jaxb2Helper.split(xmlEvents, Set.of(new QName("pojo"))); - + Flux> 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 source = Flux.just( + "", + "foo", + "foofoo", + "barbar", + ""); + + XmlEventDecoder.ReceivedByteTracker byteTracker = new XmlEventDecoder.ReceivedByteTracker(30); + Map hints = Hints.from(XmlEventDecoder.BYTE_TRACKER_HINT, byteTracker); + Flux xmlEvents = this.xmlEventDecoder.decode(source.map(this::toToDataBuffer), null, null, hints); + Flux> 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 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 { diff --git a/spring-web/src/test/java/org/springframework/http/codec/xml/XmlEventDecoderTests.java b/spring-web/src/test/java/org/springframework/http/codec/xml/XmlEventDecoderTests.java index 36425c1198c..7a97bf30f23 100644 --- a/spring-web/src/test/java/org/springframework/http/codec/xml/XmlEventDecoderTests.java +++ b/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 source = Flux.just( - "", "", "foofoo", "", "", "barbarbar", "", ""); - - Flux 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 source = Flux.concat(