chemicL/context-propagation-demo

Need suggestion!

Closed this issue · 4 comments

Hi Dariusz,
I have watched your demo on youtube and it was fantastic. I am quite new to webflux or reactor and I have a case that I need a suggestion from you. Here is what I am trying to achieve:
I am trying to write a correlation filter just like what you showed on the demo, but the difference is that I just want to put JWT subject claim into MDC context and thus in every log, of course if the web controller endpoint is secured with JWT. So first I need to extract the claim as a String and then put it into reactor context before calling chain.filter(exchange). So, in my case, filter can not get the required value right from exchange.getrequest()... I tried the following right in the filter without any success:

I am getting the right subject claim for a few logs in the terminal, but getting "defaultUser" for a quite lot more places. How should I approach to this case without any blocking call for subject claim, and staying on the pure reactive code with using suggested dependencies and beans as they are in this repo?

AtomicReference<String> userClaim = new AtomicReference<>("defaultUser");
return ReactiveSecurityContextHolder.getContext()
                .mapNotNull(this::extractSubClaim)
                .flatMap(s -> Mono.deferContextual(contextView -> {
                    userClaim.set(s);
                    log.info("USER: {}, Context MDC_USERNAME_KEY: {}", userClaim.get(), contextView.getOrEmpty(MDC_USERNAME_KEY));
                    MDC.put(MDC_USERNAME_KEY, s);
                    return Mono.just(s);
                }))
                .then()
                .then(chain.filter(exchange))
                .doOnEach(voidSignal -> {
                    try (MDC.MDCCloseable ignored = MDC.putCloseable(MDC_SECURITY_TAG_KEY, voidSignal.getContextView().get(MDC_SECURITY_TAG_KEY))) {
                        log.info("SEC key: {}", Optional.of(voidSignal.getContextView().get(MDC_SECURITY_TAG_KEY)));
                    }
                    try (MDC.MDCCloseable ignored = MDC.putCloseable(MDC_USERNAME_KEY, voidSignal.getContextView().get(MDC_USERNAME_KEY))) {
                        log.info("USR key: {}", Optional.of(voidSignal.getContextView().get(MDC_USERNAME_KEY)));
                    }
                })
                .doFinally(signalType -> MDC.clear())
                .contextWrite(context -> context.put(MDC_USERNAME_KEY, userClaim.get()))
                .contextWrite(context -> context.put(MDC_SECURITY_TAG_KEY, "SECURITY"));
private String extractSubClaim(SecurityContext securityContext) {
        final Authentication authentication = securityContext.getAuthentication();
        if (authentication instanceof JwtAuthenticationToken token) {
            final Jwt jwt = token.getToken();
            log.debug("Sub claim is: {}", jwt.getSubject());
            return jwt.getSubject();
        }
        return null;
    }

Hey! Thank you for the kind words, I'm glad the materials are helpful.
On my end, I'm not that much familiar with ReactiveSecurityContextHolder and don't have the background to generate a proper setup for the project to confirm my imagined strategy, but here's what I'd probably try:

  1. Instead of .then().then(chain.filter(exchange)) .. I'd work with the flatMap above and attach the context to the inner Mono, something along these lines:
return ReactiveSecurityContextHolder.getContext()
                .mapNotNull(this::extractSubClaim)
                .flatMap(subClaim -> chain.filter(exchange).contextWrite(ctx -> ctx.put(CLAIM_KEY, subClaim))
  1. For the above to work properly, you'd need a ThreadLocalAccessor for CLAIM_KEY and populate MDC_USERNAME_KEY ane MDC_SECURITY_TAG_KEY.
  2. Definitely avoid any external state like AtomicReference, which you populate in the flatMap and expect to access in operators further down the chain. The Context that is attached at some level of the chain using contextWrite() operator is available in all the operators above that chain. The automatic context propagation mechanism uses the registered ThreadLocalAccessors to work with corresponding keys and make that available as ThreadLocal state.

Have a look at these resources as well:

Hope this helps!

@barmaths let me know if the above works for you. I'd like to close the issue if that answers your concerns :)

Thank you very much again for your quick response too. I was already using some form of threadlocalaccessors (not the same one with general purpose SelectiveMDCAccessor) but the problem for me, I guess, was that I was using new BaseSubscriber down the reactive chain in the subscribe() method and that was causing the loss of propagated values for my logs, right? Reverting to basic doon* methods and giving up on BaseSubscriber resolved my problem, so you can close the issue right here.

Yep. Unless the BaseSubscriber carries a proper Context with it, you might not see the ThreadLocal values that would be present above the chain (above a contextWrite() operator). Thanks for the response, take care!