diff --git a/spring-messaging/src/main/java/org/springframework/messaging/simp/SimpMessageHeaderAccessor.java b/spring-messaging/src/main/java/org/springframework/messaging/simp/SimpMessageHeaderAccessor.java index b5d7d57e77..6c314f3329 100644 --- a/spring-messaging/src/main/java/org/springframework/messaging/simp/SimpMessageHeaderAccessor.java +++ b/spring-messaging/src/main/java/org/springframework/messaging/simp/SimpMessageHeaderAccessor.java @@ -1,5 +1,5 @@ /* - * Copyright 2002-2017 the original author or authors. + * Copyright 2002-2019 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. @@ -19,6 +19,7 @@ package org.springframework.messaging.simp; import java.security.Principal; import java.util.List; import java.util.Map; +import java.util.function.Consumer; import org.springframework.lang.Nullable; import org.springframework.messaging.Message; @@ -84,6 +85,10 @@ public class SimpMessageHeaderAccessor extends NativeMessageHeaderAccessor { public static final String IGNORE_ERROR = "simpIgnoreError"; + @Nullable + private Consumer userCallback; + + /** * A constructor for creating new message headers. * This constructor is protected. See factory methods in this and sub-classes. @@ -171,6 +176,9 @@ public class SimpMessageHeaderAccessor extends NativeMessageHeaderAccessor { public void setUser(@Nullable Principal principal) { setHeader(USER_HEADER, principal); + if (this.userCallback != null) { + this.userCallback.accept(principal); + } } /** @@ -181,6 +189,18 @@ public class SimpMessageHeaderAccessor extends NativeMessageHeaderAccessor { return (Principal) getHeader(USER_HEADER); } + /** + * Provide a callback to be invoked if and when {@link #setUser(Principal)} + * is called. This is used internally on the inbound channel to detect + * token-based authentications through an interceptor. + * @param callback the callback to invoke + * @since 5.1.9 + */ + public void setUserChangeCallback(Consumer callback) { + Assert.notNull(callback, "'callback' is required"); + this.userCallback = this.userCallback != null ? this.userCallback.andThen(callback) : callback; + } + @Override public String getShortLogMessage(Object payload) { if (getMessageType() == null) { diff --git a/spring-messaging/src/test/java/org/springframework/messaging/simp/SimpMessageHeaderAccessorTests.java b/spring-messaging/src/test/java/org/springframework/messaging/simp/SimpMessageHeaderAccessorTests.java index 69221a0077..9cfbc10723 100644 --- a/spring-messaging/src/test/java/org/springframework/messaging/simp/SimpMessageHeaderAccessorTests.java +++ b/spring-messaging/src/test/java/org/springframework/messaging/simp/SimpMessageHeaderAccessorTests.java @@ -16,11 +16,14 @@ package org.springframework.messaging.simp; +import java.security.Principal; import java.util.Collections; +import java.util.function.Consumer; import org.junit.Test; import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.Mockito.mock; /** * Unit tests for SimpMessageHeaderAccessor. @@ -32,7 +35,8 @@ public class SimpMessageHeaderAccessorTests { @Test public void getShortLogMessage() { - assertThat(SimpMessageHeaderAccessor.create().getShortLogMessage("p")).isEqualTo("MESSAGE session=null payload=p"); + assertThat(SimpMessageHeaderAccessor.create().getShortLogMessage("p")) + .isEqualTo("MESSAGE session=null payload=p"); } @Test @@ -44,8 +48,9 @@ public class SimpMessageHeaderAccessorTests { accessor.setUser(new TestPrincipal("user")); accessor.setSessionAttributes(Collections.singletonMap("key", "value")); - assertThat(accessor.getShortLogMessage("p")).isEqualTo(("MESSAGE destination=/destination subscriptionId=subscription " + - "session=session user=user attributes[1] payload=p")); + assertThat(accessor.getShortLogMessage("p")) + .isEqualTo(("MESSAGE destination=/destination subscriptionId=subscription " + + "session=session user=user attributes[1] payload=p")); } @Test @@ -58,9 +63,41 @@ public class SimpMessageHeaderAccessorTests { accessor.setSessionAttributes(Collections.singletonMap("key", "value")); accessor.setNativeHeader("nativeKey", "nativeValue"); - assertThat(accessor.getDetailedLogMessage("p")).isEqualTo(("MESSAGE destination=/destination subscriptionId=subscription " + - "session=session user=user attributes={key=value} nativeHeaders=" + - "{nativeKey=[nativeValue]} payload=p")); + assertThat(accessor.getDetailedLogMessage("p")) + .isEqualTo(("MESSAGE destination=/destination subscriptionId=subscription " + + "session=session user=user attributes={key=value} nativeHeaders=" + + "{nativeKey=[nativeValue]} payload=p")); + } + + @Test + public void userChangeCallback() { + UserCallback userCallback = new UserCallback(); + SimpMessageHeaderAccessor accessor = SimpMessageHeaderAccessor.create(); + accessor.setUserChangeCallback(userCallback); + + Principal user1 = mock(Principal.class); + accessor.setUser(user1); + assertThat(userCallback.getUser()).isEqualTo(user1); + + Principal user2 = mock(Principal.class); + accessor.setUser(user2); + assertThat(userCallback.getUser()).isEqualTo(user2); + } + + + private static class UserCallback implements Consumer { + + private Principal user; + + + public Principal getUser() { + return this.user; + } + + @Override + public void accept(Principal principal) { + this.user = principal; + } } } diff --git a/spring-websocket/src/main/java/org/springframework/web/socket/messaging/StompSubProtocolHandler.java b/spring-websocket/src/main/java/org/springframework/web/socket/messaging/StompSubProtocolHandler.java index 4e190f8476..e5764b1c27 100644 --- a/spring-websocket/src/main/java/org/springframework/web/socket/messaging/StompSubProtocolHandler.java +++ b/spring-websocket/src/main/java/org/springframework/web/socket/messaging/StompSubProtocolHandler.java @@ -267,9 +267,19 @@ public class StompSubProtocolHandler implements SubProtocolHandler, ApplicationE MessageHeaderAccessor.getAccessor(message, StompHeaderAccessor.class); Assert.state(headerAccessor != null, "No StompHeaderAccessor"); + StompCommand command = headerAccessor.getCommand(); + boolean isConnect = StompCommand.CONNECT.equals(command) || StompCommand.STOMP.equals(command); + headerAccessor.setSessionId(session.getId()); headerAccessor.setSessionAttributes(session.getAttributes()); headerAccessor.setUser(getUser(session)); + if (isConnect) { + headerAccessor.setUserChangeCallback(user -> { + if (user != null && user != session.getPrincipal()) { + this.stompAuthentications.put(session.getId(), user); + } + }); + } headerAccessor.setHeader(SimpMessageHeaderAccessor.HEART_BEAT_HEADER, headerAccessor.getHeartbeat()); if (!detectImmutableMessageInterceptor(outputChannel)) { headerAccessor.setImmutable(); @@ -279,8 +289,6 @@ public class StompSubProtocolHandler implements SubProtocolHandler, ApplicationE logger.trace("From client: " + headerAccessor.getShortLogMessage(message.getPayload())); } - StompCommand command = headerAccessor.getCommand(); - boolean isConnect = StompCommand.CONNECT.equals(command) || StompCommand.STOMP.equals(command); if (isConnect) { this.stats.incrementConnectCount(); } @@ -293,12 +301,6 @@ public class StompSubProtocolHandler implements SubProtocolHandler, ApplicationE boolean sent = outputChannel.send(message); if (sent) { - if (isConnect) { - Principal user = headerAccessor.getUser(); - if (user != null && user != session.getPrincipal()) { - this.stompAuthentications.put(session.getId(), user); - } - } if (this.eventPublisher != null) { Principal user = getUser(session); if (isConnect) { diff --git a/spring-websocket/src/main/java/org/springframework/web/socket/server/support/WebSocketHttpRequestHandler.java b/spring-websocket/src/main/java/org/springframework/web/socket/server/support/WebSocketHttpRequestHandler.java index 296d974dfc..4f9fa2d478 100644 --- a/spring-websocket/src/main/java/org/springframework/web/socket/server/support/WebSocketHttpRequestHandler.java +++ b/spring-websocket/src/main/java/org/springframework/web/socket/server/support/WebSocketHttpRequestHandler.java @@ -1,5 +1,5 @@ /* - * Copyright 2002-2018 the original author or authors. + * Copyright 2002-2019 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. @@ -166,7 +166,6 @@ public class WebSocketHttpRequestHandler implements HttpRequestHandler, Lifecycl } this.handshakeHandler.doHandshake(request, response, this.wsHandler, attributes); chain.applyAfterHandshake(request, response, null); - response.close(); } catch (HandshakeFailureException ex) { failure = ex; @@ -177,8 +176,10 @@ public class WebSocketHttpRequestHandler implements HttpRequestHandler, Lifecycl finally { if (failure != null) { chain.applyAfterHandshake(request, response, failure); + response.close(); throw failure; } + response.close(); } } diff --git a/spring-websocket/src/test/java/org/springframework/web/socket/messaging/StompSubProtocolHandlerTests.java b/spring-websocket/src/test/java/org/springframework/web/socket/messaging/StompSubProtocolHandlerTests.java index e7e599face..e87fb194d1 100644 --- a/spring-websocket/src/test/java/org/springframework/web/socket/messaging/StompSubProtocolHandlerTests.java +++ b/spring-websocket/src/test/java/org/springframework/web/socket/messaging/StompSubProtocolHandlerTests.java @@ -383,6 +383,15 @@ public class StompSubProtocolHandlerTests { Principal user = SimpMessageHeaderAccessor.getUser(message.getHeaders()); assertThat(user).isNotNull(); assertThat(user.getName()).isEqualTo("__pete__@gmail.com"); + + StompHeaderAccessor accessor = StompHeaderAccessor.create(StompCommand.CONNECTED); + message = MessageBuilder.createMessage(EMPTY_PAYLOAD, accessor.getMessageHeaders()); + handler.handleMessageToClient(this.session, message); + + assertThat(this.session.getSentMessages()).hasSize(1); + WebSocketMessage textMessage = this.session.getSentMessages().get(0); + assertThat(textMessage.getPayload()) + .isEqualTo("CONNECTED\n" + "user-name:__pete__@gmail.com\n" + "\n" + "\u0000"); } @Test diff --git a/spring-websocket/src/test/java/org/springframework/web/socket/server/DefaultHandshakeHandlerTests.java b/spring-websocket/src/test/java/org/springframework/web/socket/server/DefaultHandshakeHandlerTests.java index e9eb4aeffa..c0c44765ea 100644 --- a/spring-websocket/src/test/java/org/springframework/web/socket/server/DefaultHandshakeHandlerTests.java +++ b/spring-websocket/src/test/java/org/springframework/web/socket/server/DefaultHandshakeHandlerTests.java @@ -23,6 +23,7 @@ import java.util.Map; import org.junit.Test; +import org.springframework.http.HttpHeaders; import org.springframework.web.socket.AbstractHttpRequestTests; import org.springframework.web.socket.SubProtocolCapable; import org.springframework.web.socket.WebSocketExtension; @@ -51,14 +52,9 @@ public class DefaultHandshakeHandlerTests extends AbstractHttpRequestTests { public void supportedSubProtocols() { this.handshakeHandler.setSupportedProtocols("stomp", "mqtt"); given(this.upgradeStrategy.getSupportedVersions()).willReturn(new String[] {"13"}); - this.servletRequest.setMethod("GET"); - WebSocketHttpHeaders headers = new WebSocketHttpHeaders(this.request.getHeaders()); - headers.setUpgrade("WebSocket"); - headers.setConnection("Upgrade"); - headers.setSecWebSocketVersion("13"); - headers.setSecWebSocketKey("82/ZS2YHjEnUN97HLL8tbw=="); - headers.setSecWebSocketProtocol("STOMP"); + this.servletRequest.setMethod("GET"); + initHeaders(this.request.getHeaders()).setSecWebSocketProtocol("STOMP"); WebSocketHandler handler = new TextWebSocketHandler(); Map attributes = Collections.emptyMap(); @@ -77,16 +73,10 @@ public class DefaultHandshakeHandlerTests extends AbstractHttpRequestTests { given(this.upgradeStrategy.getSupportedExtensions(this.request)).willReturn(Collections.singletonList(extension1)); this.servletRequest.setMethod("GET"); - - WebSocketHttpHeaders headers = new WebSocketHttpHeaders(this.request.getHeaders()); - headers.setUpgrade("WebSocket"); - headers.setConnection("Upgrade"); - headers.setSecWebSocketVersion("13"); - headers.setSecWebSocketKey("82/ZS2YHjEnUN97HLL8tbw=="); - headers.setSecWebSocketExtensions(Arrays.asList(extension1, extension2)); + initHeaders(this.request.getHeaders()).setSecWebSocketExtensions(Arrays.asList(extension1, extension2)); WebSocketHandler handler = new TextWebSocketHandler(); - Map attributes = Collections.emptyMap(); + Map attributes = Collections.emptyMap(); this.handshakeHandler.doHandshake(this.request, this.response, handler, attributes); verify(this.upgradeStrategy).upgrade(this.request, this.response, null, @@ -98,16 +88,10 @@ public class DefaultHandshakeHandlerTests extends AbstractHttpRequestTests { given(this.upgradeStrategy.getSupportedVersions()).willReturn(new String[] {"13"}); this.servletRequest.setMethod("GET"); - - WebSocketHttpHeaders headers = new WebSocketHttpHeaders(this.request.getHeaders()); - headers.setUpgrade("WebSocket"); - headers.setConnection("Upgrade"); - headers.setSecWebSocketVersion("13"); - headers.setSecWebSocketKey("82/ZS2YHjEnUN97HLL8tbw=="); - headers.setSecWebSocketProtocol("v11.stomp"); + initHeaders(this.request.getHeaders()).setSecWebSocketProtocol("v11.stomp"); WebSocketHandler handler = new SubProtocolCapableHandler("v12.stomp", "v11.stomp"); - Map attributes = Collections.emptyMap(); + Map attributes = Collections.emptyMap(); this.handshakeHandler.doHandshake(this.request, this.response, handler, attributes); verify(this.upgradeStrategy).upgrade(this.request, this.response, "v11.stomp", @@ -119,22 +103,25 @@ public class DefaultHandshakeHandlerTests extends AbstractHttpRequestTests { given(this.upgradeStrategy.getSupportedVersions()).willReturn(new String[] {"13"}); this.servletRequest.setMethod("GET"); - - WebSocketHttpHeaders headers = new WebSocketHttpHeaders(this.request.getHeaders()); - headers.setUpgrade("WebSocket"); - headers.setConnection("Upgrade"); - headers.setSecWebSocketVersion("13"); - headers.setSecWebSocketKey("82/ZS2YHjEnUN97HLL8tbw=="); - headers.setSecWebSocketProtocol("v10.stomp"); + initHeaders(this.request.getHeaders()).setSecWebSocketProtocol("v10.stomp"); WebSocketHandler handler = new SubProtocolCapableHandler("v12.stomp", "v11.stomp"); - Map attributes = Collections.emptyMap(); + Map attributes = Collections.emptyMap(); this.handshakeHandler.doHandshake(this.request, this.response, handler, attributes); verify(this.upgradeStrategy).upgrade(this.request, this.response, null, Collections.emptyList(), null, handler, attributes); } + private WebSocketHttpHeaders initHeaders(HttpHeaders httpHeaders) { + WebSocketHttpHeaders headers = new WebSocketHttpHeaders(httpHeaders); + headers.setUpgrade("WebSocket"); + headers.setConnection("Upgrade"); + headers.setSecWebSocketVersion("13"); + headers.setSecWebSocketKey("82/ZS2YHjEnUN97HLL8tbw=="); + return headers; + } + private static class SubProtocolCapableHandler extends TextWebSocketHandler implements SubProtocolCapable { diff --git a/spring-websocket/src/test/java/org/springframework/web/socket/server/support/WebSocketHttpRequestHandlerTests.java b/spring-websocket/src/test/java/org/springframework/web/socket/server/support/WebSocketHttpRequestHandlerTests.java new file mode 100644 index 0000000000..4db37566cf --- /dev/null +++ b/spring-websocket/src/test/java/org/springframework/web/socket/server/support/WebSocketHttpRequestHandlerTests.java @@ -0,0 +1,138 @@ +/* + * Copyright 2002-2019 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. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.springframework.web.socket.server.support; + +import java.io.IOException; +import java.util.Collections; +import java.util.Map; +import javax.servlet.ServletException; + +import org.junit.Before; +import org.junit.Test; + +import org.springframework.http.server.ServerHttpRequest; +import org.springframework.http.server.ServerHttpResponse; +import org.springframework.mock.web.test.MockHttpServletRequest; +import org.springframework.mock.web.test.MockHttpServletResponse; +import org.springframework.web.socket.WebSocketHandler; +import org.springframework.web.socket.server.HandshakeFailureException; +import org.springframework.web.socket.server.HandshakeHandler; +import org.springframework.web.socket.server.HandshakeInterceptor; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.BDDMockito.when; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.verifyNoMoreInteractions; + +/** + * Unit tests for {@link WebSocketHttpRequestHandler}. + * @author Rossen Stoyanchev + * @since 5.1.9 + */ +public class WebSocketHttpRequestHandlerTests { + + private HandshakeHandler handshakeHandler; + + private WebSocketHttpRequestHandler requestHandler; + + private MockHttpServletResponse response; + + + @Before + public void setUp() { + this.handshakeHandler = mock(HandshakeHandler.class); + this.requestHandler = new WebSocketHttpRequestHandler(mock(WebSocketHandler.class), this.handshakeHandler); + this.response = new MockHttpServletResponse(); + } + + + @Test + public void success() throws ServletException, IOException { + TestInterceptor interceptor = new TestInterceptor(true); + this.requestHandler.setHandshakeInterceptors(Collections.singletonList(interceptor)); + this.requestHandler.handleRequest(new MockHttpServletRequest(), this.response); + + verify(this.handshakeHandler).doHandshake(any(), any(), any(), any()); + assertThat(this.response.getHeader("headerName")).isEqualTo("headerValue"); + } + + @Test + public void failure() { + TestInterceptor interceptor = new TestInterceptor(true); + this.requestHandler.setHandshakeInterceptors(Collections.singletonList(interceptor)); + + when(this.handshakeHandler.doHandshake(any(), any(), any(), any())) + .thenThrow(new IllegalStateException("bad state")); + + assertThatThrownBy(() -> this.requestHandler.handleRequest(new MockHttpServletRequest(), this.response)) + .isInstanceOf(HandshakeFailureException.class) + .hasRootCauseInstanceOf(IllegalStateException.class) + .hasMessageEndingWith("bad state"); + + assertThat(this.response.getHeader("headerName")).isEqualTo("headerValue"); + assertThat(this.response.getHeader("exceptionHeaderName")).isEqualTo("exceptionHeaderValue"); + } + + @Test // gh-23179 + public void handshakeNotAllowed() throws ServletException, IOException { + TestInterceptor interceptor = new TestInterceptor(false); + this.requestHandler.setHandshakeInterceptors(Collections.singletonList(interceptor)); + + this.requestHandler.handleRequest(new MockHttpServletRequest(), this.response); + + verifyNoMoreInteractions(this.handshakeHandler); + assertThat(this.response.getHeader("headerName")).isEqualTo("headerValue"); + } + + + private static class TestInterceptor implements HandshakeInterceptor { + + private final boolean allowHandshake; + + private Exception exception; + + + private TestInterceptor(boolean allowHandshake) { + this.allowHandshake = allowHandshake; + } + + + public Exception getException() { + return this.exception; + } + + + @Override + public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response, + WebSocketHandler wsHandler, Map attributes) { + + response.getHeaders().add("headerName", "headerValue"); + return this.allowHandshake; + } + + @Override + public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response, + WebSocketHandler wsHandler, Exception exception) { + + response.getHeaders().add("exceptionHeaderName", "exceptionHeaderValue"); + this.exception = exception; + } + } + +}