SSE writer respects charset in MediaType

Issue: SPR-16516

spring-web/src/main/java/org/springframework/http/codec/ServerSentEventHttpMessageWriter.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.
@@ -37,6 +37,7 @@ import org.springframework.http.ReactiveHttpOutputMessage;
 import org.springframework.http.server.reactive.ServerHttpRequest;
 import org.springframework.http.server.reactive.ServerHttpResponse;
 import org.springframework.lang.Nullable;
+import org.springframework.util.Assert;
 
 /**
  * {@code HttpMessageWriter} for {@code "text/event-stream"} responses.
@@ -48,8 +49,11 @@ import org.springframework.lang.Nullable;
  */
 public class ServerSentEventHttpMessageWriter implements HttpMessageWriter<Object> {
 
+	private static final MediaType DEFAULT_MEDIA_TYPE = new MediaType("text", "event-stream", StandardCharsets.UTF_8);
+
 	private static final List<MediaType> WRITABLE_MEDIA_TYPES = Collections.singletonList(MediaType.TEXT_EVENT_STREAM);
 
+
 	@Nullable
 	private final Encoder<?> encoder;
 
@@ -96,12 +100,15 @@ public class ServerSentEventHttpMessageWriter implements HttpMessageWriter<Objec
 	public Mono<Void> write(Publisher<?> input, ResolvableType elementType, @Nullable MediaType mediaType,
 			ReactiveHttpOutputMessage message, Map<String, Object> hints) {
 
-		message.getHeaders().setContentType(MediaType.TEXT_EVENT_STREAM);
-		return message.writeAndFlushWith(encode(input, message.bufferFactory(), elementType, hints));
+		mediaType = (mediaType != null && mediaType.getCharset() != null ? mediaType : DEFAULT_MEDIA_TYPE);
+		DataBufferFactory bufferFactory = message.bufferFactory();
+
+		message.getHeaders().setContentType(mediaType);
+		return message.writeAndFlushWith(encode(input, elementType, mediaType, bufferFactory, hints));
 	}
 
-	private Flux<Publisher<DataBuffer>> encode(Publisher<?> input, DataBufferFactory factory,
-			ResolvableType elementType, Map<String, Object> hints) {
+	private Flux<Publisher<DataBuffer>> encode(Publisher<?> input, ResolvableType elementType,
+			MediaType mediaType, DataBufferFactory factory, Map<String, Object> hints) {
 
 		Class<?> elementClass = elementType.getRawClass();
 		ResolvableType valueType = (elementClass != null && ServerSentEvent.class.isAssignableFrom(elementClass) ?
@@ -134,9 +141,9 @@ public class ServerSentEventHttpMessageWriter implements HttpMessageWriter<Objec
 				sb.append("data:");
 			}
 
-			return Flux.concat(encodeText(sb, factory),
-					encodeData(data, valueType, factory, hints),
-					encodeText("\n", factory));
+			return Flux.concat(encodeText(sb, mediaType, factory),
+					encodeData(data, valueType, mediaType, factory, hints),
+					encodeText("\n", mediaType, factory));
 		});
 	}
 
@@ -149,7 +156,7 @@ public class ServerSentEventHttpMessageWriter implements HttpMessageWriter<Objec
 
 	@SuppressWarnings("unchecked")
 	private <T> Flux<DataBuffer> encodeData(@Nullable T data, ResolvableType valueType,
-			DataBufferFactory factory, Map<String, Object> hints) {
+			MediaType mediaType, DataBufferFactory factory, Map<String, Object> hints) {
 
 		if (data == null) {
 			return Flux.empty();
@@ -157,7 +164,7 @@ public class ServerSentEventHttpMessageWriter implements HttpMessageWriter<Objec
 
 		if (data instanceof String) {
 			String text = (String) data;
-			return Flux.from(encodeText(text.replaceAll("\\n", "\ndata:") + "\n", factory));
+			return Flux.from(encodeText(text.replaceAll("\\n", "\ndata:") + "\n", mediaType, factory));
 		}
 
 		if (this.encoder == null) {
@@ -165,12 +172,13 @@ public class ServerSentEventHttpMessageWriter implements HttpMessageWriter<Objec
 		}
 
 		return ((Encoder<T>) this.encoder)
-				.encode(Mono.just(data), factory, valueType, MediaType.TEXT_EVENT_STREAM, hints)
-				.concatWith(encodeText("\n", factory));
+				.encode(Mono.just(data), factory, valueType, mediaType, hints)
+				.concatWith(encodeText("\n", mediaType, factory));
 	}
 
-	private Mono<DataBuffer> encodeText(CharSequence text, DataBufferFactory bufferFactory) {
-		byte[] bytes = text.toString().getBytes(StandardCharsets.UTF_8);
+	private Mono<DataBuffer> encodeText(CharSequence text, MediaType mediaType, DataBufferFactory bufferFactory) {
+		Assert.notNull(mediaType.getCharset(), "Expected MediaType with charset");
+		byte[] bytes = text.toString().getBytes(mediaType.getCharset());
 		DataBuffer buffer = bufferFactory.allocateBuffer(bytes.length).write(bytes);
 		return Mono.just(buffer);
 	}

spring-web/src/test/java/org/springframework/http/codec/ServerSentEventHttpMessageWriterTests.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.http.codec;
 
+import java.nio.charset.Charset;
+import java.nio.charset.StandardCharsets;
 import java.time.Duration;
 import java.util.Collections;
 import java.util.Map;
 
 import com.fasterxml.jackson.databind.ObjectMapper;
+import org.junit.Ignore;
 import org.junit.Test;
 import org.reactivestreams.Publisher;
 import reactor.core.publisher.Flux;
@@ -34,9 +37,8 @@ import org.springframework.http.codec.json.Jackson2JsonEncoder;
 import org.springframework.http.converter.json.Jackson2ObjectMapperBuilder;
 import org.springframework.mock.http.server.reactive.test.MockServerHttpResponse;
 
-import static org.junit.Assert.assertFalse;
-import static org.junit.Assert.assertTrue;
-import static org.springframework.core.ResolvableType.forClass;
+import static org.junit.Assert.*;
+import static org.springframework.core.ResolvableType.*;
 
 /**
  * Unit tests for {@link ServerSentEventHttpMessageWriter}.
@@ -45,7 +47,8 @@ import static org.springframework.core.ResolvableType.forClass;
  */
 public class ServerSentEventHttpMessageWriterTests extends AbstractDataBufferAllocatingTestCase {
 
-	public static final Map<String, Object> HINTS = Collections.emptyMap();
+	private static final Map<String, Object> HINTS = Collections.emptyMap();
+
 
 	private ServerSentEventHttpMessageWriter messageWriter =
 			new ServerSentEventHttpMessageWriter(new Jackson2JsonEncoder());
@@ -105,6 +108,18 @@ public class ServerSentEventHttpMessageWriterTests extends AbstractDataBufferAll
 				.verify();
 	}
 
+	@Test // SPR-16516
+	public void writeStringWithCustomCharset() {
+		Flux<String> source = Flux.just("\u00A3");
+		Charset charset = StandardCharsets.ISO_8859_1;
+		MediaType mediaType = new MediaType("text", "event-stream", charset);
+		MockServerHttpResponse outputMessage = new MockServerHttpResponse();
+		testWrite(source, mediaType, outputMessage, String.class);
+
+		assertEquals(mediaType, outputMessage.getHeaders().getContentType());
+		StepVerifier.create(outputMessage.getBodyAsString()).expectNext("data:\u00A3\n\n").verifyComplete();
+	}
+
 	@Test
 	public void writePojo() {
 		Flux<Pojo> source = Flux.just(new Pojo("foofoo", "barbar"), new Pojo("foofoofoo", "barbarbar"));
@@ -139,9 +154,32 @@ public class ServerSentEventHttpMessageWriterTests extends AbstractDataBufferAll
 				.verify();
 	}
 
-	private <T> void testWrite(Publisher<T> source, MockServerHttpResponse outputMessage, Class<T> clazz) {
-		this.messageWriter.write(source, forClass(clazz),
-				MediaType.TEXT_EVENT_STREAM, outputMessage, HINTS).block(Duration.ofMillis(5000));
+	@Ignore
+	@Test // SPR-16516, SPR-16539
+	public void writePojoWithCustomEncoding() {
+		Flux<Pojo> source = Flux.just(new Pojo("foo\u00A3", "bar\u00A3"));
+		Charset charset = StandardCharsets.ISO_8859_1;
+		MediaType mediaType = new MediaType("text", "event-stream", charset);
+		MockServerHttpResponse outputMessage = new MockServerHttpResponse();
+		testWrite(source, mediaType, outputMessage, Pojo.class);
+
+		assertEquals(mediaType, outputMessage.getHeaders().getContentType());
+		StepVerifier.create(outputMessage.getBodyAsString())
+				.expectNext("data:{\"foo\":\"foo\u00A3\",\"bar\":\"bar\u00A3\"}\n\n")
+				.expectComplete()
+				.verify();
+	}
+
+
+	private <T> void testWrite(Publisher<T> source, MockServerHttpResponse response, Class<T> clazz) {
+		testWrite(source, MediaType.TEXT_EVENT_STREAM, response, clazz);
+	}
+
+	private <T> void testWrite(Publisher<T> source, MediaType mediaType, MockServerHttpResponse response,
+			Class<T> clazz) {
+
+		this.messageWriter.write(source, forClass(clazz), mediaType, response, HINTS)
+				.block(Duration.ofMillis(5000));
 	}
 
 }