Skip to content

Commit f4bffea

Browse files
rstoyanchevrstoyanchev
committed
Support for @sendto and @SendToUser for same method
Issue: SPR-16891
1 parent 7a70f7c commit f4bffea

File tree

2 files changed

+94
-49
lines changed

2 files changed

+94
-49
lines changed

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

Lines changed: 54 additions & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright 2002-2017 the original author or authors.
2+
* Copyright 2002-2018 the original author or authors.
33
*
44
* Licensed under the Apache License, Version 2.0 (the "License");
55
* you may not use this file except in compliance with the License.
@@ -18,6 +18,7 @@
1818

1919
import java.lang.annotation.Annotation;
2020
import java.security.Principal;
21+
import java.util.Collections;
2122
import java.util.Map;
2223

2324
import org.springframework.core.MethodParameter;
@@ -153,11 +154,10 @@ public void handleReturnValue(@Nullable Object returnValue, MethodParameter retu
153154

154155
MessageHeaders headers = message.getHeaders();
155156
String sessionId = SimpMessageHeaderAccessor.getSessionId(headers);
156-
PlaceholderResolver varResolver = initVarResolver(headers);
157-
Object annotation = findAnnotation(returnType);
157+
DestinationHelper destinationHelper = getDestinationHelper(headers, returnType);
158158

159-
if (annotation instanceof SendToUser) {
160-
SendToUser sendToUser = (SendToUser) annotation;
159+
SendToUser sendToUser = destinationHelper.getSendToUser();
160+
if (sendToUser != null) {
161161
boolean broadcast = sendToUser.broadcast();
162162
String user = getUserName(message, headers);
163163
if (user == null) {
@@ -169,7 +169,7 @@ public void handleReturnValue(@Nullable Object returnValue, MethodParameter retu
169169
}
170170
String[] destinations = getTargetDestinations(sendToUser, message, this.defaultUserDestinationPrefix);
171171
for (String destination : destinations) {
172-
destination = this.placeholderHelper.replacePlaceholders(destination, varResolver);
172+
destination = destinationHelper.expandTemplateVars(destination);
173173
if (broadcast) {
174174
this.messagingTemplate.convertAndSendToUser(
175175
user, destination, returnValue, createHeaders(null, returnType));
@@ -180,51 +180,33 @@ public void handleReturnValue(@Nullable Object returnValue, MethodParameter retu
180180
}
181181
}
182182
}
183-
else {
184-
SendTo sendTo = (SendTo) annotation; // possibly null
183+
184+
SendTo sendTo = destinationHelper.getSendTo();
185+
if (sendTo != null || sendToUser == null) {
185186
String[] destinations = getTargetDestinations(sendTo, message, this.defaultDestinationPrefix);
186187
for (String destination : destinations) {
187-
destination = this.placeholderHelper.replacePlaceholders(destination, varResolver);
188+
destination = destinationHelper.expandTemplateVars(destination);
188189
this.messagingTemplate.convertAndSend(destination, returnValue, createHeaders(sessionId, returnType));
189190
}
190191
}
191192
}
192193

193-
@Nullable
194-
private Object findAnnotation(MethodParameter returnType) {
195-
Annotation[] anns = new Annotation[4];
196-
anns[0] = AnnotatedElementUtils.findMergedAnnotation(returnType.getExecutable(), SendToUser.class);
197-
anns[1] = AnnotatedElementUtils.findMergedAnnotation(returnType.getExecutable(), SendTo.class);
198-
anns[2] = AnnotatedElementUtils.findMergedAnnotation(returnType.getDeclaringClass(), SendToUser.class);
199-
anns[3] = AnnotatedElementUtils.findMergedAnnotation(returnType.getDeclaringClass(), SendTo.class);
200-
201-
if (anns[0] != null && !ObjectUtils.isEmpty(((SendToUser) anns[0]).value())) {
202-
return anns[0];
203-
}
204-
if (anns[1] != null && !ObjectUtils.isEmpty(((SendTo) anns[1]).value())) {
205-
return anns[1];
206-
}
207-
if (anns[2] != null && !ObjectUtils.isEmpty(((SendToUser) anns[2]).value())) {
208-
return anns[2];
209-
}
210-
if (anns[3] != null && !ObjectUtils.isEmpty(((SendTo) anns[3]).value())) {
211-
return anns[3];
212-
}
194+
private DestinationHelper getDestinationHelper(MessageHeaders headers, MethodParameter returnType) {
213195

214-
for (int i=0; i < 4; i++) {
215-
if (anns[i] != null) {
216-
return anns[i];
217-
}
196+
SendToUser m1 = AnnotatedElementUtils.findMergedAnnotation(returnType.getExecutable(), SendToUser.class);
197+
SendTo m2 = AnnotatedElementUtils.findMergedAnnotation(returnType.getExecutable(), SendTo.class);
198+
if ((m1 != null && !ObjectUtils.isEmpty(m1.value())) || (m2 != null && !ObjectUtils.isEmpty(m2.value()))) {
199+
return new DestinationHelper(headers, m1, m2);
218200
}
219201

220-
return null;
221-
}
202+
SendToUser c1 = AnnotatedElementUtils.findMergedAnnotation(returnType.getDeclaringClass(), SendToUser.class);
203+
SendTo c2 = AnnotatedElementUtils.findMergedAnnotation(returnType.getDeclaringClass(), SendTo.class);
204+
if ((c1 != null && !ObjectUtils.isEmpty(c1.value())) || (c2 != null && !ObjectUtils.isEmpty(c2.value()))) {
205+
return new DestinationHelper(headers, c1, c2);
206+
}
222207

223-
@SuppressWarnings("unchecked")
224-
private PlaceholderResolver initVarResolver(MessageHeaders headers) {
225-
String name = DestinationVariableMethodArgumentResolver.DESTINATION_TEMPLATE_VARIABLES_HEADER;
226-
Map<String, String> vars = (Map<String, String>) headers.get(name);
227-
return new DestinationVariablePlaceholderResolver(vars);
208+
return m1 != null || m2 != null ?
209+
new DestinationHelper(headers, m1, m2) : new DestinationHelper(headers, c1, c2);
228210
}
229211

230212
@Nullable
@@ -275,20 +257,43 @@ public String toString() {
275257
}
276258

277259

278-
private static class DestinationVariablePlaceholderResolver implements PlaceholderResolver {
260+
private class DestinationHelper {
261+
262+
private final PlaceholderResolver placeholderResolver;
279263

280264
@Nullable
281-
private final Map<String, String> vars;
265+
private final SendTo sendTo;
282266

283-
public DestinationVariablePlaceholderResolver(@Nullable Map<String, String> vars) {
284-
this.vars = vars;
267+
@Nullable
268+
private final SendToUser sendToUser;
269+
270+
271+
public DestinationHelper(MessageHeaders headers, @Nullable SendToUser sendToUser, @Nullable SendTo sendTo) {
272+
Map<String, String> variables = getTemplateVariables(headers);
273+
this.placeholderResolver = variables::get;
274+
this.sendTo = sendTo;
275+
this.sendToUser = sendToUser;
276+
}
277+
278+
@SuppressWarnings("unchecked")
279+
private Map<String, String> getTemplateVariables(MessageHeaders headers) {
280+
String name = DestinationVariableMethodArgumentResolver.DESTINATION_TEMPLATE_VARIABLES_HEADER;
281+
return (Map<String, String>) headers.getOrDefault(name, Collections.emptyMap());
285282
}
286283

287-
@Override
288284
@Nullable
289-
public String resolvePlaceholder(String placeholderName) {
290-
return (this.vars != null ? this.vars.get(placeholderName) : null);
285+
public SendTo getSendTo() {
286+
return this.sendTo;
291287
}
292-
}
293288

294-
}
289+
@Nullable
290+
public SendToUser getSendToUser() {
291+
return this.sendToUser;
292+
}
293+
294+
295+
public String expandTemplateVars(String destination) {
296+
return placeholderHelper.replacePlaceholders(destination, this.placeholderResolver);
297+
}
298+
}
299+
}

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

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -87,6 +87,7 @@ public class SendToMethodReturnValueHandlerTests {
8787
private MethodParameter sendToWithPlaceholdersReturnType = param("handleAndSendToWithPlaceholders");
8888
private MethodParameter sendToUserReturnType = param("handleAndSendToUser");
8989
private MethodParameter sendToUserInSessionReturnType = param("handleAndSendToUserInSession");
90+
private MethodParameter sendToSendToUserReturnType = param("handleAndSendToAndSendToUser");
9091
private MethodParameter sendToUserDefaultDestReturnType = param("handleAndSendToUserDefaultDest");
9192
private MethodParameter sendToUserInSessionDefaultDestReturnType = param("handleAndSendToUserDefaultDestInSession");
9293
private MethodParameter jsonViewReturnType = param("handleAndSendToJsonView");
@@ -355,6 +356,38 @@ public void sendToUser() throws Exception {
355356
assertEquals("/user/" + user.getName() + "/dest2", accessor.getDestination());
356357
}
357358

359+
@Test
360+
public void sendToAndSendToUser() throws Exception {
361+
given(this.messageChannel.send(any(Message.class))).willReturn(true);
362+
363+
String sessionId = "sess1";
364+
TestUser user = new TestUser();
365+
Message<?> inputMessage = createMessage(sessionId, "sub1", null, null, user);
366+
this.handler.handleReturnValue(PAYLOAD, this.sendToSendToUserReturnType, inputMessage);
367+
368+
verify(this.messageChannel, times(4)).send(this.messageCaptor.capture());
369+
370+
SimpMessageHeaderAccessor accessor = getCapturedAccessor(0);
371+
assertNull(accessor.getSessionId());
372+
assertNull(accessor.getSubscriptionId());
373+
assertEquals("/user/" + user.getName() + "/dest1", accessor.getDestination());
374+
375+
accessor = getCapturedAccessor(1);
376+
assertNull(accessor.getSessionId());
377+
assertNull(accessor.getSubscriptionId());
378+
assertEquals("/user/" + user.getName() + "/dest2", accessor.getDestination());
379+
380+
accessor = getCapturedAccessor(2);
381+
assertEquals("sess1", accessor.getSessionId());
382+
assertNull(accessor.getSubscriptionId());
383+
assertEquals("/dest1", accessor.getDestination());
384+
385+
accessor = getCapturedAccessor(3);
386+
assertEquals("sess1", accessor.getSessionId());
387+
assertNull(accessor.getSubscriptionId());
388+
assertEquals("/dest2", accessor.getDestination());
389+
}
390+
358391
@Test // SPR-12170
359392
public void sendToWithDestinationPlaceholders() throws Exception {
360393
given(this.messageChannel.send(any(Message.class))).willReturn(true);
@@ -577,6 +610,13 @@ String handleAndSendToUserInSession() {
577610
return PAYLOAD;
578611
}
579612

613+
@SendTo({"/dest1", "/dest2"})
614+
@SendToUser({"/dest1", "/dest2"})
615+
@SuppressWarnings("unused")
616+
String handleAndSendToAndSendToUser() {
617+
return PAYLOAD;
618+
}
619+
580620
@JsonView(MyJacksonView1.class)
581621
@SuppressWarnings("unused")
582622
JacksonViewBean handleAndSendToJsonView() {

0 commit comments

Comments
 (0)