com.github.netty.protocol.servlet.websocket.NettyMessageToWebSocketRunnable Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of spring-boot-protocol Show documentation
Spring Boot Protocol
The newest version!
package com.github.netty.protocol.servlet.websocket;

import com.github.netty.core.MessageToRunnable;
import com.github.netty.core.util.*;
import io.netty.channel.ChannelHandlerContext;
import io.netty.handler.codec.http.websocketx.*;

import javax.websocket.MessageHandler;
import javax.websocket.PongMessage;
import java.nio.ByteBuffer;
import java.util.Set;

/**
 * WebSocketMessageToRunnable
 * Life cycle connection
 *
 * @author wangzihao
 */
public class NettyMessageToWebSocketRunnable implements MessageToRunnable {
    private static final Recycler RECYCLER = new Recycler<>(WebsocketRunnable::new);
    private final MessageToRunnable parent;

    public NettyMessageToWebSocketRunnable(MessageToRunnable parent) {
        this.parent = parent;
    }

    @Override
    public Runnable onMessage(ChannelHandlerContext context, Object msg) {
        if (msg instanceof WebSocketFrame) {
            WebsocketRunnable task = RECYCLER.getInstance();
            task.context = context;
            task.frame = (WebSocketFrame) msg;
            return task;
        }
        if (parent != null) {
            return parent.onMessage(context, msg);
        }
        throw new IllegalStateException("[" + msg.getClass().getName() + "] Message data type that cannot be processed");
    }

    @Override
    public Runnable onClose(ChannelHandlerContext context) {
        if (parent != null) {
            return parent.onClose(context);
        }
        return null;
    }

    /**
     * Websocket task
     */
    public static class WebsocketRunnable implements Runnable, Recyclable {
        private ChannelHandlerContext context;
        private WebSocketFrame frame;

        public WebSocketSession getWebSocketSession() {
            return WebSocketSession.getSession(context.channel());
        }

        public WebSocketFrame getFrame() {
            return frame;
        }

        public void setFrame(WebSocketFrame frame) {
            this.frame = frame;
        }

        public ChannelHandlerContext getContext() {
            return context;
        }

        @Override
        public void run() {
            try {
                WebSocketSession wsSession = getWebSocketSession();
                if (wsSession == null) {
                    return;
                }

                // Close the message
                if (frame instanceof CloseWebSocketFrame) {
                    wsSession.closeByClient((CloseWebSocketFrame) frame);
                    return;
                }

                // Ping message
                if (frame instanceof PingWebSocketFrame) {
                    onWebsocketMessage(wsSession, frame, IOUtil.heap(frame.content()), PongMessage.class);
                    return;
                }

                // Binary message
                if (frame instanceof BinaryWebSocketFrame) {
                    onWebsocketMessage(wsSession, frame, IOUtil.heap(frame.content()), ByteBuffer.class);
                    return;
                }

                // String message
                if (frame instanceof TextWebSocketFrame) {
                    onWebsocketMessage(wsSession, frame, ((TextWebSocketFrame) frame).text(), String.class);
                }
            } finally {
                WebsocketRunnable.this.recycle();
            }
        }

        private void onWebsocketMessage(WebSocketSession wsSession, WebSocketFrame frame, Object message, Class messageType) {
            Set messageHandlers = wsSession.getMessageHandlers();
            for (MessageHandler handler : messageHandlers) {
                if (handler instanceof MessageHandler.Partial) {
                    MessageHandler.Partial