Refine SendTo/Subscription header propagation

See gh-36179

Signed-off-by: Junhwan Kim <musoyou1085@gmail.com>
This commit is contained in:
Junhwan Kim
2026-02-10 15:28:54 +00:00
committed by rstoyanchev
parent 35c4973504
commit fee7e6b7e2
4 changed files with 52 additions and 39 deletions
@@ -263,9 +263,10 @@ public class SendToMethodReturnValueHandler implements HandlerMethodReturnValueH
getHeaderInitializer().initHeaders(headerAccessor);
}
if (inputMessage != null && this.headerFilter != null) {
inputMessage.getHeaders().forEach((name, value) -> {
SimpMessageHeaderAccessor inputAccessor = SimpMessageHeaderAccessor.wrap(inputMessage);
inputAccessor.toNativeHeaderMap().forEach((name, values) -> {
if (this.headerFilter.test(name)) {
headerAccessor.setHeader(name, value);
headerAccessor.setNativeHeaderValues(name, values);
}
});
}
@@ -160,9 +160,10 @@ public class SubscriptionMethodReturnValueHandler implements HandlerMethodReturn
getHeaderInitializer().initHeaders(accessor);
}
if (inputMessage != null && this.headerFilter != null) {
inputMessage.getHeaders().forEach((name, value) -> {
SimpMessageHeaderAccessor inputAccessor = SimpMessageHeaderAccessor.wrap(inputMessage);
inputAccessor.toNativeHeaderMap().forEach((name, values) -> {
if (this.headerFilter.test(name)) {
accessor.setHeader(name, value);
accessor.setNativeHeaderValues(name, values);
}
});
}
@@ -298,21 +298,26 @@ public class SendToMethodReturnValueHandlerTests {
given(this.messageChannel.send(any(Message.class))).willReturn(true);
String sessionId = "sess1";
String headerName = "x-custom-header";
String headerValue = "custom-value";
Message<?> inputMessage = createMessage(sessionId, "sub1", null, null, null);
inputMessage = MessageBuilder.fromMessage(inputMessage).setHeader(headerName, headerValue).build();
String nativeHeaderName = "x-custom-header";
String nativeHeaderValue = "custom-value";
SimpMessageHeaderAccessor inputAccessor = SimpMessageHeaderAccessor.create();
inputAccessor.setSessionId(sessionId);
inputAccessor.setSubscriptionId("sub1");
inputAccessor.setNativeHeader(nativeHeaderName, nativeHeaderValue);
Message<?> inputMessage = MessageBuilder.createMessage(new byte[0], inputAccessor.getMessageHeaders());
SimpMessagingTemplate template = new SimpMessagingTemplate(this.messageChannel);
SendToMethodReturnValueHandler handler = new SendToMethodReturnValueHandler(template, true);
handler.addHeaderFilter(name -> name.equals(headerName));
handler.addHeaderFilter(name -> name.equals(nativeHeaderName));
handler.handleReturnValue(PAYLOAD, this.sendToReturnType, inputMessage);
verify(this.messageChannel, times(2)).send(this.messageCaptor.capture());
for (Message<?> sent : this.messageCaptor.getAllValues()) {
MessageHeaders headers = sent.getHeaders();
assertThat(headers.get(headerName)).isEqualTo(headerValue);
SimpMessageHeaderAccessor sentAccessor = MessageHeaderAccessor.getAccessor(sent, SimpMessageHeaderAccessor.class);
assertThat(sentAccessor).isNotNull();
assertThat(sentAccessor.getFirstNativeHeader(nativeHeaderName)).isEqualTo(nativeHeaderValue);
}
}
@@ -323,11 +328,13 @@ public class SendToMethodReturnValueHandlerTests {
String sessionId = "sess1";
String headerA = "x-header-a";
String headerB = "x-header-b";
Message<?> inputMessage = createMessage(sessionId, "sub1", null, null, null);
inputMessage = MessageBuilder.fromMessage(inputMessage)
.setHeader(headerA, "A-value")
.setHeader(headerB, "B-value")
.build();
SimpMessageHeaderAccessor inputAccessor = SimpMessageHeaderAccessor.create();
inputAccessor.setSessionId(sessionId);
inputAccessor.setSubscriptionId("sub1");
inputAccessor.setNativeHeader(headerA, "A-value");
inputAccessor.setNativeHeader(headerB, "B-value");
Message<?> inputMessage = MessageBuilder.createMessage(new byte[0], inputAccessor.getMessageHeaders());
SimpMessagingTemplate template = new SimpMessagingTemplate(this.messageChannel);
SendToMethodReturnValueHandler handler = new SendToMethodReturnValueHandler(template, true);
@@ -338,9 +345,10 @@ public class SendToMethodReturnValueHandlerTests {
verify(this.messageChannel, times(2)).send(this.messageCaptor.capture());
for (Message<?> sent : this.messageCaptor.getAllValues()) {
MessageHeaders headers = sent.getHeaders();
assertThat(headers.get(headerA)).isEqualTo("A-value");
assertThat(headers.get(headerB)).isEqualTo("B-value");
SimpMessageHeaderAccessor sentAccessor = MessageHeaderAccessor.getAccessor(sent, SimpMessageHeaderAccessor.class);
assertThat(sentAccessor).isNotNull();
assertThat(sentAccessor.getFirstNativeHeader(headerA)).isEqualTo("A-value");
assertThat(sentAccessor.getFirstNativeHeader(headerB)).isEqualTo("B-value");
}
}
@@ -191,26 +191,28 @@ class SubscriptionMethodReturnValueHandlerTests {
String sessionId = "sess1";
String subscriptionId = "subs1";
String destination = "/dest";
String customHeaderName = "x-custom-header";
String customHeaderValue = "custom-value";
String nativeHeaderName = "x-custom-header";
String nativeHeaderValue = "custom-value";
Message<?> inputMessage = MessageBuilder.withPayload(PAYLOAD)
.setHeader(SimpMessageHeaderAccessor.SESSION_ID_HEADER, sessionId)
.setHeader(SimpMessageHeaderAccessor.SUBSCRIPTION_ID_HEADER, subscriptionId)
.setHeader(SimpMessageHeaderAccessor.DESTINATION_HEADER, destination)
.setHeader(customHeaderName, customHeaderValue)
.build();
SimpMessageHeaderAccessor inputAccessor = SimpMessageHeaderAccessor.create();
inputAccessor.setSessionId(sessionId);
inputAccessor.setSubscriptionId(subscriptionId);
inputAccessor.setDestination(destination);
inputAccessor.setNativeHeader(nativeHeaderName, nativeHeaderValue);
Message<?> inputMessage = MessageBuilder.createMessage(PAYLOAD, inputAccessor.getMessageHeaders());
MessageSendingOperations template = mock();
SubscriptionMethodReturnValueHandler handler = new SubscriptionMethodReturnValueHandler(template);
handler.addHeaderFilter(name -> name.equals(customHeaderName));
handler.addHeaderFilter(name -> name.equals(nativeHeaderName));
handler.handleReturnValue(PAYLOAD, this.subscribeEventReturnType, inputMessage);
ArgumentCaptor<MessageHeaders> captor = ArgumentCaptor.forClass(MessageHeaders.class);
verify(template).convertAndSend(eq(destination), eq(PAYLOAD), captor.capture());
assertThat(captor.getValue().get(customHeaderName)).isEqualTo(customHeaderValue);
SimpMessageHeaderAccessor sentAccessor = MessageHeaderAccessor.getAccessor(captor.getValue(), SimpMessageHeaderAccessor.class);
assertThat(sentAccessor).isNotNull();
assertThat(sentAccessor.getFirstNativeHeader(nativeHeaderName)).isEqualTo(nativeHeaderValue);
}
@Test
@@ -221,13 +223,13 @@ class SubscriptionMethodReturnValueHandlerTests {
String headerA = "x-header-a";
String headerB = "x-header-b";
Message<?> inputMessage = MessageBuilder.withPayload(PAYLOAD)
.setHeader(SimpMessageHeaderAccessor.SESSION_ID_HEADER, sessionId)
.setHeader(SimpMessageHeaderAccessor.SUBSCRIPTION_ID_HEADER, subscriptionId)
.setHeader(SimpMessageHeaderAccessor.DESTINATION_HEADER, destination)
.setHeader(headerA, "A-value")
.setHeader(headerB, "B-value")
.build();
SimpMessageHeaderAccessor inputAccessor = SimpMessageHeaderAccessor.create();
inputAccessor.setSessionId(sessionId);
inputAccessor.setSubscriptionId(subscriptionId);
inputAccessor.setDestination(destination);
inputAccessor.setNativeHeader(headerA, "A-value");
inputAccessor.setNativeHeader(headerB, "B-value");
Message<?> inputMessage = MessageBuilder.createMessage(PAYLOAD, inputAccessor.getMessageHeaders());
MessageSendingOperations template = mock();
SubscriptionMethodReturnValueHandler handler = new SubscriptionMethodReturnValueHandler(template);
@@ -239,9 +241,10 @@ class SubscriptionMethodReturnValueHandlerTests {
ArgumentCaptor<MessageHeaders> captor = ArgumentCaptor.forClass(MessageHeaders.class);
verify(template).convertAndSend(eq(destination), eq(PAYLOAD), captor.capture());
MessageHeaders sentHeaders = captor.getValue();
assertThat(sentHeaders.get(headerA)).isEqualTo("A-value");
assertThat(sentHeaders.get(headerB)).isEqualTo("B-value");
SimpMessageHeaderAccessor sentAccessor = MessageHeaderAccessor.getAccessor(captor.getValue(), SimpMessageHeaderAccessor.class);
assertThat(sentAccessor).isNotNull();
assertThat(sentAccessor.getFirstNativeHeader(headerA)).isEqualTo("A-value");
assertThat(sentAccessor.getFirstNativeHeader(headerB)).isEqualTo("B-value");
}
private Message<?> createInputMessage(String sessId, String subsId, String dest, Principal principal) {