diff --git a/spring-messaging/src/main/java/org/springframework/messaging/simp/SimpMessagingTemplate.java b/spring-messaging/src/main/java/org/springframework/messaging/simp/SimpMessagingTemplate.java index a0e720edda..2b53358fd2 100644 --- a/spring-messaging/src/main/java/org/springframework/messaging/simp/SimpMessagingTemplate.java +++ b/spring-messaging/src/main/java/org/springframework/messaging/simp/SimpMessagingTemplate.java @@ -27,6 +27,7 @@ import org.springframework.messaging.core.AbstractMessageSendingTemplate; import org.springframework.messaging.core.MessagePostProcessor; import org.springframework.messaging.support.MessageBuilder; import org.springframework.messaging.support.MessageHeaderAccessor; +import org.springframework.messaging.support.MessageHeaderInitializer; import org.springframework.messaging.support.NativeMessageHeaderAccessor; import org.springframework.util.Assert; import org.springframework.util.StringUtils; @@ -48,6 +49,8 @@ public class SimpMessagingTemplate extends AbstractMessageSendingTemplateBy default this property is not set. + */ + public void setHeaderInitializer(MessageHeaderInitializer headerInitializer) { + this.headerInitializer = headerInitializer; + } + + /** + * @return the configured header initializer. + */ + public MessageHeaderInitializer getHeaderInitializer() { + return this.headerInitializer; + } + /** * @return the messageChannel */ @@ -146,10 +166,12 @@ public class SimpMessagingTemplate extends AbstractMessageSendingTemplateBy default this property is not set. + */ + public void setHeaderInitializer(MessageHeaderInitializer headerInitializer) { + this.headerInitializer = headerInitializer; + } + + /** + * @return the configured header initializer. + */ + public MessageHeaderInitializer getHeaderInitializer() { + return this.headerInitializer; + } + @Override public boolean supportsReturnType(MethodParameter returnType) { @@ -166,6 +186,9 @@ public class SendToMethodReturnValueHandler implements HandlerMethodReturnValueH private MessageHeaders createHeaders(String sessionId) { SimpMessageHeaderAccessor headerAccessor = SimpMessageHeaderAccessor.create(SimpMessageType.MESSAGE); + if (getHeaderInitializer() != null) { + getHeaderInitializer().initHeaders(headerAccessor); + } headerAccessor.setSessionId(sessionId); headerAccessor.setLeaveMutable(true); return headerAccessor.getMessageHeaders(); diff --git a/spring-messaging/src/main/java/org/springframework/messaging/simp/annotation/support/SimpAnnotationMethodMessageHandler.java b/spring-messaging/src/main/java/org/springframework/messaging/simp/annotation/support/SimpAnnotationMethodMessageHandler.java index 5d0851b6c0..c73c94dd1f 100644 --- a/spring-messaging/src/main/java/org/springframework/messaging/simp/annotation/support/SimpAnnotationMethodMessageHandler.java +++ b/spring-messaging/src/main/java/org/springframework/messaging/simp/annotation/support/SimpAnnotationMethodMessageHandler.java @@ -59,6 +59,7 @@ import org.springframework.messaging.simp.SimpMessageTypeMessageCondition; import org.springframework.messaging.simp.SimpMessagingTemplate; import org.springframework.messaging.simp.annotation.SubscribeMapping; import org.springframework.messaging.support.MessageHeaderAccessor; +import org.springframework.messaging.support.MessageHeaderInitializer; import org.springframework.stereotype.Controller; import org.springframework.util.AntPathMatcher; import org.springframework.util.Assert; @@ -95,6 +96,8 @@ public class SimpAnnotationMethodMessageHandler extends AbstractMethodMessageHan private Validator validator; + private MessageHeaderInitializer headerInitializer; + private final Object lifecycleMonitor = new Object(); private volatile boolean running = false; @@ -194,6 +197,24 @@ public class SimpAnnotationMethodMessageHandler extends AbstractMethodMessageHan this.validator = validator; } + /** + * Configure a {@link MessageHeaderInitializer} to pass on to + * {@link org.springframework.messaging.handler.invocation.HandlerMethodReturnValueHandler}s + * that send messages from controller return values. + * + *

By default this property is not set. + */ + public void setHeaderInitializer(MessageHeaderInitializer headerInitializer) { + this.headerInitializer = headerInitializer; + } + + /** + * @return the configured header initializer. + */ + public MessageHeaderInitializer getHeaderInitializer() { + return this.headerInitializer; + } + @Override public boolean isAutoStartup() { @@ -265,14 +286,21 @@ public class SimpAnnotationMethodMessageHandler extends AbstractMethodMessageHan List handlers = new ArrayList(); // Annotation-based return value types - handlers.add(new SendToMethodReturnValueHandler(this.brokerTemplate, true)); - handlers.add(new SubscriptionMethodReturnValueHandler(this.clientMessagingTemplate)); + SendToMethodReturnValueHandler sth = new SendToMethodReturnValueHandler(this.brokerTemplate, true); + sth.setHeaderInitializer(this.headerInitializer); + handlers.add(sth); + + SubscriptionMethodReturnValueHandler sh = new SubscriptionMethodReturnValueHandler(this.clientMessagingTemplate); + sh.setHeaderInitializer(this.headerInitializer); + handlers.add(sh); // custom return value types handlers.addAll(getCustomReturnValueHandlers()); // catch-all - handlers.add(new SendToMethodReturnValueHandler(this.brokerTemplate, false)); + sth = new SendToMethodReturnValueHandler(this.brokerTemplate, false); + sth.setHeaderInitializer(this.headerInitializer); + handlers.add(sth); return handlers; } diff --git a/spring-messaging/src/main/java/org/springframework/messaging/simp/annotation/support/SubscriptionMethodReturnValueHandler.java b/spring-messaging/src/main/java/org/springframework/messaging/simp/annotation/support/SubscriptionMethodReturnValueHandler.java index f17797b0e3..693e16f7af 100644 --- a/spring-messaging/src/main/java/org/springframework/messaging/simp/annotation/support/SubscriptionMethodReturnValueHandler.java +++ b/spring-messaging/src/main/java/org/springframework/messaging/simp/annotation/support/SubscriptionMethodReturnValueHandler.java @@ -26,6 +26,7 @@ import org.springframework.messaging.simp.SimpMessageHeaderAccessor; import org.springframework.messaging.simp.SimpMessageType; import org.springframework.messaging.simp.annotation.SendToUser; import org.springframework.messaging.simp.annotation.SubscribeMapping; +import org.springframework.messaging.support.MessageHeaderInitializer; import org.springframework.util.Assert; /** @@ -46,6 +47,8 @@ public class SubscriptionMethodReturnValueHandler implements HandlerMethodReturn private final MessageSendingOperations messagingTemplate; + private MessageHeaderInitializer headerInitializer; + /** * Class constructor. @@ -58,6 +61,23 @@ public class SubscriptionMethodReturnValueHandler implements HandlerMethodReturn this.messagingTemplate = messagingTemplate; } + /** + * Configure a {@link MessageHeaderInitializer} to apply to the headers of all + * messages sent to the client outbound channel. + * + *

By default this property is not set. + */ + public void setHeaderInitializer(MessageHeaderInitializer headerInitializer) { + this.headerInitializer = headerInitializer; + } + + /** + * @return the configured header initializer. + */ + public MessageHeaderInitializer getHeaderInitializer() { + return this.headerInitializer; + } + @Override public boolean supportsReturnType(MethodParameter returnType) { @@ -87,6 +107,9 @@ public class SubscriptionMethodReturnValueHandler implements HandlerMethodReturn private MessageHeaders createHeaders(String sessionId, String subscriptionId) { SimpMessageHeaderAccessor headerAccessor = SimpMessageHeaderAccessor.create(SimpMessageType.MESSAGE); + if (getHeaderInitializer() != null) { + getHeaderInitializer().initHeaders(headerAccessor); + } headerAccessor.setSessionId(sessionId); headerAccessor.setSubscriptionId(subscriptionId); headerAccessor.setLeaveMutable(true); diff --git a/spring-messaging/src/main/java/org/springframework/messaging/simp/broker/SimpleBrokerMessageHandler.java b/spring-messaging/src/main/java/org/springframework/messaging/simp/broker/SimpleBrokerMessageHandler.java index 8aa18b151e..c2b0038c63 100644 --- a/spring-messaging/src/main/java/org/springframework/messaging/simp/broker/SimpleBrokerMessageHandler.java +++ b/spring-messaging/src/main/java/org/springframework/messaging/simp/broker/SimpleBrokerMessageHandler.java @@ -25,6 +25,7 @@ import org.springframework.messaging.SubscribableChannel; import org.springframework.messaging.simp.SimpMessageHeaderAccessor; import org.springframework.messaging.simp.SimpMessageType; import org.springframework.messaging.support.MessageBuilder; +import org.springframework.messaging.support.MessageHeaderInitializer; import org.springframework.util.Assert; import org.springframework.util.MultiValueMap; @@ -48,6 +49,8 @@ public class SimpleBrokerMessageHandler extends AbstractBrokerMessageHandler { private SubscriptionRegistry subscriptionRegistry = new DefaultSubscriptionRegistry(); + private MessageHeaderInitializer headerInitializer; + /** * Create a SimpleBrokerMessageHandler instance with the given message channels @@ -94,6 +97,23 @@ public class SimpleBrokerMessageHandler extends AbstractBrokerMessageHandler { return this.subscriptionRegistry; } + /** + * Configure a {@link MessageHeaderInitializer} to apply to the headers of all + * messages sent to the client outbound channel. + * + *

By default this property is not set. + */ + public void setHeaderInitializer(MessageHeaderInitializer headerInitializer) { + this.headerInitializer = headerInitializer; + } + + /** + * @return the configured header initializer. + */ + public MessageHeaderInitializer getHeaderInitializer() { + return this.headerInitializer; + } + @Override public void startInternal() { @@ -138,6 +158,7 @@ public class SimpleBrokerMessageHandler extends AbstractBrokerMessageHandler { } else if (SimpMessageType.CONNECT.equals(messageType)) { SimpMessageHeaderAccessor accessor = SimpMessageHeaderAccessor.create(SimpMessageType.CONNECT_ACK); + initHeaders(accessor); accessor.setSessionId(sessionId); accessor.setHeader(SimpMessageHeaderAccessor.CONNECT_MESSAGE_HEADER, message); Message connectAck = MessageBuilder.createMessage(EMPTY_PAYLOAD, accessor.getMessageHeaders()); @@ -150,6 +171,12 @@ public class SimpleBrokerMessageHandler extends AbstractBrokerMessageHandler { } } + private void initHeaders(SimpMessageHeaderAccessor accessor) { + if (getHeaderInitializer() != null) { + getHeaderInitializer().initHeaders(accessor); + } + } + protected void sendMessageToSubscribers(String destination, Message message) { MultiValueMap subscriptions = this.subscriptionRegistry.findSubscriptions(message); if ((subscriptions.size() > 0) && logger.isDebugEnabled()) { @@ -159,6 +186,7 @@ public class SimpleBrokerMessageHandler extends AbstractBrokerMessageHandler { for (String sessionId : subscriptions.keySet()) { for (String subscriptionId : subscriptions.get(sessionId)) { SimpMessageHeaderAccessor headerAccessor = SimpMessageHeaderAccessor.create(SimpMessageType.MESSAGE); + initHeaders(headerAccessor); headerAccessor.setSessionId(sessionId); headerAccessor.setSubscriptionId(subscriptionId); headerAccessor.copyHeadersIfAbsent(message.getHeaders()); diff --git a/spring-messaging/src/main/java/org/springframework/messaging/simp/stomp/StompBrokerRelayMessageHandler.java b/spring-messaging/src/main/java/org/springframework/messaging/simp/stomp/StompBrokerRelayMessageHandler.java index 20b005af7d..b9745a1703 100644 --- a/spring-messaging/src/main/java/org/springframework/messaging/simp/stomp/StompBrokerRelayMessageHandler.java +++ b/spring-messaging/src/main/java/org/springframework/messaging/simp/stomp/StompBrokerRelayMessageHandler.java @@ -30,6 +30,7 @@ import org.springframework.messaging.simp.SimpMessageType; import org.springframework.messaging.simp.broker.AbstractBrokerMessageHandler; import org.springframework.messaging.support.MessageBuilder; import org.springframework.messaging.support.MessageHeaderAccessor; +import org.springframework.messaging.support.MessageHeaderInitializer; import org.springframework.messaging.tcp.FixedIntervalReconnectStrategy; import org.springframework.messaging.tcp.TcpConnection; import org.springframework.messaging.tcp.TcpConnectionHandler; @@ -111,6 +112,8 @@ public class StompBrokerRelayMessageHandler extends AbstractBrokerMessageHandler private TcpOperations tcpClient; + private MessageHeaderInitializer headerInitializer; + private final Map connectionHandlers = new ConcurrentHashMap(); @@ -323,6 +326,24 @@ public class StompBrokerRelayMessageHandler extends AbstractBrokerMessageHandler return this.tcpClient; } + /** + * Configure a {@link MessageHeaderInitializer} to apply to the headers of all + * messages created through the {@code StompBrokerRelayMessageHandler} that + * are sent to the client outbound message channel. + * + *

By default this property is not set. + */ + public void setHeaderInitializer(MessageHeaderInitializer headerInitializer) { + this.headerInitializer = headerInitializer; + } + + /** + * @return the configured header initializer. + */ + public MessageHeaderInitializer getHeaderInitializer() { + return this.headerInitializer; + } + @Override protected void startInternal() { @@ -331,7 +352,10 @@ public class StompBrokerRelayMessageHandler extends AbstractBrokerMessageHandler this.brokerChannel.subscribe(this); if (this.tcpClient == null) { - this.tcpClient = new StompTcpClientFactory().create(this.relayHost, this.relayPort); + StompDecoder decoder = new StompDecoder(); + decoder.setHeaderInitializer(getHeaderInitializer()); + StompCodec codec = new StompCodec(new StompEncoder(), decoder); + this.tcpClient = new StompTcpClientFactory().create(this.relayHost, this.relayPort, codec); } if (logger.isDebugEnabled()) { @@ -536,10 +560,13 @@ public class StompBrokerRelayMessageHandler extends AbstractBrokerMessageHandler private void sendStompErrorToClient(String errorText) { if (this.isRemoteClientSession) { - StompHeaderAccessor headers = StompHeaderAccessor.create(StompCommand.ERROR); - headers.setSessionId(this.sessionId); - headers.setMessage(errorText); - Message errorMessage = MessageBuilder.createMessage(EMPTY_PAYLOAD, headers.getMessageHeaders()); + StompHeaderAccessor headerAccessor = StompHeaderAccessor.create(StompCommand.ERROR); + if (getHeaderInitializer() != null) { + getHeaderInitializer().initHeaders(headerAccessor); + } + headerAccessor.setSessionId(this.sessionId); + headerAccessor.setMessage(errorText); + Message errorMessage = MessageBuilder.createMessage(EMPTY_PAYLOAD, headerAccessor.getMessageHeaders()); sendMessageToClient(errorMessage); } } @@ -811,8 +838,8 @@ public class StompBrokerRelayMessageHandler extends AbstractBrokerMessageHandler private static class StompTcpClientFactory { - public TcpOperations create(String relayHost, int relayPort) { - return new ReactorTcpClient(relayHost, relayPort, new StompCodec()); + public TcpOperations create(String relayHost, int relayPort, StompCodec codec) { + return new ReactorTcpClient(relayHost, relayPort, codec); } } diff --git a/spring-messaging/src/main/java/org/springframework/messaging/simp/stomp/StompCodec.java b/spring-messaging/src/main/java/org/springframework/messaging/simp/stomp/StompCodec.java index ec5569508f..7941ff8e1d 100644 --- a/spring-messaging/src/main/java/org/springframework/messaging/simp/stomp/StompCodec.java +++ b/spring-messaging/src/main/java/org/springframework/messaging/simp/stomp/StompCodec.java @@ -18,12 +18,13 @@ package org.springframework.messaging.simp.stomp; import org.springframework.messaging.Message; +import org.springframework.util.Assert; import reactor.function.Consumer; import reactor.function.Function; import reactor.io.Buffer; import reactor.io.encoding.Codec; -import java.util.List; +import java.nio.ByteBuffer; /** * A Reactor TCP {@link Codec} for sending and receiving STOMP messages. @@ -34,35 +35,68 @@ import java.util.List; */ public class StompCodec implements Codec, Message> { - private static final StompDecoder DECODER = new StompDecoder(); + private final StompDecoder stompDecoder; - private static final Function, Buffer> ENCODER_FUNCTION = new Function, Buffer>() { + private final StompEncoder stompEncoder; - private final StompEncoder encoder = new StompEncoder(); + private final Function, Buffer> encodingFunction; - @Override - public Buffer apply(Message message) { - return Buffer.wrap(this.encoder.encode(message)); - } - }; + + public StompCodec() { + this(new StompEncoder(), new StompDecoder()); + } + + public StompCodec(StompEncoder encoder, StompDecoder decoder) { + Assert.notNull(encoder, "'encoder' is required"); + Assert.notNull(decoder, "'decoder' is required"); + this.stompEncoder = encoder; + this.stompDecoder = decoder; + this.encodingFunction = new EncodingFunction(this.stompEncoder); + } @Override - public Function> decoder(final Consumer> next) { - return new Function>() { - - @Override - public Message apply(Buffer buffer) { - for (Message message : DECODER.decode(buffer.byteBuffer())) { - next.accept(message); - } - return null; - } - }; + public Function> decoder(final Consumer> messageConsumer) { + return new DecodingFunction(this.stompDecoder, messageConsumer); } @Override public Function, Buffer> encoder() { - return ENCODER_FUNCTION; + return this.encodingFunction; } + + private static class EncodingFunction implements Function, Buffer> { + + private final StompEncoder encoder; + + private EncodingFunction(StompEncoder encoder) { + this.encoder = encoder; + } + + @Override + public Buffer apply(Message message) { + byte[] bytes = this.encoder.encode(message); + return new Buffer(ByteBuffer.wrap(bytes)); + } + } + + private static class DecodingFunction implements Function> { + + private final StompDecoder decoder; + + private final Consumer> messageConsumer; + + public DecodingFunction(StompDecoder decoder, Consumer> next) { + this.decoder = decoder; + this.messageConsumer = next; + } + + @Override + public Message apply(Buffer buffer) { + for (Message message : this.decoder.decode(buffer.byteBuffer())) { + this.messageConsumer.accept(message); + } + return null; + } + } } diff --git a/spring-messaging/src/main/java/org/springframework/messaging/simp/stomp/StompDecoder.java b/spring-messaging/src/main/java/org/springframework/messaging/simp/stomp/StompDecoder.java index 13347e13ef..3afbee28e2 100644 --- a/spring-messaging/src/main/java/org/springframework/messaging/simp/stomp/StompDecoder.java +++ b/spring-messaging/src/main/java/org/springframework/messaging/simp/stomp/StompDecoder.java @@ -27,6 +27,7 @@ import org.apache.commons.logging.LogFactory; import org.springframework.messaging.Message; import org.springframework.messaging.support.MessageBuilder; +import org.springframework.messaging.support.MessageHeaderInitializer; import org.springframework.messaging.support.NativeMessageHeaderAccessor; import org.springframework.util.InvalidMimeTypeException; import org.springframework.util.MultiValueMap; @@ -52,6 +53,24 @@ public class StompDecoder { private final Log logger = LogFactory.getLog(StompDecoder.class); + private MessageHeaderInitializer headerInitializer; + + + /** + * Configure a + * {@link org.springframework.messaging.support.MessageHeaderInitializer MessageHeaderInitializer} + * to apply to the headers of {@link Message}s from decoded STOMP frames. + */ + public void setHeaderInitializer(MessageHeaderInitializer headerInitializer) { + this.headerInitializer = headerInitializer; + } + + /** + * @return the configured {@code MessageHeaderInitializer} if any. + */ + public MessageHeaderInitializer getHeaderInitializer() { + return this.headerInitializer; + } /** * Decodes one or more STOMP frames from the given {@code ByteBuffer} into a @@ -123,6 +142,7 @@ public class StompDecoder { if (buffer.remaining() > 0) { StompCommand stompCommand = StompCommand.valueOf(command); headerAccessor = StompHeaderAccessor.create(stompCommand); + initHeaders(headerAccessor); readHeaders(buffer, headerAccessor); payload = readPayload(buffer, headerAccessor); @@ -160,12 +180,18 @@ public class StompDecoder { logger.trace("Decoded heartbeat"); } StompHeaderAccessor headerAccessor = StompHeaderAccessor.createForHeartbeat(); + initHeaders(headerAccessor); headerAccessor.setLeaveMutable(true); decodedMessage = MessageBuilder.createMessage(HEARTBEAT_PAYLOAD, headerAccessor.getMessageHeaders()); } return decodedMessage; } + private void initHeaders(StompHeaderAccessor headerAccessor) { + if (getHeaderInitializer() != null) { + getHeaderInitializer().initHeaders(headerAccessor); + } + } /** * Skip one ore more EOL characters at the start of the given ByteBuffer. diff --git a/spring-messaging/src/main/java/org/springframework/messaging/simp/user/UserDestinationMessageHandler.java b/spring-messaging/src/main/java/org/springframework/messaging/simp/user/UserDestinationMessageHandler.java index e16bfdebdf..fded86db5b 100644 --- a/spring-messaging/src/main/java/org/springframework/messaging/simp/user/UserDestinationMessageHandler.java +++ b/spring-messaging/src/main/java/org/springframework/messaging/simp/user/UserDestinationMessageHandler.java @@ -31,6 +31,7 @@ import org.springframework.messaging.simp.SimpMessageHeaderAccessor; import org.springframework.messaging.simp.SimpMessageType; import org.springframework.messaging.simp.SimpMessagingTemplate; import org.springframework.messaging.support.MessageBuilder; +import org.springframework.messaging.support.MessageHeaderInitializer; import org.springframework.util.Assert; /** @@ -56,6 +57,8 @@ public class UserDestinationMessageHandler implements MessageHandler, SmartLifec private final UserDestinationResolver userDestinationResolver; + private MessageHeaderInitializer headerInitializer; + private final Object lifecycleMonitor = new Object(); private volatile boolean running = false; @@ -97,6 +100,24 @@ public class UserDestinationMessageHandler implements MessageHandler, SmartLifec return this.userDestinationResolver; } + /** + * Configure a {@link MessageHeaderInitializer} to pass on to + * {@link org.springframework.messaging.handler.invocation.HandlerMethodReturnValueHandler}s + * that send messages from controller return values. + * + *

By default this property is not set. + */ + public void setHeaderInitializer(MessageHeaderInitializer headerInitializer) { + this.headerInitializer = headerInitializer; + } + + /** + * @return the configured header initializer. + */ + public MessageHeaderInitializer getHeaderInitializer() { + return this.headerInitializer; + } + @Override public int getPhase() { @@ -154,6 +175,9 @@ public class UserDestinationMessageHandler implements MessageHandler, SmartLifec } if (SimpMessageType.MESSAGE.equals(SimpMessageHeaderAccessor.getMessageType(message.getHeaders()))) { SimpMessageHeaderAccessor headerAccessor = SimpMessageHeaderAccessor.wrap(message); + if (getHeaderInitializer() != null) { + getHeaderInitializer().initHeaders(headerAccessor); + } headerAccessor.setHeader(SimpMessageHeaderAccessor.ORIGINAL_DESTINATION, result.getSubscribeDestination()); message = MessageBuilder.createMessage(message.getPayload(), headerAccessor.getMessageHeaders()); } diff --git a/spring-messaging/src/main/java/org/springframework/messaging/support/MessageHeaderAccessor.java b/spring-messaging/src/main/java/org/springframework/messaging/support/MessageHeaderAccessor.java index d4a4f20dc9..68193fa58c 100644 --- a/spring-messaging/src/main/java/org/springframework/messaging/support/MessageHeaderAccessor.java +++ b/spring-messaging/src/main/java/org/springframework/messaging/support/MessageHeaderAccessor.java @@ -38,66 +38,74 @@ import org.springframework.util.StringUtils; /** * A base for classes providing strongly typed getters and setters as well as * behavior around specific categories of headers (e.g. STOMP headers). - * Supports creating new headers (default constructor), modifying existing headers - * (when still mutable), or copying and modifying existing headers. + * Supports creating new headers, modifying existing headers (when still mutable), + * or copying and modifying existing headers. * - *

The {@link #getMessageHeaders()} method provides access to the underlying, - * fully-prepared {@code MessageHeaders} instance that can then be used as-is - * to create a single message as follows: + *

The method {@link #getMessageHeaders()} provides access to the underlying, + * fully-prepared {@link MessageHeaders} that can then be used as-is (i.e. + * without copying) to create a single message as follows: * *

- * MessageHeaderAccessor headerAccessor = new MessageHeaderAccessor();
- * headerAccessor.set("foo", "bar");
- * Message message = MessageBuilder.createMessage("payload", headerAccessor.getMessageHeaders());
+ * MessageHeaderAccessor accessor = new MessageHeaderAccessor();
+ * accessor.set("foo", "bar");
+ * Message message = MessageBuilder.createMessage("payload", accessor.getMessageHeaders());
  * 
* - *

After the above is executed, by default the {@code MessageHeaderAccessor} - * is immutable. It is also possible to leave it mutable, for example for further - * initialization of the message in the same thread: + *

After the above, by default the {@code MessageHeaderAccessor} becomes + * immutable. However it is possible to leave it mutable for further initialization + * in the same thread, for example: * *

- * MessageHeaderAccessor headerAccessor = new MessageHeaderAccessor();
- * headerAccessor.set("foo", "bar");
- * headerAccessor.setLeaveMutable(true);
- * Message message = MessageBuilder.createMessage("payload", headerAccessor.getMessageHeaders());
+ * MessageHeaderAccessor accessor = new MessageHeaderAccessor();
+ * accessor.set("foo", "bar");
+ * accessor.setLeaveMutable(true);
+ * Message message = MessageBuilder.createMessage("payload", accessor.getMessageHeaders());
  *
  * // later on in the same thread...
  *
- * MessageHeaderAccessor headerAccessor = MessageHeaderAccessor.getAccessor(message);
- * headerAccessor.set("bar", "baz");
- * headerAccessor.setImmutable();
+ * MessageHeaderAccessor accessor = MessageHeaderAccessor.getAccessor(message);
+ * accessor.set("bar", "baz");
+ * accessor.setImmutable();
  * 
* - *

To re-obtain the {@code MessageHeaderAccessor} for a {@code MessageHeaders} - * instance, use the following: - * + *

The method {@link #toMap()} returns a copy of the underlying headers. It can + * be used to prepare multiple messages from the same {@code MessageHeaderAccessor} + * instance: *

- * MessageHeaderAccessor headerAccessor = new MessageHeaderAccessor();
- * headerAccessor.set("foo", "bar");
- * Message message = MessageBuilder.createMessage("payload", headerAccessor.getMessageHeaders());
+ * MessageHeaderAccessor accessor = new MessageHeaderAccessor();
+ * MessageBuilder builder = MessageBuilder.withPayload("payload").setHeaders(accessor);
  *
- * // later on (any thread)...
- * MessageHeaderAccessor headerAccessor = MessageHeaderAccessor.getAccessor(message);
- * headerAccessor.get("foo");
- * 
- * - *

To prepare multiple messages with the same {@code MessageHeaderAccessor} - * instance, use the code below. However note that this usage style does not allow - * re-obtaining a header accessor later on: - *

- * MessageHeaderAccessor headerAccessor = new MessageHeaderAccessor();
- * MessageBuilder builder = MessageBuilder.withPayload("payload").setHeaders(headerAccessor);
- *
- * headerAccessor.setHeader("foo", "bar1");
+ * accessor.setHeader("foo", "bar1");
  * Message message1 = builder.build();
  *
- * headerAccessor.setHeader("foo", "bar2");
+ * accessor.setHeader("foo", "bar2");
  * Message message2 = builder.build();
  *
- * headerAccessor.setHeader("foo", "bar3");
+ * accessor.setHeader("foo", "bar3");
  * Message  message3 = builder.build();
  * 
* + *

However note that with the above style, the header accessor is shared and + * cannot be re-obtained later on. Alternatively it is also possible to create + * one {@code MessageHeaderAccessor} per message: + * + *

+ * MessageHeaderAccessor accessor1 = new MessageHeaderAccessor();
+ * accessor.set("foo", "bar1");
+ * Message message1 = MessageBuilder.createMessage("payload", accessor1.getMessageHeaders());
+ *
+ * MessageHeaderAccessor accessor2 = new MessageHeaderAccessor();
+ * accessor.set("foo", "bar2");
+ * Message message2 = MessageBuilder.createMessage("payload", accessor2.getMessageHeaders());
+ *
+ * MessageHeaderAccessor accessor3 = new MessageHeaderAccessor();
+ * accessor.set("foo", "bar3");
+ * Message message3 = MessageBuilder.createMessage("payload", accessor3.getMessageHeaders());
+ * 
+ * + *

Note that the above examples aim to demonstrate the general idea of using + * header accessors. The most likely usage however is through sub-classes. + * * @author Rossen Stoyanchev * @since 4.0 */ @@ -210,7 +218,6 @@ public class MessageHeaderAccessor { * use {@link #toMap()}. */ public MessageHeaders getMessageHeaders() { - this.headers.setIdAndTimestamp(); if (!this.leaveMutable) { setImmutable(); } @@ -259,6 +266,7 @@ public class MessageHeaderAccessor { * @since 4.1 */ public void setImmutable() { + this.headers.setIdAndTimestamp(); this.headers.setImmutable(); } @@ -500,6 +508,9 @@ public class MessageHeaderAccessor { } public void setIdAndTimestamp() { + if (!isMutable()) { + return; + } if (getId() == null) { IdGenerator idGenerator = (MessageHeaderAccessor.this.idGenerator != null) ? MessageHeaderAccessor.this.idGenerator : diff --git a/spring-messaging/src/test/java/org/springframework/messaging/support/MessageHeaderAccessorTests.java b/spring-messaging/src/test/java/org/springframework/messaging/support/MessageHeaderAccessorTests.java index 52e01ee850..af7a9fc05c 100644 --- a/spring-messaging/src/test/java/org/springframework/messaging/support/MessageHeaderAccessorTests.java +++ b/spring-messaging/src/test/java/org/springframework/messaging/support/MessageHeaderAccessorTests.java @@ -201,15 +201,6 @@ public class MessageHeaderAccessorTests { assertNull(accessor.getMessageHeaders().getTimestamp()); } - @Test - public void timestampBehaviorCopyFromExistingMessage() { - MessageHeaderAccessor accessor = new MessageHeaderAccessor(); - accessor.setEnableTimestamp(true); - Message message = MessageBuilder.createMessage("payload", accessor.getMessageHeaders()); - MessageHeaderAccessor secondAccessor = new MessageHeaderAccessor(message); - assertNotNull(secondAccessor.getMessageHeaders().getTimestamp()); - } - @Test public void idGeneratorCustom() { final UUID id = new UUID(0L, 23L); @@ -229,19 +220,35 @@ public class MessageHeaderAccessorTests { assertNotNull(accessor.getMessageHeaders().getId()); } + @Test - public void idGeneratorCopyFromExistingMessage() { - final UUID id = new UUID(0L, 23L); + public void idTimestampWithMutableHeaders() { MessageHeaderAccessor accessor = new MessageHeaderAccessor(); + accessor.setIdGenerator(new IdGenerator() { + @Override + public UUID generateId() { + return MessageHeaders.ID_VALUE_NONE; + } + }); + accessor.setEnableTimestamp(false); + accessor.setLeaveMutable(true); + MessageHeaders headers = accessor.getMessageHeaders(); + + assertNull(headers.getId()); + assertNull(headers.getTimestamp()); + + final UUID id = new UUID(0L, 23L); accessor.setIdGenerator(new IdGenerator() { @Override public UUID generateId() { return id; } }); - Message message = MessageBuilder.createMessage("payload", accessor.getMessageHeaders()); - MessageHeaderAccessor secondAccessor = new MessageHeaderAccessor(message); - assertSame(id, secondAccessor.getMessageHeaders().getId()); + accessor.setEnableTimestamp(true); + accessor.setImmutable(); + + assertSame(id, accessor.getMessageHeaders().getId()); + assertNotNull(headers.getTimestamp()); } diff --git a/spring-websocket/src/main/java/org/springframework/web/socket/messaging/StompSubProtocolHandler.java b/spring-websocket/src/main/java/org/springframework/web/socket/messaging/StompSubProtocolHandler.java index 9cb3ae8384..60a8bf2aaa 100644 --- a/spring-websocket/src/main/java/org/springframework/web/socket/messaging/StompSubProtocolHandler.java +++ b/spring-websocket/src/main/java/org/springframework/web/socket/messaging/StompSubProtocolHandler.java @@ -45,6 +45,7 @@ import org.springframework.messaging.simp.user.DestinationUserNameProvider; import org.springframework.messaging.simp.user.UserSessionRegistry; import org.springframework.messaging.support.MessageBuilder; import org.springframework.messaging.support.MessageHeaderAccessor; +import org.springframework.messaging.support.MessageHeaderInitializer; import org.springframework.util.Assert; import org.springframework.web.socket.CloseStatus; import org.springframework.web.socket.TextMessage; @@ -85,13 +86,15 @@ public class StompSubProtocolHandler implements SubProtocolHandler, ApplicationE private int messageSizeLimit = 64 * 1024; + private UserSessionRegistry userSessionRegistry; + private final StompEncoder stompEncoder = new StompEncoder(); private final StompDecoder stompDecoder = new StompDecoder(); private final Map decoders = new ConcurrentHashMap(); - private UserSessionRegistry userSessionRegistry; + private MessageHeaderInitializer headerInitializer; private ApplicationEventPublisher eventPublisher; @@ -134,6 +137,25 @@ public class StompSubProtocolHandler implements SubProtocolHandler, ApplicationE return this.userSessionRegistry; } + /** + * Configure a {@link MessageHeaderInitializer} to apply to the headers of all + * messages created from decoded STOMP frames and other messages sent to the + * client inbound channel. + * + *

By default this property is not set. + */ + public void setHeaderInitializer(MessageHeaderInitializer headerInitializer) { + this.headerInitializer = headerInitializer; + this.stompDecoder.setHeaderInitializer(headerInitializer); + } + + /** + * @return the configured header initializer. + */ + public MessageHeaderInitializer getHeaderInitializer() { + return this.headerInitializer; + } + @Override public List getSupportedProtocols() { return Arrays.asList("v10.stomp", "v11.stomp", "v12.stomp"); @@ -401,9 +423,12 @@ public class StompSubProtocolHandler implements SubProtocolHandler, ApplicationE logger.debug("WebSocket session ended, sending DISCONNECT message to broker"); } - StompHeaderAccessor headers = StompHeaderAccessor.create(StompCommand.DISCONNECT); - headers.setSessionId(session.getId()); - Message message = MessageBuilder.createMessage(EMPTY_PAYLOAD, headers.getMessageHeaders()); + StompHeaderAccessor headerAccessor = StompHeaderAccessor.create(StompCommand.DISCONNECT); + if (getHeaderInitializer() != null) { + getHeaderInitializer().initHeaders(headerAccessor); + } + headerAccessor.setSessionId(session.getId()); + Message message = MessageBuilder.createMessage(EMPTY_PAYLOAD, headerAccessor.getMessageHeaders()); if (this.eventPublisher != null) { publishEvent(new SessionDisconnectEvent(this, session.getId(), closeStatus));