HttpHeadResponseDecorator sets Content-Length for Mono only

See gh-25908
This commit is contained in:
Rossen Stoyanchev 2020-10-13 23:45:51 +01:00
parent abd79d43af
commit 7e647ab1d8
2 changed files with 22 additions and 10 deletions

View File

@ -41,18 +41,22 @@ public class HttpHeadResponseDecorator extends ServerHttpResponseDecorator {
/**
* Consume and release the body without writing.
* <p>If the headers contain neither Content-Length nor Transfer-Encoding,
* count the bytes and set Content-Length.
* and the body is a {@link Mono}, count the bytes and set Content-Length.
*/
@Override
@SuppressWarnings("unchecked")
public final Mono<Void> writeWith(Publisher<? extends DataBuffer> body) {
if (shouldSetContentLength()) {
return Flux.from(body)
.reduce(0, (current, buffer) -> {
int next = current + buffer.readableByteCount();
DataBufferUtils.release(buffer);
return next;
if (shouldSetContentLength() && body instanceof Mono) {
return ((Mono<? extends DataBuffer>) body)
.doOnSuccess(buffer -> {
if (buffer != null) {
getHeaders().setContentLength(buffer.readableByteCount());
DataBufferUtils.release(buffer);
}
else {
getHeaders().setContentLength(0);
}
})
.doOnNext(length -> getHeaders().setContentLength(length))
.then();
}
else {

View File

@ -21,6 +21,7 @@ import io.netty.buffer.PooledByteBufAllocator;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.Test;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.NettyDataBufferFactory;
@ -50,10 +51,17 @@ public class HttpHeadResponseDecoratorTests {
@Test
public void write() {
public void writeWithFlux() {
Flux<DataBuffer> body = Flux.just(toDataBuffer("data1"), toDataBuffer("data2"));
this.response.writeWith(body).block();
assertThat(this.response.getHeaders().getContentLength()).isEqualTo(10);
assertThat(this.response.getHeaders().getContentLength()).isEqualTo(-1);
}
@Test
public void writeWithMono() {
Mono<DataBuffer> body = Mono.just(toDataBuffer("data1,data2"));
this.response.writeWith(body).block();
assertThat(this.response.getHeaders().getContentLength()).isEqualTo(11);
}
@Test // gh-23484