Merge branch '5.1.x'

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

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.<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;
+		}
 	}
 
 }