Added support for ResponseMessage response type in @Controller message handling methods

spring-messaging/src/main/java/org/springframework/messaging/simp/annotation/support/AbstractMethodReturnValueHandler.java

@@ -0,0 +1,147 @@
+/*
+ * 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.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.messaging.simp.annotation.support;
+
+import java.security.Principal;
+
+import org.springframework.core.MethodParameter;
+import org.springframework.messaging.Message;
+import org.springframework.messaging.MessageHeaders;
+import org.springframework.messaging.handler.DestinationPatternsMessageCondition;
+import org.springframework.messaging.handler.annotation.SendTo;
+import org.springframework.messaging.handler.invocation.HandlerMethodReturnValueHandler;
+import org.springframework.messaging.simp.SimpMessageHeaderAccessor;
+import org.springframework.messaging.simp.SimpMessageSendingOperations;
+import org.springframework.messaging.simp.SimpMessageType;
+import org.springframework.messaging.simp.annotation.SendToUser;
+import org.springframework.messaging.simp.user.DestinationUserNameProvider;
+import org.springframework.messaging.support.MessageHeaderInitializer;
+import org.springframework.util.Assert;
+
+
+/**
+ * Abstract {@link HandlerMethodReturnValueHandler}
+ * 
+ * @author Sergi Almar
+ * @since 4.1.1
+ */
+public abstract class AbstractMethodReturnValueHandler implements HandlerMethodReturnValueHandler {
+
+	protected final SimpMessageSendingOperations messagingTemplate;
+
+	protected String defaultDestinationPrefix = "/topic";
+
+	protected String defaultUserDestinationPrefix = "/queue";
+
+	protected MessageHeaderInitializer headerInitializer;
+
+	public AbstractMethodReturnValueHandler(SimpMessageSendingOperations messagingTemplate) {
+		Assert.notNull(messagingTemplate, "messagingTemplate must not be null");
+		this.messagingTemplate = messagingTemplate;
+	}
+
+	/**
+	 * Configure a default prefix to add to message destinations in cases where a method
+	 * is not annotated with {@link SendTo @SendTo} or does not specify any destinations
+	 * through the annotation's value attribute.
+	 * <p>By default, the prefix is set to "/topic".
+	 */
+	public void setDefaultDestinationPrefix(String defaultDestinationPrefix) {
+		this.defaultDestinationPrefix = defaultDestinationPrefix;
+	}
+
+	/**
+	 * Return the configured default destination prefix.
+	 * @see #setDefaultDestinationPrefix(String)
+	 */
+	public String getDefaultDestinationPrefix() {
+		return this.defaultDestinationPrefix;
+	}
+
+	/**
+	 * Configure a default prefix to add to message destinations in cases where a
+	 * method is annotated with {@link SendToUser @SendToUser} but does not specify
+	 * any destinations through the annotation's value attribute.
+	 * <p>By default, the prefix is set to "/queue".
+	 */
+	public void setDefaultUserDestinationPrefix(String prefix) {
+		this.defaultUserDestinationPrefix = prefix;
+	}
+
+	/**
+	 * Return the configured default user destination prefix.
+	 * @see #setDefaultUserDestinationPrefix(String)
+	 */
+	public String getDefaultUserDestinationPrefix() {
+		return this.defaultUserDestinationPrefix;
+	}
+
+	/**
+	 * Configure a {@link MessageHeaderInitializer} to apply to the headers of all
+	 * messages sent to the client outbound channel.
+	 *
+	 * <p>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;
+	}
+
+	public final void handleReturnValue(Object returnValue, MethodParameter returnType, Message<?> message) throws Exception {
+		if (returnValue == null) {
+			return;
+		}
+
+		handleReturnValueInternal(returnValue, returnType, message);
+	}
+
+	public abstract void handleReturnValueInternal(Object returnValue, MethodParameter returnType, Message<?> message) throws Exception ;
+
+	protected String getUserName(Message<?> message, MessageHeaders headers) {
+		Principal principal = SimpMessageHeaderAccessor.getUser(headers);
+		if (principal != null) {
+			return (principal instanceof DestinationUserNameProvider ?
+					((DestinationUserNameProvider) principal).getDestinationUserName() : principal.getName());
+		}
+		return null;
+	}
+
+	protected String[] getTargetDestinations(Message<?> message, String defaultPrefix) {
+		String name = DestinationPatternsMessageCondition.LOOKUP_DESTINATION_HEADER;
+		String destination = (String) message.getHeaders().get(name);
+		Assert.hasText(destination, "No lookup destination header in " + message);
+
+		return (destination.startsWith("/") ?
+				new String[] {defaultPrefix + destination} : new String[] {defaultPrefix + "/" + destination});
+	}
+
+	protected 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();
+	}
+}

spring-messaging/src/main/java/org/springframework/messaging/simp/annotation/support/ResponseMessageMethodReturnValueHandler.java

@@ -0,0 +1,114 @@
+/*
+ * 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.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.messaging.simp.annotation.support;
+
+import org.springframework.core.MethodParameter;
+import org.springframework.messaging.Message;
+import org.springframework.messaging.MessageHeaders;
+import org.springframework.messaging.handler.annotation.SendTo;
+import org.springframework.messaging.handler.invocation.HandlerMethodReturnValueHandler;
+import org.springframework.messaging.simp.SimpMessageHeaderAccessor;
+import org.springframework.messaging.simp.SimpMessageSendingOperations;
+import org.springframework.messaging.simp.annotation.SendToUser;
+import org.springframework.messaging.support.ResponseMessage;
+import org.springframework.util.ObjectUtils;
+
+
+/**
+ * A {@link HandlerMethodReturnValueHandler} for sending messages according to the
+ * {@link org.springframework.messaging.support.ResponseMessage} returned from message handling methods.
+ * 
+ * @author Sergi Almar
+ * @since 4.1.1
+ */
+public class ResponseMessageMethodReturnValueHandler extends AbstractMethodReturnValueHandler {
+
+	public ResponseMessageMethodReturnValueHandler(SimpMessageSendingOperations messagingTemplate) {
+		super(messagingTemplate);
+	}
+
+	@Override
+	public boolean supportsReturnType(MethodParameter returnType) {
+		return returnType.getParameterType().equals(ResponseMessage.class);
+	}
+
+	@Override
+	public void handleReturnValueInternal(Object returnValue, MethodParameter returnType, Message<?> message) throws Exception {
+		ResponseMessage<?> responseMessage = (ResponseMessage<?>) returnValue;
+
+		boolean toUniqueUser = responseMessage.getUser() != null || responseMessage.isToCurrentUser();
+
+		if(toUniqueUser) {
+			handleUserDestinations(responseMessage, message);
+		} 
+		else {
+			handleGenericDestinations(responseMessage, message);
+		}	
+	}
+
+	private void handleGenericDestinations(ResponseMessage<?> responseMessage, Message<?> message) {
+		String [] destinations = getTargetDestinations(responseMessage, message, getDefaultDestinationPrefix());
+
+		MessageHeaders headers = message.getHeaders();
+		String sessionId = SimpMessageHeaderAccessor.getSessionId(headers);
+
+		for (String destination : destinations) {
+			this.messagingTemplate.convertAndSend(destination, responseMessage.getBody(), createHeaders(sessionId));
+		}
+	}
+
+	private void handleUserDestinations(ResponseMessage<?> responseMessage, Message<?> message) {
+		String [] destinations = getTargetDestinations(responseMessage, message, getDefaultUserDestinationPrefix());
+		String user = responseMessage.getUser();
+		boolean broadcast = responseMessage.isBroadcast();
+
+		MessageHeaders headers = message.getHeaders();
+		String sessionId = SimpMessageHeaderAccessor.getSessionId(headers);
+
+		if(responseMessage.isToCurrentUser()) {
+			user = super.getUserName(message, headers);
+
+			if (user == null) {
+				if (sessionId == null) {
+					throw new MissingSessionUserException(message);
+				}
+
+				user = sessionId;
+				broadcast = false;
+			}
+		}
+
+		for (String destination : destinations) {
+			if(broadcast && responseMessage.isToCurrentUser()) {
+				this.messagingTemplate.convertAndSendToUser(user, destination, responseMessage.getBody());
+			}
+			else {
+				this.messagingTemplate.convertAndSendToUser(user, destination, responseMessage.getBody(), createHeaders(sessionId));
+			}
+		}
+	}
+
+	protected String[] getTargetDestinations(ResponseMessage<?> responseMessage, Message<?> message, String defaultPrefix) {
+		String [] destinations = responseMessage.getDestinations();
+
+		if (!ObjectUtils.isEmpty(destinations)) {
+			return destinations;
+		}
+
+		return super.getTargetDestinations(message, defaultPrefix);
+	}
+}