Spring/spring-graphql
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Add handleConnectionClosed to WebSocketInterceptor

Browse Source 
Closes gh-276
This commit is contained in:
rstoyanchev
2022-03-08 08:49:46 +00:00
parent 35359b0b24
commit ec00547130
5 changed files with 103 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
spring-graphql/src/main/java/org/springframework/graphql/web/WebSocketInterceptor.java
View File

@@ -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<Object> handleConnectionInitialization(String sessionId, Map<String, Object> connectionInitPayload) {
default Mono<Object> handleConnectionInitialization(
String sessionId, Map<String, Object> 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<String, Object> connectionInitPayload) {
}
}

21
spring-graphql/src/main/java/org/springframework/graphql/web/webflux/GraphQlWebSocketHandler.java
View File

@@ -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<Map<String, Object>> connectionInitPayloadRef = new AtomicReference<>();
Map<String, Subscription> 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<String, Object> 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<String, Object> 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)

29
spring-graphql/src/main/java/org/springframework/graphql/web/webmvc/GraphQlWebSocketHandler.java
View File

@@ -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<String, Object> 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<Map<String, Object>> connectionInitPayloadRef = new AtomicReference<>();
private final Map<String, Subscription> 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<String, Object> getConnectionInitPayload() {
return this.connectionInitPayloadRef.get();
}
synchronized boolean setConnectionInitProcessed() {
boolean previousValue = this.connectionInitProcessed;
this.connectionInitProcessed = true;
return previousValue;
boolean setConnectionInitPayload(Map<String, Object> payload) {
return this.connectionInitPayloadRef.compareAndSet(null, payload);
}
Map<String, Subscription> getSubscriptions() {
return this.subscriptions;
}

25
spring-graphql/src/test/java/org/springframework/graphql/web/webflux/GraphQlWebSocketHandlerTests.java
View File

@@ -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<String, Object> 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(

30
spring-graphql/src/test/java/org/springframework/graphql/web/webmvc/GraphQlWebSocketHandlerTests.java
View File

@@ -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<String, Object> 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 {