org.springframework.web.reactive.socket.WebSocketMessage Java Examples

@Override
public Mono<Void> handle(WebSocketSession session) {
	// pass headers along so custom headers can be sent through
	return client.execute(url, this.headers, new WebSocketHandler() {
		@Override
		public Mono<Void> handle(WebSocketSession proxySession) {
			// Use retain() for Reactor Netty
			Mono<Void> proxySessionSend = proxySession
					.send(session.receive().doOnNext(WebSocketMessage::retain));
			// .log("proxySessionSend", Level.FINE);
			Mono<Void> serverSessionSend = session.send(
					proxySession.receive().doOnNext(WebSocketMessage::retain));
			// .log("sessionSend", Level.FINE);
			return Mono.zip(proxySessionSend, serverSessionSend).then();
		}

		/**
		 * Copy subProtocols so they are available downstream.
		 * @return
		 */
		@Override
		public List<String> getSubProtocols() {
			return ProxyWebSocketHandler.this.subProtocols;
		}
	});
}
 

@Test
public void testWebSocket() {
    Flux<String> originalMessages = Flux.just("first", "second");
    List<String> responseMessages = new CopyOnWriteArrayList<>();

    disposable = client.execute(serviceUri, session -> {
        // Convert strings to WebSocket messages and send them
        Mono<Void> outputMono = session.send(originalMessages.map(session::textMessage));

        Mono<Void> inputMono = session.receive() // Receive a messages stream
            .map(WebSocketMessage::getPayloadAsText) // Extract a payload from each message
            .doOnNext(responseMessages::add) // Store the payload to a collection
            .then();

        return outputMono.then(inputMono); // Start receiving messages after sending.
    }).subscribe(); // Subscribe to the socket. Original messages will be sent and then we'll start receiving responses.

    await()
        .atMost(2, SECONDS)
        .until(() -> responseMessages, contains("FIRST", "SECOND"));
}
 

public static void main(String[] args) throws InterruptedException {

    WebSocketClient client = new ReactorNettyWebSocketClient();
    client.execute(
      URI.create("ws://localhost:8080/event-emitter"), 
      session -> session.send(
        Mono.just(session.textMessage("event-spring-reactive-client-websocket")))
        .thenMany(session.receive()
          .map(WebSocketMessage::getPayloadAsText)
          .log())
        .then())
        .block(Duration.ofSeconds(10L));
}
 

@Override
protected boolean sendMessage(WebSocketMessage message) throws IOException {
	ByteBuffer buffer = message.getPayload().asByteBuffer();
	if (WebSocketMessage.Type.TEXT.equals(message.getType())) {
		getSendProcessor().setReadyToSend(false);
		String text = new String(buffer.array(), StandardCharsets.UTF_8);
		getDelegate().getRemote().sendString(text, new SendProcessorCallback());
	}
	else if (WebSocketMessage.Type.BINARY.equals(message.getType())) {
		getSendProcessor().setReadyToSend(false);
		getDelegate().getRemote().sendBytes(buffer, new SendProcessorCallback());
	}
	else if (WebSocketMessage.Type.PING.equals(message.getType())) {
		getDelegate().getRemote().sendPing(buffer);
	}
	else if (WebSocketMessage.Type.PONG.equals(message.getType())) {
		getDelegate().getRemote().sendPong(buffer);
	}
	else {
		throw new IllegalArgumentException("Unexpected message type: " + message.getType());
	}
	return true;
}
 

private void messageHandler(WebSocketBase socket, WebSocketMessage message) {
    if (message.getType() == WebSocketMessage.Type.TEXT) {
        String payload = message.getPayloadAsText();
        socket.writeTextMessage(payload);
    } else {
        Buffer buffer = bufferConverter.toBuffer(message.getPayload());

        if (message.getType() == WebSocketMessage.Type.PING) {
            socket.writePing(buffer);
        } else if (message.getType() == WebSocketMessage.Type.PONG) {
            socket.writePong(buffer);
        } else {
            socket.writeBinaryMessage(buffer);
        }
    }
}
 

@Test
public void shouldSendAndReceiveTextMessage() {
    startServerWithoutSecurity(Handlers.class);

    AtomicReference<String> expectedMessage = new AtomicReference<>();

    getWebSocketClient()
        .execute(URI.create(WS_BASE_URL + "/echo"), session -> {
            WebSocketMessage originalMessage = session.textMessage("ping");

            Mono<Void> outputMono = session.send(Mono.just(originalMessage));
            Mono<Void> inputMono = session.receive()
                .filter(message -> message.getType().equals(WebSocketMessage.Type.TEXT))
                .map(WebSocketMessage::getPayloadAsText)
                .doOnNext(expectedMessage::set)
                .then();

            return outputMono.then(inputMono);
        }).subscribe();

    await()
        .atMost(2, SECONDS)
        .untilAtomic(expectedMessage, equalTo("ping"));
}
 

@Override
protected boolean sendMessage(WebSocketMessage message) throws IOException {
	ByteBuffer buffer = message.getPayload().asByteBuffer();
	if (WebSocketMessage.Type.TEXT.equals(message.getType())) {
		getSendProcessor().setReadyToSend(false);
		String text = new String(buffer.array(), StandardCharsets.UTF_8);
		getDelegate().getAsyncRemote().sendText(text, new SendProcessorCallback());
	}
	else if (WebSocketMessage.Type.BINARY.equals(message.getType())) {
		getSendProcessor().setReadyToSend(false);
		getDelegate().getAsyncRemote().sendBinary(buffer, new SendProcessorCallback());
	}
	else if (WebSocketMessage.Type.PING.equals(message.getType())) {
		getDelegate().getAsyncRemote().sendPing(buffer);
	}
	else if (WebSocketMessage.Type.PONG.equals(message.getType())) {
		getDelegate().getAsyncRemote().sendPong(buffer);
	}
	else {
		throw new IllegalArgumentException("Unexpected message type: " + message.getType());
	}
	return true;
}
 

@Override
protected boolean sendMessage(WebSocketMessage message) throws IOException {
	ByteBuffer buffer = message.getPayload().asByteBuffer();
	if (WebSocketMessage.Type.TEXT.equals(message.getType())) {
		getSendProcessor().setReadyToSend(false);
		String text = new String(buffer.array(), StandardCharsets.UTF_8);
		getDelegate().getAsyncRemote().sendText(text, new SendProcessorCallback());
	}
	else if (WebSocketMessage.Type.BINARY.equals(message.getType())) {
		getSendProcessor().setReadyToSend(false);
		getDelegate().getAsyncRemote().sendBinary(buffer, new SendProcessorCallback());
	}
	else if (WebSocketMessage.Type.PING.equals(message.getType())) {
		getDelegate().getAsyncRemote().sendPing(buffer);
	}
	else if (WebSocketMessage.Type.PONG.equals(message.getType())) {
		getDelegate().getAsyncRemote().sendPong(buffer);
	}
	else {
		throw new IllegalArgumentException("Unexpected message type: " + message.getType());
	}
	return true;
}
 

protected WebSocketFrame toFrame(WebSocketMessage message) {
	ByteBuf byteBuf = NettyDataBufferFactory.toByteBuf(message.getPayload());
	if (WebSocketMessage.Type.TEXT.equals(message.getType())) {
		return new TextWebSocketFrame(byteBuf);
	}
	else if (WebSocketMessage.Type.BINARY.equals(message.getType())) {
		return new BinaryWebSocketFrame(byteBuf);
	}
	else if (WebSocketMessage.Type.PING.equals(message.getType())) {
		return new PingWebSocketFrame(byteBuf);
	}
	else if (WebSocketMessage.Type.PONG.equals(message.getType())) {
		return new PongWebSocketFrame(byteBuf);
	}
	else {
		throw new IllegalArgumentException("Unexpected message type: " + message.getType());
	}
}
 

@Test
public void shouldSendPingAndReceivePong() {
    startServerWithoutSecurity(Handlers.class);

    AtomicReference<String> expectedMessage = new AtomicReference<>();

    // Ping should be handled by a server, not a handler
    getWebSocketClient()
        .execute(URI.create(BASE_URL + "/sink"), session -> {
            WebSocketMessage originalMessage = session.pingMessage(factory -> factory.wrap("ping".getBytes()));

            Mono<Void> outputMono = session.send(Mono.just(originalMessage));
            Mono<Void> inputMono = session.receive()
                .filter(message -> message.getType().equals(WebSocketMessage.Type.PONG))
                .map(WebSocketMessage::getPayloadAsText)
                .doOnNext(expectedMessage::set)
                .then();

            return outputMono.then(inputMono);
        }).subscribe();

    await()
        .atMost(2, SECONDS)
        .untilAtomic(expectedMessage, equalTo("ping"));
}
 

@Test
public void echoForHttp() throws Exception {
	int count = 100;
	Flux<String> input = Flux.range(1, count).map(index -> "msg-" + index);
	ReplayProcessor<Object> output = ReplayProcessor.create(count);

	client.execute(getHttpUrl("/echoForHttp"), session -> {
		logger.debug("Starting to send messages");
		return session
				.send(input.doOnNext(s -> logger.debug("outbound " + s))
						.map(s -> session.textMessage(s)))
				.thenMany(session.receive().take(count)
						.map(WebSocketMessage::getPayloadAsText))
				.subscribeWith(output).doOnNext(s -> logger.debug("inbound " + s))
				.then().doOnSuccess(aVoid -> logger.debug("Done with success"))
				.doOnError(ex -> logger.debug(
						"Done with " + (ex != null ? ex.getMessage() : "error")));
	}).block(Duration.ofMillis(5000));

	assertThat(output.collectList().block(Duration.ofMillis(5000)))
			.isEqualTo(input.collectList().block(Duration.ofMillis(5000)));
}
 

@Test
public void serverShouldSendFromTwoPublishers() {
    startServerWithoutSecurity(Handlers.class);

    List<String> expectedMessages = new LinkedList<>();

    getWebSocketClient()
        .execute(URI.create(BASE_URL + "/double-producer"), session -> session.receive()
            .map(WebSocketMessage::getPayloadAsText)
            .doOnNext(expectedMessages::add)
            .then()
        ).subscribe();

    await()
        .atMost(2, SECONDS)
        .until(() -> expectedMessages, contains("ping", "pong"));
}
 

@Test
public void shouldSendTextMessage() {
    TestPublisher<WebSocketMessage> source = TestPublisher.create();
    Mono<Void> result = session.send(source);

    StepVerifier.create(result)
        .expectSubscription()
        .then(() -> source.assertMinRequested(1))
        .then(() -> source.next(getWebSocketMessage(WebSocketMessage.Type.TEXT, "test1")))
        .then(() -> source.assertMinRequested(1))
        .then(() -> source.next(getWebSocketMessage(WebSocketMessage.Type.TEXT, "test2")))
        .then(() -> source.assertMinRequested(1))
        .then(source::complete)
        .verifyComplete();

    verify(mockServerWebSocket).writeTextMessage("test1");
    verify(mockServerWebSocket).writeTextMessage("test2");
}
 

@Test
public void shouldSendBinaryMessage() {
    TestPublisher<WebSocketMessage> source = TestPublisher.create();
    Mono<Void> result = session.send(source);

    StepVerifier.create(result)
        .expectSubscription()
        .then(() -> source.assertMinRequested(1))
        .then(() -> source.next(getWebSocketMessage(WebSocketMessage.Type.BINARY, "test1")))
        .then(() -> source.assertMinRequested(1))
        .then(() -> source.next(getWebSocketMessage(WebSocketMessage.Type.BINARY, "test2")))
        .then(() -> source.assertMinRequested(1))
        .then(source::complete)
        .verifyComplete();

    verify(mockServerWebSocket).writeBinaryMessage(Buffer.buffer("test1"));
    verify(mockServerWebSocket).writeBinaryMessage(Buffer.buffer("test2"));
}
 

@Test
public void shouldSendPingMessage() {
    TestPublisher<WebSocketMessage> source = TestPublisher.create();
    Mono<Void> result = session.send(source);

    StepVerifier.create(result)
        .expectSubscription()
        .then(() -> source.assertMinRequested(1))
        .then(() -> source.next(getWebSocketMessage(WebSocketMessage.Type.PING, "test1")))
        .then(() -> source.assertMinRequested(1))
        .then(() -> source.next(getWebSocketMessage(WebSocketMessage.Type.PING, "test2")))
        .then(() -> source.assertMinRequested(1))
        .then(source::complete)
        .verifyComplete();

    verify(mockServerWebSocket).writePing(Buffer.buffer("test1"));
    verify(mockServerWebSocket).writePing(Buffer.buffer("test2"));
}
 

@Test
public void shouldSendPongMessage() {
    TestPublisher<WebSocketMessage> source = TestPublisher.create();
    Mono<Void> result = session.send(source);

    StepVerifier.create(result)
        .expectSubscription()
        .then(() -> source.assertMinRequested(1))
        .then(() -> source.next(getWebSocketMessage(WebSocketMessage.Type.PONG, "test1")))
        .then(() -> source.assertMinRequested(1))
        .then(() -> source.next(getWebSocketMessage(WebSocketMessage.Type.PONG, "test2")))
        .then(() -> source.assertMinRequested(1))
        .then(source::complete)
        .verifyComplete();

    verify(mockServerWebSocket).writePong(Buffer.buffer("test1"));
    verify(mockServerWebSocket).writePong(Buffer.buffer("test2"));
}
 

@NonNull
@Override
public Mono<Void> handle(@NonNull final WebSocketSession session) {
    // pass headers along so custom headers can be sent through
    return client.execute(url, this.headers, new WebSocketHandler() {
    
        @NonNull
        @Override
        public Mono<Void> handle(@NonNull final WebSocketSession webSocketSession) {
            // Use retain() for Reactor Netty
            Mono<Void> sessionSend = webSocketSession
                    .send(session.receive().doOnNext(WebSocketMessage::retain));
            Mono<Void> serverSessionSend = session.send(
                    webSocketSession.receive().doOnNext(WebSocketMessage::retain));
            return Mono.zip(sessionSend, serverSessionSend).then();
        }
    
        @NonNull
        @Override
        public List<String> getSubProtocols() {
            return SoulWebSocketHandler.this.subProtocols;
        }
    });
}
 

@Override
protected boolean sendMessage(WebSocketMessage message) throws IOException {
	ByteBuffer buffer = message.getPayload().asByteBuffer();
	if (WebSocketMessage.Type.TEXT.equals(message.getType())) {
		getSendProcessor().setReadyToSend(false);
		String text = new String(buffer.array(), StandardCharsets.UTF_8);
		WebSockets.sendText(text, getDelegate(), new SendProcessorCallback(message.getPayload()));
	}
	else if (WebSocketMessage.Type.BINARY.equals(message.getType())) {
		getSendProcessor().setReadyToSend(false);
		WebSockets.sendBinary(buffer, getDelegate(), new SendProcessorCallback(message.getPayload()));
	}
	else if (WebSocketMessage.Type.PING.equals(message.getType())) {
		getSendProcessor().setReadyToSend(false);
		WebSockets.sendPing(buffer, getDelegate(), new SendProcessorCallback(message.getPayload()));
	}
	else if (WebSocketMessage.Type.PONG.equals(message.getType())) {
		getSendProcessor().setReadyToSend(false);
		WebSockets.sendPong(buffer, getDelegate(), new SendProcessorCallback(message.getPayload()));
	}
	else {
		throw new IllegalArgumentException("Unexpected message type: " + message.getType());
	}
	return true;
}
 

protected WebSocketFrame toFrame(WebSocketMessage message) {
	ByteBuf byteBuf = NettyDataBufferFactory.toByteBuf(message.getPayload());
	if (WebSocketMessage.Type.TEXT.equals(message.getType())) {
		return new TextWebSocketFrame(byteBuf);
	}
	else if (WebSocketMessage.Type.BINARY.equals(message.getType())) {
		return new BinaryWebSocketFrame(byteBuf);
	}
	else if (WebSocketMessage.Type.PING.equals(message.getType())) {
		return new PingWebSocketFrame(byteBuf);
	}
	else if (WebSocketMessage.Type.PONG.equals(message.getType())) {
		return new PongWebSocketFrame(byteBuf);
	}
	else {
		throw new IllegalArgumentException("Unexpected message type: " + message.getType());
	}
}
 

public static void main(String[] args) throws URISyntaxException, InterruptedException {
    WebSocketClient client = new ReactorNettyWebSocketClient();

    URI url = new URI("ws://localhost:8080/time");
    final Mono<Void> reqResp = client
            .execute(url, session -> {
                Flux<WebSocketMessage> outMessages = Flux
                        .interval(Duration.ofSeconds(1))
                        .take(10)
                        .map(x -> "Message " + x)
                        .doOnNext(x -> log.info("About to send '{}'", x))
                        .map(session::textMessage);
                Mono<Void> receiving = session
                        .receive()
                        .map(WebSocketMessage::getPayloadAsText)
                        .doOnNext(x -> log.info("Received '{}'", x))
                        .then();
                return session.send(outMessages).mergeWith(receiving).then();
            });

    final Disposable disposable = reqResp.subscribe();
    TimeUnit.SECONDS.sleep(20);
    disposable.dispose();
}
 

@Override
public Mono<Void> handle(WebSocketSession session) {
	Flux<WebSocketMessage> outMessages = session
			.receive()
			.doOnSubscribe(s -> log.info("[{}] Got new connection", session.getId()))
			.map(WebSocketMessage::getPayloadAsText)
			.doOnNext(x -> log.info("[{}] Received: '{}'", session.getId(), x))
			.map(String::toUpperCase)
			.doOnNext(x -> log.info("[{}] Sending '{}'", session.getId(), x))
			.map(session::textMessage)
			.delaySequence(Duration.ofSeconds(1));
	return session
			.send(outMessages)
			.doOnSuccess(v -> log.info("[{}] Done, terminating the connection", session.getId()));
}
 

@Test
public void shouldReceiveBinaryMessages() {
    given(mockServerWebSocket.binaryMessageHandler(any())).will(invocation -> {
        Handler<Buffer> handler = invocation.getArgument(0);
        handler.handle(Buffer.buffer("test1"));
        handler.handle(Buffer.buffer("test2"));
        return mockServerWebSocket;
    });

    StepVerifier.create(session.receive())
        .expectNext(getWebSocketMessage(WebSocketMessage.Type.BINARY, "test1"))
        .expectNext(getWebSocketMessage(WebSocketMessage.Type.BINARY, "test2"))
        .verifyComplete();
}
 

@Test
public void shouldReceiveTextMessages() {
    given(mockServerWebSocket.textMessageHandler(any())).will(invocation -> {
        Handler<String> handler = invocation.getArgument(0);
        handler.handle("test1");
        handler.handle("test2");
        return mockServerWebSocket;
    });

    StepVerifier.create(session.receive())
        .expectNext(getWebSocketMessage(WebSocketMessage.Type.TEXT, "test1"))
        .expectNext(getWebSocketMessage(WebSocketMessage.Type.TEXT, "test2"))
        .verifyComplete();
}
 

@Test
public void shouldSendAndReceiveMultiFrameBinaryMessage() {
    Properties properties = new Properties();
    properties.setProperty("vertx.http.client.maxWebsocketFrameSize", "5");
    properties.setProperty("vertx.http.server.maxWebsocketFrameSize", "5");
    startServerWithoutSecurity(properties, Handlers.class);

    AtomicReference<String> expectedMessage = new AtomicReference<>();

    getWebSocketClient()
        .execute(URI.create(BASE_URL + "/echo"), session -> {
            WebSocketMessage originalMessage =
                session.binaryMessage(factory -> factory.wrap("ping pong".getBytes()));

            Mono<Void> outputMono = session.send(Mono.just(originalMessage));
            Mono<Void> inputMono = session.receive()
                .filter(message -> message.getType().equals(WebSocketMessage.Type.BINARY))
                .map(WebSocketMessage::getPayloadAsText)
                .doOnNext(expectedMessage::set)
                .then();

            return outputMono.then(inputMono);
        }).subscribe();

    await()
        .atMost(2, SECONDS)
        .untilAtomic(expectedMessage, equalTo("ping pong"));
}
 

@Override
protected Mono<Void> doHandle(WebSocketSession session) {
	//end::code[]
	return session.receive()
		.log(session.getId()
			+ "-inbound-incoming-chat-message")
		.map(WebSocketMessage::getPayloadAsText)
		.log(session.getId()
			+ "-inbound-convert-to-text")
		.flatMap(message -> broadcast(message, session))
		.log(session.getId()
			+ "-inbound-broadcast-to-broker")
		.then();
}
 

@Override                                                          // (2)
public Mono<Void> handle(WebSocketSession session) {               //
    return session                                                 // (3)
        .receive()                                                 // (4)
        .map(WebSocketMessage::getPayloadAsText)                   // (5)
        .map(tm -> "Echo: " + tm)                                  // (6)
        .map(session::textMessage)                                 // (7)
        .as(session::send);                                        // (8)
}
 

@Override
protected boolean write(WebSocketMessage message) throws IOException {
	if (logger.isTraceEnabled()) {
		logger.trace(getLogPrefix() + "Sending " + message);
	}
	else if (rsWriteLogger.isTraceEnabled()) {
		rsWriteLogger.trace(getLogPrefix() + "Sending " + message);
	}
	// In case of IOException, onError handling should call discardData(WebSocketMessage)..
	return sendMessage(message);
}
 

void handleMessage(WebSocketMessage message) {
	if (logger.isTraceEnabled()) {
		logger.trace(getLogPrefix() + "Received " + message);
	}
	else if (rsReadLogger.isTraceEnabled()) {
		rsReadLogger.trace(getLogPrefix() + "Received " + message);
	}
	if (!this.pendingMessages.offer(message)) {
		discardData();
		throw new IllegalStateException(
				"Too many messages. Please ensure WebSocketSession.receive() is subscribed to.");
	}
	onDataAvailable();
}
 

@Test
@Ignore
public void customHeader() throws Exception {
	HttpHeaders headers = new HttpHeaders();
	headers.add("my-header", "my-value");
	MonoProcessor<Object> output = MonoProcessor.create();

	client.execute(getUrl("/custom-header"), headers, session -> session.receive()
			.map(WebSocketMessage::getPayloadAsText).subscribeWith(output).then())
			.block(Duration.ofMillis(5000));

	assertThat(output.block(Duration.ofMillis(5000))).isEqualTo("my-header:my-value");
}
 

@Override
public Mono<Void> send(Publisher<WebSocketMessage> messages) {
	if (this.sendCalled.compareAndSet(false, true)) {
		WebSocketSendProcessor sendProcessor = new WebSocketSendProcessor();
		this.sendProcessor = sendProcessor;
		return Mono.from(subscriber -> {
				messages.subscribe(sendProcessor);
				sendProcessor.subscribe(subscriber);
		});
	}
	else {
		return Mono.error(new IllegalStateException("send() has already been called"));
	}
}