diff --git a/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/filter/AdaptCachedBodyGlobalFilter.java b/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/filter/AdaptCachedBodyGlobalFilter.java index 16760f23..a7f758b2 100644 --- a/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/filter/AdaptCachedBodyGlobalFilter.java +++ b/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/filter/AdaptCachedBodyGlobalFilter.java @@ -39,6 +39,7 @@ public class AdaptCachedBodyGlobalFilter implements GlobalFilter, Ordered { return body; } }; + exchange.getAttributes().remove(CACHED_REQUEST_BODY_KEY); return chain.filter(exchange.mutate().request(decorator).build()); } diff --git a/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/handler/predicate/ReadBodyPredicateFactory.java b/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/handler/predicate/ReadBodyPredicateFactory.java index 9218f1f0..05133916 100644 --- a/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/handler/predicate/ReadBodyPredicateFactory.java +++ b/spring-cloud-gateway-core/src/main/java/org/springframework/cloud/gateway/handler/predicate/ReadBodyPredicateFactory.java @@ -17,43 +17,28 @@ package org.springframework.cloud.gateway.handler.predicate; -import java.util.Arrays; -import java.util.Collections; -import java.util.HashMap; -import java.util.HashSet; import java.util.List; import java.util.Map; -import java.util.Set; -import java.util.function.BiConsumer; -import java.util.function.BinaryOperator; -import java.util.function.Function; import java.util.function.Predicate; -import java.util.function.Supplier; -import java.util.stream.Collector; import org.apache.commons.logging.Log; import org.apache.commons.logging.LogFactory; -import org.springframework.cloud.gateway.support.BodyInserterContext; import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; -import org.springframework.cloud.gateway.support.CachedBodyOutputMessage; import org.springframework.cloud.gateway.handler.AsyncPredicate; -import org.springframework.core.ResolvableType; import org.springframework.core.io.buffer.DataBuffer; -import org.springframework.http.HttpHeaders; -import org.springframework.http.ReactiveHttpInputMessage; +import org.springframework.core.io.buffer.DataBufferUtils; import org.springframework.http.codec.HttpMessageReader; -import org.springframework.web.reactive.function.BodyInserter; -import org.springframework.web.reactive.function.BodyInserters; +import org.springframework.http.server.reactive.ServerHttpRequest; +import org.springframework.http.server.reactive.ServerHttpRequestDecorator; import org.springframework.web.reactive.function.server.HandlerStrategies; +import org.springframework.web.reactive.function.server.ServerRequest; import org.springframework.web.server.ServerWebExchange; import static org.springframework.cloud.gateway.filter.AdaptCachedBodyGlobalFilter.CACHED_REQUEST_BODY_KEY; -import static org.springframework.cloud.gateway.handler.predicate.ReadBodyPredicateFactory.DataBufferMapCollector.BODY_ONE; -import static org.springframework.cloud.gateway.handler.predicate.ReadBodyPredicateFactory.DataBufferMapCollector.BODY_TWO; /** * This predicate is BETA and may be subject to change in a future release. @@ -95,90 +80,37 @@ public class ReadBodyPredicateFactory } return Mono.just(false); } else { - return exchange.getRequest().getBody().collect(new DataBufferMapCollector()).flatMap(dataBufferMap -> { - BodyInserter bodyInserter = BodyInserters.fromPublisher(dataBufferMap.get(BODY_ONE), DataBuffer.class); - CachedBodyOutputMessage outputMessage = new CachedBodyOutputMessage(exchange, - exchange.getRequest().getHeaders()); - return bodyInserter.insert(outputMessage, new BodyInserterContext()) - // .log("modify_request", Level.INFO) - .then(Mono.defer(() -> { - ResolvableType type = ResolvableType.forClass(inClass); - for (HttpMessageReader messageReader : messageReaders) { - if (messageReader.canRead(type, exchange.getRequest().getHeaders().getContentType())) { - ReactiveHttpInputMessage inputMessage = new ReadBodyReactiveHttpInputMessage(dataBufferMap.get(BODY_TWO), - exchange.getRequest().getHeaders()); - Function mapper = (bodyObj) -> { - exchange.getAttributes().put(CACHE_REQUEST_BODY_OBJECT_KEY, bodyObj); - exchange.getAttributes().put(CACHED_REQUEST_BODY_KEY, - outputMessage.getBody()); - boolean test = config.predicate.test(bodyObj); - return Mono.just(test); - }; - return messageReader.readMono(type, inputMessage, Collections.EMPTY_MAP).flatMap(mapper); - } - } - return Mono.just(false); - })); + //Join all the DataBuffers so we have a single DataBuffer for the body + return DataBufferUtils.join(exchange.getRequest().getBody()) + .flatMap(dataBuffer -> { + //Update the retain counts so we can read the body twice, once to parse into an object + //that we can test the predicate against and a second time when the HTTP client sends + //the request downstream + //Note: if we end up reading the body twice we will run into a problem, but as of right + //now there is no good use case for doing this + DataBufferUtils.retain(dataBuffer); + //Make a slice for each read so each read has its own read/write indexes + Flux cachedFlux = Flux.defer(() -> Flux.just(dataBuffer.slice(0, dataBuffer.readableByteCount()))); + + ServerHttpRequest mutatedRequest = new ServerHttpRequestDecorator(exchange.getRequest()) { + @Override + public Flux getBody() { + return cachedFlux; + } + }; + return ServerRequest.create(exchange.mutate().request(mutatedRequest).build(), messageReaders) + .bodyToMono(inClass) + .doOnNext(objectValue -> { + exchange.getAttributes().put(CACHE_REQUEST_BODY_OBJECT_KEY, objectValue); + exchange.getAttributes().put(CACHED_REQUEST_BODY_KEY, cachedFlux); + }) + .map(objectValue -> config.predicate.test(objectValue)); + }); - }); } }; } - /** - * This {@link Collector} is meant to collect the {@code Flux} from the request body into a {@link Map} - * which contains two copy of the body, one under the key {@code orig} and the other under the {@key copy}. - */ - class DataBufferMapCollector implements Collector>, Map>> { - public static final String BODY_ONE = "bodyOne"; - public static final String BODY_TWO = "bodyTwo"; - private final Set CHARACTERISTICS = new HashSet<>(Arrays.asList( - Characteristics.IDENTITY_FINISH)); - - @Override - public Supplier>> supplier() { - return () -> new HashMap>(); - } - - @Override - public BiConsumer>, DataBuffer> accumulator() { - return (dataBufferMap, dataBuffer) -> { - accumulate(BODY_ONE, dataBufferMap, dataBuffer); - accumulate(BODY_TWO, dataBufferMap, dataBuffer); - }; - } - - private void accumulate(String key, Map> dataBufferMap, DataBuffer dataBuffer) { - if (dataBufferMap.get(key) == null) { - dataBufferMap.put(key, Flux.just(copy(dataBuffer))); - } else { - dataBufferMap.put(key, dataBufferMap.get(key).mergeWith(Flux.just(copy(dataBuffer)))); - } - } - - @Override - public BinaryOperator>> combiner() { - return (map1, map2) -> { - map2.forEach((k, v) -> map1.merge(k, v, (v1, v2) -> v1.mergeWith(v2))); - return map1; - }; - } - - @Override - public Function>, Map>> finisher() { - return Function.identity(); - } - - @Override - public Set characteristics() { - return CHARACTERISTICS; - } - - private DataBuffer copy(DataBuffer dataBuffer) { - return dataBuffer.factory().allocateBuffer().write(dataBuffer.asByteBuffer()); - } - } - @Override @SuppressWarnings("unchecked") public Predicate apply(Config config) { @@ -186,26 +118,6 @@ public class ReadBodyPredicateFactory "ReadBodyPredicateFactory is only async."); } - static class ReadBodyReactiveHttpInputMessage implements ReactiveHttpInputMessage { - private Flux body; - private HttpHeaders httpHeaders; - - public ReadBodyReactiveHttpInputMessage(Flux body, HttpHeaders headers) { - this.body = body; - this.httpHeaders = headers; - } - - @Override - public Flux getBody() { - return body; - } - - @Override - public HttpHeaders getHeaders() { - return httpHeaders; - } - } - public static class Config { private Class inClass; private Predicate predicate;