Support to customize rather than replace default codecs

See gh-26212

Author: rstoyanchev

spring-web/src/main/java/org/springframework/http/codec/CodecConfigurer.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2020 the original author or authors.
+ * Copyright 2002-2021 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.
@@ -192,6 +192,17 @@ interface DefaultCodecs {
 		 */
 		void kotlinSerializationJsonEncoder(Encoder<?> encoder);
 
+		/**
+		 * Register a consumer to apply to default config instances. This can be
+		 * used to configure rather than replace a specific codec or multiple
+		 * codecs. The consumer is applied to every default {@link Encoder},
+		 * {@link Decoder}, {@link HttpMessageReader} and {@link HttpMessageWriter}
+		 * instance.
+		 * @param codecConsumer the consumer to apply
+		 * @since 5.3.4
+		 */
+		void configureDefaultCodec(Consumer<Object> codecConsumer);
+
 		/**
 		 * Configure a limit on the number of bytes that can be buffered whenever
 		 * the input stream needs to be aggregated. This can be a result of

spring-web/src/main/java/org/springframework/http/codec/json/Jackson2CodecSupport.java

@@ -84,7 +84,7 @@ public abstract class Jackson2CodecSupport {
 
 	protected final Log logger = HttpLogging.forLogName(getClass());
 
-	private final ObjectMapper defaultObjectMapper;
+	private ObjectMapper defaultObjectMapper;
 
 	@Nullable
 	private Map<Class<?>, Map<MimeType, ObjectMapper>> objectMapperRegistrations;
@@ -103,6 +103,19 @@ protected Jackson2CodecSupport(ObjectMapper objectMapper, MimeType... mimeTypes)
 	}
 
 
+	/**
+	 * Configure the default ObjectMapper instance to use.
+	 * @param objectMapper the ObjectMapper instance
+	 * @since 5.3.4
+	 */
+	public void setObjectMapper(ObjectMapper objectMapper) {
+		Assert.notNull(objectMapper, "ObjectMapper must not be null");
+		this.defaultObjectMapper = objectMapper;
+	}
+
+	/**
+	 * Return the {@link #setObjectMapper configured} default ObjectMapper.
+	 */
 	public ObjectMapper getObjectMapper() {
 		return this.defaultObjectMapper;
 	}

spring-web/src/main/java/org/springframework/http/codec/support/BaseDefaultCodecs.java

@@ -20,6 +20,7 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
+import java.util.function.Consumer;
 
 import org.springframework.core.SpringProperties;
 import org.springframework.core.codec.AbstractDataBufferDecoder;
@@ -46,6 +47,7 @@
 import org.springframework.http.codec.ResourceHttpMessageReader;
 import org.springframework.http.codec.ResourceHttpMessageWriter;
 import org.springframework.http.codec.ServerSentEventHttpMessageReader;
+import org.springframework.http.codec.ServerSentEventHttpMessageWriter;
 import org.springframework.http.codec.json.AbstractJackson2Decoder;
 import org.springframework.http.codec.json.Jackson2JsonDecoder;
 import org.springframework.http.codec.json.Jackson2JsonEncoder;
@@ -139,6 +141,9 @@ class BaseDefaultCodecs implements CodecConfigurer.DefaultCodecs, CodecConfigure
 	@Nullable
 	private Encoder<?> kotlinSerializationJsonEncoder;
 
+	@Nullable
+	private Consumer<Object> codecConsumer;
+
 	@Nullable
 	private Integer maxInMemorySize;
 
@@ -196,6 +201,7 @@ protected BaseDefaultCodecs(BaseDefaultCodecs other) {
 		this.jaxb2Encoder = other.jaxb2Encoder;
 		this.kotlinSerializationJsonDecoder = other.kotlinSerializationJsonDecoder;
 		this.kotlinSerializationJsonEncoder = other.kotlinSerializationJsonEncoder;
+		this.codecConsumer = other.codecConsumer;
 		this.maxInMemorySize = other.maxInMemorySize;
 		this.enableLoggingRequestDetails = other.enableLoggingRequestDetails;
 		this.registerDefaults = other.registerDefaults;
@@ -265,6 +271,14 @@ public void kotlinSerializationJsonEncoder(Encoder<?> encoder) {
 		initObjectWriters();
 	}
 
+	@Override
+	public void configureDefaultCodec(Consumer<Object> codecConsumer) {
+		this.codecConsumer = (this.codecConsumer != null ?
+				this.codecConsumer.andThen(codecConsumer) : codecConsumer);
+		initReaders();
+		initWriters();
+	}
+
 	@Override
 	public void maxInMemorySize(int byteCount) {
 		if (!ObjectUtils.nullSafeEquals(this.maxInMemorySize, byteCount)) {
@@ -359,6 +373,9 @@ private void initCodec(@Nullable Object codec) {
 		if (codec instanceof DecoderHttpMessageReader) {
 			codec = ((DecoderHttpMessageReader) codec).getDecoder();
 		}
+		else if (codec instanceof EncoderHttpMessageWriter) {
+			codec = ((EncoderHttpMessageWriter<?>) codec).getEncoder();
+		}
 
 		if (codec == null) {
 			return;
@@ -394,7 +411,6 @@ private void initCodec(@Nullable Object codec) {
 			}
 			if (codec instanceof ServerSentEventHttpMessageReader) {
 				((ServerSentEventHttpMessageReader) codec).setMaxInMemorySize(size);
-				initCodec(((ServerSentEventHttpMessageReader) codec).getDecoder());
 			}
 			if (codec instanceof DefaultPartHttpMessageReader) {
 				((DefaultPartHttpMessageReader) codec).setMaxInMemorySize(size);
@@ -430,12 +446,23 @@ private void initCodec(@Nullable Object codec) {
 			}
 		}
 
+		if (this.codecConsumer != null) {
+			this.codecConsumer.accept(codec);
+		}
+
+		// Recurse for nested codecs
 		if (codec instanceof MultipartHttpMessageReader) {
 			initCodec(((MultipartHttpMessageReader) codec).getPartReader());
 		}
 		else if (codec instanceof MultipartHttpMessageWriter) {
 			initCodec(((MultipartHttpMessageWriter) codec).getFormWriter());
 		}
+		else if (codec instanceof ServerSentEventHttpMessageReader) {
+			initCodec(((ServerSentEventHttpMessageReader) codec).getDecoder());
+		}
+		else if (codec instanceof ServerSentEventHttpMessageWriter) {
+			initCodec(((ServerSentEventHttpMessageWriter) codec).getEncoder());
+		}
 	}
 
 	/**
@@ -521,22 +548,21 @@ protected void initTypedWriters() {
 	/**
 	 * Return "base" typed writers only, i.e. common to client and server.
 	 */
-	@SuppressWarnings("unchecked")
 	final List<HttpMessageWriter<?>> getBaseTypedWriters() {
 		if (!this.registerDefaults) {
 			return Collections.emptyList();
 		}
 		List<HttpMessageWriter<?>> writers = new ArrayList<>();
-		writers.add(new EncoderHttpMessageWriter<>(new ByteArrayEncoder()));
-		writers.add(new EncoderHttpMessageWriter<>(new ByteBufferEncoder()));
-		writers.add(new EncoderHttpMessageWriter<>(new DataBufferEncoder()));
+		addCodec(writers, new EncoderHttpMessageWriter<>(new ByteArrayEncoder()));
+		addCodec(writers, new EncoderHttpMessageWriter<>(new ByteBufferEncoder()));
+		addCodec(writers, new EncoderHttpMessageWriter<>(new DataBufferEncoder()));
 		if (nettyByteBufPresent) {
-			writers.add(new EncoderHttpMessageWriter<>(new NettyByteBufEncoder()));
+			addCodec(writers, new EncoderHttpMessageWriter<>(new NettyByteBufEncoder()));
 		}
-		writers.add(new ResourceHttpMessageWriter());
-		writers.add(new EncoderHttpMessageWriter<>(CharSequenceEncoder.textPlainOnly()));
+		addCodec(writers, new ResourceHttpMessageWriter());
+		addCodec(writers, new EncoderHttpMessageWriter<>(CharSequenceEncoder.textPlainOnly()));
 		if (protobufPresent) {
-			writers.add(new ProtobufHttpMessageWriter(this.protobufEncoder != null ?
+			addCodec(writers, new ProtobufHttpMessageWriter(this.protobufEncoder != null ?
 					(ProtobufEncoder) this.protobufEncoder : new ProtobufEncoder()));
 		}
 		return writers;
@@ -574,17 +600,17 @@ protected void initObjectWriters() {
 	final List<HttpMessageWriter<?>> getBaseObjectWriters() {
 		List<HttpMessageWriter<?>> writers = new ArrayList<>();
 		if (kotlinSerializationJsonPresent) {
-			writers.add(new EncoderHttpMessageWriter<>(getKotlinSerializationJsonEncoder()));
+			addCodec(writers, new EncoderHttpMessageWriter<>(getKotlinSerializationJsonEncoder()));
 		}
 		if (jackson2Present) {
-			writers.add(new EncoderHttpMessageWriter<>(getJackson2JsonEncoder()));
+			addCodec(writers, new EncoderHttpMessageWriter<>(getJackson2JsonEncoder()));
 		}
 		if (jackson2SmilePresent) {
-			writers.add(new EncoderHttpMessageWriter<>(this.jackson2SmileEncoder != null ?
+			addCodec(writers, new EncoderHttpMessageWriter<>(this.jackson2SmileEncoder != null ?
 					(Jackson2SmileEncoder) this.jackson2SmileEncoder : new Jackson2SmileEncoder()));
 		}
 		if (jaxb2Present && !shouldIgnoreXml) {
-			writers.add(new EncoderHttpMessageWriter<>(this.jaxb2Encoder != null ?
+			addCodec(writers, new EncoderHttpMessageWriter<>(this.jaxb2Encoder != null ?
 					(Jaxb2XmlEncoder) this.jaxb2Encoder : new Jaxb2XmlEncoder()));
 		}
 		return writers;

spring-web/src/test/java/org/springframework/http/codec/support/ClientCodecConfigurerTests.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2020 the original author or authors.
+ * Copyright 2002-2021 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.
@@ -23,6 +23,7 @@
 import java.util.List;
 import java.util.concurrent.atomic.AtomicInteger;
 
+import com.fasterxml.jackson.databind.ObjectMapper;
 import org.junit.jupiter.api.Test;
 import reactor.core.publisher.Flux;
 
@@ -52,6 +53,7 @@
 import org.springframework.http.codec.ResourceHttpMessageReader;
 import org.springframework.http.codec.ResourceHttpMessageWriter;
 import org.springframework.http.codec.ServerSentEventHttpMessageReader;
+import org.springframework.http.codec.json.Jackson2CodecSupport;
 import org.springframework.http.codec.json.Jackson2JsonDecoder;
 import org.springframework.http.codec.json.Jackson2JsonEncoder;
 import org.springframework.http.codec.json.Jackson2SmileDecoder;
@@ -121,12 +123,29 @@ public void defaultWriters() {
 	}
 
 	@Test
-	public void jackson2EncoderOverride() {
+	public void jackson2CodecCustomizations() {
 		Jackson2JsonDecoder decoder = new Jackson2JsonDecoder();
+		Jackson2JsonEncoder encoder = new Jackson2JsonEncoder();
 		this.configurer.defaultCodecs().jackson2JsonDecoder(decoder);
+		this.configurer.defaultCodecs().jackson2JsonEncoder(encoder);
+
+		ObjectMapper objectMapper = new ObjectMapper();
+		this.configurer.defaultCodecs().configureDefaultCodec(codec -> {
+			if (codec instanceof Jackson2CodecSupport) {
+				((Jackson2CodecSupport) codec).setObjectMapper(objectMapper);
+			}
+		});
 
 		List<HttpMessageReader<?>> readers = this.configurer.getReaders();
+		Jackson2JsonDecoder actualDecoder = findCodec(readers, Jackson2JsonDecoder.class);
+		assertThat(actualDecoder).isSameAs(decoder);
+		assertThat(actualDecoder.getObjectMapper()).isSameAs(objectMapper);
 		assertThat(findCodec(readers, ServerSentEventHttpMessageReader.class).getDecoder()).isSameAs(decoder);
+
+		List<HttpMessageWriter<?>> writers = this.configurer.getWriters();
+		Jackson2JsonEncoder actualEncoder = findCodec(writers, Jackson2JsonEncoder.class);
+		assertThat(actualEncoder).isSameAs(encoder);
+		assertThat(actualEncoder.getObjectMapper()).isSameAs(objectMapper);
 	}
 
 	@Test
@@ -237,7 +256,19 @@ private Encoder<?> getNextEncoder(List<HttpMessageWriter<?>> writers) {
 
 	@SuppressWarnings("unchecked")
 	private <T> T findCodec(List<?> codecs, Class<T> type) {
-		return (T) codecs.stream().filter(type::isInstance).findFirst().get();
+		return (T) codecs.stream()
+				.map(c -> {
+					if (c instanceof EncoderHttpMessageWriter) {
+						return ((EncoderHttpMessageWriter<?>) c).getEncoder();
+					}
+					else if (c instanceof DecoderHttpMessageReader) {
+						return ((DecoderHttpMessageReader<?>) c).getDecoder();
+					}
+					else {
+						return c;
+					}
+				})
+				.filter(type::isInstance).findFirst().get();
 	}
 
 	@SuppressWarnings("unchecked")