dev.snowdrop.vertx.http.common.VertxWebSocketSession Maven / Gradle / Ivy

Go to download
package dev.snowdrop.vertx.http.common;

import dev.snowdrop.vertx.http.utils.BufferConverter;
import io.vertx.core.buffer.Buffer;
import io.vertx.core.http.WebSocketBase;
import org.reactivestreams.Publisher;
import org.reactivestreams.Subscriber;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.util.ObjectUtils;
import org.springframework.web.reactive.socket.CloseStatus;
import org.springframework.web.reactive.socket.HandshakeInfo;
import org.springframework.web.reactive.socket.WebSocketMessage;
import org.springframework.web.reactive.socket.adapter.AbstractWebSocketSession;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

public class VertxWebSocketSession extends AbstractWebSocketSession {

    private final Logger logger = LoggerFactory.getLogger(getClass());

    private final BufferConverter bufferConverter;

    private final int requestLimit;

    public VertxWebSocketSession(WebSocketBase delegate, HandshakeInfo handshakeInfo, BufferConverter bufferConverter,
        int maxWebSocketFrameSize, int maxWebSocketMessageSize) {
        super(delegate, ObjectUtils.getIdentityHexString(delegate), handshakeInfo,
            bufferConverter.getDataBufferFactory());
        this.bufferConverter = bufferConverter;
        if (maxWebSocketMessageSize < 1 || maxWebSocketFrameSize < 1) {
            throw new IllegalArgumentException("Max web socket frame and message sizes cannot be less than 1");
        }
        this.requestLimit = maxWebSocketMessageSize / maxWebSocketFrameSize + 1;
    }

    @Override
    public Flux receive() {
        return Flux.create(sink -> {
                logger.debug("{}Connecting to a web socket read stream", getLogPrefix());
                WebSocketBase socket = getDelegate();
                socket.pause()
                    .textMessageHandler(payload -> {
                        logger.debug("{}Received text '{}' from a web socket read stream", getLogPrefix(), payload);
                        sink.next(textMessage(payload));
                    })
                    .binaryMessageHandler(payload -> {
                        logger.debug("{}Received binary '{}' from a web socket read stream", getLogPrefix(), payload);
                        sink.next(binaryMessage(payload));
                    })
                    .pongHandler(payload -> {
                        logger.debug("{}Received pong '{}' from a web socket read stream", getLogPrefix(), payload);
                        sink.next(pongMessage(payload));
                    })
                    .exceptionHandler(throwable -> {
                        logger.debug("{}Received exception '{}' from a web socket read stream", getLogPrefix(), throwable);
                        sink.error(throwable);
                    })
                    .endHandler(e -> {
                        logger.debug("{}Web socket read stream ended", getLogPrefix());
                        sink.complete();
                    });
                sink.onRequest(i -> {
                    logger.debug("{}Fetching '{}' entries from a web socket read stream", getLogPrefix(), i);
                    socket.fetch(i);
                });
            }
        );
    }

    @Override
    public Mono send(Publisher messages) {
        return Mono.create(sink -> {
            logger.debug("{}Subscribing to messages publisher", getLogPrefix());
            Subscriber subscriber =
                new WriteStreamSubscriber.Builder()
                    .writeStream(getDelegate())
                    .nextHandler(this::messageHandler)
                    .endHook(sink)
                    .requestLimit(requestLimit)
                    .build();
            messages.subscribe(subscriber);
        });
    }

    @Override
    public boolean isOpen() {
        return !getDelegate().isClosed();
    }

    @Override
    public Mono close(CloseStatus status) {
        logger.debug("{}Closing web socket with status '{}'", getLogPrefix(), status);
        return Mono.create(sink -> getDelegate()
            .closeHandler(e -> {
                logger.debug("{}Web socket closed", getLogPrefix());
                sink.success();
            })
            .close((short) status.getCode(), status.getReason()));
    }

    @Override
    public Mono closeStatus() {
        Short code = getDelegate().closeStatusCode();
        if (code == null) {
            return Mono.empty();
        }

        String reason = getDelegate().closeReason();
        if (reason == null) {
            return Mono.just(new CloseStatus(code));
        }

        return Mono.just(new CloseStatus(code, reason));
    }

    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);
            }
        }
    }

    private WebSocketMessage binaryMessage(Buffer payloadBuffer) {
        DataBuffer payload = bufferConverter.toDataBuffer(payloadBuffer);
        return new WebSocketMessage(WebSocketMessage.Type.BINARY, payload);
    }

    private WebSocketMessage pongMessage(Buffer payloadBuffer) {
        DataBuffer payload = bufferConverter.toDataBuffer(payloadBuffer);
        return new WebSocketMessage(WebSocketMessage.Type.PONG, payload);
    }
}