Polishing contribution

See gh-790

Author: rstoyanchev

spring-graphql/src/main/java/org/springframework/graphql/data/method/annotation/support/AuthenticationPrincipalArgumentResolver.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2022 the original author or authors.
+ * Copyright 2002-2023 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.
@@ -96,26 +96,20 @@ private static AuthenticationPrincipal findMethodAnnotation(MethodParameter para
 
 	@Override
 	public Object resolveArgument(MethodParameter parameter, DataFetchingEnvironment environment) throws Exception {
-		return getCurrentAuthentication(parameter.isOptional())
-				.flatMap(auth -> Mono.justOrEmpty(resolvePrincipal(parameter, auth.getPrincipal())))
-				.transform((argument) -> isParameterMonoAssignable(parameter) ? Mono.just(argument) : argument);
+		return getCurrentAuthentication(parameter)
+				.mapNotNull(auth -> resolvePrincipal(parameter, auth.getPrincipal()))
+				.transform((argument) -> isPublisherOrMono(parameter) ? Mono.just(argument) : argument);
 	}
 
-	private static boolean isParameterMonoAssignable(MethodParameter parameter) {
+	private static boolean isPublisherOrMono(MethodParameter parameter) {
 		Class<?> type = parameter.getParameterType();
 		return (Publisher.class.equals(type) || Mono.class.equals(type));
 	}
 
 	@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));
+	private Mono<Authentication> getCurrentAuthentication(MethodParameter parameter) {
+		Object value = PrincipalMethodArgumentResolver.resolveAuthentication(parameter);
+		return (value instanceof Authentication auth ? Mono.just(auth) : (Mono<Authentication>) value);
 	}
 
 	@Nullable
@@ -144,7 +138,7 @@ private boolean isInvalidType(MethodParameter parameter, @Nullable Object princi
 			return false;
 		}
 		Class<?> typeToCheck = parameter.getParameterType();
-		if (isParameterMonoAssignable(parameter)) {
+		if (isPublisherOrMono(parameter)) {
 			Class<?> genericType = parameter.nested().getNestedParameterType();
 			if (genericType.equals(Object.class)) {
 				return false;

spring-graphql/src/main/java/org/springframework/graphql/data/method/annotation/support/BatchLoaderHandlerMethod.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2022 the original author or authors.
+ * Copyright 2002-2023 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.
@@ -145,7 +145,7 @@ else if ("kotlin.coroutines.Continuation".equals(parameterType.getName())) {
 			return null;
 		}
 		else if (springSecurityPresent && Principal.class.isAssignableFrom(parameter.getParameterType())) {
-			return PrincipalMethodArgumentResolver.doResolve(parameter.isOptional());
+			return PrincipalMethodArgumentResolver.resolveAuthentication(parameter);
 		}
 		else {
 			throw new IllegalStateException(formatArgumentError(parameter, "Unexpected argument type."));

spring-graphql/src/main/java/org/springframework/graphql/data/method/annotation/support/PrincipalMethodArgumentResolver.java

@@ -16,15 +16,13 @@
 package org.springframework.graphql.data.method.annotation.support;
 
 import java.security.Principal;
-import java.util.function.Function;
 
 import graphql.schema.DataFetchingEnvironment;
 
 import org.springframework.core.MethodParameter;
 import org.springframework.graphql.data.method.HandlerMethodArgumentResolver;
 import org.springframework.security.authentication.AuthenticationCredentialsNotFoundException;
 import org.springframework.security.core.Authentication;
-import org.springframework.security.core.AuthenticationException;
 import org.springframework.security.core.context.ReactiveSecurityContextHolder;
 import org.springframework.security.core.context.SecurityContext;
 import org.springframework.security.core.context.SecurityContextHolder;
@@ -44,7 +42,7 @@
 public class PrincipalMethodArgumentResolver implements HandlerMethodArgumentResolver {
 
 	/**
-	 * Return "true" if the argument is {@link Principal} or a sub-type.
+	 * Return "true" if the argument is {@link Principal} or a subtype.
 	 */
 	@Override
 	public boolean supportsParameter(MethodParameter parameter) {
@@ -54,29 +52,24 @@ public boolean supportsParameter(MethodParameter parameter) {
 
 	@Override
 	public Object resolveArgument(MethodParameter parameter, DataFetchingEnvironment environment) {
-		return doResolve(parameter.isOptional());
+		return resolveAuthentication(parameter);
 	}
 
-	static Object doResolve(boolean optional) {
-		Authentication authentication = SecurityContextHolder.getContext().getAuthentication();
-
-		if (authentication != null) {
-			return authentication;
+	static Object resolveAuthentication(MethodParameter parameter) {
+		Authentication auth = SecurityContextHolder.getContext().getAuthentication();
+		if (auth != null) {
+			return auth;
 		}
 
-		return ReactiveSecurityContextHolder.getContext()
-				.switchIfEmpty(optional ? Mono.empty() : Mono.error(new AuthenticationCredentialsNotFoundException("SecurityContext not available")))
-				.handle((context, sink) -> {
-					Authentication auth = context.getAuthentication();
+		Mono<Authentication> authMono =
+				ReactiveSecurityContextHolder.getContext().mapNotNull(SecurityContext::getAuthentication);
+
+		if (!parameter.isOptional()) {
+			authMono = authMono.switchIfEmpty(
+					Mono.error(new AuthenticationCredentialsNotFoundException("No Authentication")));
+		}
 
-					if (auth != null) {
-						sink.next(auth);
-					} else if (!optional) {
-						sink.error(new AuthenticationCredentialsNotFoundException("An Authentication object was not found in the SecurityContext"));
-					} else {
-						sink.complete();
-					}
-				});
+		return authMono;
 	}
 
 }

spring-graphql/src/test/java/org/springframework/graphql/data/method/annotation/support/SchemaMappingPrincipalMethodArgumentResolverTests.java

@@ -134,7 +134,7 @@ void nonNullPrincipalRequiresSecurityContext() {
 
 			assertThat(responseHelper.errorCount()).isEqualTo(1);
 			assertThat(responseHelper.error(0).errorType()).isEqualTo("UNAUTHORIZED");
-			assertThat(responseHelper.error(0).message()).isEqualTo("Resolved error: SecurityContext not available");
+			assertThat(responseHelper.error(0).message()).isEqualTo("Resolved error: No Authentication");
 		}
 
 		@Test
@@ -154,7 +154,7 @@ void nonNullPrincipalRequiresAuthentication() {
 
 			assertThat(responseHelper.errorCount()).isEqualTo(1);
 			assertThat(responseHelper.error(0).errorType()).isEqualTo("UNAUTHORIZED");
-			assertThat(responseHelper.error(0).message()).isEqualTo("Resolved error: An Authentication object was not found in the SecurityContext");
+			assertThat(responseHelper.error(0).message()).isEqualTo("Resolved error: No Authentication");
 		}
 
 		@Test