From 0bad69d5fbab74cd1afded96fbf5ed287ece7f18 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?S=C3=A9bastien=20Deleuze?= <sdeleuze@vmware.com>
Date: Fri, 7 Apr 2023 11:25:47 +0200
Subject: [PATCH] Fix SSE with indenting serializer in WebMvc.fn

This commit ensures that HTTP headers like "text/event-stream"
are correctly forwarded to the converter used in
SseServerResponse for proper pretty print handling.

Close gh-30302
---
 .../servlet/function/SseServerResponse.java   | 14 +++++++---
 .../function/SseServerResponseTests.java      | 28 +++++++++++++++++++
 2 files changed, 38 insertions(+), 4 deletions(-)

diff --git a/spring-webmvc/src/main/java/org/springframework/web/servlet/function/SseServerResponse.java b/spring-webmvc/src/main/java/org/springframework/web/servlet/function/SseServerResponse.java
index e078a66d110a..ffc74fde6369 100644
--- a/spring-webmvc/src/main/java/org/springframework/web/servlet/function/SseServerResponse.java
+++ b/spring-webmvc/src/main/java/org/springframework/web/servlet/function/SseServerResponse.java
@@ -48,6 +48,7 @@
  * <a href="https://www.w3.org/TR/eventsource/">Server-Sent Events</a>.
  *
  * @author Arjen Poutsma
+ * @author Sebastien Deleuze
  * @since 5.3.2
  */
 final class SseServerResponse extends AbstractServerResponse {
@@ -90,7 +91,7 @@ protected ModelAndView writeToInternal(HttpServletRequest request, HttpServletRe
 		}
 
 		DefaultAsyncServerResponse.writeAsync(request, response, result);
-		this.sseConsumer.accept(new DefaultSseBuilder(response, context, result));
+		this.sseConsumer.accept(new DefaultSseBuilder(response, context, result, this.headers()));
 		return null;
 	}
 
@@ -113,15 +114,19 @@ private static final class DefaultSseBuilder implements SseBuilder {
 
 		private final List<HttpMessageConverter<?>> messageConverters;
 
+		private final HttpHeaders httpHeaders;
+
 		private final StringBuilder builder = new StringBuilder();
 
 		private boolean sendFailed;
 
 
-		public DefaultSseBuilder(HttpServletResponse response, Context context, DeferredResult<?> deferredResult) {
+		public DefaultSseBuilder(HttpServletResponse response, Context context, DeferredResult<?> deferredResult,
+				HttpHeaders httpHeaders) {
 			this.outputMessage = new ServletServerHttpResponse(response);
 			this.deferredResult = deferredResult;
 			this.messageConverters = context.messageConverters();
+			this.httpHeaders = httpHeaders;
 		}
 
 		@Override
@@ -206,7 +211,7 @@ private void writeObject(Object data) throws IOException {
 				for (HttpMessageConverter<?> converter : this.messageConverters) {
 					if (converter.canWrite(dataClass, MediaType.APPLICATION_JSON)) {
 						HttpMessageConverter<Object> objectConverter = (HttpMessageConverter<Object>) converter;
-						ServerHttpResponse response = new MutableHeadersServerHttpResponse(this.outputMessage);
+						ServerHttpResponse response = new MutableHeadersServerHttpResponse(this.outputMessage, this.httpHeaders);
 						objectConverter.write(data, MediaType.APPLICATION_JSON, response);
 						this.outputMessage.getBody().write(NL_NL);
 						this.outputMessage.flush();
@@ -276,9 +281,10 @@ private static final class MutableHeadersServerHttpResponse extends DelegatingSe
 
 			private final HttpHeaders mutableHeaders = new HttpHeaders();
 
-			public MutableHeadersServerHttpResponse(ServerHttpResponse delegate) {
+			public MutableHeadersServerHttpResponse(ServerHttpResponse delegate, HttpHeaders headers) {
 				super(delegate);
 				this.mutableHeaders.putAll(delegate.getHeaders());
+				this.mutableHeaders.putAll(headers);
 			}
 
 			@Override
diff --git a/spring-webmvc/src/test/java/org/springframework/web/servlet/function/SseServerResponseTests.java b/spring-webmvc/src/test/java/org/springframework/web/servlet/function/SseServerResponseTests.java
index 0b9aff9d27d6..6e87a9eed366 100644
--- a/spring-webmvc/src/test/java/org/springframework/web/servlet/function/SseServerResponseTests.java
+++ b/spring-webmvc/src/test/java/org/springframework/web/servlet/function/SseServerResponseTests.java
@@ -33,6 +33,7 @@
 
 /**
  * @author Arjen Poutsma
+ * @author Sebastien Deleuze
  */
 class SseServerResponseTests {
 
@@ -89,6 +90,33 @@ void sendObject() throws Exception {
 		assertThat(this.mockResponse.getContentAsString()).isEqualTo(expected);
 	}
 
+	@Test
+	void sendObjectWithPrettyPrint() throws Exception {
+		Person person = new Person("John Doe", 42);
+		ServerResponse response = ServerResponse.sse(sse -> {
+			try {
+				sse.send(person);
+			}
+			catch (IOException ex) {
+				throw new UncheckedIOException(ex);
+			}
+		});
+
+		MappingJackson2HttpMessageConverter converter = new MappingJackson2HttpMessageConverter();
+		converter.setPrettyPrint(true);
+		ServerResponse.Context context = () -> Collections.singletonList(converter);
+
+		ModelAndView mav = response.writeTo(this.mockRequest, this.mockResponse, context);
+		assertThat(mav).isNull();
+
+		String expected = "data:{\n" +
+				"data:  \"name\" : \"John Doe\",\n" +
+				"data:  \"age\" : 42\n" +
+				"data:}\n" +
+				"\n";
+		assertThat(this.mockResponse.getContentAsString()).isEqualTo(expected);
+	}
+
 	@Test
 	void builder() throws Exception {
 		ServerResponse response = ServerResponse.sse(sse -> {