use single method to resolve authentication

Author: koenpunt

spring-graphql/src/main/java/org/springframework/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