io.vertx.core.http.WebSocket Java Examples

/**
 * This does not set up a handler on the websocket
 */
private WebSocket setupRawWebsocketClient(String serverPath)
  throws InterruptedException {
  String requestURI = serverPath + "/websocket";

  AtomicReference<WebSocket> openedWebSocketReference = new AtomicReference<>();
  CountDownLatch openSocketCountDown = new CountDownLatch(1);
  client.webSocket(requestURI, onSuccess(ws -> {
    openedWebSocketReference.set(ws);
    openSocketCountDown.countDown();
    ws.endHandler(v -> testComplete());
    ws.exceptionHandler(this::fail);
  }));

  openSocketCountDown.await(5, TimeUnit.SECONDS);
  return openedWebSocketReference.get();
}
 

private void report(WebSocket ws) {
  BlockStats head = newHead.getAndSet(null);
  if (head != null) {
    writeCommand(ws, "block", "block", head);
  }
  Integer count = newTxCount.getAndSet(null);
  if (count != null) {
    writeCommand(ws, "pending", "stats", Collections.singletonMap("pending", count));
  }
  NodeStats nodeStats = newNodeStats.getAndSet(null);
  if (nodeStats != null) {
    writeCommand(ws, "stats", "stats", nodeStats);
  }
  List<BlockStats> newBlocks = newHistory.getAndSet(null);
  if (newBlocks != null && !newBlocks.isEmpty()) {
    writeCommand(ws, "history", "history", newBlocks);
  }
}
 

/**
 * This sets up a handler on the websocket
 */
private WebSocket setupSockJsClient(String serverPath, List<Buffer> receivedMessagesCollector)
  throws InterruptedException {
  String requestURI = serverPath + "/000/000/websocket";

  AtomicReference<WebSocket> openedWebSocketReference = new AtomicReference<>();
  CountDownLatch openSocketCountDown = new CountDownLatch(1);
  client.webSocket(requestURI, onSuccess(ws -> {
    openedWebSocketReference.set(ws);
    ws.handler(replyBuffer -> {
      log.debug("Client received " + replyBuffer);
      String textReply = replyBuffer.toString();
      if ("o".equals(textReply)) {
        openSocketCountDown.countDown();
      } else {
        receivedMessagesCollector.add(replyBuffer);
      }
    });
    ws.endHandler(v -> testComplete());
    ws.exceptionHandler(this::fail);
  }));

  openSocketCountDown.await(5, TimeUnit.SECONDS);
  return openedWebSocketReference.get();
}
 

@Test
public void testTextFrameSockJs() throws InterruptedException {
  String serverPath = "/text-sockjs";
  setupSockJsServer(serverPath, this::echoRequest);

  List<Buffer> receivedMessages = new ArrayList<>();
  WebSocket openedWebSocket = setupSockJsClient(serverPath, receivedMessages);
  String messageToSend = "[\"testMessage\"]";
  openedWebSocket.writeFrame(WebSocketFrame.textFrame(messageToSend, true));

  await(5, TimeUnit.SECONDS);

  assertEquals("Client should have received 2 messages: the reply and the close.", 2, receivedMessages.size());
  Buffer expectedReply = Buffer.buffer("a" + messageToSend);
  assertEquals("Client reply should have matched request", expectedReply, receivedMessages.get(0));
  assertEquals("Final message should have been a close", SOCKJS_CLOSE_REPLY, receivedMessages.get(1));
}
 

@Test
public void testTextFrameRawWebSocket() throws InterruptedException {
  String serverPath = "/textecho";
  setupSockJsServer(serverPath, this::echoRequest);

  String message = "hello";
  AtomicReference<String> receivedReply = new AtomicReference<>();
  WebSocket ws = setupRawWebsocketClient(serverPath);

  ws.handler(replyBuffer -> receivedReply.set(replyBuffer.toString()));

  ws.writeFrame(WebSocketFrame.textFrame(message, true));

  await(5, TimeUnit.SECONDS);

  assertEquals("Client reply should have matched request", message, receivedReply.get());
}
 

private void handleEmitEvent(ArrayNode event, WebSocket ws) {
  String command = event.get(0).textValue();
  switch (command) {
    case "node-pong":
      logger.debug("Received a pong {}", event.get(1));
      if (!waitingOnPong.compareAndSet(true, false)) {
        logger.warn("Received pong when we didn't expect one");
      } else {
        long start = event.get(1).get("clientTime").longValue();
        long latency = (Instant.now().toEpochMilli() - start) / (2 * 1000);
        writeCommand(ws, "latency", "latency", latency);
      }
      break;
    case "history":
      logger.debug("History request {}", event.get(1));
      requestHistory(event.get(1));
      break;
    default:
      logger.warn("Unexpected message {}", command);

  }
}
 

@Test
public void testConnection() {
  AtomicReference<Throwable> error = new AtomicReference<>();
  AtomicReference<Buffer> frame = new AtomicReference<>();
  AtomicReference<WebSocket> socket = new AtomicReference<>();

  vertx.createHttpClient().webSocket(options, ar -> {
    if (ar.succeeded()) {
      WebSocket ws = ar.result();
      socket.set(ws);
      ws.exceptionHandler(error::set)
        .handler(frame::set)
        .write(new Frame(Frame.Command.CONNECT, Headers.create("accept-version", "1.2,1.1,1.0",
          "heart-beat", "10000,10000"), null).toBuffer());
    }
  });

  await().atMost(10, TimeUnit.SECONDS).until(() -> error.get() == null && frame.get() != null);
  assertThat(frame.get().toString()).startsWith("CONNECTED")
      .contains("server:vertx-stomp", "heart-beat:", "session:", "version:1.2");
  socket.get().close();
}
 

@Test
public void testSendingAMessage() {
  AtomicReference<Throwable> error = new AtomicReference<>();
  AtomicReference<Frame> frame = new AtomicReference<>();
  AtomicReference<WebSocket> socket = new AtomicReference<>();

  AtomicReference<StompClientConnection> client = new AtomicReference<>();

  clients.add(StompClient.create(vertx).connect(61613, "localhost", connection -> {
    connection.result().subscribe("foo", frame::set, r -> {
      client.set(connection.result());
    });
  }));

  await().atMost(10, TimeUnit.SECONDS).until(() -> client.get() != null);
  await().atMost(10, TimeUnit.SECONDS).until(() -> server.stompHandler().getDestination("foo") != null);


  vertx.createHttpClient().webSocket(options, ar -> {
    if (ar.succeeded()) {
      WebSocket ws = ar.result();
      socket.set(ws);
      ws.exceptionHandler(error::set)
        .handler(buffer -> {
          if (buffer.toString().startsWith("CONNECTED")) {
            ws.write(
              new Frame(Frame.Command.SEND, Headers.create("header", "value", "destination", "foo"), Buffer
                .buffer("hello")).toBuffer());
          }
        })
        .write(new Frame(Frame.Command.CONNECT, Headers.create("accept-version", "1.2,1.1,1.0",
          "heart-beat", "10000,10000"), null).toBuffer());
    }
  });

  await().atMost(10, TimeUnit.SECONDS).until(() -> error.get() == null && frame.get() != null);
  assertThat(frame.get().toString()).startsWith("MESSAGE")
      .contains("destination:foo", "header:value", "\nhello");
  socket.get().close();
}
 

@Override
public void handle(WebSocket webSocket) {
    connecting = false;
    connected = true;
    logger.info("Bridge Client - connected to server [" + remoteBridgeHost + ":" + remoteBridgePort + "]");

    webSocket.write(Buffer.buffer( tenant + "\n" ));
    webSocket.write(Buffer.buffer( "START SESSION" + "\n" ));
    webSocket.pause();

    final EventBusWebsocketBridge ebnb = new EventBusWebsocketBridge(webSocket, vertx.eventBus(), address);
    webSocket.closeHandler(aVoid -> {
        logger.error("Bridge Client - closed connection from server [" + remoteBridgeHost + ":" + remoteBridgePort + "]" + webSocket.textHandlerID());
        ebnb.stop();
        connected = false;
    });
    webSocket.exceptionHandler(throwable -> {
        logger.error("Bridge Client - Exception: " + throwable.getMessage(), throwable);
        ebnb.stop();
        connected = false;
    });
    ebnb.setTenant(tenant);
    ebnb.start();

    logger.info("Bridge Client - bridgeUUID: "+ ebnb.getBridgeUUID());
    webSocket.resume();
}
 

@Test
public void testSplitLargeReplySockJs() throws InterruptedException {
  String serverPath = "/large-reply-sockjs";

  String largeMessage = TestUtils.randomAlphaString(65536 * 2);
  Buffer largeReplyBuffer = Buffer.buffer(largeMessage);

  setupSockJsServer(serverPath, (sock, requestBuffer) -> {
    sock.write(largeReplyBuffer);
    sock.close();
  });

  List<Buffer> receivedMessages = new ArrayList<>();
  WebSocket openedWebSocket = setupSockJsClient(serverPath, receivedMessages);

  String messageToSend = "[\"hello\"]";
  openedWebSocket.writeFrame(WebSocketFrame.textFrame(messageToSend, true));

  await(5, TimeUnit.SECONDS);

  int receivedReplyCount = receivedMessages.size();
  assertTrue("Should have received > 2 reply frame, actually received " + receivedReplyCount, receivedReplyCount > 2);

  Buffer expectedReplyBuffer = Buffer.buffer("a[\"").appendBuffer(largeReplyBuffer).appendBuffer(Buffer.buffer("\"]"));
  Buffer clientReplyBuffer = combineReplies(receivedMessages.subList(0, receivedMessages.size() - 1));
  assertEquals(String.format("Combined reply on client (length %s) should equal message from server (%s)",
    clientReplyBuffer.length(), expectedReplyBuffer.length()),
    expectedReplyBuffer, clientReplyBuffer);

  Buffer finalMessage = receivedMessages.get(receivedMessages.size() - 1);
  assertEquals("Final message should have been a close", SOCKJS_CLOSE_REPLY, finalMessage);
}
 

@Test
public void testCombineTextFrameSockJs() throws InterruptedException {
  String serverPath = "/text-combine-sockjs";
  setupSockJsServer(serverPath, this::echoRequest);

  List<Buffer> receivedMessages = new ArrayList<>();
  WebSocket openedWebSocket = setupSockJsClient(serverPath, receivedMessages);

  Buffer largeMessage = Buffer.buffer("[\"" + TestUtils.randomAlphaString(30) + "\"]");
  WebSocketFrame frame1 = new WebSocketFrameImpl(FrameType.TEXT, largeMessage.slice(0, 10).getByteBuf(), false);
  WebSocketFrame frame2 = WebSocketFrame.continuationFrame(largeMessage.slice(10, 20), false);
  WebSocketFrame frame3 = WebSocketFrame.continuationFrame(largeMessage.slice(20, largeMessage.length()), true);

  log.debug("Client sending " + frame1.textData());
  openedWebSocket.writeFrame(frame1);
  log.debug("Client sending " + frame2.textData());
  openedWebSocket.writeFrame(frame2);
  log.debug("Client sending " + frame3.textData());
  openedWebSocket.writeFrame(frame3);

  await(5, TimeUnit.SECONDS);

  assertEquals("Client should have received 2 messages: the reply and the close.", 2, receivedMessages.size());
  Buffer expectedReply = Buffer.buffer("a" + largeMessage.toString());
  assertEquals("Client reply should have matched request", expectedReply, receivedMessages.get(0));
  assertEquals("Final message should have been a close", SOCKJS_CLOSE_REPLY, receivedMessages.get(1));
}
 

@Test
public void testSplitLargeReplyRawWebSocket() throws InterruptedException {
  String serverPath = "/split";

  String largeReply = TestUtils.randomAlphaString(65536 * 5);
  Buffer largeReplyBuffer = Buffer.buffer(largeReply);

  setupSockJsServer(serverPath, (sock, requestBuffer) -> {
    sock.write(largeReplyBuffer);
    sock.close();
  });

  Buffer totalReplyBuffer = Buffer.buffer(largeReplyBuffer.length());
  AtomicInteger receivedReplies = new AtomicInteger(0);
  WebSocket ws = setupRawWebsocketClient(serverPath);
  ws.handler(replyBuffer -> {
    totalReplyBuffer.appendBuffer(replyBuffer);
    receivedReplies.incrementAndGet();

  });

  ws.writeFrame(WebSocketFrame.binaryFrame(Buffer.buffer("hello"), true));

  await(5, TimeUnit.SECONDS);

  int receivedReplyCount = receivedReplies.get();
  assertEquals("Combined reply on client should equal message from server", largeReplyBuffer, totalReplyBuffer);
  assertTrue("Should have received > 1 reply frame, actually received " + receivedReplyCount, receivedReplyCount > 1);
}
 

/**
 * Writing multiple continuation frames from the client side should result in a single message on the server side
 * after the frames are re-combined
 */
@Test
public void testCombineBinaryContinuationFramesRawWebSocket() throws InterruptedException {
  String serverPath = "/combine";

  AtomicReference<Buffer> serverReceivedMessage = new AtomicReference<>();
  setupSockJsServer(serverPath, (sock, requestBuffer) -> {
    serverReceivedMessage.set(requestBuffer);
    sock.write(Buffer.buffer("reply"));
    sock.close();
  });


  Buffer largeMessage = Buffer.buffer(TestUtils.randomAlphaString(30));
  WebSocketFrame frame1 = WebSocketFrame.binaryFrame(largeMessage.slice(0, 10), false);
  WebSocketFrame frame2 = WebSocketFrame.continuationFrame(largeMessage.slice(10, 20), false);
  WebSocketFrame frame3 = WebSocketFrame.continuationFrame(largeMessage.slice(20, largeMessage.length()), true);

  WebSocket ws = setupRawWebsocketClient(serverPath);
  ws.writeFrame(frame1);
  ws.writeFrame(frame2);
  ws.writeFrame(frame3);

  await(5, TimeUnit.SECONDS);

  assertEquals("Server did not combine continuation frames correctly", largeMessage, serverReceivedMessage.get());
}
 

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.textMessageHandler(text -> {
      JsonObject obj = new JsonObject(text);
      int current = counter.getAndIncrement();
      if (current == 0) {
        assertTrue(id.compareAndSet(null, obj.getString("id")));
        assertEquals(DATA, ApolloWSMessageType.from(obj.getString("type")));
        int val = obj.getJsonObject("payload").getJsonObject("data").getJsonObject("staticCounter").getInteger("count");
        assertEquals(STATIC_COUNT, val);
        complete();
      } else if (current == 1) {
        assertEquals(id.get(), obj.getString("id"));
        assertEquals(COMPLETE, ApolloWSMessageType.from(obj.getString("type")));
        complete();
      } else {
        fail();
      }
    });

    JsonObject message = new JsonObject()
      .put("payload", new JsonObject()
        .put("query", "query Query { staticCounter { count } }"))
      .put("type", "start")
      .put("id", "1");
    sender.accept(websocket, message);
  }));
  await();
}
 

protected static void sendTypeToBridge(WebSocket ws, String type, String address, String msg) {
  JsonObject json = new JsonObject();
  json.put("type", type);
  json.put("address", address);
  json.put("body", msg);
  ws.write(Buffer.buffer(json.toString()));
}
 

@Test
public void testPingFromServer() {
  AtomicReference<Throwable> error = new AtomicReference<>();
  AtomicReference<WebSocket> socket = new AtomicReference<>();
  AtomicReference<Boolean> flag = new AtomicReference<>();

  AtomicReference<StompClientConnection> client = new AtomicReference<>();

  clients.add(StompClient.create(vertx).connect(61613, "localhost", connection -> {
    client.set(connection.result());
  }));

  await().atMost(10, TimeUnit.SECONDS).until(() -> client.get() != null);

  vertx.createHttpClient().webSocket(options, ar -> {
    if (ar.succeeded()) {
      WebSocket ws = ar.result();
      socket.set(ws);
      ws.exceptionHandler(error::set)
        .handler(buffer -> {
          vertx.setTimer(1000, id -> {
            flag.set(true);
          });
        })
        .write(new Frame(Frame.Command.CONNECT, Headers.create("accept-version", "1.2,1.1,1.0",
          "heart-beat", "100,0"), null).toBuffer());
    }
  });

  await().atMost(10, TimeUnit.SECONDS).until(() -> error.get() == null && flag.get() != null);
  socket.get().close();
}
 

private void sendHeartbeat(WebSocket ws) {
  try {
    ws.writeFrame(new WebSocketFrameImpl(FrameType.PING));
    EventManager.post(new ConnSuccEvent());
  } catch (IllegalStateException e) {
    EventManager.post(new ConnFailEvent("heartbeat fail, " + e.getMessage()));
    LOGGER.error("heartbeat fail", e);
  }
}
 

private void reportPeriodically(WebSocket ws) {
  TimeoutStream reportingStream = vertx.periodicStream(REPORTING_PERIOD).handler(ev -> {
    report(ws);
  });
  TimeoutStream pingStream = vertx.periodicStream(PING_PERIOD).handler(ev -> {
    writePing(ws);
  });
  ws.closeHandler(h -> {
    reportingStream.cancel();
    pingStream.cancel();
    attemptConnect(null);
  });
}
 

private void startHeartBeatThread(WebSocket ws) {
  heartbeatTask = Executors.newScheduledThreadPool(1);
  heartbeatTask.scheduleWithFixedDelay(() -> sendHeartbeat(ws),
      HEARTBEAT_INTERVAL,
      HEARTBEAT_INTERVAL,
      TimeUnit.MILLISECONDS);
}
 

private void writeCommand(WebSocket ws, String command, Object payload) {
  try {
    String message =
        mapper.writer().writeValueAsString(Collections.singletonMap("emit", Arrays.asList(command, payload)));
    logger.debug("Sending {} message {}", command, message);
    ws.writeTextMessage(message);
  } catch (JsonProcessingException e) {
    throw new UncheckedIOException(e);
  }
}
 

@Test
public void webSocketDoesNotToHandlePingPayloadAsJsonRpcRequest(final TestContext context) {
  final Async async = context.async();

  httpClient.webSocket(
      "/",
      result -> {
        WebSocket websocket = result.result();

        websocket.handler(
            buffer -> {
              final String payload = buffer.toString();
              if (!payload.equals("foo")) {
                context.fail("Only expected PONG response with same payload as PING request");
              }
            });

        websocket.closeHandler(
            h -> {
              verifyNoInteractions(webSocketRequestHandlerSpy);
              async.complete();
            });

        websocket.writeFrame(WebSocketFrame.pingFrame(Buffer.buffer("foo")));
        websocket.close();
      });

  async.awaitSuccess(VERTX_AWAIT_TIMEOUT_MILLIS);
}
 

@BeforeEach
public void setUp() {
    httpClientMock = Mockito.mock(HttpClient.class);
    String url = "http://nem.com:3000/";
    namespaceRepository = Mockito.mock(NamespaceRepository.class);
    listener = new ListenerVertx(httpClientMock, url, namespaceRepository);
    jsonHelper = listener.getJsonHelper();
    webSocketMock = Mockito.mock(WebSocket.class);
}
 

@Before
@SuppressWarnings("unchecked")
public void setUp() {
    // Configure mock http client to invoke passed web socket handler
    given(mockHttpClient.websocket(anyInt(), any(), anyString(), any(), any(Handler.class), any(Handler.class)))
        .will(answer -> {
            Handler<WebSocket> handler = answer.getArgument(4);
            handler.handle(mockWebSocket);
            return mockHttpClient;
        });
    given(mockVertx.createHttpClient(any(HttpClientOptions.class))).willReturn(mockHttpClient);

    webSocketClient = new VertxWebSocketClient(mockVertx);
}
 

private void writePing(WebSocket ws) {
  waitingOnPong.set(true);
  writeCommand(ws, "node-ping", "clientTime", Instant.now().toEpochMilli());
}
 

@Override
public Object connected(DefaultClientEndpointMetric endpointMetric, DefaultHttpSocketMetric socketMetric,
    WebSocket webSocket) {
  return null;
}
 

private void writeCommand(WebSocket ws, String command, String key, Object payload) {
  Map<String, Object> body = new HashMap<>();
  body.put("id", id);
  body.put(key, payload);
  writeCommand(ws, command, body);
}
 

private void webSocketConnection(final WebSocket connection) {
  this.connection = connection;
}
 

private VertxWebSocketSession initSession(URI uri, WebSocket socket) {
    // Vert.x handshake doesn't return headers so passing an empty collection
    HandshakeInfo handshakeInfo = new HandshakeInfo(uri, new HttpHeaders(), Mono.empty(), socket.subProtocol());

    return new VertxWebSocketSession(socket, handshakeInfo, bufferConverter);
}
 

@Override
public String connected(WebSocket webSocket) {
  String remote = webSocket.remoteAddress().toString();
  wsConnections.get(local, remote).increment();
  return remote;
}
 

@Override
public WebSocketMetric connected(WebSocket webSocket) {
  return clientReporter.createWebSocketMetric();
}