Spring/spring-framework
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Support for maxInMemorySize in SSE reader

Browse Source 
Closes gh-24312
This commit is contained in:
Rossen Stoyanchev
2020-01-13 14:56:48 +00:00
parent a741ae422b
commit cbc57460b7
5 changed files with 138 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files

76
spring-web/src/main/java/org/springframework/http/codec/ServerSentEventHttpMessageReader.java
View File

@@ -1,5 +1,5 @@
/*
* Copyright 2002-2019 the original author or authors.
* Copyright 2002-2020 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
@@ -31,6 +31,7 @@ import org.springframework.core.codec.Decoder;
import org.springframework.core.codec.StringDecoder;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferFactory;
import org.springframework.core.io.buffer.DataBufferLimitException;
import org.springframework.core.io.buffer.DefaultDataBufferFactory;
import org.springframework.http.MediaType;
import org.springframework.http.ReactiveHttpInputMessage;
@@ -48,14 +49,16 @@ public class ServerSentEventHttpMessageReader implements HttpMessageReader<Objec
private static final DataBufferFactory bufferFactory = new DefaultDataBufferFactory();
private static final StringDecoder stringDecoder = StringDecoder.textPlainOnly();
private static final ResolvableType STRING_TYPE = ResolvableType.forClass(String.class);
@Nullable
private final Decoder<?> decoder;
private final StringDecoder lineDecoder = StringDecoder.textPlainOnly();
/**
* Constructor without a {@code Decoder}. In this mode only {@code String}
@@ -82,6 +85,29 @@ public class ServerSentEventHttpMessageReader implements HttpMessageReader<Objec
return this.decoder;
}
/**
* Configure a limit on the maximum number of bytes per SSE event which are
* buffered before the event is parsed.
* <p>Note that the {@link #getDecoder() data decoder}, if provided, must
* also be customized accordingly to raise the limit if necessary in order
* to be able to parse the data portion of the event.
* <p>By default this is set to 256K.
* @param byteCount the max number of bytes to buffer, or -1 for unlimited
* @since 5.1.13
*/
public void setMaxInMemorySize(int byteCount) {
this.lineDecoder.setMaxInMemorySize(byteCount);
}
/**
* Return the {@link #setMaxInMemorySize configured} byte count limit.
* @since 5.1.13
*/
public int getMaxInMemorySize() {
return this.lineDecoder.getMaxInMemorySize();
}
@Override
public List<MediaType> getReadableMediaTypes() {
return Collections.singletonList(MediaType.TEXT_EVENT_STREAM);
@@ -101,12 +127,15 @@ public class ServerSentEventHttpMessageReader implements HttpMessageReader<Objec
public Flux<Object> read(
ResolvableType elementType, ReactiveHttpInputMessage message, Map<String, Object> hints) {
LimitTracker limitTracker = new LimitTracker();
boolean shouldWrap = isServerSentEvent(elementType);
ResolvableType valueType = (shouldWrap ? elementType.getGeneric() : elementType);
return stringDecoder.decode(message.getBody(), STRING_TYPE, null, hints)
return this.lineDecoder.decode(message.getBody(), STRING_TYPE, null, hints)
.doOnNext(limitTracker::afterLineParsed)
.bufferUntil(String::isEmpty)
.concatMap(lines -> Mono.justOrEmpty(buildEvent(lines, valueType, shouldWrap, hints)));
.map(lines -> buildEvent(lines, valueType, shouldWrap, hints));
}
@Nullable
@@ -172,16 +201,47 @@ public class ServerSentEventHttpMessageReader implements HttpMessageReader<Objec
public Mono<Object> readMono(
ResolvableType elementType, ReactiveHttpInputMessage message, Map<String, Object> hints) {
// We're ahead of String + "*/*"
// Let's see if we can aggregate the output (lest we time out)...
// In order of readers, we're ahead of String + "*/*"
// If this is called, simply delegate to StringDecoder
if (elementType.resolve() == String.class) {
Flux<DataBuffer> body = message.getBody();
return stringDecoder.decodeToMono(body, elementType, null, null).cast(Object.class);
return this.lineDecoder.decodeToMono(body, elementType, null, null).cast(Object.class);
}
return Mono.error(new UnsupportedOperationException(
"ServerSentEventHttpMessageReader only supports reading stream of events as a Flux"));
}
private class LimitTracker {
private int accumulated = 0;
public void afterLineParsed(String line) {
if (getMaxInMemorySize() < 0) {
return;
}
if (line.isEmpty()) {
this.accumulated = 0;
}
if (line.length() > Integer.MAX_VALUE - this.accumulated) {
raiseLimitException();
}
else {
this.accumulated += line.length();
if (this.accumulated > getMaxInMemorySize()) {
raiseLimitException();
}
}
}
private void raiseLimitException() {
// Do not release here, it's likely down via doOnDiscard..
throw new DataBufferLimitException(
"Exceeded limit on max bytes to buffer : " + getMaxInMemorySize());
}
}
}

9
spring-web/src/main/java/org/springframework/http/codec/support/BaseDefaultCodecs.java
View File

@@ -1,5 +1,5 @@
/*
* Copyright 2002-2019 the original author or authors.
* Copyright 2002-2020 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
@@ -238,9 +238,6 @@ class BaseDefaultCodecs implements CodecConfigurer.DefaultCodecs, CodecConfigure
if (codec instanceof DecoderHttpMessageReader) {
codec = ((DecoderHttpMessageReader) codec).getDecoder();
}
else if (codec instanceof ServerSentEventHttpMessageReader) {
codec = ((ServerSentEventHttpMessageReader) codec).getDecoder();
}
if (codec == null) {
return;
@@ -269,6 +266,10 @@ class BaseDefaultCodecs implements CodecConfigurer.DefaultCodecs, CodecConfigure
if (codec instanceof FormHttpMessageReader) {
((FormHttpMessageReader) codec).setMaxInMemorySize(size);
}
if (codec instanceof ServerSentEventHttpMessageReader) {
((ServerSentEventHttpMessageReader) codec).setMaxInMemorySize(size);
initCodec(((ServerSentEventHttpMessageReader) codec).getDecoder());
}
if (synchronossMultipartPresent) {
if (codec instanceof SynchronossPartHttpMessageReader) {
((SynchronossPartHttpMessageReader) codec).setMaxInMemorySize(size);