BufferedImage converter writes Content-Type again

Issue: SPR-13906

spring-web/src/main/java/org/springframework/http/converter/BufferedImageHttpMessageConverter.java

@@ -38,7 +38,6 @@ import javax.imageio.stream.ImageOutputStream;
 import javax.imageio.stream.MemoryCacheImageInputStream;
 import javax.imageio.stream.MemoryCacheImageOutputStream;
 
-import org.springframework.http.HttpHeaders;
 import org.springframework.http.HttpInputMessage;
 import org.springframework.http.HttpOutputMessage;
 import org.springframework.http.MediaType;
@@ -209,29 +208,35 @@ public class BufferedImageHttpMessageConverter implements HttpMessageConverter<B
 			final HttpOutputMessage outputMessage)
 			throws IOException, HttpMessageNotWritableException {
 
+		final MediaType selectedContentType = getContentType(contentType);
+		outputMessage.getHeaders().setContentType(selectedContentType);
+
 		if (outputMessage instanceof StreamingHttpOutputMessage) {
 			StreamingHttpOutputMessage streamingOutputMessage = (StreamingHttpOutputMessage) outputMessage;
 			streamingOutputMessage.setBody(new StreamingHttpOutputMessage.Body() {
 				@Override
 				public void writeTo(OutputStream outputStream) throws IOException {
-					writeInternal(image, contentType, outputMessage.getHeaders(), outputStream);
+					writeInternal(image, selectedContentType, outputStream);
 				}
 			});
 		}
 		else {
-			writeInternal(image, contentType, outputMessage.getHeaders(), outputMessage.getBody());
+			writeInternal(image, selectedContentType, outputMessage.getBody());
 		}
 	}
 
-	private void writeInternal(BufferedImage image, MediaType contentType, HttpHeaders headers,
-			OutputStream body) throws IOException, HttpMessageNotWritableException {
-
+	private MediaType getContentType(MediaType contentType) {
 		if (contentType == null || contentType.isWildcardType() || contentType.isWildcardSubtype()) {
 			contentType = getDefaultContentType();
 		}
-		Assert.notNull(contentType,
-				"Count not determine Content-Type, set one using the 'defaultContentType' property");
-		headers.setContentType(contentType);
+		Assert.notNull(contentType, "Could not select Content-Type. " +
+				"Please specify one through the 'defaultContentType' property.");
+		return contentType;
+	}
+
+	private void writeInternal(BufferedImage image, MediaType contentType, OutputStream body)
+			throws IOException, HttpMessageNotWritableException {
+
 		ImageOutputStream imageOutputStream = null;
 		ImageWriter imageWriter = null;
 		try {

spring-web/src/test/java/org/springframework/http/MockHttpOutputMessage.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2014 the original author or authors.
+ * Copyright 2002-2016 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.
@@ -21,10 +21,11 @@ import java.io.IOException;
 import java.io.OutputStream;
 import java.nio.charset.Charset;
 
-import static org.mockito.Mockito.*;
+import static org.mockito.Mockito.spy;
 
 /**
  * @author Arjen Poutsma
+ * @author Rossen Stoyanchev
  */
 public class MockHttpOutputMessage implements HttpOutputMessage {
 
@@ -34,19 +35,31 @@ public class MockHttpOutputMessage implements HttpOutputMessage {
 
 	private boolean headersWritten = false;
 
+	private final HttpHeaders writtenHeaders = new HttpHeaders();
+
+
 	@Override
 	public HttpHeaders getHeaders() {
 		return (this.headersWritten ? HttpHeaders.readOnlyHttpHeaders(this.headers) : this.headers);
 	}
 
+	/**
+	 * Return a copy of the actual headers written at the time of the call to
+	 * getBody, i.e. ignoring any further changes that may have been made to
+	 * the underlying headers, e.g. via a previously obtained instance.
+	 */
+	public HttpHeaders getWrittenHeaders() {
+		return writtenHeaders;
+	}
+
 	@Override
 	public OutputStream getBody() throws IOException {
-		this.headersWritten = true;
+		writeHeaders();
 		return body;
 	}
 
 	public byte[] getBodyAsBytes() {
-		this.headersWritten = true;
+		writeHeaders();
 		return body.toByteArray();
 	}
 
@@ -54,4 +67,13 @@ public class MockHttpOutputMessage implements HttpOutputMessage {
 		byte[] bytes = getBodyAsBytes();
 		return new String(bytes, charset);
 	}
+
+	private void writeHeaders() {
+		if (this.headersWritten) {
+			return;
+		}
+		this.headersWritten = true;
+		this.writtenHeaders.putAll(this.headers);
+	}
+
 }

spring-web/src/test/java/org/springframework/http/converter/BufferedImageHttpMessageConverterTests.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2013 the original author or authors.
+ * Copyright 2002-2016 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.
@@ -33,6 +33,11 @@ import org.springframework.util.FileCopyUtils;
 
 import static org.junit.Assert.*;
 
+/**
+ * Unit tests for BufferedImageHttpMessageConverter.
+ * @author Arjen Poutsma
+ * @author Rossen Stoyanchev
+ */
 public class BufferedImageHttpMessageConverterTests {
 
 	private BufferedImageHttpMessageConverter converter;
@@ -73,7 +78,7 @@ public class BufferedImageHttpMessageConverterTests {
 		MockHttpOutputMessage outputMessage = new MockHttpOutputMessage();
 		MediaType contentType = new MediaType("image", "png");
 		converter.write(body, contentType, outputMessage);
-		assertEquals("Invalid content type", contentType, outputMessage.getHeaders().getContentType());
+		assertEquals("Invalid content type", contentType, outputMessage.getWrittenHeaders().getContentType());
 		assertTrue("Invalid size", outputMessage.getBodyAsBytes().length > 0);
 		BufferedImage result = ImageIO.read(new ByteArrayInputStream(outputMessage.getBodyAsBytes()));
 		assertEquals("Invalid height", 500, result.getHeight());
@@ -88,7 +93,7 @@ public class BufferedImageHttpMessageConverterTests {
 		BufferedImage body = ImageIO.read(logo.getFile());
 		MockHttpOutputMessage outputMessage = new MockHttpOutputMessage();
 		converter.write(body, new MediaType("*", "*"), outputMessage);
-		assertEquals("Invalid content type", contentType, outputMessage.getHeaders().getContentType());
+		assertEquals("Invalid content type", contentType, outputMessage.getWrittenHeaders().getContentType());
 		assertTrue("Invalid size", outputMessage.getBodyAsBytes().length > 0);
 		BufferedImage result = ImageIO.read(new ByteArrayInputStream(outputMessage.getBodyAsBytes()));
 		assertEquals("Invalid height", 500, result.getHeight());