Spring/spring-framework
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Merge branch '5.1.x'

Browse Source
This commit is contained in:
Rossen Stoyanchev
2019-07-03 17:25:12 +01:00
parent 03a342378a 4e6e47b726
commit 3d913b8134
7 changed files with 242 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files

22
spring-messaging/src/main/java/org/springframework/messaging/simp/SimpMessageHeaderAccessor.java
View File

@@ -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<Principal> 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<Principal> 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) {

49
spring-messaging/src/test/java/org/springframework/messaging/simp/SimpMessageHeaderAccessorTests.java
View File

@@ -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.<String, Object>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.<String, Object>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<Principal> {
private Principal user;
public Principal getUser() {
return this.user;
}
@Override
public void accept(Principal principal) {
this.user = principal;
}
}
}

18
spring-websocket/src/main/java/org/springframework/web/socket/messaging/StompSubProtocolHandler.java
View File

@@ -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) {

5
spring-websocket/src/main/java/org/springframework/web/socket/server/support/WebSocketHttpRequestHandler.java
View File

@@ -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();
}
}

9
spring-websocket/src/test/java/org/springframework/web/socket/messaging/StompSubProtocolHandlerTests.java
View File

@@ -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

49
spring-websocket/src/test/java/org/springframework/web/socket/server/DefaultHandshakeHandlerTests.java
View File

@@ -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<String, Object> 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<String, Object> attributes = Collections.<String, Object>emptyMap();
Map<String, Object> 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<String, Object> attributes = Collections.<String, Object>emptyMap();
Map<String, Object> 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<String, Object> attributes = Collections.<String, Object>emptyMap();
Map<String, Object> 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 {

138
spring-websocket/src/test/java/org/springframework/web/socket/server/support/WebSocketHttpRequestHandlerTests.java Normal file
View File

@@ -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<String, Object> 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;
}
}
}