How to use

body

method

in

org.springframework.web.reactive.function.client.ClientResponse$Builder

private static Function<ClientResponse, Mono<ClientResponse>> convertClientResponse(Function<Flux<DataBuffer>, Flux<DataBuffer>> bodConverter,
                                          MediaType contentType) {
  return response -> {
    ClientResponse convertedResponse = ClientResponse.from(response).headers(headers -> {
      headers.replace(HttpHeaders.CONTENT_TYPE, singletonList(contentType.toString()));
      headers.remove(HttpHeaders.CONTENT_LENGTH);
    }).body(response.bodyToFlux(DataBuffer.class).transform(bodConverter)).build();
    return Mono.just(convertedResponse);
  };
}

/**
 * Consume up to the specified number of bytes from the response body and
 * cancel if any more data arrives.
 * <p>Internally delegates to {@link DataBufferUtils#takeUntilByteCount}.
 * @param maxByteCount the limit as number of bytes
 * @return the filter to limit the response size with
 * @since 5.1
 */
public static ExchangeFilterFunction limitResponseSize(long maxByteCount) {
  return (request, next) ->
      next.exchange(request).map(response -> {
        Flux<DataBuffer> body = response.body(BodyExtractors.toDataBuffers());
        body = DataBufferUtils.takeUntilByteCount(body, maxByteCount);
        return ClientResponse.from(response).body(body).build();
      });
}

@Test
public void from() {
  Flux<DataBuffer> otherBody = Flux.just("foo", "bar")
      .map(s -> s.getBytes(StandardCharsets.UTF_8))
      .map(dataBufferFactory::wrap);
  ClientResponse other = ClientResponse.create(HttpStatus.BAD_REQUEST, ExchangeStrategies.withDefaults())
      .header("foo", "bar")
      .cookie("baz", "qux")
      .body(otherBody)
      .build();
  Flux<DataBuffer> body = Flux.just("baz")
      .map(s -> s.getBytes(StandardCharsets.UTF_8))
      .map(dataBufferFactory::wrap);
  ClientResponse result = ClientResponse.from(other)
      .headers(httpHeaders -> httpHeaders.set("foo", "baar"))
      .cookies(cookies -> cookies.set("baz", ResponseCookie.from("baz", "quux").build()))
      .body(body)
      .build();
  assertEquals(HttpStatus.BAD_REQUEST, result.statusCode());
  assertEquals(1, result.headers().asHttpHeaders().size());
  assertEquals("baar", result.headers().asHttpHeaders().getFirst("foo"));
  assertEquals(1, result.cookies().size());
  assertEquals("quux", result.cookies().getFirst("baz").getValue());
  StepVerifier.create(result.bodyToFlux(String.class))
      .expectNext("baz")
      .verifyComplete();
}

@Override
public Mono<ClientResponse> filter(ClientRequest request, ExchangeFunction next) {
  URI originalUrl = request.url();
  String serviceId = originalUrl.getHost();
  if(serviceId == null) {
    String msg = String.format("Request URI does not contain a valid hostname: %s", originalUrl.toString());
    logger.warn(msg);
    return Mono.just(ClientResponse.create(HttpStatus.BAD_REQUEST).body(msg).build());
  }
  //TODO: reactive lb client
  ServiceInstance instance = this.loadBalancerClient.choose(serviceId);
  if(instance == null) {
    String msg = String.format("Load balancer does not contain an instance for the service %s", serviceId);
    logger.warn(msg);
    return Mono.just(ClientResponse.create(HttpStatus.SERVICE_UNAVAILABLE).body(msg).build());
  }
  URI uri = this.loadBalancerClient.reconstructURI(instance, originalUrl);
  ClientRequest newRequest = ClientRequest.method(request.method(), uri)
      .headers(headers -> headers.addAll(request.headers()))
      .cookies(cookies -> cookies.addAll(request.cookies()))
      .attributes(attributes -> attributes.putAll(request.attributes()))
      .body(request.body())
      .build();
  return next.exchange(newRequest);
}

@Override
public void onNext(ClientResponse response) {
  this.done = true;
  try {
    // decorate response body
    this.actual
        .onNext(ClientResponse.from(response)
            .body(response.bodyToFlux(DataBuffer.class)
                .transform(this.scopePassingTransformer))
            .build());
  }
  finally {
    terminateSpan(response, null);
  }
}

@Override
public Mono<ClientResponse> filter(ClientRequest request, ExchangeFunction next) {
  URI originalUrl = request.url();
  String serviceId = originalUrl.getHost();
  if(serviceId == null) {
    String msg = String.format("Request URI does not contain a valid hostname: %s", originalUrl.toString());
    logger.warn(msg);
    return Mono.just(ClientResponse.create(HttpStatus.BAD_REQUEST).body(msg).build());
  }
  //TODO: reactive lb client
  ServiceInstance instance = this.loadBalancerClient.choose(serviceId);
  if(instance == null) {
    String msg = String.format("Load balancer does not contain an instance for the service %s", serviceId);
    logger.warn(msg);
    return Mono.just(ClientResponse.create(HttpStatus.SERVICE_UNAVAILABLE).body(msg).build());
  }
  URI uri = this.loadBalancerClient.reconstructURI(instance, originalUrl);
  ClientRequest newRequest = ClientRequest.method(request.method(), uri)
      .headers(headers -> headers.addAll(request.headers()))
      .cookies(cookies -> cookies.addAll(request.cookies()))
      .attributes(attributes -> attributes.putAll(request.attributes()))
      .body(request.body())
      .build();
  return next.exchange(newRequest);
}

@Override
public void onNext(ClientResponse response) {
  this.done = true;
  try {
    // decorate response body
    this.actual
        .onNext(ClientResponse.from(response)
            .body(response.bodyToFlux(DataBuffer.class)
                .transform(this.scopePassingTransformer))
            .build());
  }
  finally {
    terminateSpan(response, null);
  }
}

@Test
public void normal() {
  Flux<DataBuffer> body = Flux.just("baz")
      .map(s -> s.getBytes(StandardCharsets.UTF_8))
      .map(dataBufferFactory::wrap);
  ClientResponse response = ClientResponse.create(HttpStatus.BAD_GATEWAY, ExchangeStrategies.withDefaults())
      .header("foo", "bar")
      .cookie("baz", "qux")
      .body(body)
      .build();
  assertEquals(HttpStatus.BAD_GATEWAY, response.statusCode());
  HttpHeaders responseHeaders = response.headers().asHttpHeaders();
  assertEquals("bar", responseHeaders.getFirst("foo"));
  assertNotNull("qux", response.cookies().getFirst("baz"));
  assertEquals("qux", response.cookies().getFirst("baz").getValue());
  StepVerifier.create(response.bodyToFlux(String.class))
      .expectNext("baz")
      .verifyComplete();
}

@Test
public void limitResponseSize() {
  DefaultDataBufferFactory bufferFactory = new DefaultDataBufferFactory();
  DataBuffer b1 = dataBuffer("foo", bufferFactory);
  DataBuffer b2 = dataBuffer("bar", bufferFactory);
  DataBuffer b3 = dataBuffer("baz", bufferFactory);
  ClientRequest request = ClientRequest.create(HttpMethod.GET, DEFAULT_URL).build();
  ClientResponse response = ClientResponse.create(HttpStatus.OK).body(Flux.just(b1, b2, b3)).build();
  Mono<ClientResponse> result = ExchangeFilterFunctions.limitResponseSize(5)
      .filter(request, req -> Mono.just(response));
  StepVerifier.create(result.flatMapMany(res -> res.body(BodyExtractors.toDataBuffers())))
      .consumeNextWith(buffer -> assertEquals("foo", string(buffer)))
      .consumeNextWith(buffer -> assertEquals("ba", string(buffer)))
      .expectComplete()
      .verify();
}