From b9348bb89c6116c6c57b68c34b2ff28b5b3d4ce3 Mon Sep 17 00:00:00 2001
From: Brian Clozel <bclozel@pivotal.io>
Date: Mon, 22 Sep 2014 14:02:35 +0200
Subject: [PATCH] Fix Protobuf support - HTTP headers already written

Prior to this commit, the `ProtobufHttpMessageConverter` would call
`outputMessage.getBody()` at the beginning of the `writeInternal`
method, thus writing HTTP headers. Since this method is trying to write
"x-protobuf" headers after that, protobuf support wasn't working
properly for the default "x-protobuf" media type.

Thanks Toshiaki Maki for the repro project!

Also fixed:
* improve `MockHttpOutputMessage` behavior to reproduce the read-only
state of HTTP headers once they've been written.

Issue: SPR-12229
---
 .../protobuf/ProtobufHttpMessageConverter.java        | 11 ++++++-----
 .../springframework/http/MockHttpOutputMessage.java   |  8 ++++++--
 2 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/spring-web/src/main/java/org/springframework/http/converter/protobuf/ProtobufHttpMessageConverter.java b/spring-web/src/main/java/org/springframework/http/converter/protobuf/ProtobufHttpMessageConverter.java
index faf2173e82..b8c3d16b88 100644
--- a/spring-web/src/main/java/org/springframework/http/converter/protobuf/ProtobufHttpMessageConverter.java
+++ b/spring-web/src/main/java/org/springframework/http/converter/protobuf/ProtobufHttpMessageConverter.java
@@ -164,19 +164,18 @@ public class ProtobufHttpMessageConverter extends AbstractHttpMessageConverter<M
 
 		MediaType contentType = outputMessage.getHeaders().getContentType();
 		Charset charset = getCharset(contentType);
-		OutputStreamWriter writer = new OutputStreamWriter(outputMessage.getBody(), charset);
 
 		if (MediaType.TEXT_HTML.isCompatibleWith(contentType)) {
-			HtmlFormat.print(message, writer);
+			HtmlFormat.print(message, new OutputStreamWriter(outputMessage.getBody(), charset));
 		}
 		else if (MediaType.APPLICATION_JSON.isCompatibleWith(contentType)) {
-			JsonFormat.print(message, writer);
+			JsonFormat.print(message, new OutputStreamWriter(outputMessage.getBody(), charset));
 		}
 		else if (MediaType.TEXT_PLAIN.isCompatibleWith(contentType)) {
-			TextFormat.print(message, writer);
+			TextFormat.print(message, new OutputStreamWriter(outputMessage.getBody(), charset));
 		}
 		else if (MediaType.APPLICATION_XML.isCompatibleWith(contentType)) {
-			XmlFormat.print(message, writer);
+			XmlFormat.print(message, new OutputStreamWriter(outputMessage.getBody(), charset));
 		}
 		else if (PROTOBUF.isCompatibleWith(contentType)) {
 			setProtoHeader(outputMessage, message);
@@ -191,6 +190,8 @@ public class ProtobufHttpMessageConverter extends AbstractHttpMessageConverter<M
 	/**
 	 * Set the "X-Protobuf-*" HTTP headers when responding with a message of
 	 * content type "application/x-protobuf"
+	 * <p><b>Note:</b> <code>outputMessage.getBody()</code> should not have been called
+	 * before because it writes HTTP headers (making them read only).</p>
 	 */
 	private void setProtoHeader(HttpOutputMessage response, Message message) {
 		response.getHeaders().set(X_PROTOBUF_SCHEMA_HEADER, message.getDescriptorForType().getFile().getName());
diff --git a/spring-web/src/test/java/org/springframework/http/MockHttpOutputMessage.java b/spring-web/src/test/java/org/springframework/http/MockHttpOutputMessage.java
index dbb39e62b9..9ce157e4b9 100644
--- a/spring-web/src/test/java/org/springframework/http/MockHttpOutputMessage.java
+++ b/spring-web/src/test/java/org/springframework/http/MockHttpOutputMessage.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2013 the original author or authors.
+ * Copyright 2002-2014 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.
@@ -32,17 +32,21 @@ public class MockHttpOutputMessage implements HttpOutputMessage {
 
 	private final ByteArrayOutputStream body = spy(new ByteArrayOutputStream());
 
+	private boolean headersWritten = false;
+
 	@Override
 	public HttpHeaders getHeaders() {
-		return headers;
+		return (this.headersWritten ? HttpHeaders.readOnlyHttpHeaders(this.headers) : this.headers);
 	}
 
 	@Override
 	public OutputStream getBody() throws IOException {
+		this.headersWritten = true;
 		return body;
 	}
 
 	public byte[] getBodyAsBytes() {
+		this.headersWritten = true;
 		return body.toByteArray();
 	}