diff --git a/spring-messaging/src/main/java/org/springframework/messaging/simp/annotation/support/SendToMethodReturnValueHandler.java b/spring-messaging/src/main/java/org/springframework/messaging/simp/annotation/support/SendToMethodReturnValueHandler.java index 257cb2e7939..db524ca50da 100644 --- a/spring-messaging/src/main/java/org/springframework/messaging/simp/annotation/support/SendToMethodReturnValueHandler.java +++ b/spring-messaging/src/main/java/org/springframework/messaging/simp/annotation/support/SendToMethodReturnValueHandler.java @@ -263,9 +263,10 @@ private MessageHeaders createHeaders(@Nullable String sessionId, MethodParameter 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); } }); } diff --git a/spring-messaging/src/main/java/org/springframework/messaging/simp/annotation/support/SubscriptionMethodReturnValueHandler.java b/spring-messaging/src/main/java/org/springframework/messaging/simp/annotation/support/SubscriptionMethodReturnValueHandler.java index c7ba0a7de47..f133daaa2b1 100644 --- a/spring-messaging/src/main/java/org/springframework/messaging/simp/annotation/support/SubscriptionMethodReturnValueHandler.java +++ b/spring-messaging/src/main/java/org/springframework/messaging/simp/annotation/support/SubscriptionMethodReturnValueHandler.java @@ -160,9 +160,10 @@ private MessageHeaders createHeaders(@Nullable String sessionId, String subscrip 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); } }); } diff --git a/spring-messaging/src/test/java/org/springframework/messaging/simp/annotation/support/SendToMethodReturnValueHandlerTests.java b/spring-messaging/src/test/java/org/springframework/messaging/simp/annotation/support/SendToMethodReturnValueHandlerTests.java index 86ded07f1af..b3d67cbc26b 100644 --- a/spring-messaging/src/test/java/org/springframework/messaging/simp/annotation/support/SendToMethodReturnValueHandlerTests.java +++ b/spring-messaging/src/test/java/org/springframework/messaging/simp/annotation/support/SendToMethodReturnValueHandlerTests.java @@ -298,21 +298,26 @@ void sendToWithHeaderFilterSinglePredicate() throws Exception { 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 @@ void sendToWithHeaderFilterMultiplePredicates() throws Exception { 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 @@ void sendToWithHeaderFilterMultiplePredicates() throws Exception { 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"); } } diff --git a/spring-messaging/src/test/java/org/springframework/messaging/simp/annotation/support/SubscriptionMethodReturnValueHandlerTests.java b/spring-messaging/src/test/java/org/springframework/messaging/simp/annotation/support/SubscriptionMethodReturnValueHandlerTests.java index b795b4e1636..3ed01bbe38f 100644 --- a/spring-messaging/src/test/java/org/springframework/messaging/simp/annotation/support/SubscriptionMethodReturnValueHandlerTests.java +++ b/spring-messaging/src/test/java/org/springframework/messaging/simp/annotation/support/SubscriptionMethodReturnValueHandlerTests.java @@ -191,26 +191,28 @@ void testHeaderFilterSinglePredicate() throws Exception { 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 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 @@ void testHeaderFilterMultiplePredicates() throws Exception { 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 @@ void testHeaderFilterMultiplePredicates() throws Exception { ArgumentCaptor 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) {