From ec005471306f0f559c2a5c1afc8e529daf476ce1 Mon Sep 17 00:00:00 2001 From: rstoyanchev Date: Tue, 8 Mar 2022 08:49:46 +0000 Subject: [PATCH] Add handleConnectionClosed to WebSocketInterceptor Closes gh-276 --- .../graphql/web/WebSocketInterceptor.java | 15 +++++++++- .../web/webflux/GraphQlWebSocketHandler.java | 21 +++++++++---- .../web/webmvc/GraphQlWebSocketHandler.java | 29 +++++++++++------- .../webflux/GraphQlWebSocketHandlerTests.java | 25 ++++++++++++++++ .../webmvc/GraphQlWebSocketHandlerTests.java | 30 +++++++++++++++++++ 5 files changed, 103 insertions(+), 17 deletions(-) diff --git a/spring-graphql/src/main/java/org/springframework/graphql/web/WebSocketInterceptor.java b/spring-graphql/src/main/java/org/springframework/graphql/web/WebSocketInterceptor.java index 42591042..52aa91e9 100644 --- a/spring-graphql/src/main/java/org/springframework/graphql/web/WebSocketInterceptor.java +++ b/spring-graphql/src/main/java/org/springframework/graphql/web/WebSocketInterceptor.java @@ -43,7 +43,9 @@ public interface WebSocketInterceptor extends WebInterceptor { * @param connectionInitPayload the payload from the {@code "connection_init"} message * @return the payload for the {@code "connection_ack"}, or empty */ - default Mono handleConnectionInitialization(String sessionId, Map connectionInitPayload) { + default Mono handleConnectionInitialization( + String sessionId, Map connectionInitPayload) { + return Mono.empty(); } @@ -62,4 +64,15 @@ public interface WebSocketInterceptor extends WebInterceptor { return Mono.empty(); } + /** + * Invoked when the WebSocket session is closed, from either side. + * @param sessionId the id of the WebSocket session + * @param statusCode the WebSocket "close" status code + * @param connectionInitPayload the payload from the {@code "connect_init"} + * message received at the start of the connection + */ + default void handleConnectionClosed( + String sessionId, int statusCode, Map connectionInitPayload) { + } + } diff --git a/spring-graphql/src/main/java/org/springframework/graphql/web/webflux/GraphQlWebSocketHandler.java b/spring-graphql/src/main/java/org/springframework/graphql/web/webflux/GraphQlWebSocketHandler.java index 1313b812..982a5759 100644 --- a/spring-graphql/src/main/java/org/springframework/graphql/web/webflux/GraphQlWebSocketHandler.java +++ b/spring-graphql/src/main/java/org/springframework/graphql/web/webflux/GraphQlWebSocketHandler.java @@ -22,7 +22,7 @@ import java.util.Collections; import java.util.List; import java.util.Map; import java.util.concurrent.ConcurrentHashMap; -import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicReference; import graphql.ExecutionResult; import org.apache.commons.logging.Log; @@ -105,23 +105,34 @@ public class GraphQlWebSocketHandler implements WebSocketHandler { } // Session state - AtomicBoolean connectionInitProcessed = new AtomicBoolean(); + AtomicReference> connectionInitPayloadRef = new AtomicReference<>(); Map subscriptions = new ConcurrentHashMap<>(); Mono.delay(this.initTimeoutDuration) .then(Mono.defer(() -> - connectionInitProcessed.compareAndSet(false, true) ? + connectionInitPayloadRef.compareAndSet(null, Collections.emptyMap()) ? session.close(GraphQlStatus.INIT_TIMEOUT_STATUS) : Mono.empty())) .subscribe(); + session.closeStatus() + .doOnSuccess(closeStatus -> { + Map connectionInitPayload = connectionInitPayloadRef.get(); + if (connectionInitPayload == null) { + return; + } + int statusCode = (closeStatus != null ? closeStatus.getCode() : 1005); + this.webSocketInterceptor.handleConnectionClosed(session.getId(), statusCode, connectionInitPayload); + }) + .subscribe(); + return session.send(session.receive().flatMap(webSocketMessage -> { GraphQlWebSocketMessage message = this.codecDelegate.decode(webSocketMessage); String id = message.getId(); Map payload = message.getPayloadOrDefault(Collections.emptyMap()); switch (message.getType()) { case "subscribe": - if (!connectionInitProcessed.get()) { + if (connectionInitPayloadRef.get() == null) { return GraphQlStatus.close(session, GraphQlStatus.UNAUTHORIZED_STATUS); } if (id == null) { @@ -146,7 +157,7 @@ public class GraphQlWebSocketHandler implements WebSocketHandler { } return Flux.empty(); case "connection_init": - if (!connectionInitProcessed.compareAndSet(false, true)) { + if (!connectionInitPayloadRef.compareAndSet(null, payload)) { return GraphQlStatus.close(session, GraphQlStatus.TOO_MANY_INIT_REQUESTS_STATUS); } return this.webSocketInterceptor.handleConnectionInitialization(session.getId(), payload) diff --git a/spring-graphql/src/main/java/org/springframework/graphql/web/webmvc/GraphQlWebSocketHandler.java b/spring-graphql/src/main/java/org/springframework/graphql/web/webmvc/GraphQlWebSocketHandler.java index e4919e9c..27212cbf 100644 --- a/spring-graphql/src/main/java/org/springframework/graphql/web/webmvc/GraphQlWebSocketHandler.java +++ b/spring-graphql/src/main/java/org/springframework/graphql/web/webmvc/GraphQlWebSocketHandler.java @@ -28,6 +28,7 @@ import java.util.Collections; import java.util.List; import java.util.Map; import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.atomic.AtomicReference; import graphql.ExecutionResult; import graphql.GraphQLError; @@ -52,6 +53,7 @@ import org.springframework.http.HttpInputMessage; import org.springframework.http.HttpOutputMessage; import org.springframework.http.converter.GenericHttpMessageConverter; import org.springframework.http.converter.HttpMessageConverter; +import org.springframework.lang.Nullable; import org.springframework.util.Assert; import org.springframework.util.CollectionUtils; import org.springframework.web.socket.CloseStatus; @@ -127,7 +129,7 @@ public class GraphQlWebSocketHandler extends TextWebSocketHandler implements Sub Mono.delay(this.initTimeoutDuration) .then(Mono.fromRunnable(() -> { - if (sessionState.isConnectionInitNotProcessed()) { + if (sessionState.setConnectionInitPayload(Collections.emptyMap())) { GraphQlStatus.closeSession(session, GraphQlStatus.INIT_TIMEOUT_STATUS); } })) @@ -143,7 +145,7 @@ public class GraphQlWebSocketHandler extends TextWebSocketHandler implements Sub SessionState sessionState = getSessionInfo(session); switch (message.getType()) { case "subscribe": - if (sessionState.isConnectionInitNotProcessed()) { + if (sessionState.getConnectionInitPayload() == null) { GraphQlStatus.closeSession(session, GraphQlStatus.UNAUTHORIZED_STATUS); return; } @@ -174,7 +176,7 @@ public class GraphQlWebSocketHandler extends TextWebSocketHandler implements Sub } return; case "connection_init": - if (sessionState.setConnectionInitProcessed()) { + if (!sessionState.setConnectionInitPayload(payload)) { GraphQlStatus.closeSession(session, GraphQlStatus.TOO_MANY_INIT_REQUESTS_STATUS); return; } @@ -274,9 +276,14 @@ public class GraphQlWebSocketHandler extends TextWebSocketHandler implements Sub @Override public void afterConnectionClosed(WebSocketSession session, CloseStatus closeStatus) { - SessionState info = this.sessionInfoMap.remove(session.getId()); + String id = session.getId(); + SessionState info = this.sessionInfoMap.remove(id); if (info != null) { info.dispose(); + Map connectionInitPayload = info.getConnectionInitPayload(); + if (connectionInitPayload != null) { + this.webSocketInterceptor.handleConnectionClosed(id, closeStatus.getCode(), connectionInitPayload); + } } } @@ -345,7 +352,7 @@ public class GraphQlWebSocketHandler extends TextWebSocketHandler implements Sub private static class SessionState { - private boolean connectionInitProcessed; + private final AtomicReference> connectionInitPayloadRef = new AtomicReference<>(); private final Map subscriptions = new ConcurrentHashMap<>(); @@ -355,16 +362,16 @@ public class GraphQlWebSocketHandler extends TextWebSocketHandler implements Sub this.scheduler = Schedulers.newSingle("GraphQL-WsSession-" + sessionId); } - boolean isConnectionInitNotProcessed() { - return !this.connectionInitProcessed; + @Nullable + Map getConnectionInitPayload() { + return this.connectionInitPayloadRef.get(); } - synchronized boolean setConnectionInitProcessed() { - boolean previousValue = this.connectionInitProcessed; - this.connectionInitProcessed = true; - return previousValue; + boolean setConnectionInitPayload(Map payload) { + return this.connectionInitPayloadRef.compareAndSet(null, payload); } + Map getSubscriptions() { return this.subscriptions; } diff --git a/spring-graphql/src/test/java/org/springframework/graphql/web/webflux/GraphQlWebSocketHandlerTests.java b/spring-graphql/src/test/java/org/springframework/graphql/web/webflux/GraphQlWebSocketHandlerTests.java index 02594e86..d51c5fad 100644 --- a/spring-graphql/src/test/java/org/springframework/graphql/web/webflux/GraphQlWebSocketHandlerTests.java +++ b/spring-graphql/src/test/java/org/springframework/graphql/web/webflux/GraphQlWebSocketHandlerTests.java @@ -22,6 +22,7 @@ import java.util.ArrayList; import java.util.Collections; import java.util.List; import java.util.Map; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.function.BiConsumer; import org.assertj.core.api.InstanceOfAssertFactories; @@ -155,6 +156,30 @@ public class GraphQlWebSocketHandlerTests extends WebSocketHandlerTestSupport { .verifyComplete(); } + @Test + void connectionClosedHandling() { + + CloseStatus closeStatus = CloseStatus.PROTOCOL_ERROR; + AtomicBoolean called = new AtomicBoolean(); + + TestWebSocketSession session = handle( + Flux.just(toWebSocketMessage("{\"type\":\"connection_init\",\"payload\":{\"key\":\"A\"}}")), + new WebSocketInterceptor() { + + @Override + public void handleConnectionClosed(String sessionId, int status, Map payload) { + called.set(true); + assertThat(sessionId).isEqualTo("1"); + assertThat(status).isEqualTo(closeStatus.getCode()); + assertThat(payload).hasSize(1).containsEntry("key", "A"); + } + }); + + StepVerifier.create(session.getOutput()).expectNextCount(1).verifyComplete(); + StepVerifier.create(session.close(closeStatus)).verifyComplete(); + assertThat(called).isTrue(); + } + @Test void connectionInitRejected() { TestWebSocketSession session = handle( diff --git a/spring-graphql/src/test/java/org/springframework/graphql/web/webmvc/GraphQlWebSocketHandlerTests.java b/spring-graphql/src/test/java/org/springframework/graphql/web/webmvc/GraphQlWebSocketHandlerTests.java index 7d7141df..3674adc1 100644 --- a/spring-graphql/src/test/java/org/springframework/graphql/web/webmvc/GraphQlWebSocketHandlerTests.java +++ b/spring-graphql/src/test/java/org/springframework/graphql/web/webmvc/GraphQlWebSocketHandlerTests.java @@ -24,6 +24,7 @@ import java.util.ArrayList; import java.util.Collections; import java.util.List; import java.util.Map; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.function.BiConsumer; import java.util.function.Consumer; @@ -160,6 +161,35 @@ public class GraphQlWebSocketHandlerTests extends WebSocketHandlerTestSupport { .verifyComplete(); } + @Test + void connectionClosedHandling() throws Exception { + + CloseStatus closeStatus = CloseStatus.PROTOCOL_ERROR; + AtomicBoolean called = new AtomicBoolean(); + + WebSocketInterceptor interceptor = new WebSocketInterceptor() { + + @Override + public void handleConnectionClosed(String sessionId, int status, Map payload) { + called.set(true); + assertThat(sessionId).isEqualTo("1"); + assertThat(status).isEqualTo(closeStatus.getCode()); + assertThat(payload).hasSize(1).containsEntry("key", "A"); + } + }; + + GraphQlWebSocketHandler handler = initWebSocketHandler(interceptor); + handle(handler, new TextMessage("{\"type\":\"connection_init\",\"payload\":{\"key\":\"A\"}}")); + + StepVerifier.create(session.getOutput()) + .expectNextCount(1) + .then(this.session::close) // Complete output Flux + .verifyComplete(); + + handler.afterConnectionClosed(this.session, closeStatus); + assertThat(called).isTrue(); + } + @Test void connectionInitRejected() throws Exception {