From 7e647ab1d831170b912ea29add4e79839535ec93 Mon Sep 17 00:00:00 2001
From: Rossen Stoyanchev <rstoyanchev@vmware.com>
Date: Tue, 13 Oct 2020 23:45:51 +0100
Subject: [PATCH] HttpHeadResponseDecorator sets Content-Length for Mono only

See gh-25908
---
 .../reactive/HttpHeadResponseDecorator.java   | 20 +++++++++++--------
 .../HttpHeadResponseDecoratorTests.java       | 12 +++++++++--
 2 files changed, 22 insertions(+), 10 deletions(-)

diff --git a/spring-web/src/main/java/org/springframework/http/server/reactive/HttpHeadResponseDecorator.java b/spring-web/src/main/java/org/springframework/http/server/reactive/HttpHeadResponseDecorator.java
index b52a38f722..04277cdead 100644
--- a/spring-web/src/main/java/org/springframework/http/server/reactive/HttpHeadResponseDecorator.java
+++ b/spring-web/src/main/java/org/springframework/http/server/reactive/HttpHeadResponseDecorator.java
@@ -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 {
diff --git a/spring-web/src/test/java/org/springframework/http/server/reactive/HttpHeadResponseDecoratorTests.java b/spring-web/src/test/java/org/springframework/http/server/reactive/HttpHeadResponseDecoratorTests.java
index d98b19356f..479bc3fc11 100644
--- a/spring-web/src/test/java/org/springframework/http/server/reactive/HttpHeadResponseDecoratorTests.java
+++ b/spring-web/src/test/java/org/springframework/http/server/reactive/HttpHeadResponseDecoratorTests.java
@@ -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