Kotlin coroutine context lost in suspend functions with method security annotations

[calebdelnay]

Describe the bug
Working on a new backend for a product that makes heavy use of Spring with Kotlin, WebFlux, Spring Method Security, and Spring Cloud Sleuth.

Ran into an issue where Sleuth trace and span IDs were disappearing from log messages in the middle of handling requests.

After quite a bit of debugging, narrowed it down to what appears to be an issue related to reactive method security.

The context of a Kotlin coroutine is lost when execution proceeds into a suspend function annotated with @PreAuthorize or @PostAuthorize method security annotations.

Suspend functions that are not annotated do correctly preserve, propagate, and can access the coroutine context.

Spring Cloud Sleuth bridges between Reactor context and Kotlin coroutine context, so that is why it was being affected.

Was able to reproduce the issue without anything related to Sleuth.

Original related implementation looks to be: #9586

To Reproduce
Populate context in a @RestController class' suspend function and call into a @Service class' suspend function that is annotated with @PreAuthorize. The coroutine context will be missing. When calling into a method that isn't annotated, the coroutine context is available.

Example code below. See the sample code for a complete example.

@RestController
class ExampleController(private val exampleService: ExampleService) {

	@GetMapping("/test")
	suspend fun endpoint(): String {
		withContext(CoroutineName("methodWithoutSecurity")) {
			exampleService.methodWithoutSecurity()
		}

		withContext(CoroutineName("methodWithPreAuthorizeSecurity")) {
			exampleService.methodWithPreAuthorizeSecurity()
		}

		return "Test"
	}

}

@Service
class ExampleService {

	suspend fun methodWithoutSecurity() {
		// Name will be: CoroutineName("methodWithoutSecurity")
		val name = coroutineContext[CoroutineName.Key]
		println("methodWithoutSecurity: name is: $name")
	}

	@PreAuthorize("true")
	suspend fun methodWithPreAuthorizeSecurity() {
		// Name will be: null
		val name = coroutineContext[CoroutineName.Key]
		println("methodWithPreAuthorizeSecurity: name is: $name")
	}

}

Expected behavior
That suspend function coroutine context is preserved and available for suspend functions annotated with method security annotations.

Sample
https://gist.github.com/calebdelnay/7abc79922418cf914ef28eaa0d3ae368

[eleftherias]

Thanks for submitting this issue @calebdelnay.
This looks to have the same root cause as gh-10252.

We are planning a fix for a future release, but the suggested workaround for now is to use the Reactor operators Mono / Flux instead of suspending functions when combining them with @PreAuthorize or @PostAuthorize.

[calebdelnay]

@eleftherias Thanks for taking a look!

[ydolzhenko]

Anyone has any update on this?

[bdalenoord]

We've also run into this issue, and after some digging, I figured out that the ReactorContext (https://kotlinlang.org/api/kotlinx.coroutines/kotlinx-coroutines-reactor/kotlinx.coroutines.reactor/-reactor-context/) is preserved after the PrePostAdviceReactiveMethodInterceptor.

By leveraging that ReactorContext, I'm able to push my custom context's value into that context. This is not quite optimal as we've now got a util-function that first attempts to look up our custom Coroutine Context and if that fails, it looks for the ReactorContext and attempts to extract the custom value from it instead. It allows us to keep using suspend funs though, which is nice enough, and the util-function should be more or less easily removable later.

[mwalkerr]

@bdalenoord we're experiencing a similar issue. Would you be able to share a code snippet or some pseudo-code demonstrating how you're accessing the ReactorContext?

[bdalenoord]

@mwalkerr we've got a helper function for our specific Coroutine context. Let's say we've got a SomeCoroutineContext containing a Some, and if that's not present, the ReactorContext will contain the Some instead. Our helper-function would look like this:

suspend fun getSomeFromCoroutineOrReactorContext(): Some? {
    val someContext = coroutineContext[SomeCoroutineContext]
    logger.trace("Found${if (someContext == null) " no" else ""} SomeCoroutineContext")

    val reactorContextFallback = suspend {
         coroutineContext[ReactorContext]?.context?.get(Some::class.java)
    }

    return someContext?.some ?: reactorContextFallback()
}

[jzheaux]

Please see my comment in #10252. Given that, I'll mark this as closed as of 6.2.0.

I can also confirm that when I take the sample provided in the OP and update it to the latest, it demonstrates the expected behavior.