content-length support in EncoderHttpMessageWriter

EncoderHttpMessageWriter checks explicitly for Mono publishers and sets
the content length, if it is known for the given data item.

Issue: SPR-16542

Author: rstoyanchev

spring-core/src/main/java/org/springframework/core/codec/ByteArrayEncoder.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2017 the original author or authors.
+ * Copyright 2002-2018 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.
@@ -55,4 +55,8 @@ public Flux<DataBuffer> encode(Publisher<? extends byte[]> inputStream,
 		return Flux.from(inputStream).map(bufferFactory::wrap);
 	}
 
+	@Override
+	public Long getContentLength(byte[] bytes, @Nullable MimeType mimeType) {
+		return (long) bytes.length;
+	}
 }

spring-core/src/main/java/org/springframework/core/codec/ByteBufferEncoder.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2017 the original author or authors.
+ * Copyright 2002-2018 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.
@@ -56,4 +56,8 @@ public Flux<DataBuffer> encode(Publisher<? extends ByteBuffer> inputStream,
 		return Flux.from(inputStream).map(bufferFactory::wrap);
 	}
 
+	@Override
+	public Long getContentLength(ByteBuffer byteBuffer, @Nullable MimeType mimeType) {
+		return (long) byteBuffer.array().length;
+	}
 }

spring-core/src/main/java/org/springframework/core/codec/CharSequenceEncoder.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2017 the original author or authors.
+ * Copyright 2002-2018 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.
@@ -62,20 +62,30 @@ public Flux<DataBuffer> encode(Publisher<? extends CharSequence> inputStream,
 			DataBufferFactory bufferFactory, ResolvableType elementType,
 			@Nullable MimeType mimeType, @Nullable Map<String, Object> hints) {
 
+		Charset charset = getCharset(mimeType);
+
+		return Flux.from(inputStream).map(charSequence -> {
+			CharBuffer charBuffer = CharBuffer.wrap(charSequence);
+			ByteBuffer byteBuffer = charset.encode(charBuffer);
+			return bufferFactory.wrap(byteBuffer);
+		});
+	}
+
+	private Charset getCharset(@Nullable MimeType mimeType) {
 		Charset charset;
 		if (mimeType != null && mimeType.getCharset() != null) {
 			charset = mimeType.getCharset();
 		}
 		else {
 			 charset = DEFAULT_CHARSET;
 		}
-		return Flux.from(inputStream).map(charSequence -> {
-			CharBuffer charBuffer = CharBuffer.wrap(charSequence);
-			ByteBuffer byteBuffer = charset.encode(charBuffer);
-			return bufferFactory.wrap(byteBuffer);
-		});
+		return charset;
 	}
 
+	@Override
+	public Long getContentLength(CharSequence data, @Nullable MimeType mimeType) {
+		return (long) data.toString().getBytes(getCharset(mimeType)).length;
+	}
 
 	/**
 	 * Create a {@code CharSequenceEncoder} that supports only "text/plain".

spring-core/src/main/java/org/springframework/core/codec/DataBufferEncoder.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2017 the original author or authors.
+ * Copyright 2002-2018 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.
@@ -55,4 +55,8 @@ public Flux<DataBuffer> encode(Publisher<? extends DataBuffer> inputStream,
 		return Flux.from(inputStream);
 	}
 
+	@Override
+	public Long getContentLength(DataBuffer dataBuffer, MimeType mimeType) {
+		return (long) dataBuffer.readableByteCount();
+	}
 }

spring-core/src/main/java/org/springframework/core/codec/Encoder.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2016 the original author or authors.
+ * Copyright 2002-2018 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.
@@ -67,6 +67,17 @@ public interface Encoder<T> {
 	Flux<DataBuffer> encode(Publisher<? extends T> inputStream, DataBufferFactory bufferFactory,
 			ResolvableType elementType, @Nullable MimeType mimeType, @Nullable Map<String, Object> hints);
 
+	/**
+	 * Return the length for the given item, if known.
+	 * @param t the item to check
+	 * @return the length in bytes, or {@code null} if not known.
+	 * @since 5.0.5
+	 */
+	@Nullable
+	default Long getContentLength(T t, @Nullable MimeType mimeType) {
+		return null;
+	}
+
 	/**
 	 * Return the list of mime types this encoder supports.
 	 */

spring-core/src/main/java/org/springframework/core/codec/ResourceEncoder.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2017 the original author or authors.
+ * Copyright 2002-2018 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.
@@ -16,11 +16,14 @@
 
 package org.springframework.core.codec;
 
+import java.io.IOException;
 import java.util.Map;
+import java.util.OptionalLong;
 
 import reactor.core.publisher.Flux;
 
 import org.springframework.core.ResolvableType;
+import org.springframework.core.io.InputStreamResource;
 import org.springframework.core.io.Resource;
 import org.springframework.core.io.buffer.DataBuffer;
 import org.springframework.core.io.buffer.DataBufferFactory;
@@ -68,4 +71,17 @@ protected Flux<DataBuffer> encode(Resource resource, DataBufferFactory dataBuffe
 		return DataBufferUtils.read(resource, dataBufferFactory, this.bufferSize);
 	}
 
+	@Override
+	public Long getContentLength(Resource resource, @Nullable MimeType mimeType) {
+		// Don't consume InputStream...
+		if (InputStreamResource.class != resource.getClass()) {
+			try {
+				return resource.contentLength();
+			}
+			catch (IOException ignored) {
+			}
+		}
+		return null;
+	}
+
 }

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2017 the original author or authors.
+ * Copyright 2002-2018 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.
@@ -28,6 +28,7 @@
 import org.springframework.core.ResolvableType;
 import org.springframework.core.codec.Encoder;
 import org.springframework.core.io.buffer.DataBuffer;
+import org.springframework.http.HttpHeaders;
 import org.springframework.http.MediaType;
 import org.springframework.http.ReactiveHttpOutputMessage;
 import org.springframework.http.server.reactive.ServerHttpRequest;
@@ -91,11 +92,25 @@ public boolean canWrite(ResolvableType elementType, @Nullable MediaType mediaTyp
 		return this.encoder.canEncode(elementType, mediaType);
 	}
 
+	@SuppressWarnings("unchecked")
 	@Override
 	public Mono<Void> write(Publisher<? extends T> inputStream, ResolvableType elementType,
 			@Nullable MediaType mediaType, ReactiveHttpOutputMessage message, Map<String, Object> hints) {
 
 		MediaType contentType = updateContentType(message, mediaType);
+		HttpHeaders headers = message.getHeaders();
+
+		if (headers.getContentLength() < 0 && !headers.containsKey(HttpHeaders.TRANSFER_ENCODING)) {
+			if (inputStream instanceof Mono) {
+				// This works because we don't actually commit until after the first signal...
+				inputStream = ((Mono<T>) inputStream).doOnNext(data -> {
+					Long contentLength = this.encoder.getContentLength(data, contentType);
+					if (contentLength != null) {
+						headers.setContentLength(contentLength);
+					}
+				});
+			}
+		}
 
 		Flux<DataBuffer> body = this.encoder.encode(
 				inputStream, message.bufferFactory(), elementType, contentType, hints);

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2017 the original author or authors.
+ * Copyright 2002-2018 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.
@@ -22,7 +22,6 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
-import java.util.OptionalLong;
 
 import org.reactivestreams.Publisher;
 import reactor.core.publisher.Flux;
@@ -32,7 +31,6 @@
 import org.springframework.core.codec.ResourceDecoder;
 import org.springframework.core.codec.ResourceEncoder;
 import org.springframework.core.codec.ResourceRegionEncoder;
-import org.springframework.core.io.InputStreamResource;
 import org.springframework.core.io.Resource;
 import org.springframework.core.io.buffer.DataBuffer;
 import org.springframework.core.io.buffer.DataBufferFactory;
@@ -49,7 +47,7 @@
 import org.springframework.lang.Nullable;
 import org.springframework.util.MimeTypeUtils;
 
-import static java.util.Collections.emptyMap;
+import static java.util.Collections.*;
 
 /**
  * {@code HttpMessageWriter} that can write a {@link Resource}.
@@ -121,7 +119,10 @@ private Mono<Void> writeResource(Resource resource, ResolvableType type, @Nullab
 		headers.setContentType(resourceMediaType);
 
 		if (headers.getContentLength() < 0) {
-			lengthOf(resource).ifPresent(headers::setContentLength);
+			Long contentLength = this.encoder.getContentLength(resource, mediaType);
+			if (contentLength != null) {
+				headers.setContentLength(contentLength);
+			}
 		}
 
 		return zeroCopy(resource, null, message)
@@ -140,18 +141,6 @@ private static MediaType getResourceMediaType(@Nullable MediaType mediaType, Res
 		return MediaTypeFactory.getMediaType(resource).orElse(MediaType.APPLICATION_OCTET_STREAM);
 	}
 
-	private static OptionalLong lengthOf(Resource resource) {
-		// Don't consume InputStream...
-		if (InputStreamResource.class != resource.getClass()) {
-			try {
-				return OptionalLong.of(resource.contentLength());
-			}
-			catch (IOException ignored) {
-			}
-		}
-		return OptionalLong.empty();
-	}
-
 	private static Optional<Mono<Void>> zeroCopy(Resource resource, @Nullable ResourceRegion region,
 			ReactiveHttpOutputMessage message) {
 
@@ -205,13 +194,14 @@ public Mono<Void> write(Publisher<? extends Resource> inputStream, @Nullable Res
 			if (regions.size() == 1){
 				ResourceRegion region = regions.get(0);
 				headers.setContentType(resourceMediaType);
-				lengthOf(resource).ifPresent(length -> {
+				Long contentLength = this.encoder.getContentLength(resource, mediaType);
+				if (contentLength != null) {
 					long start = region.getPosition();
 					long end = start + region.getCount() - 1;
-					end = Math.min(end, length - 1);
-					headers.add("Content-Range", "bytes " + start + '-' + end + '/' + length);
+					end = Math.min(end, contentLength - 1);
+					headers.add("Content-Range", "bytes " + start + '-' + end + '/' + contentLength);
 					headers.setContentLength(end - start + 1);
-				});
+				}
 				return writeSingleRegion(region, response);
 			}
 			else {

spring-webflux/src/test/java/org/springframework/web/reactive/result/method/annotation/RequestMappingMessageConversionIntegrationTests.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2017 the original author or authors.
+ * Copyright 2002-2018 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.
@@ -59,10 +59,9 @@
 import org.springframework.web.bind.annotation.RestController;
 import org.springframework.web.reactive.config.EnableWebFlux;
 
-import static java.util.Arrays.asList;
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertTrue;
-import static org.springframework.http.MediaType.APPLICATION_XML;
+import static java.util.Arrays.*;
+import static org.junit.Assert.*;
+import static org.springframework.http.MediaType.*;
 
 /**
  * {@code @RequestMapping} integration tests focusing on serialization and
@@ -87,7 +86,6 @@ protected ApplicationContext initApplicationContext() {
 		return wac;
 	}
 
-
 	@Test
 	public void byteBufferResponseBodyWithPublisher() throws Exception {
 		Person expected = new Person("Robert");
@@ -100,6 +98,14 @@ public void byteBufferResponseBodyWithFlux() throws Exception {
 		assertEquals(expected, performGet("/raw-response/flux", new HttpHeaders(), String.class).getBody());
 	}
 
+	@Test
+	public void byteBufferResponseBodyWithMono() throws Exception {
+		String expected = "Hello!";
+		ResponseEntity<String> responseEntity = performGet("/raw-response/mono", new HttpHeaders(), String.class);
+		assertEquals(6, responseEntity.getHeaders().getContentLength());
+		assertEquals(expected, responseEntity.getBody());
+	}
+
 	@Test
 	public void byteBufferResponseBodyWithObservable() throws Exception {
 		String expected = "Hello!";
@@ -422,6 +428,11 @@ public Flux<ByteBuffer> getFlux() {
 			return Flux.just(ByteBuffer.wrap("Hello!".getBytes()));
 		}
 
+		@GetMapping("/mono")
+		public Mono<ByteBuffer> getMonoString() {
+			return Mono.just(ByteBuffer.wrap("Hello!".getBytes()));
+		}
+
 		@GetMapping("/observable")
 		public Observable<ByteBuffer> getObservable() {
 			return Observable.just(ByteBuffer.wrap("Hello!".getBytes()));