Complete fix for "connection_init" not having id

Closes gh-27
This commit is contained in:
Rossen Stoyanchev
2021-01-17 22:04:58 +00:00
parent edcd801c8c
commit bbb0d98511
2 changed files with 57 additions and 30 deletions

View File

@@ -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 <T> WebSocketMessage encode(
WebSocketSession session, String id, MessageType messageType, @Nullable Object payload) {
WebSocketSession session, @Nullable String id, MessageType messageType, @Nullable Object payload) {
Map<String, Object> 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);

View File

@@ -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<WebSocketMessage> 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<WebSocketMessage> 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<WebSocketMessage> 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<WebSocketMessage> 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<WebSocketMessage> 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<WebSocketMessage> input = Flux.just(toWebSocketMessage(
"{\"id\":\"" + SUBSCRIPTION_ID + "\",\"type\":\"connection_init\"}"),
Flux<WebSocketMessage> 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<WebSocketMessage> 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<WebInterceptor> 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<String, Object> map = decode(message);
assertThat(map).containsEntry("type", messageType);
if (!messageType.equals("connection_ack")) {
assertThat(map).containsEntry("id", SUBSCRIPTION_ID);
}
}