diff --git a/spring-graphql-web/src/main/java/org/springframework/graphql/webflux/GraphQLWebSocketHandler.java b/spring-graphql-web/src/main/java/org/springframework/graphql/webflux/GraphQLWebSocketHandler.java index ca89adaa..6abbb308 100644 --- a/spring-graphql-web/src/main/java/org/springframework/graphql/webflux/GraphQLWebSocketHandler.java +++ b/spring-graphql-web/src/main/java/org/springframework/graphql/webflux/GraphQLWebSocketHandler.java @@ -1,5 +1,5 @@ /* - * Copyright 2002-2020 the original author or authors. + * Copyright 2002-2021 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. @@ -15,6 +15,14 @@ */ package org.springframework.graphql.webflux; +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; @@ -23,6 +31,9 @@ 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.ParameterizedTypeReference; import org.springframework.core.ResolvableType; import org.springframework.core.codec.Decoder; @@ -46,16 +57,6 @@ import org.springframework.web.reactive.socket.HandshakeInfo; import org.springframework.web.reactive.socket.WebSocketHandler; import org.springframework.web.reactive.socket.WebSocketMessage; import org.springframework.web.reactive.socket.WebSocketSession; -import reactor.core.publisher.Flux; -import reactor.core.publisher.Mono; - -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; /** * WebSocketHandler for GraphQL based on @@ -144,6 +145,9 @@ public class GraphQLWebSocketHandler implements WebSocketHandler { if (!initialized.get()) { return GraphQLStatus.unauthorized(session); } + if (id == null) { + return GraphQLStatus.invalidMessage(session); + } HandshakeInfo handshakeInfo = session.getHandshakeInfo(); WebSocketInput input = new WebSocketInput(handshakeInfo, id, getPayload(map)); if (logger.isDebugEnabled()) { @@ -152,16 +156,18 @@ public class GraphQLWebSocketHandler implements WebSocketHandler { return executionChain.execute(input) .flatMapMany(output -> handleWebOutput(session, input.id(), output)); case COMPLETE: - Subscription subscription = this.subscriptions.remove(id); - if (subscription != null) { - subscription.cancel(); + if (id != null) { + 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)); + return Flux.just(encode(session, null, MessageType.CONNECTION_ACK, null)); default: return GraphQLStatus.invalidMessage(session); } @@ -232,10 +238,12 @@ public class GraphQLWebSocketHandler implements WebSocketHandler { @SuppressWarnings("unchecked") private WebSocketMessage encode( - WebSocketSession session, String id, MessageType messageType, @Nullable Object payload) { + WebSocketSession session, @Nullable String id, MessageType messageType, @Nullable Object payload) { Map payloadMap = new HashMap<>(3); - payloadMap.put("id", id); + if (id != null) { + payloadMap.put("id", id); + } payloadMap.put("type", messageType.getType()); if (payload != null) { payloadMap.put("payload", payload); diff --git a/spring-graphql-web/src/test/java/org/springframework/graphql/webflux/GraphQLWebSocketHandlerTests.java b/spring-graphql-web/src/test/java/org/springframework/graphql/webflux/GraphQLWebSocketHandlerTests.java index 7204317a..b760e13d 100644 --- a/spring-graphql-web/src/test/java/org/springframework/graphql/webflux/GraphQLWebSocketHandlerTests.java +++ b/spring-graphql-web/src/test/java/org/springframework/graphql/webflux/GraphQLWebSocketHandlerTests.java @@ -1,5 +1,5 @@ /* - * Copyright 2002-2020 the original author or authors. + * Copyright 2002-2021 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. @@ -44,7 +44,6 @@ import org.springframework.core.io.buffer.DefaultDataBufferFactory; import org.springframework.graphql.GraphQLDataFetchers; import org.springframework.graphql.WebInterceptor; import org.springframework.graphql.WebOutput; -import org.springframework.graphql.webflux.GraphQLWebSocketHandler; import org.springframework.http.HttpHeaders; import org.springframework.http.codec.ServerCodecConfigurer; import org.springframework.http.codec.json.Jackson2JsonDecoder; @@ -99,7 +98,7 @@ public class GraphQLWebSocketHandlerTests { "}"; Flux input = Flux.just( - toWebSocketMessage("{\"id\":\"" + SUBSCRIPTION_ID + "\",\"type\":\"connection_init\"}"), + toWebSocketMessage("{\"type\":\"connection_init\"}"), toWebSocketMessage(bookQuery)); TestWebSocketSession session = new TestWebSocketSession(input); @@ -123,7 +122,7 @@ public class GraphQLWebSocketHandlerTests { @Test void subscription() throws Exception { Flux input = Flux.just( - toWebSocketMessage("{\"id\":\"" + SUBSCRIPTION_ID + "\",\"type\":\"connection_init\"}"), + toWebSocketMessage("{\"type\":\"connection_init\"}"), toWebSocketMessage(BOOK_SEARCH_QUERY)); TestWebSocketSession session = new TestWebSocketSession(input); @@ -151,7 +150,7 @@ public class GraphQLWebSocketHandlerTests { @Test void unauthorizedWithoutMessageType() throws Exception { Flux input = Flux.just( - toWebSocketMessage("{\"id\":\"" + SUBSCRIPTION_ID + "\",\"type\":\"connection_init\"}"), + toWebSocketMessage("{\"type\":\"connection_init\"}"), toWebSocketMessage("{\"id\":\"" + SUBSCRIPTION_ID + "\"}")); // No message type TestWebSocketSession session = new TestWebSocketSession(input); @@ -166,6 +165,24 @@ public class GraphQLWebSocketHandlerTests { .verifyComplete(); } + @Test + void invalidMessageWithoutId() throws Exception { + Flux input = Flux.just( + toWebSocketMessage("{\"type\":\"connection_init\"}"), + toWebSocketMessage("{\"type\":\"subscribe\"}")); // No message id + + 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))); @@ -180,8 +197,8 @@ public class GraphQLWebSocketHandlerTests { @Test void tooManyConnectionInitRequests() throws Exception { Flux input = Flux.just( - toWebSocketMessage("{\"id\":\"" + SUBSCRIPTION_ID + "\",\"type\":\"connection_init\"}"), - toWebSocketMessage("{\"id\":\"" + SUBSCRIPTION_ID + "\",\"type\":\"connection_init\"}")); + toWebSocketMessage("{\"type\":\"connection_init\"}"), + toWebSocketMessage("{\"type\":\"connection_init\"}")); TestWebSocketSession session = new TestWebSocketSession(input); initWebSocketHandler().handle(session).block(); @@ -207,8 +224,8 @@ public class GraphQLWebSocketHandlerTests { @Test void subscriptionExists() throws Exception { - Flux input = Flux.just(toWebSocketMessage( - "{\"id\":\"" + SUBSCRIPTION_ID + "\",\"type\":\"connection_init\"}"), + Flux input = Flux.just( + toWebSocketMessage("{\"type\":\"connection_init\"}"), toWebSocketMessage(BOOK_SEARCH_QUERY), toWebSocketMessage(BOOK_SEARCH_QUERY)); @@ -233,7 +250,7 @@ public class GraphQLWebSocketHandlerTests { @Test void clientCompletion() throws Exception { Sinks.Many input = Sinks.many().unicast().onBackpressureBuffer(); - input.tryEmitNext(toWebSocketMessage("{\"id\":\"" + SUBSCRIPTION_ID + "\",\"type\":\"connection_init\"}")); + input.tryEmitNext(toWebSocketMessage("{\"type\":\"connection_init\"}")); input.tryEmitNext(toWebSocketMessage(BOOK_SEARCH_QUERY)); List interceptors = Collections.singletonList(new TakeOneAndNeverCompleteInterceptor()); @@ -293,9 +310,11 @@ public class GraphQLWebSocketHandlerTests { } private void assertMessageType(WebSocketMessage message, String messageType) { - assertThat(decode(message)) - .containsEntry("id", SUBSCRIPTION_ID) - .containsEntry("type", messageType); + Map map = decode(message); + assertThat(map).containsEntry("type", messageType); + if (!messageType.equals("connection_ack")) { + assertThat(map).containsEntry("id", SUBSCRIPTION_ID); + } }