WebFilter in WebFlux application

Question

I have a Spring Boot WebFlux application using Spring Boot 2.0.0.M5/2.0.0.BUILD-SNAPSHOT. I have a requirement to add trace-ids to all logs.

In order to get this to work in a WebFlux application, I tried using the WebFilter approach described here and here

@Component
public class TraceIdFilter implements WebFilter {

@Override
public Mono<Void> filter(ServerWebExchange exchange, WebFilterChain chain) {
    return chain.filter(exchange).subscriberContext((Context context) ->
        context.put(AuditContext.class, getAuditContext(exchange.getRequest().getHeaders()))
    );
}

My controller

@GetMapping(value = "/some_mapping")
public Mono<ResponseEntity<WrappedResponse>> getResource(@PathVariable("resourceId") String id) {
    Mono.subscriberContext().flatMap(context -> {
        AuditContext auditContext = context.get(AuditContext.class);
        ...
    });

The problem I have is that the filter method never gets executed, and the context is not set. I have confirmed that the Webfilter is loaded on startup. Is there anything else needed to get the filter to work?

Answer 1

I had a ton of issues figuring this out so hopefully it helps someone. My use case was to validate a signature on the request. This required me to parse the request body for PUT/POST's. The other major use case I see is logging so the below will be helpful too.

MiddlewareAuthenticator.java

@Component
public class MiddlewareAuthenticator implements WebFilter { 

    @Autowired private RequestValidationService requestValidationService;

@Override
public Mono<Void> filter(ServerWebExchange serverWebExchange, WebFilterChain chain) {
  return HEALTH_ENDPOINTS
      .matches(serverWebExchange)
      .flatMap(
          matches -> {
            if (matches.isMatch()) {
              return chain.filter(serverWebExchange);
            } else {
              return requestValidationService
                  .validate(serverWebExchange, 
                       new BiPredicate<ServerWebExchange, String> { 
                         @Override
                         public boolean test(ServerWebExchange e, String body) {
                             /** application logic can go here. few points:
                              1. I used a BiPredicate because I just need a true or false if the request should be passed to the controller. 
                              2. If you want todo other mutations you could swap the predicate to a normal function and return a mutated ServerWebExchange. 
                              3. I pass body separately here to ensure safety of accessing the request body and not having to rewrap the ServerWebExchange. A side affect of this though is any mutations to the String body do not affect downstream.
                              **/
                              return true;
                            }

                      })
                 .flatMap((ServerWebExchange r) -> chain.filter(r));
            }});
}

RequestValidationService.java

@Service
public class RequestValidationService {
private DataBuffer stringBuffer(String value) {
  byte[] bytes = value.getBytes(StandardCharsets.UTF_8);

  NettyDataBufferFactory nettyDataBufferFactory =
      new NettyDataBufferFactory(ByteBufAllocator.DEFAULT);
  DataBuffer buffer = nettyDataBufferFactory.allocateBuffer(bytes.length);
  buffer.write(bytes);
  return buffer;
}

private String bodyToString(InputStream bodyBytes) {
  byte[] currArr = null;
  try {
    currArr = bodyBytes.readAllBytes();
    bodyBytes.read(currArr);
  } catch (IOException ioe) {
    throw new RuntimeException("could not parse body");
  }

  if (currArr.length == 0) {
    return null;
  }

  return new String(currArr, StandardCharsets.UTF_8);
}

private ServerHttpRequestDecorator requestWrapper(ServerHttpRequest request, String bodyStr) {
  URI uri = request.getURI();
  ServerHttpRequest newRequest = request.mutate().uri(uri).build();
  final DataBuffer bodyDataBuffer = stringBuffer(bodyStr);
  Flux<DataBuffer> newBodyFlux = Flux.just(bodyDataBuffer);
  ServerHttpRequestDecorator requestDecorator =
      new ServerHttpRequestDecorator(newRequest) {
        @Override
        public Flux<DataBuffer> getBody() {
          return newBodyFlux;
        }
      };

  return requestDecorator;
}

private InputStream newInputStream() {
  return new InputStream() {
    public int read() {
      return -1;
    }
  };
}

private InputStream processRequestBody(InputStream s, DataBuffer d) {
  SequenceInputStream seq = new SequenceInputStream(s, d.asInputStream());
  return seq;
}

private Mono<ServerWebExchange> processInputStream(
    InputStream aggregatedBodyBytes,
    ServerWebExchange exchange,
    BiPredicate<ServerHttpRequest, String> predicate) {

  ServerHttpRequest request = exchange.getRequest();
  HttpHeaders headers = request.getHeaders();

  String bodyStr = bodyToString(aggregatedBodyBytes);

  ServerWebExchange mutatedExchange = exchange;

  // if the body exists on the request we need to mutate the ServerWebExchange to not
  // reparse the body because DataBuffers can only be read once;
  if (bodyStr != null) {
    mutatedExchange = exchange.mutate().request(requestWrapper(request, bodyStr)).build();
  }

  ServerHttpRequest mutatedRequest = mutatedExchange.getRequest();

  if (predicate.test(mutatedRequest, bodyStr)) {
    return Mono.just(mutatedExchange);
  }

  return Mono.error(new RuntimeException("invalid signature"));
}

/*
 * Because the DataBuffer is in a Flux we must reduce it to a Mono type via Flux.reduce
 * This covers large payloads or requests bodies that get sent in multiple byte chunks
 * and need to be concatentated.
 *
 * 1. The reduce is initialized with a newInputStream
 * 2. processRequestBody is called on each step of the Flux where a step is a body byte
 *    chunk. The method processRequestBody casts the Inbound DataBuffer to a InputStream
 *    and concats the new InputStream with the existing one
 * 3. Once the Flux is complete flatMap is executed with the resulting InputStream which is
 *    passed with the ServerWebExchange to processInputStream which will do the request validation
 */
public Mono<ServerWebExchange> validate(
    ServerWebExchange exchange, BiPredicate<ServerHttpRequest, String> p) {
  Flux<DataBuffer> body = exchange.getRequest().getBody();

  return body.reduce(newInputStream(), this::processRequestBody)
      .flatMap((InputStream b) -> processInputStream(b, exchange, p));
}

}

BiPredicate docs: https://docs.oracle.com/javase/8/docs/api/java/util/function/BiPredicate.html