Support for MediaType mappings in ResourceWebHandler

Closes gh-26170

Author: rstoyanchev

spring-webflux/src/main/java/org/springframework/web/reactive/config/ResourceHandlerRegistration.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2016 the original author or authors.
+ * Copyright 2002-2020 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.
@@ -18,12 +18,16 @@
 
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 
 import org.springframework.cache.Cache;
 import org.springframework.core.io.Resource;
 import org.springframework.core.io.ResourceLoader;
 import org.springframework.http.CacheControl;
+import org.springframework.http.MediaType;
+import org.springframework.http.MediaTypeFactory;
 import org.springframework.lang.Nullable;
 import org.springframework.util.Assert;
 import org.springframework.web.reactive.resource.ResourceWebHandler;
@@ -50,6 +54,10 @@ public class ResourceHandlerRegistration {
 
 	private boolean useLastModified = true;
 
+	@Nullable
+	private Map<String, MediaType> mediaTypes;
+
+
 
 	/**
 	 * Create a {@link ResourceHandlerRegistration} instance.
@@ -146,6 +154,23 @@ public ResourceChainRegistration resourceChain(boolean cacheResources, Cache cac
 		return this.resourceChainRegistration;
 	}
 
+	/**
+	 * Add mappings between file extensions extracted from the filename of static
+	 * {@link Resource}s and the media types to use for the response.
+	 * <p>Use of this method is typically not necessary since mappings can be
+	 * also determined via {@link MediaTypeFactory#getMediaType(Resource)}.
+	 * @param mediaTypes media type mappings
+	 * @since 5.3.2
+	 */
+	public void setMediaTypes(Map<String, MediaType> mediaTypes) {
+		if (this.mediaTypes == null) {
+			this.mediaTypes = new HashMap<>(mediaTypes.size());
+		}
+		this.mediaTypes.clear();
+		this.mediaTypes.putAll(mediaTypes);
+	}
+
+
 	/**
 	 * Returns the URL path patterns for the resource handler.
 	 */
@@ -168,6 +193,9 @@ protected ResourceWebHandler getRequestHandler() {
 			handler.setCacheControl(this.cacheControl);
 		}
 		handler.setUseLastModified(this.useLastModified);
+		if (this.mediaTypes != null) {
+			handler.setMediaTypes(this.mediaTypes);
+		}
 		return handler;
 	}
 

spring-webflux/src/main/java/org/springframework/web/reactive/resource/ResourceWebHandler.java

@@ -23,7 +23,10 @@
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.EnumSet;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Locale;
+import java.util.Map;
 import java.util.Set;
 import java.util.stream.Collectors;
 
@@ -111,6 +114,9 @@ public class ResourceWebHandler implements WebHandler, InitializingBean {
 	@Nullable
 	private ResourceHttpMessageWriter resourceHttpMessageWriter;
 
+	@Nullable
+	private Map<String, MediaType> mediaTypes;
+
 	@Nullable
 	private ResourceLoader resourceLoader;
 
@@ -230,6 +236,30 @@ public ResourceHttpMessageWriter getResourceHttpMessageWriter() {
 		return this.resourceHttpMessageWriter;
 	}
 
+	/**
+	 * Add mappings between file extensions extracted from the filename of static
+	 * {@link Resource}s and the media types to use for the response.
+	 * <p>Use of this method is typically not necessary since mappings can be
+	 * also determined via {@link MediaTypeFactory#getMediaType(Resource)}.
+	 * @param mediaTypes media type mappings
+	 * @since 5.3.2
+	 */
+	public void setMediaTypes(Map<String, MediaType> mediaTypes) {
+		if (this.mediaTypes == null) {
+			this.mediaTypes = new HashMap<>(mediaTypes.size());
+		}
+		mediaTypes.forEach((ext, type) ->
+				this.mediaTypes.put(ext.toLowerCase(Locale.ENGLISH), type));
+	}
+
+	/**
+	 * Return the {@link #setMediaTypes(Map) configured} media type mappings.
+	 * @since 5.3.2
+	 */
+	public Map<String, MediaType> getMediaTypes() {
+		return (this.mediaTypes != null ? this.mediaTypes : Collections.emptyMap());
+	}
+
 	/**
 	 * Provide the ResourceLoader to load {@link #setLocationValues(List)
 	 * location values} with.
@@ -374,7 +404,7 @@ public Mono<Void> handle(ServerWebExchange exchange) {
 						}
 
 						// Check the media type for the resource
-						MediaType mediaType = MediaTypeFactory.getMediaType(resource).orElse(null);
+						MediaType mediaType = getMediaType(resource);
 						setHeaders(exchange, resource, mediaType);
 
 						// Content phase
@@ -535,6 +565,25 @@ protected boolean isInvalidPath(String path) {
 		return false;
 	}
 
+	@Nullable
+	private MediaType getMediaType(Resource resource) {
+		MediaType mediaType = null;
+		String filename = resource.getFilename();
+		if (!CollectionUtils.isEmpty(this.mediaTypes)) {
+			String ext = StringUtils.getFilenameExtension(filename);
+			if (ext != null) {
+				mediaType = this.mediaTypes.get(ext.toLowerCase(Locale.ENGLISH));
+			}
+		}
+		if (mediaType == null) {
+			List<MediaType> mediaTypes = MediaTypeFactory.getMediaTypes(filename);
+			if (!CollectionUtils.isEmpty(mediaTypes)) {
+				mediaType = mediaTypes.get(0);
+			}
+		}
+		return mediaType;
+	}
+
 	/**
 	 * Set headers on the response. Called for both GET and HEAD requests.
 	 * @param exchange current exchange

spring-webflux/src/test/java/org/springframework/web/reactive/config/ResourceHandlerRegistryTests.java

@@ -17,6 +17,7 @@
 package org.springframework.web.reactive.config;
 
 import java.time.Duration;
+import java.util.Collections;
 import java.util.List;
 import java.util.concurrent.TimeUnit;
 
@@ -28,6 +29,7 @@
 import org.springframework.cache.concurrent.ConcurrentMapCache;
 import org.springframework.context.support.GenericApplicationContext;
 import org.springframework.http.CacheControl;
+import org.springframework.http.MediaType;
 import org.springframework.http.server.PathContainer;
 import org.springframework.web.reactive.HandlerMapping;
 import org.springframework.web.reactive.handler.SimpleUrlHandlerMapping;
@@ -99,6 +101,16 @@ public void cacheControl() {
 				.isEqualTo(CacheControl.noCache().cachePrivate().getHeaderValue());
 	}
 
+	@Test
+	public void mediaTypes() {
+		MediaType mediaType = MediaType.parseMediaType("foo/bar");
+		this.registration.setMediaTypes(Collections.singletonMap("bar", mediaType));
+		ResourceWebHandler requestHandler = this.registration.getRequestHandler();
+
+		assertThat(requestHandler.getMediaTypes()).size().isEqualTo(1);
+		assertThat(requestHandler.getMediaTypes()).containsEntry("bar", mediaType);
+	}
+
 	@Test
 	public void order() {
 		assertThat(this.registry.getHandlerMapping().getOrder()).isEqualTo(Integer.MAX_VALUE -1);

spring-webflux/src/test/java/org/springframework/web/reactive/resource/ResourceWebHandlerTests.java

@@ -214,6 +214,24 @@ public void getResourceFromSubDirectoryOfAlternatePath() {
 		assertResponseBody(exchange, "function foo() { console.log(\"hello world\"); }");
 	}
 
+	@Test
+	public void getResourceWithRegisteredMediaType() throws Exception {
+		MediaType mediaType = new MediaType("foo", "bar");
+
+		ResourceWebHandler handler = new ResourceWebHandler();
+		handler.setLocations(Collections.singletonList(new ClassPathResource("test/", getClass())));
+		handler.setMediaTypes(Collections.singletonMap("bar", mediaType));
+		handler.afterPropertiesSet();
+
+		MockServerWebExchange exchange = MockServerWebExchange.from(MockServerHttpRequest.get(""));
+		setPathWithinHandlerMapping(exchange, "foo.bar");
+		handler.handle(exchange).block(TIMEOUT);
+
+		HttpHeaders headers = exchange.getResponse().getHeaders();
+		assertThat(headers.getContentType()).isEqualTo(mediaType);
+		assertResponseBody(exchange, "foo bar foo bar foo bar");
+	}
+
 	@Test // SPR-14577
 	public void getMediaTypeWithFavorPathExtensionOff() throws Exception {
 		List<Resource> paths = Collections.singletonList(new ClassPathResource("test/", getClass()));

spring-webflux/src/test/resources/org/springframework/web/reactive/resource/test/foo.bar

@@ -0,0 +1 @@
+foo bar foo bar foo bar