HttpHeaderResponseDecorator checks for "Transfer-Encoding"

This commit extends the fix from b86c11cc9b
by checking for both existing Content-Length and Transfer-Encoding.

Closes gh-25908
This commit is contained in:
Rossen Stoyanchev 2020-10-13 22:59:45 +01:00
parent 7b6293fa05
commit 1d96f6a266
2 changed files with 30 additions and 19 deletions

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"); * Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License. * you may not use this file except in compliance with the License.
@ -16,8 +16,6 @@
package org.springframework.http.server.reactive; package org.springframework.http.server.reactive;
import java.util.function.BiFunction;
import org.reactivestreams.Publisher; import org.reactivestreams.Publisher;
import reactor.core.publisher.Flux; import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono; import reactor.core.publisher.Mono;
@ -41,24 +39,30 @@ public class HttpHeadResponseDecorator extends ServerHttpResponseDecorator {
/** /**
* Apply {@link Flux#reduce(Object, BiFunction) reduce} on the body, count * Consume and release the body without writing.
* the number of bytes produced, release data buffers without writing, and * <p>If the headers contain neither Content-Length nor Transfer-Encoding,
* set the {@literal Content-Length} header. * count the bytes and set Content-Length.
*/ */
@Override @Override
public final Mono<Void> writeWith(Publisher<? extends DataBuffer> body) { public final Mono<Void> writeWith(Publisher<? extends DataBuffer> body) {
return Flux.from(body) if (shouldSetContentLength()) {
.reduce(0, (current, buffer) -> { return Flux.from(body)
int next = current + buffer.readableByteCount(); .reduce(0, (current, buffer) -> {
DataBufferUtils.release(buffer); int next = current + buffer.readableByteCount();
return next; DataBufferUtils.release(buffer);
}) return next;
.doOnNext(length -> { })
if (length > 0 || getHeaders().getFirst(HttpHeaders.CONTENT_LENGTH) == null) { .doOnNext(length -> getHeaders().setContentLength(length))
getHeaders().setContentLength(length); .then();
} }
}) else {
.then(); return Flux.from(body).then();
}
}
private boolean shouldSetContentLength() {
return (getHeaders().getFirst(HttpHeaders.CONTENT_LENGTH) == null &&
getHeaders().getFirst(HttpHeaders.TRANSFER_ENCODING) == null);
} }
/** /**

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"); * Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License. * you may not use this file except in compliance with the License.
@ -25,6 +25,7 @@ import reactor.core.publisher.Flux;
import org.springframework.core.io.buffer.DataBuffer; import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.NettyDataBufferFactory; import org.springframework.core.io.buffer.NettyDataBufferFactory;
import org.springframework.core.testfixture.io.buffer.LeakAwareDataBufferFactory; import org.springframework.core.testfixture.io.buffer.LeakAwareDataBufferFactory;
import org.springframework.http.HttpHeaders;
import org.springframework.web.testfixture.http.server.reactive.MockServerHttpResponse; import org.springframework.web.testfixture.http.server.reactive.MockServerHttpResponse;
import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThat;
@ -63,6 +64,12 @@ public class HttpHeadResponseDecoratorTests {
assertThat(this.response.getHeaders().getContentLength()).isEqualTo(length); assertThat(this.response.getHeaders().getContentLength()).isEqualTo(length);
} }
@Test // gh-25908
public void writeWithGivenTransferEncoding() {
this.response.getHeaders().add(HttpHeaders.TRANSFER_ENCODING, "chunked");
this.response.writeWith(Flux.empty()).block();
assertThat(this.response.getHeaders().getContentLength()).isEqualTo(-1);
}
private DataBuffer toDataBuffer(String s) { private DataBuffer toDataBuffer(String s) {
DataBuffer buffer = this.bufferFactory.allocateBuffer(); DataBuffer buffer = this.bufferFactory.allocateBuffer();