From 2e32a06a558d527f3789a14e4a134aeafa2fa9a3 Mon Sep 17 00:00:00 2001 From: Rossen Stoyanchev Date: Sun, 3 Jan 2021 21:24:55 +0000 Subject: [PATCH] Update to latest GraphQL over WebSocket protocol --- samples/webflux-websocket/build.gradle | 20 + .../sample/graphql/SampleApplication.java | 28 ++ .../spring/sample/graphql/SampleWiring.java | 23 ++ .../src/main/resources/application.properties | 5 + .../src/main/resources/schema.graphqls | 11 + .../src/main/resources/static/index.html | 55 +++ settings.gradle | 1 + .../boot/graphql/GraphQLProperties.java | 16 + .../WebFluxGraphQLAutoConfiguration.java | 7 +- .../WebFluxGraphQLWebSocketHandler.java | 269 ++++++++++--- .../org/springframework/graphql/WebInput.java | 4 + .../graphql/WebInterceptorExecutionChain.java | 10 +- .../springframework/graphql/WebOutput.java | 23 +- .../graphql/WebSocketInput.java | 60 +++ .../WebFluxApplicationContextTests.java | 127 +----- .../WebFluxGraphQLWebSocketHandlerTests.java | 363 ++++++++++++++++++ 16 files changed, 858 insertions(+), 164 deletions(-) create mode 100644 samples/webflux-websocket/build.gradle create mode 100644 samples/webflux-websocket/src/main/java/io/spring/sample/graphql/SampleApplication.java create mode 100644 samples/webflux-websocket/src/main/java/io/spring/sample/graphql/SampleWiring.java create mode 100644 samples/webflux-websocket/src/main/resources/application.properties create mode 100644 samples/webflux-websocket/src/main/resources/schema.graphqls create mode 100644 samples/webflux-websocket/src/main/resources/static/index.html create mode 100644 spring-graphql-web/src/main/java/org/springframework/graphql/WebSocketInput.java create mode 100644 spring-graphql-web/src/test/java/org/springframework/graphql/WebFluxGraphQLWebSocketHandlerTests.java diff --git a/samples/webflux-websocket/build.gradle b/samples/webflux-websocket/build.gradle new file mode 100644 index 00000000..f722c6e3 --- /dev/null +++ b/samples/webflux-websocket/build.gradle @@ -0,0 +1,20 @@ +plugins { + id 'org.springframework.boot' version '2.4.0-SNAPSHOT' + id 'io.spring.dependency-management' version '1.0.10.RELEASE' + id 'java' +} +group = 'com.example' +version = '0.0.1-SNAPSHOT' +description = "GraphQL over WebSocket sample" +sourceCompatibility = '1.8' + +dependencies { + implementation project(':spring-graphql-web') + implementation 'org.springframework.boot:spring-boot-starter-webflux' + implementation 'org.springframework.boot:spring-boot-starter-actuator' + developmentOnly 'org.springframework.boot:spring-boot-devtools' + testImplementation 'org.springframework.boot:spring-boot-starter-test' +} +test { + useJUnitPlatform() +} \ No newline at end of file diff --git a/samples/webflux-websocket/src/main/java/io/spring/sample/graphql/SampleApplication.java b/samples/webflux-websocket/src/main/java/io/spring/sample/graphql/SampleApplication.java new file mode 100644 index 00000000..19225fa6 --- /dev/null +++ b/samples/webflux-websocket/src/main/java/io/spring/sample/graphql/SampleApplication.java @@ -0,0 +1,28 @@ +/* + * Copyright 2002-2020 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.spring.sample.graphql; + +import org.springframework.boot.SpringApplication; +import org.springframework.boot.autoconfigure.SpringBootApplication; + +@SpringBootApplication +public class SampleApplication { + + public static void main(String[] args) { + SpringApplication.run(SampleApplication.class, args); + } +} diff --git a/samples/webflux-websocket/src/main/java/io/spring/sample/graphql/SampleWiring.java b/samples/webflux-websocket/src/main/java/io/spring/sample/graphql/SampleWiring.java new file mode 100644 index 00000000..14be4617 --- /dev/null +++ b/samples/webflux-websocket/src/main/java/io/spring/sample/graphql/SampleWiring.java @@ -0,0 +1,23 @@ +package io.spring.sample.graphql; + +import java.time.Duration; + +import graphql.schema.idl.RuntimeWiring; +import reactor.core.publisher.Flux; + +import org.springframework.boot.graphql.RuntimeWiringCustomizer; +import org.springframework.stereotype.Component; + +@Component +public class SampleWiring implements RuntimeWiringCustomizer { + + @Override + public void customize(RuntimeWiring.Builder builder) { + builder.type("Query", wiringBuilder -> wiringBuilder.dataFetcher("hello", + env -> "Hello world!")); + builder.type("Subscription", wiringBuilder -> wiringBuilder.dataFetcher("greetings", + env -> Flux.just("Hi", "Bonjour", "Hola", "Cio", "Zdravo") + .delayElements(Duration.ofMillis(500)))); + } + +} diff --git a/samples/webflux-websocket/src/main/resources/application.properties b/samples/webflux-websocket/src/main/resources/application.properties new file mode 100644 index 00000000..256081aa --- /dev/null +++ b/samples/webflux-websocket/src/main/resources/application.properties @@ -0,0 +1,5 @@ +management.endpoints.web.exposure.include=health,metrics,info +logging.level.org.springframework.web=debug +logging.level.org.springframework.http=debug +logging.level.org.springframework.graphql=debug +logging.level.reactor.netty=debug \ No newline at end of file diff --git a/samples/webflux-websocket/src/main/resources/schema.graphqls b/samples/webflux-websocket/src/main/resources/schema.graphqls new file mode 100644 index 00000000..8716049a --- /dev/null +++ b/samples/webflux-websocket/src/main/resources/schema.graphqls @@ -0,0 +1,11 @@ +schema { + query: Query + subscription : Subscription +} + +type Query { + hello: String +} +type Subscription { + greetings: String +} \ No newline at end of file diff --git a/samples/webflux-websocket/src/main/resources/static/index.html b/samples/webflux-websocket/src/main/resources/static/index.html new file mode 100644 index 00000000..f9654872 --- /dev/null +++ b/samples/webflux-websocket/src/main/resources/static/index.html @@ -0,0 +1,55 @@ + + + + + GraphQL over WebSocket + + + + + + \ No newline at end of file diff --git a/settings.gradle b/settings.gradle index 13576afa..cb60eb3d 100644 --- a/settings.gradle +++ b/settings.gradle @@ -16,3 +16,4 @@ pluginManagement { rootProject.name = 'spring-graphql' include 'spring-graphql-web' include 'samples:webmvc-http-endpoint' +include 'samples:webflux-websocket' diff --git a/spring-graphql-web/src/main/java/org/springframework/boot/graphql/GraphQLProperties.java b/spring-graphql-web/src/main/java/org/springframework/boot/graphql/GraphQLProperties.java index 5031c619..4c8de061 100644 --- a/spring-graphql-web/src/main/java/org/springframework/boot/graphql/GraphQLProperties.java +++ b/spring-graphql-web/src/main/java/org/springframework/boot/graphql/GraphQLProperties.java @@ -15,6 +15,8 @@ */ package org.springframework.boot.graphql; +import java.time.Duration; + import org.springframework.boot.context.properties.ConfigurationProperties; @ConfigurationProperties(prefix = "spring.graphql") @@ -35,6 +37,13 @@ public class GraphQLProperties { */ private String webSocketPath = path + "/websocket"; + /** + * For the GraphQL over WebSocket endpoint, this is time within which the + * initial {@code CONNECTION_INIT} type message must be received. + */ + private Duration connectionInitTimeoutDuration = Duration.ofSeconds(60); + + public String getPath() { return path; } @@ -59,4 +68,11 @@ public class GraphQLProperties { this.schemaLocation = schemaLocation; } + public Duration getConnectionInitTimeoutDuration() { + return this.connectionInitTimeoutDuration; + } + + public void setConnectionInitTimeoutDuration(Duration connectionInitTimeoutDuration) { + this.connectionInitTimeoutDuration = connectionInitTimeoutDuration; + } } diff --git a/spring-graphql-web/src/main/java/org/springframework/boot/graphql/WebFluxGraphQLAutoConfiguration.java b/spring-graphql-web/src/main/java/org/springframework/boot/graphql/WebFluxGraphQLAutoConfiguration.java index d1e6a770..876aa86f 100644 --- a/spring-graphql-web/src/main/java/org/springframework/boot/graphql/WebFluxGraphQLAutoConfiguration.java +++ b/spring-graphql-web/src/main/java/org/springframework/boot/graphql/WebFluxGraphQLAutoConfiguration.java @@ -57,9 +57,12 @@ public class WebFluxGraphQLAutoConfiguration { @Bean @ConditionalOnMissingBean public WebFluxGraphQLWebSocketHandler graphQLWebSocketHandler( - GraphQL.Builder graphQLBuilder, ServerCodecConfigurer configurer) { + GraphQL.Builder graphQLBuilder, GraphQLProperties properties, ServerCodecConfigurer configurer) { - return new WebFluxGraphQLWebSocketHandler(graphQLBuilder.build(), Collections.emptyList(), configurer); + return new WebFluxGraphQLWebSocketHandler( + graphQLBuilder.build(), Collections.emptyList(), + configurer, properties.getConnectionInitTimeoutDuration() + ); } @Bean diff --git a/spring-graphql-web/src/main/java/org/springframework/graphql/WebFluxGraphQLWebSocketHandler.java b/spring-graphql-web/src/main/java/org/springframework/graphql/WebFluxGraphQLWebSocketHandler.java index 14f5d0b7..f3867006 100644 --- a/spring-graphql-web/src/main/java/org/springframework/graphql/WebFluxGraphQLWebSocketHandler.java +++ b/spring-graphql-web/src/main/java/org/springframework/graphql/WebFluxGraphQLWebSocketHandler.java @@ -15,18 +15,25 @@ */ package org.springframework.graphql; +import java.time.Duration; import java.util.Collections; +import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.atomic.AtomicBoolean; +import graphql.ErrorType; import graphql.ExecutionResult; import graphql.GraphQL; +import graphql.GraphqlErrorBuilder; import org.apache.commons.logging.Log; import org.apache.commons.logging.LogFactory; import org.reactivestreams.Publisher; +import org.reactivestreams.Subscription; +import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; -import org.springframework.core.ResolvableType; import org.springframework.core.codec.Decoder; import org.springframework.core.codec.Encoder; import org.springframework.core.io.buffer.DataBuffer; @@ -35,8 +42,11 @@ import org.springframework.http.MediaType; import org.springframework.http.codec.DecoderHttpMessageReader; import org.springframework.http.codec.EncoderHttpMessageWriter; import org.springframework.http.codec.ServerCodecConfigurer; +import org.springframework.lang.Nullable; +import org.springframework.util.Assert; import org.springframework.util.CollectionUtils; import org.springframework.util.MimeTypeUtils; +import org.springframework.web.reactive.socket.CloseStatus; import org.springframework.web.reactive.socket.HandshakeInfo; import org.springframework.web.reactive.socket.WebSocketHandler; import org.springframework.web.reactive.socket.WebSocketMessage; @@ -50,33 +60,45 @@ public class WebFluxGraphQLWebSocketHandler implements WebSocketHandler { private static final Log logger = LogFactory.getLog(WebFluxGraphQLWebSocketHandler.class); - private static final ResolvableType MAP_TYPE = ResolvableType.forClass(Map.class); + private static final List SUB_PROTOCOL_LIST = Collections.singletonList("graphql-transport-ws"); private final WebInterceptorExecutionChain executionChain; - private final Decoder jsonDecoder; + private final Duration initTimeoutDuration; - private final Encoder jsonEncoder; + private final Decoder decoder; + private final Encoder encoder; + + private final Map subscriptions = new ConcurrentHashMap<>(); + + + @Override + public List getSubProtocols() { + return SUB_PROTOCOL_LIST; + } /** * Create a new instance. * @param graphQL the GraphQL instance to use for query execution * @param interceptors 0 or more interceptors to customize input and output * @param configurer codec configurer for JSON encoding and decoding + * @param initTimeoutDuration the time within which the +* {@code CONNECTION_INIT} type message must be received. */ public WebFluxGraphQLWebSocketHandler(GraphQL graphQL, List interceptors, - ServerCodecConfigurer configurer) { + ServerCodecConfigurer configurer, Duration initTimeoutDuration) { this.executionChain = new WebInterceptorExecutionChain(graphQL, interceptors); - this.jsonDecoder = initDecoder(configurer); - this.jsonEncoder = initEncoder(configurer); + this.initTimeoutDuration = initTimeoutDuration; + this.decoder = initDecoder(configurer); + this.encoder = initEncoder(configurer); } private static Decoder initDecoder(ServerCodecConfigurer configurer) { return configurer.getReaders().stream() - .filter(reader -> reader.canRead(MAP_TYPE, MediaType.APPLICATION_JSON)) + .filter(reader -> reader.canRead(WebInput.MAP_RESOLVABLE_TYPE, MediaType.APPLICATION_JSON)) .map(reader -> ((DecoderHttpMessageReader) reader).getDecoder()) .findFirst() .orElseThrow(() -> new IllegalArgumentException("No JSON Decoder")); @@ -84,7 +106,7 @@ public class WebFluxGraphQLWebSocketHandler implements WebSocketHandler { private static Encoder initEncoder(ServerCodecConfigurer configurer) { return configurer.getWriters().stream() - .filter(writer -> writer.canWrite(MAP_TYPE, MediaType.APPLICATION_JSON)) + .filter(writer -> writer.canWrite(WebInput.MAP_RESOLVABLE_TYPE, MediaType.APPLICATION_JSON)) .map(writer -> ((EncoderHttpMessageWriter) writer).getEncoder()) .findFirst() .orElseThrow(() -> new IllegalArgumentException("No JSON Encoder")); @@ -92,54 +114,211 @@ public class WebFluxGraphQLWebSocketHandler implements WebSocketHandler { @Override - @SuppressWarnings("unchecked") public Mono handle(WebSocketSession session) { - return session.send(session.receive() - .concatMap(message -> { + AtomicBoolean initialized = new AtomicBoolean(); + Mono.delay(this.initTimeoutDuration) + .then(Mono.defer(() -> + initialized.compareAndSet(false, true) ? + GraphQLStatus.initTimeout(session) : + Mono.empty())) + .subscribe(); + + Flux responseFlux = session.receive() + .flatMap(message -> { Map map = decode(message); - HandshakeInfo handshakeInfo = session.getHandshakeInfo(); - WebInput webInput = new WebInput(handshakeInfo.getUri(), handshakeInfo.getHeaders(), map); - if (logger.isDebugEnabled()) { - logger.debug("Executing: " + webInput); + String id = (String) map.get("id"); + MessageType messageType = MessageType.resolve((String) map.get("type")); + if (id == null || messageType == null) { + return GraphQLStatus.invalidMessage(session); } - return executionChain.execute(webInput); - }) - .concatMap(output -> { - if (!CollectionUtils.isEmpty(output.getErrors())) { - throw new IllegalStateException( - "Execution failed: " + output.getErrors()); + switch (messageType) { + case SUBSCRIBE: + if (!initialized.get()) { + return GraphQLStatus.unauthorized(session); + } + HandshakeInfo handshakeInfo = session.getHandshakeInfo(); + WebSocketInput input = new WebSocketInput(handshakeInfo, id, getPayload(map)); + if (logger.isDebugEnabled()) { + logger.debug("Executing: " + input); + } + return executionChain.execute(input) + .flatMapMany(output -> handleWebOutput(session, input.id(), output)); + case COMPLETE: + Subscription subscription = this.subscriptions.remove(id); + if (subscription != null) { + subscription.cancel(); + } + return Flux.empty(); + case CONNECTION_INIT: + if (!initialized.compareAndSet(false, true)) { + return GraphQLStatus.tooManyInitRequests(session); + } + return Flux.just(encode(session, id, MessageType.CONNECTION_ACK, null)); + default: + return GraphQLStatus.invalidMessage(session); } - if (!(output.getData() instanceof Publisher)) { - throw new IllegalStateException( - "Expected Publisher: " + output.toSpecification()); - } - if (logger.isDebugEnabled()) { - logger.debug("Execution complete, subscribing for events."); - } - return (Publisher) output.getData(); - }) - .map(result -> { - Object data = result.getData(); - return encode(session, data); - }) - ); + }); + + return session.send(responseFlux); } @SuppressWarnings({"unchecked", "ConstantConditions"}) private Map decode(WebSocketMessage message) { - DataBuffer buffer = message.getPayload(); - return (Map) jsonDecoder.decode( - DataBufferUtils.retain(buffer), WebInput.MAP_RESOLVABLE_TYPE, null, Collections.emptyMap()); + DataBuffer buffer = DataBufferUtils.retain(message.getPayload()); + return (Map) decoder.decode(buffer, WebInput.MAP_RESOLVABLE_TYPE, null, null); } @SuppressWarnings("unchecked") - private WebSocketMessage encode(WebSocketSession session, Object data) { - DataBuffer buffer = ((Encoder) jsonEncoder).encodeValue((T) data, - session.bufferFactory(), - ResolvableType.forInstance(data), - MimeTypeUtils.APPLICATION_JSON, - Collections.emptyMap()); + private static Map getPayload(Map message) { + Map payload = (Map) message.get("payload"); + Assert.notNull(payload, "No \"payload\" in message: " + message); + return payload; + } + + @SuppressWarnings("unchecked") + private Flux handleWebOutput(WebSocketSession session, String id, WebOutput output) { + if (logger.isDebugEnabled()) { + logger.debug("Execution result ready" + + (!CollectionUtils.isEmpty(output.getErrors()) ? + " with errors: " + output.getErrors() : "") + "."); + } + + Flux outputFlux; + if (output.getData() instanceof Publisher) { + // Subscription + outputFlux = Flux.from((Publisher) output.getData()) + .doOnSubscribe(subscription -> { + Subscription previous = this.subscriptions.putIfAbsent(id, subscription); + if (previous != null) { + throw new SubscriptionExistsException(); + } + }); + } + else { + // Query + outputFlux = (CollectionUtils.isEmpty(output.getErrors()) ? + Flux.just(output) : + Flux.error(new IllegalStateException("Execution failed: " + output.getErrors()))); + } + + return outputFlux + .map(result -> { + Map dataMap = result.toSpecification(); + return encode(session, id, MessageType.NEXT, dataMap); + }) + .concatWith(Mono.defer(() -> Mono.just(encode(session, id, MessageType.COMPLETE, null)))) + .onErrorResume(ex -> { + if (ex instanceof SubscriptionExistsException) { + return GraphQLStatus.subscriptionExists(session, id); + } + ErrorType errorType = ErrorType.DataFetchingException; + String message = ex.getMessage(); + Map errorMap = GraphqlErrorBuilder.newError() + .errorType(errorType) + .message(message) + .build() + .toSpecification(); + return Mono.just(encode(session, id, MessageType.ERROR, errorMap)); + }); + } + + @SuppressWarnings("unchecked") + private WebSocketMessage encode( + WebSocketSession session, String id, MessageType messageType, @Nullable Object payload) { + + Map payloadMap = new HashMap<>(3); + payloadMap.put("id", id); + payloadMap.put("type", messageType.getType()); + if (payload != null) { + payloadMap.put("payload", payload); + } + + DataBuffer buffer = ((Encoder) encoder).encodeValue( + (T) payloadMap, session.bufferFactory(), WebInput.MAP_RESOLVABLE_TYPE, + MimeTypeUtils.APPLICATION_JSON, null); + return new WebSocketMessage(WebSocketMessage.Type.TEXT, buffer); } + + private enum MessageType { + + CONNECTION_INIT("connection_init"), + CONNECTION_ACK("connection_ack"), + SUBSCRIBE("subscribe"), + NEXT("next"), + ERROR("error"), + COMPLETE("complete"); + + + private static final Map messageTypes = new HashMap<>(6); + + static { + for (MessageType messageType : MessageType.values()) { + messageTypes.put(messageType.getType(), messageType); + } + } + + + private final String type; + + MessageType(String type) { + this.type = type; + } + + public String getType() { + return this.type; + } + + @Nullable + public static MessageType resolve(@Nullable String type) { + return (type != null ? messageTypes.get(type) : null); + } + } + + + private static class GraphQLStatus { + + private static final CloseStatus INVALID_MESSAGE_STATUS = + new CloseStatus(4400, "Invalid message"); + + private static final CloseStatus UNAUTHORIZED_STATUS = + new CloseStatus(4401, "Unauthorized"); + + private static final CloseStatus INIT_TIMEOUT_STATUS = + new CloseStatus(4408, "Connection initialisation timeout"); + + private static final CloseStatus TOO_MANY_INIT_REQUESTS_STATUS = + new CloseStatus(4429, "Too many initialisation requests"); + + + public static Flux invalidMessage(WebSocketSession session) { + return closeSession(session, INVALID_MESSAGE_STATUS); + } + + public static Flux unauthorized(WebSocketSession session) { + return closeSession(session, UNAUTHORIZED_STATUS); + } + + public static Mono initTimeout(WebSocketSession session) { + return session.close(INIT_TIMEOUT_STATUS); + } + + public static Flux subscriptionExists(WebSocketSession session, String id) { + return closeSession(session, new CloseStatus(4409, "Subscriber for " + id + " already exists")); + } + + public static Flux tooManyInitRequests(WebSocketSession session) { + return closeSession(session, TOO_MANY_INIT_REQUESTS_STATUS); + } + + private static Flux closeSession(WebSocketSession session, CloseStatus status) { + return session.close(status).thenMany(Mono.empty()); + } + } + + + private static class SubscriptionExistsException extends RuntimeException { + } + } diff --git a/spring-graphql-web/src/main/java/org/springframework/graphql/WebInput.java b/spring-graphql-web/src/main/java/org/springframework/graphql/WebInput.java index f064ed4a..0713e8a9 100644 --- a/spring-graphql-web/src/main/java/org/springframework/graphql/WebInput.java +++ b/spring-graphql-web/src/main/java/org/springframework/graphql/WebInput.java @@ -25,6 +25,7 @@ import org.springframework.core.ParameterizedTypeReference; import org.springframework.core.ResolvableType; import org.springframework.http.HttpHeaders; import org.springframework.lang.Nullable; +import org.springframework.util.Assert; import org.springframework.util.CollectionUtils; import org.springframework.util.StringUtils; import org.springframework.web.server.ServerWebInputException; @@ -58,6 +59,9 @@ public class WebInput { @SuppressWarnings("unchecked") public WebInput(URI uri, HttpHeaders headers, Map body) { + Assert.notNull(uri, "URI is required'"); + Assert.notNull(body, "HttpHeaders is required'"); + Assert.notNull(body, "'body' is required'"); this.uri = UriComponentsBuilder.fromUri(uri).build(true); this.headers = headers; this.query = getAndValidateQuery(body); diff --git a/spring-graphql-web/src/main/java/org/springframework/graphql/WebInterceptorExecutionChain.java b/spring-graphql-web/src/main/java/org/springframework/graphql/WebInterceptorExecutionChain.java index 57a94152..4bc6019e 100644 --- a/spring-graphql-web/src/main/java/org/springframework/graphql/WebInterceptorExecutionChain.java +++ b/spring-graphql-web/src/main/java/org/springframework/graphql/WebInterceptorExecutionChain.java @@ -48,10 +48,10 @@ class WebInterceptorExecutionChain { } - public Mono execute(WebInput webInput) { - return createInputChain(webInput).flatMap(executionInput -> { + public Mono execute(WebInput input) { + return createInputChain(input).flatMap(executionInput -> { CompletableFuture future = this.graphQL.executeAsync(executionInput); - return createOutputChain(Mono.fromFuture(future)); + return createOutputChain(input, Mono.fromFuture(future)); }); } @@ -63,8 +63,8 @@ class WebInterceptorExecutionChain { return preHandleMono; } - private Mono createOutputChain(Mono resultMono) { - Mono outputMono = resultMono.map((ExecutionResult executionResult) -> new WebOutput(executionResult, null)); + private Mono createOutputChain(WebInput input, Mono resultMono) { + Mono outputMono = resultMono.map((ExecutionResult result) -> new WebOutput(input, result, null)); for (int i = this.interceptors.size() - 1 ; i >= 0; i--) { WebInterceptor interceptor = this.interceptors.get(i); outputMono = outputMono.flatMap(interceptor::postHandle); diff --git a/spring-graphql-web/src/main/java/org/springframework/graphql/WebOutput.java b/spring-graphql-web/src/main/java/org/springframework/graphql/WebOutput.java index 5fe937a1..14c109c2 100644 --- a/spring-graphql-web/src/main/java/org/springframework/graphql/WebOutput.java +++ b/spring-graphql-web/src/main/java/org/springframework/graphql/WebOutput.java @@ -26,6 +26,7 @@ import graphql.GraphQLError; import org.springframework.http.HttpHeaders; import org.springframework.lang.Nullable; +import org.springframework.util.Assert; /** @@ -34,6 +35,8 @@ import org.springframework.lang.Nullable; */ public class WebOutput implements ExecutionResult { + private final WebInput input; + private final ExecutionResult executionResult; @Nullable @@ -43,12 +46,22 @@ public class WebOutput implements ExecutionResult { /** * Create an instance that wraps the given {@link ExecutionResult}. */ - public WebOutput(ExecutionResult executionResult, @Nullable HttpHeaders headers) { + public WebOutput(WebInput input, ExecutionResult executionResult, @Nullable HttpHeaders headers) { + Assert.notNull(input, "WebInput is required."); + Assert.notNull(executionResult, "ExecutionResult is required."); + this.input = input; this.executionResult = executionResult; this.headers = headers; } + /** + * Return the associated {@link WebInput} used for the execution. + */ + public WebInput getWebInput() { + return this.input; + } + @Nullable @Override public T getData() { @@ -95,6 +108,8 @@ public class WebOutput implements ExecutionResult { public static class Builder { + private final WebInput input; + @Nullable private Object data; @@ -107,7 +122,8 @@ public class WebOutput implements ExecutionResult { private HttpHeaders headers; - public Builder(WebOutput output) { + private Builder(WebOutput output) { + this.input = output.getWebInput(); this.data = output.getData(); this.errors = output.getErrors(); this.extensions = output.getExtensions(); @@ -158,7 +174,8 @@ public class WebOutput implements ExecutionResult { } public WebOutput build() { - return new WebOutput(new ExecutionResultImpl(this.data, this.errors, this.extensions), this.headers); + ExecutionResult result = new ExecutionResultImpl(this.data, this.errors, this.extensions); + return new WebOutput(this.input, result, this.headers); } } diff --git a/spring-graphql-web/src/main/java/org/springframework/graphql/WebSocketInput.java b/spring-graphql-web/src/main/java/org/springframework/graphql/WebSocketInput.java new file mode 100644 index 00000000..06a8b95d --- /dev/null +++ b/spring-graphql-web/src/main/java/org/springframework/graphql/WebSocketInput.java @@ -0,0 +1,60 @@ +/* + * Copyright 2002-2020 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.springframework.graphql; + +import java.util.Map; + +import org.springframework.web.reactive.socket.HandshakeInfo; + +/** + * Extension of {@link WebInput} that contains a GraphQL subscription received + * over a WebSocket connection. + */ +public class WebSocketInput extends WebInput { + + private final HandshakeInfo handshakeInfo; + + private final String id; + + + public WebSocketInput(HandshakeInfo handshakeInfo, String id, Map payload) { + super(handshakeInfo.getUri(), handshakeInfo.getHeaders(), payload); + this.handshakeInfo = handshakeInfo; + this.id = id; + } + + + /** + * Return information about the WebSocket handshake. + */ + public HandshakeInfo handshakeInfo() { + return this.handshakeInfo; + } + + /** + * Return the id that will correlate server responses to client requests + * within a multiplexed WebSocket connection. + */ + public String id() { + return this.id; + } + + @Override + public String toString() { + return "id='" + id() + "', " + super.toString(); + } + +} diff --git a/spring-graphql-web/src/test/java/org/springframework/boot/graphql/WebFluxApplicationContextTests.java b/spring-graphql-web/src/test/java/org/springframework/boot/graphql/WebFluxApplicationContextTests.java index 14cc55ad..5c38da82 100644 --- a/spring-graphql-web/src/test/java/org/springframework/boot/graphql/WebFluxApplicationContextTests.java +++ b/spring-graphql-web/src/test/java/org/springframework/boot/graphql/WebFluxApplicationContextTests.java @@ -1,17 +1,24 @@ +/* + * Copyright 2002-2020 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ package org.springframework.boot.graphql; - -import java.net.URI; -import java.nio.charset.StandardCharsets; import java.util.Collections; -import java.util.Map; import java.util.function.Consumer; import org.junit.jupiter.api.Test; -import org.reactivestreams.Publisher; -import reactor.core.publisher.Flux; -import reactor.core.publisher.Mono; -import reactor.test.StepVerifier; import org.springframework.boot.autoconfigure.AutoConfigurations; import org.springframework.boot.autoconfigure.http.codec.CodecsAutoConfiguration; @@ -23,24 +30,11 @@ import org.springframework.boot.test.context.runner.ReactiveWebApplicationContex import org.springframework.context.ApplicationContext; import org.springframework.context.annotation.Bean; import org.springframework.context.annotation.Configuration; -import org.springframework.core.ResolvableType; -import org.springframework.core.io.buffer.DataBuffer; -import org.springframework.core.io.buffer.DataBufferUtils; -import org.springframework.core.io.buffer.DefaultDataBufferFactory; import org.springframework.graphql.GraphQLDataFetchers; -import org.springframework.graphql.WebFluxGraphQLHandler; -import org.springframework.graphql.WebFluxGraphQLWebSocketHandler; -import org.springframework.http.HttpHeaders; import org.springframework.http.MediaType; -import org.springframework.http.codec.json.Jackson2JsonDecoder; import org.springframework.test.web.reactive.server.WebTestClient; -import org.springframework.web.reactive.socket.CloseStatus; -import org.springframework.web.reactive.socket.HandshakeInfo; -import org.springframework.web.reactive.socket.WebSocketMessage; -import org.springframework.web.reactive.socket.adapter.AbstractWebSocketSession; import static graphql.schema.idl.TypeRuntimeWiring.newTypeWiring; -import static org.assertj.core.api.Assertions.assertThat; class WebFluxApplicationContextTests { @@ -82,43 +76,6 @@ class WebFluxApplicationContextTests { testWithWebClient(client -> client.post().uri("").bodyValue(":)").exchange().expectStatus().isBadRequest()); } - @Test - void subscription() { - testWithApplicationContext(context -> { - String query = - "{ \"query\": \"" + - " subscription TestSubscription {" + - " bookSearch(minPages: 200) {" + - " id" + - " name" + - " pageCount" + - " author" + - " }" + - "}" + - "\"}"; - - DataBuffer buffer = DefaultDataBufferFactory.sharedInstance.wrap(query.getBytes(StandardCharsets.UTF_8)); - Flux input = Flux.just(new WebSocketMessage(WebSocketMessage.Type.TEXT, buffer)); - TestWebSocketSession session = new TestWebSocketSession("1", URI.create(BASE_URL), input); - - context.getBean(WebFluxGraphQLWebSocketHandler.class).handle(session).block(); - - StepVerifier.create(session.getOutput()) - .consumeNextWith(message -> assertThat(extractBook(message)).containsEntry("id", "book-2")) - .consumeNextWith(message -> assertThat(extractBook(message)).containsEntry("id", "book-3")) - .consumeNextWith(message -> assertThat(extractBook(message)).containsEntry("id", "book-3")) - .verifyComplete(); - }); - } - - @SuppressWarnings({"unchecked", "ConstantConditions"}) - private Map extractBook(WebSocketMessage message) { - Map map = (Map) new Jackson2JsonDecoder().decode( - DataBufferUtils.retain(message.getPayload()), - ResolvableType.forClass(Map.class), null, Collections.emptyMap()); - return (Map) map.get("bookSearch"); - } - private void testWithWebClient(Consumer consumer) { testWithApplicationContext(context -> { WebTestClient client = WebTestClient.bindToApplicationContext(context) @@ -149,57 +106,9 @@ class WebFluxApplicationContextTests { @Bean public RuntimeWiringCustomizer bookDataFetcher() { - return (runtimeWiring) -> { - runtimeWiring.type(newTypeWiring("Query") - .dataFetcher("bookById", GraphQLDataFetchers.getBookByIdDataFetcher())); - runtimeWiring.type(newTypeWiring("Subscription") - .dataFetcher("bookSearch", GraphQLDataFetchers.getBooksOnSale())); - }; - } - } - - - private static class TestWebSocketSession extends AbstractWebSocketSession { - - private final Flux input; - - private Flux output; - - public TestWebSocketSession(String id, URI uri, Flux input) { - super(new Object(), id, - new HandshakeInfo(uri, new HttpHeaders(), Mono.empty(), null), - DefaultDataBufferFactory.sharedInstance); - this.input = input; - } - - @Override - public Flux receive() { - return this.input; - } - - @Override - public Mono send(Publisher messages) { - this.output = Flux.from(messages); - return Mono.empty(); - } - - public Flux getOutput() { - return this.output; - } - - @Override - public boolean isOpen() { - throw new java.lang.UnsupportedOperationException(); - } - - @Override - public Mono close(CloseStatus status) { - throw new java.lang.UnsupportedOperationException(); - } - - @Override - public Mono closeStatus() { - throw new java.lang.UnsupportedOperationException(); + return (runtimeWiring) -> + runtimeWiring.type(newTypeWiring("Query") + .dataFetcher("bookById", GraphQLDataFetchers.getBookByIdDataFetcher())); } } diff --git a/spring-graphql-web/src/test/java/org/springframework/graphql/WebFluxGraphQLWebSocketHandlerTests.java b/spring-graphql-web/src/test/java/org/springframework/graphql/WebFluxGraphQLWebSocketHandlerTests.java new file mode 100644 index 00000000..a983f2c1 --- /dev/null +++ b/spring-graphql-web/src/test/java/org/springframework/graphql/WebFluxGraphQLWebSocketHandlerTests.java @@ -0,0 +1,363 @@ +/* + * Copyright 2002-2020 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.springframework.graphql; + +import java.io.File; +import java.net.URI; +import java.nio.charset.StandardCharsets; +import java.time.Duration; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.function.BiConsumer; + +import graphql.GraphQL; +import graphql.schema.GraphQLSchema; +import graphql.schema.idl.RuntimeWiring; +import graphql.schema.idl.SchemaGenerator; +import graphql.schema.idl.SchemaParser; +import graphql.schema.idl.TypeDefinitionRegistry; +import org.junit.jupiter.api.Test; +import org.reactivestreams.Publisher; +import reactor.core.publisher.Flux; +import reactor.core.publisher.Mono; +import reactor.core.publisher.Sinks; +import reactor.test.StepVerifier; + +import org.springframework.core.io.buffer.DataBuffer; +import org.springframework.core.io.buffer.DataBufferUtils; +import org.springframework.core.io.buffer.DefaultDataBufferFactory; +import org.springframework.http.HttpHeaders; +import org.springframework.http.codec.ServerCodecConfigurer; +import org.springframework.http.codec.json.Jackson2JsonDecoder; +import org.springframework.lang.Nullable; +import org.springframework.util.ResourceUtils; +import org.springframework.web.reactive.socket.CloseStatus; +import org.springframework.web.reactive.socket.HandshakeInfo; +import org.springframework.web.reactive.socket.WebSocketMessage; +import org.springframework.web.reactive.socket.adapter.AbstractWebSocketSession; + +import static graphql.schema.idl.TypeRuntimeWiring.newTypeWiring; +import static org.assertj.core.api.Assertions.as; +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.InstanceOfAssertFactories.map; + +/** + * Unit tests for {@link WebFluxGraphQLWebSocketHandler}. + */ +public class WebFluxGraphQLWebSocketHandlerTests { + + private static final Jackson2JsonDecoder decoder = new Jackson2JsonDecoder(); + + private static final String SUBSCRIPTION_ID = "123"; + + private static final String BOOK_SEARCH_QUERY = "{" + + "\"id\":\"" + SUBSCRIPTION_ID + "\"," + + "\"type\":\"subscribe\"," + + "\"payload\":{\"query\": \"" + + " subscription TestSubscription {" + + " bookSearch(minPages: 200) {" + + " id" + + " name" + + " pageCount" + + " author" + + " }}\"}" + + "}"; + + + @Test + void query() throws Exception { + String bookQuery = "{" + + "\"id\":\"" + SUBSCRIPTION_ID + "\"," + + "\"type\":\"subscribe\"," + + "\"payload\":{\"query\": \"" + + " query TestQuery {" + + " bookById(id: \\\"book-1\\\"){ " + + " id" + + " name" + + " pageCount" + + " author" + + " }}\"}" + + "}"; + + Flux input = Flux.just( + toWebSocketMessage("{\"id\":\"" + SUBSCRIPTION_ID + "\",\"type\":\"connection_init\"}"), + toWebSocketMessage(bookQuery)); + + TestWebSocketSession session = new TestWebSocketSession(input); + initWebSocketHandler().handle(session).block(); + + StepVerifier.create(session.getOutput()) + .consumeNextWith(message -> assertMessageType(message, "connection_ack")) + .consumeNextWith(message -> + assertThat(decode(message)) + .hasSize(3) + .containsEntry("id", SUBSCRIPTION_ID) + .containsEntry("type", "next") + .extractingByKey("payload", as(map(String.class, Object.class))) + .extractingByKey("data", as(map(String.class, Object.class))) + .extractingByKey("bookById", as(map(String.class, Object.class))) + .containsEntry("name", "GraphQL for beginners")) + .consumeNextWith(message -> assertMessageType(message, "complete")) + .verifyComplete(); + } + + @Test + void subscription() throws Exception { + Flux input = Flux.just( + toWebSocketMessage("{\"id\":\"" + SUBSCRIPTION_ID + "\",\"type\":\"connection_init\"}"), + toWebSocketMessage(BOOK_SEARCH_QUERY)); + + TestWebSocketSession session = new TestWebSocketSession(input); + initWebSocketHandler().handle(session).block(); + + BiConsumer bookPayloadAssertion = (message, bookId) -> + assertThat(decode(message)) + .hasSize(3) + .containsEntry("id", SUBSCRIPTION_ID) + .containsEntry("type", "next") + .extractingByKey("payload", as(map(String.class, Object.class))) + .extractingByKey("data", as(map(String.class, Object.class))) + .extractingByKey("bookSearch", as(map(String.class, Object.class))) + .containsEntry("id", bookId); + + StepVerifier.create(session.getOutput()) + .consumeNextWith(message -> assertMessageType(message, "connection_ack")) + .consumeNextWith(message -> bookPayloadAssertion.accept(message, "book-2")) + .consumeNextWith(message -> bookPayloadAssertion.accept(message, "book-3")) + .consumeNextWith(message -> bookPayloadAssertion.accept(message, "book-3")) + .consumeNextWith(message -> assertMessageType(message, "complete")) + .verifyComplete(); + } + + @Test + void unauthorizedWithoutMessageType() throws Exception { + Flux input = Flux.just( + toWebSocketMessage("{\"id\":\"" + SUBSCRIPTION_ID + "\",\"type\":\"connection_init\"}"), + toWebSocketMessage("{\"id\":\"" + SUBSCRIPTION_ID + "\"}")); // No message type + + TestWebSocketSession session = new TestWebSocketSession(input); + initWebSocketHandler().handle(session).block(); + + StepVerifier.create(session.getOutput()) + .consumeNextWith(message -> assertMessageType(message, "connection_ack")) + .verifyComplete(); + + StepVerifier.create(session.closeStatus()) + .expectNext(new CloseStatus(4400, "Invalid message")) + .verifyComplete(); + } + + @Test + void unauthorizedWithoutConnectionInit() throws Exception { + TestWebSocketSession session = new TestWebSocketSession(Flux.just(toWebSocketMessage(BOOK_SEARCH_QUERY))); + initWebSocketHandler().handle(session).block(); + + StepVerifier.create(session.getOutput()).verifyComplete(); + StepVerifier.create(session.closeStatus()) + .expectNext(new CloseStatus(4401, "Unauthorized")) + .verifyComplete(); + } + + @Test + void tooManyConnectionInitRequests() throws Exception { + Flux input = Flux.just( + toWebSocketMessage("{\"id\":\"" + SUBSCRIPTION_ID + "\",\"type\":\"connection_init\"}"), + toWebSocketMessage("{\"id\":\"" + SUBSCRIPTION_ID + "\",\"type\":\"connection_init\"}")); + + TestWebSocketSession session = new TestWebSocketSession(input); + initWebSocketHandler().handle(session).block(); + + StepVerifier.create(session.getOutput()) + .consumeNextWith(message -> assertMessageType(message, "connection_ack")) + .verifyComplete(); + + StepVerifier.create(session.closeStatus()) + .expectNext(new CloseStatus(4429, "Too many initialisation requests")) + .verifyComplete(); + } + + @Test + void connectionInitTimeout() throws Exception { + TestWebSocketSession session = new TestWebSocketSession(Flux.empty()); + initWebSocketHandler(Collections.emptyList(), Duration.ofMillis(50)).handle(session).block(); + + StepVerifier.create(session.closeStatus()) + .expectNext(new CloseStatus(4408, "Connection initialisation timeout")) + .verifyComplete(); + } + + @Test + void subscriptionExists() throws Exception { + Flux input = Flux.just(toWebSocketMessage( + "{\"id\":\"" + SUBSCRIPTION_ID + "\",\"type\":\"connection_init\"}"), + toWebSocketMessage(BOOK_SEARCH_QUERY), + toWebSocketMessage(BOOK_SEARCH_QUERY)); + + List interceptors = Collections.singletonList(new TakeOneAndNeverCompleteInterceptor()); + + TestWebSocketSession session = new TestWebSocketSession(input); + initWebSocketHandler(interceptors, null).handle(session).block(); + + // Collect messages until session closed + List> messages = new ArrayList<>(); + session.getOutput().subscribe(message -> messages.add(decode(message))); + + StepVerifier.create(session.closeStatus()) + .expectNext(new CloseStatus(4409, "Subscriber for " + SUBSCRIPTION_ID + " already exists")) + .verifyComplete(); + + assertThat(messages.size()).isEqualTo(2); + assertThat(messages.get(0).get("type")).isEqualTo("connection_ack"); + assertThat(messages.get(1).get("type")).isEqualTo("next"); + } + + @Test + void clientCompletion() throws Exception { + Sinks.Many input = Sinks.many().unicast().onBackpressureBuffer(); + input.tryEmitNext(toWebSocketMessage("{\"id\":\"" + SUBSCRIPTION_ID + "\",\"type\":\"connection_init\"}")); + input.tryEmitNext(toWebSocketMessage(BOOK_SEARCH_QUERY)); + + List interceptors = Collections.singletonList(new TakeOneAndNeverCompleteInterceptor()); + TestWebSocketSession session = new TestWebSocketSession(input.asFlux()); + initWebSocketHandler(interceptors, null).handle(session).block(); + + String completeMessage = "{\"id\":\"" + SUBSCRIPTION_ID + "\",\"type\":\"complete\"}"; + + StepVerifier.create(session.getOutput()) + .consumeNextWith(message -> assertMessageType(message, "connection_ack")) + .consumeNextWith(message -> assertMessageType(message, "next")) + .then(() -> input.tryEmitNext(toWebSocketMessage(completeMessage))) + .as("Second subscription with same id is possible only if the first was properly removed") + .then(() -> input.tryEmitNext(toWebSocketMessage(BOOK_SEARCH_QUERY))) + .consumeNextWith(message -> assertMessageType(message, "next")) + .then(() -> input.tryEmitNext(toWebSocketMessage(completeMessage))) + .verifyTimeout(Duration.ofMillis(500)); + } + + private WebFluxGraphQLWebSocketHandler initWebSocketHandler() throws Exception { + return initWebSocketHandler(Collections.emptyList(), Duration.ofSeconds(69)); + } + + private WebFluxGraphQLWebSocketHandler initWebSocketHandler( + @Nullable List interceptors, @Nullable Duration initTimeoutDuration) throws Exception { + + GraphQL graphQL = initGraphQL(); + return new WebFluxGraphQLWebSocketHandler(graphQL, + (interceptors != null ? interceptors : Collections.emptyList()), + ServerCodecConfigurer.create(), + (initTimeoutDuration != null ? initTimeoutDuration : Duration.ofSeconds(60))); + } + + private static GraphQL initGraphQL() throws Exception { + File schemaFile = ResourceUtils.getFile("classpath:books/schema.graphqls"); + TypeDefinitionRegistry typeDefinitionRegistry = new SchemaParser().parse(schemaFile); + + RuntimeWiring.Builder builder = RuntimeWiring.newRuntimeWiring(); + builder.type(newTypeWiring("Query").dataFetcher("bookById", GraphQLDataFetchers.getBookByIdDataFetcher())); + builder.type(newTypeWiring("Subscription").dataFetcher("bookSearch", GraphQLDataFetchers.getBooksOnSale())); + RuntimeWiring runtimeWiring = builder.build(); + + GraphQLSchema schema = new SchemaGenerator().makeExecutableSchema(typeDefinitionRegistry, runtimeWiring); + return GraphQL.newGraphQL(schema).build(); + } + + private static WebSocketMessage toWebSocketMessage(String data) { + DataBuffer buffer = DefaultDataBufferFactory.sharedInstance.wrap(data.getBytes(StandardCharsets.UTF_8)); + return new WebSocketMessage(WebSocketMessage.Type.TEXT, buffer); + } + + @SuppressWarnings({"unchecked", "ConstantConditions"}) + private Map decode(WebSocketMessage message) { + return (Map) decoder.decode( + DataBufferUtils.retain(message.getPayload()), + WebInput.MAP_RESOLVABLE_TYPE, null, Collections.emptyMap()); + } + + private void assertMessageType(WebSocketMessage message, String messageType) { + assertThat(decode(message)) + .containsEntry("id", SUBSCRIPTION_ID) + .containsEntry("type", messageType); + } + + + private static class TestWebSocketSession extends AbstractWebSocketSession { + + private final Flux input; + + private Flux output = Flux.empty(); + + private final Sinks.One closeStatusSink = Sinks.one(); + + + public TestWebSocketSession(Flux input) { + this("1", URI.create("https://example.org/graphql"), input); + } + + public TestWebSocketSession(String id, URI uri, Flux input) { + super(new Object(), id, + new HandshakeInfo(uri, new HttpHeaders(), Mono.empty(), null), + DefaultDataBufferFactory.sharedInstance); + this.input = input; + } + + + @Override + public Flux receive() { + return this.input; + } + + @Override + public Mono send(Publisher messages) { + this.output = Flux.from(messages); + return Mono.empty(); + } + + public Flux getOutput() { + return this.output; + } + + @Override + public boolean isOpen() { + throw new java.lang.UnsupportedOperationException(); + } + + @Override + public Mono close(CloseStatus status) { + this.closeStatusSink.tryEmitValue(status); + return Mono.empty(); + } + + @Override + public Mono closeStatus() { + return this.closeStatusSink.asMono(); + } + } + + + private static class TakeOneAndNeverCompleteInterceptor implements WebInterceptor { + + @Override + public Mono postHandle(WebOutput output) { + return Mono.just(output.transform(builder -> { + Publisher publisher = output.getData(); + assertThat(publisher).isNotNull(); + builder.data(Flux.from(publisher).take(1).concatWith(Flux.never())); + })); + } + } +}