From 9ac5eb0f55c5b49797949ae8186906f8f5dea2a7 Mon Sep 17 00:00:00 2001
From: Brian Clozel <bclozel@vmware.com>
Date: Wed, 25 Oct 2023 16:01:15 +0200
Subject: [PATCH] Support ETag generation on ResourceWebHandler

This commit replicates the ETag generation option now available on
`ResourceHttpRequestHandler` but for its WebFlux counterpart.

See gh-29031
---
 .../reactive/resource/ResourceWebHandler.java | 31 ++++++++++++-
 .../resource/ResourceWebHandlerTests.java     | 43 ++++++++++++++++++-
 2 files changed, 72 insertions(+), 2 deletions(-)

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 aa3e2f556a..8c6819517d 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
@@ -27,6 +27,7 @@ import java.util.List;
 import java.util.Locale;
 import java.util.Map;
 import java.util.Set;
+import java.util.function.Function;
 
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
@@ -123,6 +124,9 @@ public class ResourceWebHandler implements WebHandler, InitializingBean {
 
 	private boolean useLastModified = true;
 
+	@Nullable
+	private Function<Resource, String> etagGenerator;
+
 	private boolean optimizeLocations = false;
 
 
@@ -275,6 +279,29 @@ public class ResourceWebHandler implements WebHandler, InitializingBean {
 		return this.useLastModified;
 	}
 
+	/**
+	 * Configure a generator function that will be used to create the ETag information,
+	 * given a {@link Resource} that is about to be written to the response.
+	 * <p>This function should return a String that will be used as an argument in
+	 * {@link ServerWebExchange#checkNotModified(String)}, or {@code null} if no value
+	 * can be generated for the given resource.
+	 * @param etagGenerator the HTTP ETag generator function to use.
+	 * @since 6.1
+	 */
+	public void setEtagGenerator(@Nullable Function<Resource, String> etagGenerator) {
+		this.etagGenerator = etagGenerator;
+	}
+
+	/**
+	 * Return the HTTP ETag generator function to be used when serving resources.
+	 * @return the HTTP ETag generator function
+	 * @since 6.1
+	 */
+	@Nullable
+	public Function<Resource, String> getEtagGenerator() {
+		return this.etagGenerator;
+	}
+
 	/**
 	 * Set whether to optimize the specified locations through an existence
 	 * check on startup, filtering non-existing directories upfront so that
@@ -418,7 +445,9 @@ public class ResourceWebHandler implements WebHandler, InitializingBean {
 						}
 
 						// Header phase
-						if (isUseLastModified() && exchange.checkNotModified(Instant.ofEpochMilli(resource.lastModified()))) {
+						String eTagValue = (this.getEtagGenerator() != null) ? this.getEtagGenerator().apply(resource) : null;
+						Instant lastModified = isUseLastModified() ? Instant.ofEpochMilli(resource.lastModified()) : Instant.MIN;
+						if (exchange.checkNotModified(eTagValue, lastModified)) {
 							logger.trace(exchange.getLogPrefix() + "Resource not modified");
 							return Mono.empty();
 						}
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 9c09a5d5c4..2bb6f4b48c 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
@@ -500,7 +500,48 @@ class ResourceWebHandlerTests {
 		}
 
 		@Test
-			// SPR-14005
+		void shouldRespondWithNotModifiedWhenEtag() throws Exception {
+			this.handler.setEtagGenerator(resource -> "testEtag");
+			this.handler.afterPropertiesSet();
+			MockServerWebExchange exchange = MockServerWebExchange.from(
+					MockServerHttpRequest.get("").ifNoneMatch( "\"testEtag\""));
+
+			setPathWithinHandlerMapping(exchange, "foo.css");
+			setBestMachingPattern(exchange, "/**");
+			this.handler.handle(exchange).block(TIMEOUT);
+			assertThat(exchange.getResponse().getStatusCode()).isEqualTo(HttpStatus.NOT_MODIFIED);
+		}
+
+		@Test
+		void shouldRespondWithModifiedResourceWhenEtagNoMatch() throws Exception {
+			this.handler.setEtagGenerator(resource -> "noMatch");
+			this.handler.afterPropertiesSet();
+			MockServerWebExchange exchange = MockServerWebExchange.from(
+					MockServerHttpRequest.get("").ifNoneMatch( "\"testEtag\""));
+
+			setPathWithinHandlerMapping(exchange, "foo.css");
+			setBestMachingPattern(exchange, "/**");
+			this.handler.handle(exchange).block(TIMEOUT);
+			assertThat((Object) exchange.getResponse().getStatusCode()).isNull();
+			assertResponseBody(exchange, "h1 { color:red; }");
+		}
+
+		@Test
+		void shouldRespondWithNotModifiedWhenEtagAndLastModified() throws Exception {
+			this.handler.setEtagGenerator(resource -> "testEtag");
+			this.handler.afterPropertiesSet();
+			MockServerWebExchange exchange = MockServerWebExchange.from(
+					MockServerHttpRequest.get("")
+							.ifModifiedSince(resourceLastModified("test/foo.css"))
+							.ifNoneMatch( "\"testEtag\""));
+
+			setPathWithinHandlerMapping(exchange, "foo.css");
+			setBestMachingPattern(exchange, "/**");
+			this.handler.handle(exchange).block(TIMEOUT);
+			assertThat(exchange.getResponse().getStatusCode()).isEqualTo(HttpStatus.NOT_MODIFIED);
+		}
+
+		@Test  // SPR-14005
 		void doOverwriteExistingCacheControlHeaders() throws Exception {
 			this.handler.setCacheControl(CacheControl.maxAge(3600, TimeUnit.SECONDS));
 			this.handler.afterPropertiesSet();