use single method to resolve authentication

spring-projects · Sep 14, 2023 · 1f1ac3c · 1f1ac3c
1 parent 808de81
commit 1f1ac3c
Showing 1 changed file with 12 additions and 4 deletions.
diff --git a/...ework/graphql/data/method/annotation/support/AuthenticationPrincipalArgumentResolver.java b/...ework/graphql/data/method/annotation/support/AuthenticationPrincipalArgumentResolver.java
@@ -16,6 +16,7 @@
 package org.springframework.graphql.data.method.annotation.support;
 
 import java.lang.annotation.Annotation;
+import java.util.Map;
 
 import graphql.schema.DataFetchingEnvironment;
 import org.reactivestreams.Publisher;
@@ -96,7 +97,7 @@ private static AuthenticationPrincipal findMethodAnnotation(MethodParameter para
 
 	@Override
 	public Object resolveArgument(MethodParameter parameter, DataFetchingEnvironment environment) throws Exception {
-		return getCurrentAuthentication()
+		return getCurrentAuthentication(parameter.isOptional())
 				.flatMap(auth -> Mono.justOrEmpty(resolvePrincipal(parameter, auth.getPrincipal())))
 				.transform((argument) -> isParameterMonoAssignable(parameter) ? Mono.just(argument) : argument);
 	}
@@ -106,9 +107,16 @@ private static boolean isParameterMonoAssignable(MethodParameter parameter) {
 		return (Publisher.class.equals(type) || Mono.class.equals(type));
 	}
 
-	private Mono<Authentication> getCurrentAuthentication() {
-		return Mono.justOrEmpty(SecurityContextHolder.getContext().getAuthentication())
-				.switchIfEmpty(ReactiveSecurityContextHolder.getContext().map(SecurityContext::getAuthentication));
+	@SuppressWarnings("unchecked")
+	private Mono<Authentication> getCurrentAuthentication(boolean optional) {
+		Object principal = PrincipalMethodArgumentResolver.doResolve(optional);
+		if (principal instanceof Authentication) {
+			return Mono.just((Authentication) principal);
+		}
+		else if (principal instanceof Mono) {
+			return (Mono<Authentication>) principal;
+		}
+		return Mono.error(new IllegalStateException("Unexpected return value: " + principal));
 	}
 
 	@Nullable