How to response custom json body on unauthorized requests while implementing custom authentication manager in webflux

Question

I was trying to implement custom JWT token authentication while i am also handling global exceptions to customize response body for each type of exceptions. Everything is working fine except I would like to return custom json response when an unauthorized request is received instead of just 401 status code.

Below is my implementation for JwtServerAuthenticationConverter and JwtAuthenticationManager.

@Component
public class JwtServerAuthenticationConverter implements ServerAuthenticationConverter {

    private static final String AUTH_HEADER_VALUE_PREFIX = "Bearer ";

    @Override
    public Mono<Authentication> convert(ServerWebExchange exchange) {

        return Mono.justOrEmpty(exchange)
                .flatMap(serverWebExchange -> Mono.justOrEmpty(
                        serverWebExchange
                                .getRequest()
                                .getHeaders()
                                .getFirst(HttpHeaders.AUTHORIZATION)
                        )
                )
                .filter(header -> !header.trim().isEmpty() && header.trim().startsWith(AUTH_HEADER_VALUE_PREFIX))
                .map(header -> header.substring(AUTH_HEADER_VALUE_PREFIX.length()))
                .map(token -> new UsernamePasswordAuthenticationToken(token, token))
                ;
    }
}

@Component
public class JwtAuthenticationManager implements ReactiveAuthenticationManager {

    private final JWTConfig jwtConfig;
    private final ObjectMapper objectMapper;

    public JwtAuthenticationManager(JWTConfig jwtConfig, ObjectMapper objectMapper) {
        this.jwtConfig = jwtConfig;
        this.objectMapper = objectMapper;
    }

    @Override
    public Mono<Authentication> authenticate(Authentication authentication) {

        return Mono.just(authentication)
                .map(auth -> JWTHelper.loadAllClaimsFromToken(auth.getCredentials().toString(), jwtConfig.getSecret()))
                .onErrorResume(throwable -> Mono.error(new JwtException("Unauthorized")))
                .map(claims -> objectMapper.convertValue(claims, JWTUserDetails.class))
                .map(jwtUserDetails ->
                        new UsernamePasswordAuthenticationToken(
                                jwtUserDetails,
                                authentication.getCredentials(),
                                jwtUserDetails.getGrantedAuthorities()
                        )
                )
                ;
    }
}

And below is my global exception handling which is working absolutely fine except the case where webflux return 401 from JwtServerAuthenticationConverter convert method.

@Configuration
@Order(-2)
public class ExceptionHandler implements WebExceptionHandler {

    @Override
    public Mono<Void> handle(ServerWebExchange exchange, Throwable ex) {

        exchange.getResponse().getHeaders().set("Content-Type", MediaType.APPLICATION_JSON_VALUE);

        return buildErrorResponse(ex)
                .flatMap(
                        r -> r.writeTo(exchange, new HandlerStrategiesResponseContext(HandlerStrategies.withDefaults()))
                );
    }

    private Mono<ServerResponse> buildErrorResponse(Throwable ex) {

        if (ex instanceof RequestEntityValidationException) {

            return ServerResponse.badRequest().contentType(MediaType.APPLICATION_JSON).body(
                    Mono.just(new ErrorResponse(ex.getMessage())),
                    ErrorResponse.class
            );

        } else if (ex instanceof ResponseStatusException) {
            ResponseStatusException exception = (ResponseStatusException) ex;

            if (exception.getStatus().value() == 404) {
                return ServerResponse.status(HttpStatus.NOT_FOUND).contentType(MediaType.APPLICATION_JSON).body(
                        Mono.just(new ErrorResponse("Resource not found - 404")),
                        ErrorResponse.class
                );
            } else if (exception.getStatus().value() == 400) {
                return ServerResponse.status(HttpStatus.BAD_REQUEST).contentType(MediaType.APPLICATION_JSON).body(
                        Mono.just(new ErrorResponse("Unable to parse request body - 400")),
                        ErrorResponse.class
                );
            }

        } else if (ex instanceof JwtException) {

            return ServerResponse.status(HttpStatus.UNAUTHORIZED).contentType(MediaType.APPLICATION_JSON).body(
                    Mono.just(new ErrorResponse(ex.getMessage())),
                    ErrorResponse.class
            );
        }

        ex.printStackTrace();
        return ServerResponse.status(HttpStatus.INTERNAL_SERVER_ERROR).contentType(MediaType.APPLICATION_JSON).body(
                Mono.just(new ErrorResponse("Internal server error - 500")),
                ErrorResponse.class
        );

    }
}

@RequiredArgsConstructor
class HandlerStrategiesResponseContext implements ServerResponse.Context {

    private final HandlerStrategies handlerStrategies;

    @Override
    public List<HttpMessageWriter<?>> messageWriters() {
        return this.handlerStrategies.messageWriters();
    }

    @Override
    public List<ViewResolver> viewResolvers() {
        return this.handlerStrategies.viewResolvers();
    }
}

@Configuration
@EnableWebFluxSecurity
public class SecurityConfig {

    @Bean
    public SecurityWebFilterChain securityWebFilterChain(
            ServerHttpSecurity http,
            ReactiveAuthenticationManager jwtAuthenticationManager,
            ServerAuthenticationConverter jwtAuthenticationConverter
    ) {

        AuthenticationWebFilter authenticationWebFilter = new AuthenticationWebFilter(jwtAuthenticationManager);
        authenticationWebFilter.setServerAuthenticationConverter(jwtAuthenticationConverter);

        return http
                .authorizeExchange()
                .pathMatchers("/auth/login", "/auth/logout").permitAll()
                .anyExchange().authenticated()
                .and()
                .addFilterAt(authenticationWebFilter, SecurityWebFiltersOrder.AUTHENTICATION)
                .httpBasic()
                .disable()
                .csrf()
                .disable()
                .formLogin()
                .disable()
                .logout()
                .disable()
                .build();
    }

    @Bean
    public BCryptPasswordEncoder bCryptPasswordEncoder() {
        return new BCryptPasswordEncoder();
    }
}

So when i am hitting it with an invalid JWT token in header. This got handled by my ExceptioHandler class and I got below output which is great.

But when i hit it with empty jwt token I got this.

Now i would like to return the same body which i am returning in the case of invalid JWT token. but the problem is when empty token is provided its not even falling in handle method of ExceptionHandler class. thats why its not in my control like i did for JwtException in the same class. How could i do that please help?

Answer 1

I sort it out myself. webflux provides ServerAuthenticationFailureHandler to handle custom response for that but unfortunately ServerAuthenticationFailureHandler not works and its a known issue so i created a failure route and write my custom response in it and setup login page.

.formLogin()
.loginPage("/auth/failed")
.and()

.andRoute(path("/auth/failed").and(accept(MediaType.APPLICATION_JSON)), (serverRequest) ->
        ServerResponse
                .status(HttpStatus.UNAUTHORIZED)
                .body(
                        Mono.just(new ErrorResponse("Unauthorized")),
                        ErrorResponse.class
                )
);