Repository: vert-x3/vertx-web
Issue: ApolloWS handler does not support messages sent in multiple frames

GraphQL Subscriptions websocket uses a handler instead of textMessageHandler resulting in lost messages since handler will be called for all binary / text / continuation frames.

In our case we should add a textMessageHandler since we want to concatenate messages when a continuation frame is received

https://github.com/vert-x3/vertx-web/blob/13c10f88ec3b20e53950fd63ed740c9dc317f65a/vertx-web-graphql/src/main/java/io/vertx/ext/web/handler/graphql/impl/ApolloWSConnectionHandler.java#L71

Changed files:
vertx-web-graphql/src/test/java/io/vertx/ext/web/handler/graphql/ApolloWSHandlerTest.java
vertx-web-graphql/src/main/java/io/vertx/ext/web/handler/graphql/impl/ApolloWSConnectionHandler.java

diff --git a/vertx-web-graphql/src/main/java/io/vertx/ext/web/handler/graphql/impl/ApolloWSConnectionHandler.java b/vertx-web-graphql/src/main/java/io/vertx/ext/web/handler/graphql/impl/ApolloWSConnectionHandler.java
index abdd7e4ed..7394d8f96 100644
--- a/vertx-web-graphql/src/main/java/io/vertx/ext/web/handler/graphql/impl/ApolloWSConnectionHandler.java
+++ b/vertx-web-graphql/src/main/java/io/vertx/ext/web/handler/graphql/impl/ApolloWSConnectionHandler.java
@@ -19,6 +19,7 @@ package io.vertx.ext.web.handler.graphql.impl;
 import graphql.ExecutionInput;
 import graphql.ExecutionResult;
 import io.vertx.core.Handler;
+import io.vertx.core.buffer.Buffer;
 import io.vertx.core.http.ServerWebSocket;
 import io.vertx.core.impl.ContextInternal;
 import io.vertx.core.impl.logging.Logger;
@@ -68,50 +69,52 @@ class ApolloWSConnectionHandler {
       ch.handle(serverWebSocket);
     }
 
-    serverWebSocket.handler(buffer -> {
-      JsonObject content = buffer.toJsonObject();
-      String opId = content.getString("id");
-      ApolloWSMessageType type = from(content.getString("type"));
+    serverWebSocket.binaryMessageHandler(this::handleBinaryMessage);
+    serverWebSocket.textMessageHandler(this::handleTextMessage);
+    serverWebSocket.closeHandler(this::close);
+  }
 
-      if (type == null) {
-        sendMessage(opId, ERROR, "Unknown message type: " + content.getString("type"));
-        return;
-      }
+  private void handleBinaryMessage(Buffer buffer) {
+    handleMessage(new JsonObject(buffer));
+  }
 
-      ApolloWSMessage message = new ApolloWSMessageImpl(serverWebSocket, type, content);
+  private void handleTextMessage(String text) {
+    handleMessage(new JsonObject(text));
+  }
 
-      Handler<ApolloWSMessage> mh = apolloWSHandler.getMessageHandler();
-      if (mh != null) {
-        mh.handle(message);
-      }
+  private void handleMessage(JsonObject jsonObject) {
+    String opId = jsonObject.getString("id");
+    ApolloWSMessageType type = from(jsonObject.getString("type"));
 
-      switch (type) {
-        case CONNECTION_INIT:
-          connect();
-          break;
-        case CONNECTION_TERMINATE:
-          serverWebSocket.close();
-          break;
-        case START:
-          start(message);
-          break;
-        case STOP:
-          stop(opId);
-          break;
-        default:
-          sendMessage(opId, ERROR, "Unsupported message type: " + type);
-          break;
-      }
-    });
+    if (type == null) {
+      sendMessage(opId, ERROR, "Unknown message type: " + jsonObject.getString("type"));
+      return;
+    }
 
-    serverWebSocket.closeHandler(v -> {
-      subscriptions.values().forEach(Subscription::cancel);
+    ApolloWSMessage message = new ApolloWSMessageImpl(serverWebSocket, type, jsonObject);
 
-      Handler<ServerWebSocket> eh = apolloWSHandler.getEndHandler();
-      if (eh != null) {
-        eh.handle(serverWebSocket);
-      }
-    });
+    Handler<ApolloWSMessage> mh = apolloWSHandler.getMessageHandler();
+    if (mh != null) {
+      mh.handle(message);
+    }
+
+    switch (type) {
+      case CONNECTION_INIT:
+        connect();
+        break;
+      case CONNECTION_TERMINATE:
+        serverWebSocket.close();
+        break;
+      case START:
+        start(message);
+        break;
+      case STOP:
+        stop(opId);
+        break;
+      default:
+        sendMessage(opId, ERROR, "Unsupported message type: " + type);
+        break;
+    }
   }
 
   private void connect() {
@@ -250,4 +253,13 @@ class ApolloWSConnectionHandler {
     }
     serverWebSocket.writeTextMessage(message.toString());
   }
+
+  private void close(Void v) {
+    subscriptions.values().forEach(Subscription::cancel);
+
+    Handler<ServerWebSocket> eh = apolloWSHandler.getEndHandler();
+    if (eh != null) {
+      eh.handle(serverWebSocket);
+    }
+  }
 }
diff --git a/vertx-web-graphql/src/test/java/io/vertx/ext/web/handler/graphql/ApolloWSHandlerTest.java b/vertx-web-graphql/src/test/java/io/vertx/ext/web/handler/graphql/ApolloWSHandlerTest.java
index d282ab59e..d09c4789d 100644
--- a/vertx-web-graphql/src/test/java/io/vertx/ext/web/handler/graphql/ApolloWSHandlerTest.java
+++ b/vertx-web-graphql/src/test/java/io/vertx/ext/web/handler/graphql/ApolloWSHandlerTest.java
@@ -25,7 +25,10 @@ import graphql.schema.idl.SchemaParser;
 import graphql.schema.idl.TypeDefinitionRegistry;
 import io.vertx.core.AsyncResult;
 import io.vertx.core.Handler;
+import io.vertx.core.buffer.Buffer;
 import io.vertx.core.http.HttpClientOptions;
+import io.vertx.core.http.WebSocket;
+import io.vertx.core.http.WebSocketFrame;
 import io.vertx.core.json.JsonObject;
 import io.vertx.core.net.NetClientOptions;
 import io.vertx.core.net.NetServer;
@@ -40,6 +43,7 @@ import java.util.Map;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.BiConsumer;
 import java.util.stream.IntStream;
 
 import static graphql.schema.idl.RuntimeWiring.newRuntimeWiring;
@@ -132,8 +136,8 @@ public class ApolloWSHandlerTest extends WebTestBase {
 
       AtomicReference<String> id = new AtomicReference<>();
       AtomicInteger counter = new AtomicInteger();
-      websocket.handler(buffer -> {
-        JsonObject obj = buffer.toJsonObject();
+      websocket.textMessageHandler(text -> {
+        JsonObject obj = new JsonObject(text);
         int current = counter.getAndIncrement();
         if (current >= 0 && current <= MAX_COUNT) {
           if (current == 0) {
@@ -166,14 +170,30 @@ public class ApolloWSHandlerTest extends WebTestBase {
 
   @Test
   public void testQueryWsCall() {
+    testQueryWsCall((webSocket, message) -> webSocket.write(message.toBuffer()));
+  }
+
+  @Test
+  public void testQueryWsCallMultipleFrames() {
+    testQueryWsCall((webSocket, message) -> {
+      Buffer buffer = message.toBuffer();
+      int part = buffer.length() / 3;
+      if (part == 0) fail("Cannot perform test");
+      webSocket.writeFrame(WebSocketFrame.binaryFrame(buffer.getBuffer(0, part), false));
+      webSocket.writeFrame(WebSocketFrame.continuationFrame(buffer.getBuffer(part, 2 * part), false));
+      webSocket.writeFrame(WebSocketFrame.continuationFrame(buffer.getBuffer(2 * part, buffer.length()), true));
+    });
+  }
+
+  private void testQueryWsCall(BiConsumer<WebSocket, JsonObject> sender) {
     waitFor(2);
     client.webSocket("/graphql", onSuccess(websocket -> {
       websocket.exceptionHandler(this::fail);
 
       AtomicReference<String> id = new AtomicReference<>();
       AtomicInteger counter = new AtomicInteger();
-      websocket.handler(buffer -> {
-        JsonObject obj = buffer.toJsonObject();
+      websocket.textMessageHandler(text -> {
+        JsonObject obj = new JsonObject(text);
         int current = counter.getAndIncrement();
         if (current == 0) {
           assertTrue(id.compareAndSet(null, obj.getString("id")));
@@ -195,7 +215,7 @@ public class ApolloWSHandlerTest extends WebTestBase {
           .put("query", "query Query { staticCounter { count } }"))
         .put("type", "start")
         .put("id", "1");
-      websocket.write(message.toBuffer());
+      sender.accept(websocket, message);
     }));
     await();
   }
@@ -224,9 +244,9 @@ public class ApolloWSHandlerTest extends WebTestBase {
       websocket.exceptionHandler(this::fail);
 
       AtomicInteger counter = new AtomicInteger(0);
-      websocket.handler(buffer -> {
+      websocket.textMessageHandler(text -> {
         try {
-          JsonObject obj = buffer.toJsonObject();
+          JsonObject obj = new JsonObject(text);
 
           if (counter.getAndIncrement() == 0) {
             assertEquals(ApolloWSMessageType.CONNECTION_ACK.getText(), obj.getString("type"));
@@ -261,7 +281,7 @@ public class ApolloWSHandlerTest extends WebTestBase {
       websocket.exceptionHandler(this::fail);
 
       AtomicInteger counter = new AtomicInteger();
-      websocket.handler(buffer -> {
+      websocket.textMessageHandler(text -> {
         if (counter.getAndIncrement() == MAX_COUNT) {
           if (subscriptionRef.get() == null) {
             fail("Expected a live subscription");