Improve header processing in SimpMessagingTemplate

Headers provided to the SimpMessagingTemplate's convertAndSend methods
are now automatically moved into the "nativeHeaders" sub-map. This
ensures the headers will go out with the STOMP message and be received
by subscribers.

Issue: SPR-11387

spring-messaging/src/main/java/org/springframework/messaging/core/AbstractMessageSendingTemplate.java

@@ -129,6 +129,7 @@ public abstract class AbstractMessageSendingTemplate<D> implements MessageSendin
 	public void convertAndSend(D destination, Object payload, Map<String, Object> headers,
 			MessagePostProcessor postProcessor) throws MessagingException {
 
+		headers = processHeadersToSend(headers);
 		MessageHeaders messageHeaders = (headers != null) ? new MessageHeaders(headers) : null;
 		Message<?> message = this.converter.toMessage(payload, messageHeaders);
 
@@ -145,4 +146,16 @@ public abstract class AbstractMessageSendingTemplate<D> implements MessageSendin
 		this.send(destination, message);
 	}
 
+	/**
+	 * Provides access to the map of headers before a send operation.
+	 * Implementations can modify the headers by returning a different map.
+	 * This implementation returns the map that was passed in (i.e. without any changes).
+	 *
+	 * @param headers the headers to send, possibly {@code null}
+	 * @return the actual headers to send
+	 */
+	protected Map<String, Object> processHeadersToSend(Map<String, Object> headers) {
+		return headers;
+	}
+
 }

spring-messaging/src/main/java/org/springframework/messaging/simp/SimpMessagingTemplate.java

@@ -16,6 +16,7 @@
 
 package org.springframework.messaging.simp;
 
+import java.util.HashMap;
 import java.util.Map;
 
 import org.springframework.messaging.Message;
@@ -25,7 +26,10 @@ import org.springframework.messaging.MessagingException;
 import org.springframework.messaging.core.AbstractMessageSendingTemplate;
 import org.springframework.messaging.core.MessagePostProcessor;
 import org.springframework.messaging.support.MessageBuilder;
+import org.springframework.messaging.support.NativeMessageHeaderAccessor;
 import org.springframework.util.Assert;
+import org.springframework.util.LinkedMultiValueMap;
+import org.springframework.util.MultiValueMap;
 import org.springframework.util.StringUtils;
 
 /**
@@ -160,4 +164,36 @@ public class SimpMessagingTemplate extends AbstractMessageSendingTemplate<String
 		super.convertAndSend(this.userDestinationPrefix + user + destination, payload, headers, postProcessor);
 	}
 
+	/**
+	 * Creates a new map and puts the given headers under the key
+	 * {@link org.springframework.messaging.support.NativeMessageHeaderAccessor#NATIVE_HEADERS NATIVE_HEADERS}.
+	 * Effectively this treats all given headers as headers to be sent out to the
+	 * external source.
+	 * <p>
+	 * If the given headers already contain the key
+	 * {@link org.springframework.messaging.support.NativeMessageHeaderAccessor#NATIVE_HEADERS NATIVE_HEADERS}
+	 * then the same header map is returned (i.e. without any changes).
+	 */
+	@Override
+	protected Map<String, Object> processHeadersToSend(Map<String, Object> headers) {
+
+		if (headers == null) {
+			return null;
+		}
+		else if (headers.containsKey(NativeMessageHeaderAccessor.NATIVE_HEADERS)) {
+			return headers;
+		}
+		else {
+			MultiValueMap<String, String> nativeHeaders = new LinkedMultiValueMap<String, String>(headers.size());
+			for (String key : headers.keySet()) {
+				Object value = headers.get(key);
+				nativeHeaders.set(key, (value != null ? value.toString() : null));
+			}
+
+			headers = new HashMap<String, Object>(1);
+			headers.put(NativeMessageHeaderAccessor.NATIVE_HEADERS, nativeHeaders);
+			return headers;
+		}
+	}
+
 }

spring-messaging/src/test/java/org/springframework/messaging/simp/SimpMessagingTemplateTests.java

@@ -20,10 +20,12 @@ import org.junit.Before;
 import org.junit.Test;
 import org.springframework.messaging.Message;
 import org.springframework.messaging.StubMessageChannel;
+import org.springframework.messaging.support.NativeMessageHeaderAccessor;
 
-import java.util.List;
+import java.util.*;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNull;
 
 /**
  * Unit tests for {@link org.springframework.messaging.simp.SimpMessagingTemplate}.
@@ -43,6 +45,7 @@ public class SimpMessagingTemplateTests {
 		this.messagingTemplate = new SimpMessagingTemplate(messageChannel);
 	}
 
+
 	@Test
 	public void convertAndSendToUser() {
 		this.messagingTemplate.convertAndSendToUser("joe", "/queue/foo", "data");
@@ -69,4 +72,32 @@ public class SimpMessagingTemplateTests {
 		assertEquals("/user/http:%2F%2Fjoe.openid.example.org%2F/queue/foo", headers.getDestination());
 	}
 
+	@Test
+	public void convertAndSendWithCustomHeader() {
+		Map<String, Object> headers = Collections.singletonMap("key", "value");
+		this.messagingTemplate.convertAndSend("/foo", "data", headers);
+
+		List<Message<byte[]>> messages = this.messageChannel.getMessages();
+		Message<byte[]> message = messages.get(0);
+		SimpMessageHeaderAccessor resultHeaders = SimpMessageHeaderAccessor.wrap(message);
+
+		assertNull(resultHeaders.toMap().get("key"));
+		assertEquals(Arrays.asList("value"), resultHeaders.getNativeHeader("key"));
+	}
+
+	@Test
+	public void convertAndSendWithCustomHeaderNonNative() {
+		Map<String, Object> headers = new HashMap<String, Object>();
+		headers.put("key", "value");
+		headers.put(NativeMessageHeaderAccessor.NATIVE_HEADERS, Collections.emptyMap());
+		this.messagingTemplate.convertAndSend("/foo", "data", headers);
+
+		List<Message<byte[]>> messages = this.messageChannel.getMessages();
+		Message<byte[]> message = messages.get(0);
+		SimpMessageHeaderAccessor resultHeaders = SimpMessageHeaderAccessor.wrap(message);
+
+		assertEquals("value", resultHeaders.toMap().get("key"));
+		assertNull(resultHeaders.getNativeHeader("key"));
+	}
+
 }