Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit fee7e6b

Browse files
catturtle123rstoyanchev
catturtle123
authored and
rstoyanchev
committed
Refine SendTo/Subscription header propagation
See gh-36179 Signed-off-by: Junhwan Kim <[email protected]>
1 parent 35c4973 commit fee7e6b

4 files changed

Lines changed: 52 additions & 39 deletions

File tree

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

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -263,9 +263,10 @@ private MessageHeaders createHeaders(@Nullable String sessionId, MethodParameter
263263
getHeaderInitializer().initHeaders(headerAccessor);
264264
}
265265
if (inputMessage != null && this.headerFilter != null) {
266-
inputMessage.getHeaders().forEach((name, value) -> {
266+
SimpMessageHeaderAccessor inputAccessor = SimpMessageHeaderAccessor.wrap(inputMessage);
267+
inputAccessor.toNativeHeaderMap().forEach((name, values) -> {
267268
if (this.headerFilter.test(name)) {
268-
headerAccessor.setHeader(name, value);
269+
headerAccessor.setNativeHeaderValues(name, values);
269270
}
270271
});
271272
}

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

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -160,9 +160,10 @@ private MessageHeaders createHeaders(@Nullable String sessionId, String subscrip
160160
getHeaderInitializer().initHeaders(accessor);
161161
}
162162
if (inputMessage != null && this.headerFilter != null) {
163-
inputMessage.getHeaders().forEach((name, value) -> {
163+
SimpMessageHeaderAccessor inputAccessor = SimpMessageHeaderAccessor.wrap(inputMessage);
164+
inputAccessor.toNativeHeaderMap().forEach((name, values) -> {
164165
if (this.headerFilter.test(name)) {
165-
accessor.setHeader(name, value);
166+
accessor.setNativeHeaderValues(name, values);
166167
}
167168
});
168169
}

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

Lines changed: 23 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -298,21 +298,26 @@ void sendToWithHeaderFilterSinglePredicate() throws Exception {
298298
given(this.messageChannel.send(any(Message.class))).willReturn(true);
299299

300300
String sessionId = "sess1";
301-
String headerName = "x-custom-header";
302-
String headerValue = "custom-value";
303-
Message<?> inputMessage = createMessage(sessionId, "sub1", null, null, null);
304-
inputMessage = MessageBuilder.fromMessage(inputMessage).setHeader(headerName, headerValue).build();
301+
String nativeHeaderName = "x-custom-header";
302+
String nativeHeaderValue = "custom-value";
303+
304+
SimpMessageHeaderAccessor inputAccessor = SimpMessageHeaderAccessor.create();
305+
inputAccessor.setSessionId(sessionId);
306+
inputAccessor.setSubscriptionId("sub1");
307+
inputAccessor.setNativeHeader(nativeHeaderName, nativeHeaderValue);
308+
Message<?> inputMessage = MessageBuilder.createMessage(new byte[0], inputAccessor.getMessageHeaders());
305309

306310
SimpMessagingTemplate template = new SimpMessagingTemplate(this.messageChannel);
307311
SendToMethodReturnValueHandler handler = new SendToMethodReturnValueHandler(template, true);
308-
handler.addHeaderFilter(name -> name.equals(headerName));
312+
handler.addHeaderFilter(name -> name.equals(nativeHeaderName));
309313

310314
handler.handleReturnValue(PAYLOAD, this.sendToReturnType, inputMessage);
311315

312316
verify(this.messageChannel, times(2)).send(this.messageCaptor.capture());
313317
for (Message<?> sent : this.messageCaptor.getAllValues()) {
314-
MessageHeaders headers = sent.getHeaders();
315-
assertThat(headers.get(headerName)).isEqualTo(headerValue);
318+
SimpMessageHeaderAccessor sentAccessor = MessageHeaderAccessor.getAccessor(sent, SimpMessageHeaderAccessor.class);
319+
assertThat(sentAccessor).isNotNull();
320+
assertThat(sentAccessor.getFirstNativeHeader(nativeHeaderName)).isEqualTo(nativeHeaderValue);
316321
}
317322
}
318323

@@ -323,11 +328,13 @@ void sendToWithHeaderFilterMultiplePredicates() throws Exception {
323328
String sessionId = "sess1";
324329
String headerA = "x-header-a";
325330
String headerB = "x-header-b";
326-
Message<?> inputMessage = createMessage(sessionId, "sub1", null, null, null);
327-
inputMessage = MessageBuilder.fromMessage(inputMessage)
328-
.setHeader(headerA, "A-value")
329-
.setHeader(headerB, "B-value")
330-
.build();
331+
332+
SimpMessageHeaderAccessor inputAccessor = SimpMessageHeaderAccessor.create();
333+
inputAccessor.setSessionId(sessionId);
334+
inputAccessor.setSubscriptionId("sub1");
335+
inputAccessor.setNativeHeader(headerA, "A-value");
336+
inputAccessor.setNativeHeader(headerB, "B-value");
337+
Message<?> inputMessage = MessageBuilder.createMessage(new byte[0], inputAccessor.getMessageHeaders());
331338

332339
SimpMessagingTemplate template = new SimpMessagingTemplate(this.messageChannel);
333340
SendToMethodReturnValueHandler handler = new SendToMethodReturnValueHandler(template, true);
@@ -338,9 +345,10 @@ void sendToWithHeaderFilterMultiplePredicates() throws Exception {
338345

339346
verify(this.messageChannel, times(2)).send(this.messageCaptor.capture());
340347
for (Message<?> sent : this.messageCaptor.getAllValues()) {
341-
MessageHeaders headers = sent.getHeaders();
342-
assertThat(headers.get(headerA)).isEqualTo("A-value");
343-
assertThat(headers.get(headerB)).isEqualTo("B-value");
348+
SimpMessageHeaderAccessor sentAccessor = MessageHeaderAccessor.getAccessor(sent, SimpMessageHeaderAccessor.class);
349+
assertThat(sentAccessor).isNotNull();
350+
assertThat(sentAccessor.getFirstNativeHeader(headerA)).isEqualTo("A-value");
351+
assertThat(sentAccessor.getFirstNativeHeader(headerB)).isEqualTo("B-value");
344352
}
345353
}
346354

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

Lines changed: 23 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -191,26 +191,28 @@ void testHeaderFilterSinglePredicate() throws Exception {
191191
String sessionId = "sess1";
192192
String subscriptionId = "subs1";
193193
String destination = "/dest";
194-
String customHeaderName = "x-custom-header";
195-
String customHeaderValue = "custom-value";
194+
String nativeHeaderName = "x-custom-header";
195+
String nativeHeaderValue = "custom-value";
196196

197-
Message<?> inputMessage = MessageBuilder.withPayload(PAYLOAD)
198-
.setHeader(SimpMessageHeaderAccessor.SESSION_ID_HEADER, sessionId)
199-
.setHeader(SimpMessageHeaderAccessor.SUBSCRIPTION_ID_HEADER, subscriptionId)
200-
.setHeader(SimpMessageHeaderAccessor.DESTINATION_HEADER, destination)
201-
.setHeader(customHeaderName, customHeaderValue)
202-
.build();
197+
SimpMessageHeaderAccessor inputAccessor = SimpMessageHeaderAccessor.create();
198+
inputAccessor.setSessionId(sessionId);
199+
inputAccessor.setSubscriptionId(subscriptionId);
200+
inputAccessor.setDestination(destination);
201+
inputAccessor.setNativeHeader(nativeHeaderName, nativeHeaderValue);
202+
Message<?> inputMessage = MessageBuilder.createMessage(PAYLOAD, inputAccessor.getMessageHeaders());
203203

204204
MessageSendingOperations template = mock();
205205
SubscriptionMethodReturnValueHandler handler = new SubscriptionMethodReturnValueHandler(template);
206-
handler.addHeaderFilter(name -> name.equals(customHeaderName));
206+
handler.addHeaderFilter(name -> name.equals(nativeHeaderName));
207207

208208
handler.handleReturnValue(PAYLOAD, this.subscribeEventReturnType, inputMessage);
209209

210210
ArgumentCaptor<MessageHeaders> captor = ArgumentCaptor.forClass(MessageHeaders.class);
211211
verify(template).convertAndSend(eq(destination), eq(PAYLOAD), captor.capture());
212212

213-
assertThat(captor.getValue().get(customHeaderName)).isEqualTo(customHeaderValue);
213+
SimpMessageHeaderAccessor sentAccessor = MessageHeaderAccessor.getAccessor(captor.getValue(), SimpMessageHeaderAccessor.class);
214+
assertThat(sentAccessor).isNotNull();
215+
assertThat(sentAccessor.getFirstNativeHeader(nativeHeaderName)).isEqualTo(nativeHeaderValue);
214216
}
215217

216218
@Test
@@ -221,13 +223,13 @@ void testHeaderFilterMultiplePredicates() throws Exception {
221223
String headerA = "x-header-a";
222224
String headerB = "x-header-b";
223225

224-
Message<?> inputMessage = MessageBuilder.withPayload(PAYLOAD)
225-
.setHeader(SimpMessageHeaderAccessor.SESSION_ID_HEADER, sessionId)
226-
.setHeader(SimpMessageHeaderAccessor.SUBSCRIPTION_ID_HEADER, subscriptionId)
227-
.setHeader(SimpMessageHeaderAccessor.DESTINATION_HEADER, destination)
228-
.setHeader(headerA, "A-value")
229-
.setHeader(headerB, "B-value")
230-
.build();
226+
SimpMessageHeaderAccessor inputAccessor = SimpMessageHeaderAccessor.create();
227+
inputAccessor.setSessionId(sessionId);
228+
inputAccessor.setSubscriptionId(subscriptionId);
229+
inputAccessor.setDestination(destination);
230+
inputAccessor.setNativeHeader(headerA, "A-value");
231+
inputAccessor.setNativeHeader(headerB, "B-value");
232+
Message<?> inputMessage = MessageBuilder.createMessage(PAYLOAD, inputAccessor.getMessageHeaders());
231233

232234
MessageSendingOperations template = mock();
233235
SubscriptionMethodReturnValueHandler handler = new SubscriptionMethodReturnValueHandler(template);
@@ -239,9 +241,10 @@ void testHeaderFilterMultiplePredicates() throws Exception {
239241
ArgumentCaptor<MessageHeaders> captor = ArgumentCaptor.forClass(MessageHeaders.class);
240242
verify(template).convertAndSend(eq(destination), eq(PAYLOAD), captor.capture());
241243

242-
MessageHeaders sentHeaders = captor.getValue();
243-
assertThat(sentHeaders.get(headerA)).isEqualTo("A-value");
244-
assertThat(sentHeaders.get(headerB)).isEqualTo("B-value");
244+
SimpMessageHeaderAccessor sentAccessor = MessageHeaderAccessor.getAccessor(captor.getValue(), SimpMessageHeaderAccessor.class);
245+
assertThat(sentAccessor).isNotNull();
246+
assertThat(sentAccessor.getFirstNativeHeader(headerA)).isEqualTo("A-value");
247+
assertThat(sentAccessor.getFirstNativeHeader(headerB)).isEqualTo("B-value");
245248
}
246249

247250
private Message<?> createInputMessage(String sessId, String subsId, String dest, Principal principal) {

0 commit comments

Comments
 (0)