From 9adfa5e8b0f0cd65a0b14740e0a0f9832d80edcb Mon Sep 17 00:00:00 2001
From: Arjen Poutsma <poutsmaa@vmware.com>
Date: Thu, 7 Apr 2022 13:51:21 +0200
Subject: [PATCH] Add HEAD support in MVC/WebFlux Resource handling

This commit introduces explicit HEAD support in Spring
MVC's ResourceHttpRequestHandler and WebFlux's ResourceWebHandler,
adding just headers but no body.

Closes gh-28291
---
 .../http/codec/ResourceHttpMessageWriter.java | 33 ++++++++++++-------
 .../ResourceHttpMessageConverter.java         |  8 +++++
 .../reactive/resource/ResourceWebHandler.java | 15 ++++++---
 .../resource/ResourceWebHandlerTests.java     |  5 ++-
 .../resource/ResourceHttpRequestHandler.java  |  9 ++++-
 .../ResourceHttpRequestHandlerTests.java      |  3 +-
 6 files changed, 54 insertions(+), 19 deletions(-)

diff --git a/spring-web/src/main/java/org/springframework/http/codec/ResourceHttpMessageWriter.java b/spring-web/src/main/java/org/springframework/http/codec/ResourceHttpMessageWriter.java
index af456989e0..5765fe218c 100644
--- a/spring-web/src/main/java/org/springframework/http/codec/ResourceHttpMessageWriter.java
+++ b/spring-web/src/main/java/org/springframework/http/codec/ResourceHttpMessageWriter.java
@@ -116,8 +116,27 @@ public class ResourceHttpMessageWriter implements HttpMessageWriter<Resource> {
 	private Mono<Void> writeResource(Resource resource, ResolvableType type, @Nullable MediaType mediaType,
 			ReactiveHttpOutputMessage message, Map<String, Object> hints) {
 
+		addHeaders(message, resource, mediaType, hints);
+
+		return zeroCopy(resource, null, message, hints)
+				.orElseGet(() -> {
+					Mono<Resource> input = Mono.just(resource);
+					DataBufferFactory factory = message.bufferFactory();
+					Flux<DataBuffer> body = this.encoder.encode(input, factory, type, message.getHeaders().getContentType(), hints);
+					if (logger.isDebugEnabled()) {
+						body = body.doOnNext(buffer -> Hints.touchDataBuffer(buffer, hints, logger));
+					}
+					return message.writeWith(body);
+				});
+	}
+
+	/**
+	 * Adds the default headers for the given resource to the given message.
+	 * @since 6.0
+	 */
+	public void addHeaders(ReactiveHttpOutputMessage message, Resource resource, @Nullable MediaType contentType, Map<String, Object> hints) {
 		HttpHeaders headers = message.getHeaders();
-		MediaType resourceMediaType = getResourceMediaType(mediaType, resource, hints);
+		MediaType resourceMediaType = getResourceMediaType(contentType, resource, hints);
 		headers.setContentType(resourceMediaType);
 
 		if (headers.getContentLength() < 0) {
@@ -126,17 +145,7 @@ public class ResourceHttpMessageWriter implements HttpMessageWriter<Resource> {
 				headers.setContentLength(length);
 			}
 		}
-
-		return zeroCopy(resource, null, message, hints)
-				.orElseGet(() -> {
-					Mono<Resource> input = Mono.just(resource);
-					DataBufferFactory factory = message.bufferFactory();
-					Flux<DataBuffer> body = this.encoder.encode(input, factory, type, resourceMediaType, hints);
-					if (logger.isDebugEnabled()) {
-						body = body.doOnNext(buffer -> Hints.touchDataBuffer(buffer, hints, logger));
-					}
-					return message.writeWith(body);
-				});
+		headers.set(HttpHeaders.ACCEPT_RANGES, "bytes");
 	}
 
 	private static MediaType getResourceMediaType(
diff --git a/spring-web/src/main/java/org/springframework/http/converter/ResourceHttpMessageConverter.java b/spring-web/src/main/java/org/springframework/http/converter/ResourceHttpMessageConverter.java
index 17aa290ca9..e49bdd2551 100644
--- a/spring-web/src/main/java/org/springframework/http/converter/ResourceHttpMessageConverter.java
+++ b/spring-web/src/main/java/org/springframework/http/converter/ResourceHttpMessageConverter.java
@@ -122,6 +122,14 @@ public class ResourceHttpMessageConverter extends AbstractHttpMessageConverter<R
 		return (contentLength < 0 ? null : contentLength);
 	}
 
+	/**
+	 * Adds the default headers for the given resource to the given message.
+	 * @since 6.0
+	 */
+	public void addDefaultHeaders(HttpOutputMessage message, Resource resource, @Nullable MediaType contentType) throws IOException {
+		addDefaultHeaders(message.getHeaders(), resource, contentType);
+	}
+
 	@Override
 	protected void writeInternal(Resource resource, HttpOutputMessage outputMessage)
 			throws IOException, HttpMessageNotWritableException {
diff --git a/spring-webflux/src/main/java/org/springframework/web/reactive/resource/ResourceWebHandler.java b/spring-webflux/src/main/java/org/springframework/web/reactive/resource/ResourceWebHandler.java
index b83c697e71..4e8439cc1a 100644
--- a/spring-webflux/src/main/java/org/springframework/web/reactive/resource/ResourceWebHandler.java
+++ b/spring-webflux/src/main/java/org/springframework/web/reactive/resource/ResourceWebHandler.java
@@ -438,10 +438,17 @@ public class ResourceWebHandler implements WebHandler, InitializingBean {
 						// Content phase
 						ResourceHttpMessageWriter writer = getResourceHttpMessageWriter();
 						Assert.state(writer != null, "No ResourceHttpMessageWriter");
-						return writer.write(Mono.just(resource),
-								null, ResolvableType.forClass(Resource.class), mediaType,
-								exchange.getRequest(), exchange.getResponse(),
-								Hints.from(Hints.LOG_PREFIX_HINT, exchange.getLogPrefix()));
+						if (HttpMethod.HEAD == httpMethod) {
+							writer.addHeaders(exchange.getResponse(), resource, mediaType,
+									Hints.from(Hints.LOG_PREFIX_HINT, exchange.getLogPrefix()));
+							return exchange.getResponse().setComplete();
+						}
+						else {
+							return writer.write(Mono.just(resource),
+									null, ResolvableType.forClass(Resource.class), mediaType,
+									exchange.getRequest(), exchange.getResponse(),
+									Hints.from(Hints.LOG_PREFIX_HINT, exchange.getLogPrefix()));
+						}
 					}
 					catch (IOException ex) {
 						return Mono.error(ex);
diff --git a/spring-webflux/src/test/java/org/springframework/web/reactive/resource/ResourceWebHandlerTests.java b/spring-webflux/src/test/java/org/springframework/web/reactive/resource/ResourceWebHandlerTests.java
index e896404ac8..9a37835461 100644
--- a/spring-webflux/src/test/java/org/springframework/web/reactive/resource/ResourceWebHandlerTests.java
+++ b/spring-webflux/src/test/java/org/springframework/web/reactive/resource/ResourceWebHandlerTests.java
@@ -120,7 +120,10 @@ public class ResourceWebHandlerTests {
 		assertThat(headers.containsKey("Last-Modified")).isTrue();
 		assertThat(resourceLastModifiedDate("test/foo.css") / 1000).isEqualTo(headers.getLastModified() / 1000);
 		assertThat(headers.getFirst("Accept-Ranges")).isEqualTo("bytes");
-		assertThat(headers.get("Accept-Ranges").size()).isEqualTo(1);
+		assertThat(headers.get("Accept-Ranges")).hasSize(1);
+
+		StepVerifier.create(exchange.getResponse().getBody())
+				.verifyComplete();
 	}
 
 	@Test
diff --git a/spring-webmvc/src/main/java/org/springframework/web/servlet/resource/ResourceHttpRequestHandler.java b/spring-webmvc/src/main/java/org/springframework/web/servlet/resource/ResourceHttpRequestHandler.java
index 032698a268..cde1998390 100644
--- a/spring-webmvc/src/main/java/org/springframework/web/servlet/resource/ResourceHttpRequestHandler.java
+++ b/spring-webmvc/src/main/java/org/springframework/web/servlet/resource/ResourceHttpRequestHandler.java
@@ -585,7 +585,14 @@ public class ResourceHttpRequestHandler extends WebContentGenerator
 		ServletServerHttpResponse outputMessage = new ServletServerHttpResponse(response);
 		if (request.getHeader(HttpHeaders.RANGE) == null) {
 			Assert.state(this.resourceHttpMessageConverter != null, "Not initialized");
-			this.resourceHttpMessageConverter.write(resource, mediaType, outputMessage);
+
+			if (HttpMethod.HEAD.matches(request.getMethod())) {
+				this.resourceHttpMessageConverter.addDefaultHeaders(outputMessage, resource, mediaType);
+				outputMessage.flush();
+			}
+			else {
+				this.resourceHttpMessageConverter.write(resource, mediaType, outputMessage);
+			}
 		}
 		else {
 			Assert.state(this.resourceRegionHttpMessageConverter != null, "Not initialized");
diff --git a/spring-webmvc/src/test/java/org/springframework/web/servlet/resource/ResourceHttpRequestHandlerTests.java b/spring-webmvc/src/test/java/org/springframework/web/servlet/resource/ResourceHttpRequestHandlerTests.java
index 3dffcd5c4a..b6889bf206 100644
--- a/spring-webmvc/src/test/java/org/springframework/web/servlet/resource/ResourceHttpRequestHandlerTests.java
+++ b/spring-webmvc/src/test/java/org/springframework/web/servlet/resource/ResourceHttpRequestHandlerTests.java
@@ -118,7 +118,8 @@ public class ResourceHttpRequestHandlerTests {
 		assertThat(this.response.containsHeader("Last-Modified")).isTrue();
 		assertThat(this.response.getDateHeader("Last-Modified") / 1000).isEqualTo(resourceLastModified("test/foo.css") / 1000);
 		assertThat(this.response.getHeader("Accept-Ranges")).isEqualTo("bytes");
-		assertThat(this.response.getHeaders("Accept-Ranges").size()).isEqualTo(1);
+		assertThat(this.response.getHeaders("Accept-Ranges")).hasSize(1);
+		assertThat(this.response.getContentAsByteArray()).isEmpty();
 	}
 
 	@Test