use single method to resolve authentication

spring-projects · Sep 14, 2023 · cb7b8dc · cb7b8dc
1 parent 808de81
commit cb7b8dc
Showing 1 changed file with 11 additions and 4 deletions.
diff --git a/...ework/graphql/data/method/annotation/support/AuthenticationPrincipalArgumentResolver.java b/...ework/graphql/data/method/annotation/support/AuthenticationPrincipalArgumentResolver.java
@@ -96,7 +96,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 +106,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