From b86c11cc9b00693fc3446724ffd37ae379189b7f Mon Sep 17 00:00:00 2001
From: Rossen Stoyanchev <rstoyanchev@pivotal.io>
Date: Wed, 21 Aug 2019 02:28:19 +0300
Subject: [PATCH] Respect existing content-length for HTTP HEAD

Closes gh-23484
---
 .../reactive/HttpHeadResponseDecorator.java   |  9 ++-
 .../HttpHeadResponseDecoratorTests.java       | 73 +++++++++++++++++++
 2 files changed, 80 insertions(+), 2 deletions(-)
 create mode 100644 spring-web/src/test/java/org/springframework/http/server/reactive/HttpHeadResponseDecoratorTests.java

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 4c8d0b0f06c..d387abd329b 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
@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2018 the original author or authors.
+ * Copyright 2002-2019 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.
@@ -24,6 +24,7 @@ import reactor.core.publisher.Mono;
 
 import org.springframework.core.io.buffer.DataBuffer;
 import org.springframework.core.io.buffer.DataBufferUtils;
+import org.springframework.http.HttpHeaders;
 
 /**
  * {@link ServerHttpResponse} decorator for HTTP HEAD requests.
@@ -52,7 +53,11 @@ public class HttpHeadResponseDecorator extends ServerHttpResponseDecorator {
 					DataBufferUtils.release(buffer);
 					return next;
 				})
-				.doOnNext(count -> getHeaders().setContentLength(count))
+				.doOnNext(length -> {
+					if (length > 0 || getHeaders().getFirst(HttpHeaders.CONTENT_LENGTH) == null) {
+						getHeaders().setContentLength(length);
+					}
+				})
 				.then();
 	}
 
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
new file mode 100644
index 00000000000..8231bc4789c
--- /dev/null
+++ b/spring-web/src/test/java/org/springframework/http/server/reactive/HttpHeadResponseDecoratorTests.java
@@ -0,0 +1,73 @@
+/*
+ * Copyright 2002-2019 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.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.springframework.http.server.reactive;
+
+import java.nio.charset.StandardCharsets;
+
+import io.netty.buffer.PooledByteBufAllocator;
+import org.junit.After;
+import org.junit.Test;
+import reactor.core.publisher.Flux;
+
+import org.springframework.core.io.buffer.DataBuffer;
+import org.springframework.core.io.buffer.LeakAwareDataBufferFactory;
+import org.springframework.core.io.buffer.NettyDataBufferFactory;
+import org.springframework.mock.http.server.reactive.test.MockServerHttpResponse;
+
+import static org.junit.Assert.assertEquals;
+
+/**
+ * Unit tests for {@link HttpHeadResponseDecorator}.
+ * @author Rossen Stoyanchev
+ */
+public class HttpHeadResponseDecoratorTests {
+
+	private final LeakAwareDataBufferFactory bufferFactory =
+			new LeakAwareDataBufferFactory(new NettyDataBufferFactory(PooledByteBufAllocator.DEFAULT));
+
+	private final ServerHttpResponse response =
+			new HttpHeadResponseDecorator(new MockServerHttpResponse(this.bufferFactory));
+
+
+	@After
+	public void tearDown() {
+		this.bufferFactory.checkForLeaks();
+	}
+
+
+	@Test
+	public void write() {
+		Flux<DataBuffer> body = Flux.just(toDataBuffer("data1"), toDataBuffer("data2"));
+		response.writeWith(body).block();
+		assertEquals(10, response.getHeaders().getContentLength());
+	}
+
+	@Test // gh-23484
+	public void writeWithGivenContentLength() {
+		int length = 15;
+		this.response.getHeaders().setContentLength(length);
+		this.response.writeWith(Flux.empty()).block();
+		assertEquals(length, this.response.getHeaders().getContentLength());
+	}
+
+
+	private DataBuffer toDataBuffer(String s) {
+		DataBuffer buffer = this.bufferFactory.allocateBuffer();
+		buffer.write(s.getBytes(StandardCharsets.UTF_8));
+		return buffer;
+	}
+
+}