Merge pull request #487 from cloudmark/SPR-11506

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2013 the original author or authors.
+ * 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.
@@ -26,17 +26,16 @@ import org.springframework.messaging.Message;
 import org.springframework.messaging.simp.SimpMessageHeaderAccessor;
 
 /**
- * Annotation that can be used on methods processing an input message to indicate that the
- * method's return value should be converted to a {@link Message} and sent to the
- * specified destination with the prefix <code>"/user/{username}"</code> automatically
- * prepended with the user information expected to be the input message header
- * {@link SimpMessageHeaderAccessor#USER_HEADER}. Such user destinations may need to be
- * further resolved to actual destinations.
+ * Annotation that indicates the return value of a message-handling method should
+ * be sent as a {@link org.springframework.messaging.Message} to the specified
+ * destination(s) prepended with {@code "/user/{username}"} where the user
+ * name is extracted from the headers of the input message being handled.
  *
  * @author Rossen Stoyanchev
  * @since 4.0
- * @see org.springframework.messaging.handler.annotation.SendTo
+ * @see org.springframework.messaging.simp.annotation.support.SendToMethodReturnValueHandler
  * @see org.springframework.messaging.simp.user.UserDestinationMessageHandler
+ * @see org.springframework.messaging.simp.SimpMessageHeaderAccessor#getUser()
  */
 @Target(ElementType.METHOD)
 @Retention(RetentionPolicy.RUNTIME)
@@ -44,8 +43,20 @@ import org.springframework.messaging.simp.SimpMessageHeaderAccessor;
 public @interface SendToUser {
 
 	/**
-	 * The destination for a message based on the return value of a method.
+	 * One or more destinations to send a message to. If left unspecified, a
+	 * default destination is selected based on the destination of the input
+	 * message being handled.
+	 * @see org.springframework.messaging.simp.annotation.support.SendToMethodReturnValueHandler
 	 */
 	String[] value() default {};
 
+	/**
+	 * Whether messages should be sent to all sessions associated with the user
+	 * or only to the session of the input message being handled.
+	 *
+	 * <p>By default this is set to {@code true} in which case messages are
+	 * broadcast to all sessions.
+     */
+    boolean broadcast() default true;
+
 }

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

@@ -16,9 +16,6 @@
 
 package org.springframework.messaging.simp.annotation.support;
 
-import java.lang.annotation.Annotation;
-import java.security.Principal;
-
 import org.springframework.core.MethodParameter;
 import org.springframework.core.annotation.AnnotationUtils;
 import org.springframework.messaging.Message;
@@ -36,6 +33,9 @@ import org.springframework.messaging.support.MessageHeaderInitializer;
 import org.springframework.util.Assert;
 import org.springframework.util.ObjectUtils;
 
+import java.lang.annotation.Annotation;
+import java.security.Principal;
+
 /**
  * A {@link HandlerMethodReturnValueHandler} for sending to destinations specified in a
  * {@link SendTo} or {@link SendToUser} method-level annotations.
@@ -148,7 +148,12 @@ public class SendToMethodReturnValueHandler implements HandlerMethodReturnValueH
 			String user = getUserName(message, headers);
 			String[] destinations = getTargetDestinations(sendToUser, message, this.defaultUserDestinationPrefix);
 			for (String destination : destinations) {
-				this.messagingTemplate.convertAndSendToUser(user, destination, returnValue, createHeaders(sessionId));
+				if (sendToUser.broadcast()) {
+					this.messagingTemplate.convertAndSendToUser(user, destination, returnValue);
+				}
+				else {
+					this.messagingTemplate.convertAndSendToUser(user, destination, returnValue, createHeaders(sessionId));
+				}
 			}
 			return;
 		}

spring-messaging/src/main/java/org/springframework/messaging/simp/user/DefaultUserDestinationResolver.java

@@ -123,6 +123,7 @@ public class DefaultUserDestinationResolver implements UserDestinationResolver {
 		SimpMessageType messageType = SimpMessageHeaderAccessor.getMessageType(headers);
 		String destination = SimpMessageHeaderAccessor.getDestination(headers);
 		Principal principal = SimpMessageHeaderAccessor.getUser(headers);
+		String sessionId = SimpMessageHeaderAccessor.getSessionId(headers);
 
 		String destinationWithoutPrefix;
 		String subscribeDestination;
@@ -137,7 +138,6 @@ public class DefaultUserDestinationResolver implements UserDestinationResolver {
 				logger.error("Ignoring message, no principal info available");
 				return null;
 			}
-			String sessionId = SimpMessageHeaderAccessor.getSessionId(headers);
 			if (sessionId == null) {
 				logger.error("Ignoring message, no session id available");
 				return null;
@@ -158,7 +158,8 @@ public class DefaultUserDestinationResolver implements UserDestinationResolver {
 			subscribeDestination = this.destinationPrefix.substring(0, startIndex-1) + destinationWithoutPrefix;
 			user = destination.substring(startIndex, endIndex);
 			user = StringUtils.replace(user, "%2F", "/");
-			sessionIds = this.userSessionRegistry.getSessionIds(user);
+			sessionIds = (sessionId != null ?
+					Collections.singleton(sessionId) : this.userSessionRegistry.getSessionIds(user));
 		}
 		else {
 			if (logger.isTraceEnabled()) {

spring-messaging/src/test/java/org/springframework/messaging/simp/annotation/support/SendToMethodReturnValueHandlerTests.java

@@ -74,7 +74,9 @@ public class SendToMethodReturnValueHandlerTests {
 	private MethodParameter sendToReturnType;
 	private MethodParameter sendToDefaultDestReturnType;
 	private MethodParameter sendToUserReturnType;
+	private MethodParameter sendToUserSingleSessionReturnType;
 	private MethodParameter sendToUserDefaultDestReturnType;
+	private MethodParameter sendToUserSingleSessionDefaultDestReturnType;
 
 
 	@Before
@@ -101,8 +103,14 @@ public class SendToMethodReturnValueHandlerTests {
 		method = this.getClass().getDeclaredMethod("handleAndSendToUser");
 		this.sendToUserReturnType = new MethodParameter(method, -1);
 
+		method = this.getClass().getDeclaredMethod("handleAndSendToUserSingleSession");
+		this.sendToUserSingleSessionReturnType = new MethodParameter(method, -1);
+
 		method = this.getClass().getDeclaredMethod("handleAndSendToUserDefaultDestination");
 		this.sendToUserDefaultDestReturnType = new MethodParameter(method, -1);
+
+		method = this.getClass().getDeclaredMethod("handleAndSendToUserDefaultDestinationSingleSession");
+		this.sendToUserSingleSessionDefaultDestReturnType = new MethodParameter(method, -1);
 	}
 
 
@@ -211,6 +219,31 @@ public class SendToMethodReturnValueHandlerTests {
 
 		verify(this.messageChannel, times(2)).send(this.messageCaptor.capture());
 
+		Message<?> message = this.messageCaptor.getAllValues().get(0);
+		SimpMessageHeaderAccessor headers = SimpMessageHeaderAccessor.wrap(message);
+		assertNull(headers.getSessionId());
+		assertNull(headers.getSubscriptionId());
+		assertEquals("/user/" + user.getName() + "/dest1", headers.getDestination());
+
+		message = this.messageCaptor.getAllValues().get(1);
+		headers = SimpMessageHeaderAccessor.wrap(message);
+		assertNull(headers.getSessionId());
+		assertNull(headers.getSubscriptionId());
+		assertEquals("/user/" + user.getName() + "/dest2", headers.getDestination());
+	}
+
+	@Test
+	public void sendToUserSingleSession() throws Exception {
+
+		when(this.messageChannel.send(any(Message.class))).thenReturn(true);
+
+		String sessionId = "sess1";
+		TestUser user = new TestUser();
+		Message<?> inputMessage = createInputMessage(sessionId, "sub1", null, null, user);
+		this.handler.handleReturnValue(PAYLOAD, this.sendToUserSingleSessionReturnType, inputMessage);
+
+		verify(this.messageChannel, times(2)).send(this.messageCaptor.capture());
+
 		Message<?> message = this.messageCaptor.getAllValues().get(0);
 		SimpMessageHeaderAccessor headers = SimpMessageHeaderAccessor.wrap(message);
 		assertEquals(sessionId, headers.getSessionId());
@@ -257,6 +290,25 @@ public class SendToMethodReturnValueHandlerTests {
 
 		verify(this.messageChannel, times(1)).send(this.messageCaptor.capture());
 
+		Message<?> message = this.messageCaptor.getAllValues().get(0);
+		SimpMessageHeaderAccessor headers = SimpMessageHeaderAccessor.wrap(message);
+		assertNull(headers.getSessionId());
+		assertNull(headers.getSubscriptionId());
+		assertEquals("/user/" + user.getName() + "/queue/dest", headers.getDestination());
+	}
+
+	@Test
+	public void sendToUserDefaultDestinationSingleSession() throws Exception {
+
+		when(this.messageChannel.send(any(Message.class))).thenReturn(true);
+
+		String sessionId = "sess1";
+		TestUser user = new TestUser();
+		Message<?> inputMessage = createInputMessage(sessionId, "sub1", "/app", "/dest", user);
+		this.handler.handleReturnValue(PAYLOAD, this.sendToUserSingleSessionDefaultDestReturnType, inputMessage);
+
+		verify(this.messageChannel, times(1)).send(this.messageCaptor.capture());
+
 		Message<?> message = this.messageCaptor.getAllValues().get(0);
 		SimpMessageHeaderAccessor headers = SimpMessageHeaderAccessor.wrap(message);
 		assertEquals(sessionId, headers.getSessionId());
@@ -276,16 +328,8 @@ public class SendToMethodReturnValueHandlerTests {
 
 		handler.handleReturnValue(PAYLOAD, this.sendToUserDefaultDestReturnType, inputMessage);
 
-		ArgumentCaptor<MessageHeaders> captor = ArgumentCaptor.forClass(MessageHeaders.class);
-		verify(messagingTemplate).convertAndSendToUser(eq("joe"), eq("/queue/dest"), eq(PAYLOAD), captor.capture());
-
-		SimpMessageHeaderAccessor headerAccessor =
-				MessageHeaderAccessor.getAccessor(captor.getValue(), SimpMessageHeaderAccessor.class);
-
-		assertNotNull(headerAccessor);
-		assertTrue(headerAccessor.isMutable());
-		assertEquals("sess1", headerAccessor.getSessionId());
-		assertNull("Subscription id should not be copied", headerAccessor.getSubscriptionId());
+		verify(messagingTemplate).convertAndSendToUser(eq("joe"), eq("/queue/dest"), eq(PAYLOAD));
+		verifyNoMoreInteractions(messagingTemplate);
 	}
 
 
@@ -324,28 +368,45 @@ public class SendToMethodReturnValueHandlerTests {
 		}
 	}
 
+	@SuppressWarnings("unused")
 	public String handleNoAnnotations() {
 		return PAYLOAD;
 	}
 
+	@SuppressWarnings("unused")
 	@SendTo
 	public String handleAndSendToDefaultDestination() {
 		return PAYLOAD;
 	}
 
+	@SuppressWarnings("unused")
 	@SendTo({"/dest1", "/dest2"})
 	public String handleAndSendTo() {
 		return PAYLOAD;
 	}
 
+	@SuppressWarnings("unused")
 	@SendToUser
 	public String handleAndSendToUserDefaultDestination() {
 		return PAYLOAD;
 	}
 
+	@SuppressWarnings("unused")
+	@SendToUser(broadcast=false)
+	public String handleAndSendToUserDefaultDestinationSingleSession() {
+		return PAYLOAD;
+	}
+
+	@SuppressWarnings("unused")
 	@SendToUser({"/dest1", "/dest2"})
 	public String handleAndSendToUser() {
 		return PAYLOAD;
 	}
 
+	@SuppressWarnings("unused")
+	@SendToUser(value={"/dest1", "/dest2"}, broadcast=false)
+	public String handleAndSendToUserSingleSession() {
+		return PAYLOAD;
+	}
+
 }

spring-messaging/src/test/java/org/springframework/messaging/simp/user/DefaultUserDestinationResolverTests.java

@@ -109,7 +109,7 @@ public class DefaultUserDestinationResolverTests {
 		String userName = "http://joe.openid.example.org/";
 		this.registry.registerSessionId(userName, "openid123");
 		String destination = "/user/" + StringUtils.replace(userName, "/", "%2F") + "/queue/foo";
-		Message<?> message = createMessage(SimpMessageType.MESSAGE, this.user, SESSION_ID, destination);
+		Message<?> message = createMessage(SimpMessageType.MESSAGE, this.user, null, destination);
 		UserDestinationResult actual = this.resolver.resolveDestination(message);
 
 		assertEquals(1, actual.getTargetDestinations().size());