Add coroutine support to pre/post authorize

Closes gh-8143

core/spring-security-core.gradle

@@ -26,6 +26,7 @@ dependencies {
 	optional 'org.aspectj:aspectjrt'
 	optional 'org.springframework:spring-jdbc'
 	optional 'org.springframework:spring-tx'
+	optional 'org.jetbrains.kotlinx:kotlinx-coroutines-reactor'
 
 	testImplementation powerMock2Dependencies
 	testImplementation 'commons-collections:commons-collections'

core/src/main/java/org/springframework/security/access/prepost/PrePostAdviceReactiveMethodInterceptor.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2018 the original author or authors.
+ * Copyright 2002-2021 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.
@@ -19,6 +19,9 @@ package org.springframework.security.access.prepost;
 import java.lang.reflect.Method;
 import java.util.Collection;
 
+import kotlin.coroutines.Continuation;
+import kotlinx.coroutines.reactive.AwaitKt;
+import kotlinx.coroutines.reactive.ReactiveFlowKt;
 import org.aopalliance.intercept.MethodInterceptor;
 import org.aopalliance.intercept.MethodInvocation;
 import org.reactivestreams.Publisher;
@@ -26,6 +29,11 @@ import reactor.core.Exceptions;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 
+import org.springframework.core.CoroutinesUtils;
+import org.springframework.core.KotlinDetector;
+import org.springframework.core.MethodParameter;
+import org.springframework.core.ReactiveAdapter;
+import org.springframework.core.ReactiveAdapterRegistry;
 import org.springframework.security.access.AccessDeniedException;
 import org.springframework.security.access.ConfigAttribute;
 import org.springframework.security.access.method.MethodSecurityMetadataSource;
@@ -38,9 +46,11 @@ import org.springframework.util.Assert;
 
 /**
  * A {@link MethodInterceptor} that supports {@link PreAuthorize} and
- * {@link PostAuthorize} for methods that return {@link Mono} or {@link Flux}
+ * {@link PostAuthorize} for methods that return {@link Mono} or {@link Flux} and Kotlin
+ * coroutine functions.
  *
  * @author Rob Winch
+ * @author Eleftheria Stein
  * @since 5.0
  */
 public class PrePostAdviceReactiveMethodInterceptor implements MethodInterceptor {
@@ -54,6 +64,10 @@ public class PrePostAdviceReactiveMethodInterceptor implements MethodInterceptor
 
 	private final PostInvocationAuthorizationAdvice postAdvice;
 
+	private static final String COROUTINES_FLOW_CLASS_NAME = "kotlinx.coroutines.flow.Flow";
+
+	private static final int RETURN_TYPE_METHOD_PARAMETER_INDEX = -1;
+
 	/**
 	 * Creates a new instance
 	 * @param attributeSource the {@link MethodSecurityMetadataSource} to use
@@ -75,10 +89,18 @@ public class PrePostAdviceReactiveMethodInterceptor implements MethodInterceptor
 	public Object invoke(final MethodInvocation invocation) {
 		Method method = invocation.getMethod();
 		Class<?> returnType = method.getReturnType();
-		Assert.state(Publisher.class.isAssignableFrom(returnType),
+
+		boolean isSuspendingFunction = KotlinDetector.isSuspendingFunction(method);
+		boolean hasFlowReturnType = COROUTINES_FLOW_CLASS_NAME
+				.equals(new MethodParameter(method, RETURN_TYPE_METHOD_PARAMETER_INDEX).getParameterType().getName());
+		boolean hasReactiveReturnType = Publisher.class.isAssignableFrom(returnType) || isSuspendingFunction
+				|| hasFlowReturnType;
+
+		Assert.state(hasReactiveReturnType,
 				() -> "The returnType " + returnType + " on " + method
 						+ " must return an instance of org.reactivestreams.Publisher "
-						+ "(i.e. Mono / Flux) in order to support Reactor Context");
+						+ "(i.e. Mono / Flux) or the function must be a Kotlin coroutine "
+						+ "function in order to support Reactor Context");
 		Class<?> targetClass = invocation.getThis().getClass();
 		Collection<ConfigAttribute> attributes = this.attributeSource.getAttributes(method, targetClass);
 		PreInvocationAttribute preAttr = findPreInvocationAttribute(attributes);
@@ -98,6 +120,30 @@ public class PrePostAdviceReactiveMethodInterceptor implements MethodInterceptor
 			return toInvoke.flatMapMany((auth) -> PrePostAdviceReactiveMethodInterceptor.<Flux<?>>proceed(invocation)
 					.map((r) -> (attr != null) ? this.postAdvice.after(auth, invocation, attr, r) : r));
 		}
+		if (hasFlowReturnType) {
+			Publisher<?> publisher;
+			if (isSuspendingFunction) {
+				publisher = CoroutinesUtils.invokeSuspendingFunction(invocation.getMethod(), invocation.getThis(),
+						invocation.getArguments());
+			}
+			else {
+				ReactiveAdapter adapter = ReactiveAdapterRegistry.getSharedInstance().getAdapter(returnType);
+				Assert.state(adapter != null, () -> "The returnType " + returnType + " on " + method
+						+ " must have a org.springframework.core.ReactiveAdapter registered");
+				publisher = adapter.toPublisher(PrePostAdviceReactiveMethodInterceptor.flowProceed(invocation));
+			}
+			Flux<?> response = toInvoke.flatMapMany((auth) -> Flux.from(publisher)
+					.map((r) -> (attr != null) ? this.postAdvice.after(auth, invocation, attr, r) : r));
+			return KotlinDelegate.asFlow(response);
+		}
+		if (isSuspendingFunction) {
+			Mono<?> response = toInvoke.flatMap((auth) -> Mono
+					.from(CoroutinesUtils.invokeSuspendingFunction(invocation.getMethod(), invocation.getThis(),
+							invocation.getArguments()))
+					.map((r) -> (attr != null) ? this.postAdvice.after(auth, invocation, attr, r) : r));
+			return KotlinDelegate.awaitSingleOrNull(response,
+					invocation.getArguments()[invocation.getArguments().length - 1]);
+		}
 		return toInvoke.flatMapMany(
 				(auth) -> Flux.from(PrePostAdviceReactiveMethodInterceptor.<Publisher<?>>proceed(invocation))
 						.map((r) -> (attr != null) ? this.postAdvice.after(auth, invocation, attr, r) : r));
@@ -112,6 +158,15 @@ public class PrePostAdviceReactiveMethodInterceptor implements MethodInterceptor
 		}
 	}
 
+	private static Object flowProceed(final MethodInvocation invocation) {
+		try {
+			return invocation.proceed();
+		}
+		catch (Throwable throwable) {
+			throw Exceptions.propagate(throwable);
+		}
+	}
+
 	private static PostInvocationAttribute findPostInvocationAttribute(Collection<ConfigAttribute> config) {
 		for (ConfigAttribute attribute : config) {
 			if (attribute instanceof PostInvocationAttribute) {
@@ -130,4 +185,19 @@ public class PrePostAdviceReactiveMethodInterceptor implements MethodInterceptor
 		return null;
 	}
 
+	/**
+	 * Inner class to avoid a hard dependency on Kotlin at runtime.
+	 */
+	private static class KotlinDelegate {
+
+		private static Object asFlow(Publisher<?> publisher) {
+			return ReactiveFlowKt.asFlow(publisher);
+		}
+
+		private static Object awaitSingleOrNull(Publisher<?> publisher, Object continuation) {
+			return AwaitKt.awaitSingleOrNull(publisher, (Continuation<Object>) continuation);
+		}
+
+	}
+
 }