Add argument resolver for SecurityContext

Closes gh-13425

messaging/src/main/java/org/springframework/security/messaging/handler/invocation/reactive/CurrentSecurityContextArgumentResolver.java

@@ -118,7 +118,21 @@ public class CurrentSecurityContextArgumentResolver implements HandlerMethodArgu
 
 	@Override
 	public boolean supportsParameter(MethodParameter parameter) {
-		return findMethodAnnotation(CurrentSecurityContext.class, parameter) != null;
+		return isMonoSecurityContext(parameter)
+				|| findMethodAnnotation(CurrentSecurityContext.class, parameter) != null;
+	}
+
+	private boolean isMonoSecurityContext(MethodParameter parameter) {
+		boolean isParameterPublisher = Publisher.class.isAssignableFrom(parameter.getParameterType());
+		if (isParameterPublisher) {
+			ResolvableType resolvableType = ResolvableType.forMethodParameter(parameter);
+			Class<?> genericType = resolvableType.resolveGeneric(0);
+			if (genericType == null) {
+				return false;
+			}
+			return SecurityContext.class.isAssignableFrom(genericType);
+		}
+		return false;
 	}
 
 	@Override
@@ -136,6 +150,14 @@ public class CurrentSecurityContextArgumentResolver implements HandlerMethodArgu
 
 	private Object resolveSecurityContext(MethodParameter parameter, Object securityContext) {
 		CurrentSecurityContext contextAnno = findMethodAnnotation(CurrentSecurityContext.class, parameter);
+		if (contextAnno != null) {
+			return resolveSecurityContextFromAnnotation(contextAnno, parameter, securityContext);
+		}
+		return securityContext;
+	}
+
+	private Object resolveSecurityContextFromAnnotation(CurrentSecurityContext contextAnno, MethodParameter parameter,
+			Object securityContext) {
 		String expressionToParse = contextAnno.expression();
 		if (StringUtils.hasLength(expressionToParse)) {
 			StandardEvaluationContext context = new StandardEvaluationContext();

messaging/src/test/java/org/springframework/security/messaging/handler/invocation/reactive/CurrentSecurityContextArgumentResolverTests.java

@@ -46,6 +46,24 @@ public class CurrentSecurityContextArgumentResolverTests {
 		assertThat(this.resolver.supportsParameter(arg0("currentSecurityContextOnMonoSecurityContext"))).isTrue();
 	}
 
+	@Test
+	public void supportsParameterWhenMonoSecurityContextNoAnnotationThenTrue() {
+		assertThat(this.resolver.supportsParameter(arg0("currentSecurityContextOnMonoSecurityContextNoAnnotation")))
+			.isTrue();
+	}
+
+	@Test
+	public void supportsParameterWhenMonoCustomSecurityContextNoAnnotationThenTrue() {
+		assertThat(
+				this.resolver.supportsParameter(arg0("currentCustomSecurityContextOnMonoSecurityContextNoAnnotation")))
+			.isTrue();
+	}
+
+	@Test
+	public void supportsParameterWhenNoSecurityContextNoAnnotationThenFalse() {
+		assertThat(this.resolver.supportsParameter(arg0("currentSecurityContextOnMonoStringNoAnnotation"))).isFalse();
+	}
+
 	@Test
 	public void resolveArgumentWhenAuthenticationPrincipalAndEmptyContextThenNull() {
 		Object result = this.resolver.resolveArgument(arg0("currentSecurityContextOnMonoSecurityContext"), null)
@@ -67,6 +85,18 @@ public class CurrentSecurityContextArgumentResolverTests {
 	private void currentSecurityContextOnMonoSecurityContext(@CurrentSecurityContext Mono<SecurityContext> context) {
 	}
 
+	@SuppressWarnings("unused")
+	private void currentSecurityContextOnMonoSecurityContextNoAnnotation(Mono<SecurityContext> context) {
+	}
+
+	@SuppressWarnings("unused")
+	private void currentCustomSecurityContextOnMonoSecurityContextNoAnnotation(Mono<CustomSecurityContext> context) {
+	}
+
+	@SuppressWarnings("unused")
+	private void currentSecurityContextOnMonoStringNoAnnotation(Mono<String> context) {
+	}
+
 	@Test
 	public void supportsParameterWhenCurrentUserThenTrue() {
 		assertThat(this.resolver.supportsParameter(arg0("currentUserOnMonoUserDetails"))).isTrue();
@@ -110,6 +140,41 @@ public class CurrentSecurityContextArgumentResolverTests {
 	private void monoUserDetails(Mono<UserDetails> user) {
 	}
 
+	@Test
+	public void supportsParameterWhenSecurityContextNotAnnotatedThenTrue() {
+		assertThat(this.resolver.supportsParameter(arg0("monoSecurityContext"))).isTrue();
+	}
+
+	@Test
+	public void resolveArgumentWhenMonoSecurityContextNoAnnotationThenFound() {
+		Authentication authentication = TestAuthentication.authenticatedUser();
+		Mono<SecurityContext> result = (Mono<SecurityContext>) this.resolver
+			.resolveArgument(arg0("monoSecurityContext"), null)
+			.contextWrite(ReactiveSecurityContextHolder.withAuthentication(authentication))
+			.block();
+		assertThat(result.block().getAuthentication().getPrincipal()).isEqualTo(authentication.getPrincipal());
+	}
+
+	@SuppressWarnings("unused")
+	private void monoSecurityContext(Mono<SecurityContext> securityContext) {
+	}
+
+	@Test
+	public void resolveArgumentWhenMonoCustomSecurityContextNoAnnotationThenFound() {
+		Authentication authentication = TestAuthentication.authenticatedUser();
+		CustomSecurityContext securityContext = new CustomSecurityContext();
+		securityContext.setAuthentication(authentication);
+		Mono<CustomSecurityContext> result = (Mono<CustomSecurityContext>) this.resolver
+			.resolveArgument(arg0("monoCustomSecurityContext"), null)
+			.contextWrite(ReactiveSecurityContextHolder.withSecurityContext(Mono.just(securityContext)))
+			.block();
+		assertThat(result.block().getAuthentication().getPrincipal()).isEqualTo(authentication.getPrincipal());
+	}
+
+	@SuppressWarnings("unused")
+	private void monoCustomSecurityContext(Mono<CustomSecurityContext> securityContext) {
+	}
+
 	private MethodParameter arg0(String methodName) {
 		ResolvableMethod method = ResolvableMethod.on(getClass()).named(methodName).method();
 		return new SynthesizingMethodParameter(method.method(), 0);
@@ -121,4 +186,20 @@ public class CurrentSecurityContextArgumentResolverTests {
 
 	}
 
+	static class CustomSecurityContext implements SecurityContext {
+
+		private Authentication authentication;
+
+		@Override
+		public Authentication getAuthentication() {
+			return this.authentication;
+		}
+
+		@Override
+		public void setAuthentication(Authentication authentication) {
+			this.authentication = authentication;
+		}
+
+	}
+
 }