com.xzixi.websocket.interceptablewebsocket.extension.InterceptableStompSubProtocolHandler Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of interceptable-websocket Show documentation
支持消息拦截的websocket
The newest version!
package com.xzixi.websocket.interceptablewebsocket.extension;

import com.xzixi.websocket.interceptablewebsocket.interceptor.FromClientExecutionChain;
import com.xzixi.websocket.interceptablewebsocket.interceptor.FromClientInterceptor;
import com.xzixi.websocket.interceptablewebsocket.interceptor.ToClientExecutionChain;
import com.xzixi.websocket.interceptablewebsocket.interceptor.ToClientInterceptor;
import com.xzixi.websocket.interceptablewebsocket.util.MessageFromClient;
import org.apache.commons.lang3.StringUtils;
import org.springframework.messaging.Message;
import org.springframework.messaging.MessageChannel;
import org.springframework.messaging.simp.stomp.StompHeaderAccessor;
import org.springframework.util.Assert;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.messaging.StompSubProtocolHandler;

import java.security.Principal;
import java.util.ArrayList;
import java.util.List;
import java.util.Queue;
import java.util.concurrent.LinkedTransferQueue;

/**
 * 代替{@link StompSubProtocolHandler}
 * 增加了对拦截器的支持
 * @author 薛凌康
 */
public class InterceptableStompSubProtocolHandler extends StompSubProtocolHandler {

    private List fromClientInterceptors = new ArrayList<>();

    private List toClientInterceptors = new ArrayList<>();

    @Override
    public void handleMessageFromClient(WebSocketSession session, WebSocketMessage webSocketMessage, MessageChannel outputChannel) {
        FromClientExecutionChain chain = new FromClientExecutionChain(fromClientInterceptors);
        Principal principal = session.getPrincipal();
        MessageFromClient message = getMessageFromClient(webSocketMessage);
        if (chain.applyPreHandle(session, principal, message, outputChannel, this)) {
            super.handleMessageFromClient(session, webSocketMessage, outputChannel);
            chain.applyPostHandle(session, principal, message, outputChannel, this);
        }
    }

    @Override
    public void handleMessageToClient(WebSocketSession session, Message message) {
        ToClientExecutionChain chain = new ToClientExecutionChain(toClientInterceptors);
        Principal principal = session.getPrincipal();
        StompHeaderAccessor accessor = StompHeaderAccessor.wrap(message);
        Object payload = message.getPayload();
        if (chain.applyPreHandle(session, principal, accessor, payload, this)) {
            super.handleMessageToClient(session, message);
            chain.applyPostHandle(session, principal, accessor, payload, this);
        }
    }

    public void addFromClientInterceptor(FromClientInterceptor interceptor) {
        Assert.notNull(interceptor, "interceptor不能为null");
        this.fromClientInterceptors.add(interceptor);
    }

    public void addToClientInterceptor(ToClientInterceptor interceptor) {
        Assert.notNull(interceptor, "interceptor不能为null");
        this.toClientInterceptors.add(interceptor);
    }

    private static final String SEND = "SEND";
    private static final String ID = "id:";
    private static final String DESTINATION = "destination:";
    private static final String CONTENT_LENGTH = "content-length:";

    private MessageFromClient getMessageFromClient(WebSocketMessage webSocketMessage) {
        if (webSocketMessage instanceof TextMessage) {
            MessageFromClient message = new MessageFromClient();
            String payload = ((TextMessage) webSocketMessage).getPayload();
            String[] arr = payload.split("\n");
            Queue queue = new LinkedTransferQueue<>();
            for (String str: arr) {
                String strTrim = str.trim();
                if (StringUtils.isEmpty(strTrim)) {
                    continue;
                }
                queue.offer(strTrim);
            }
            String type = queue.poll();
            message.setType(type);
            int last = 0;
            if (SEND.equals(type)) {
                last = 1;
            }
            while (queue.size()>last) {
                String param = queue.poll();
                if (param.startsWith(ID)) {
                    message.setSubId(param.split(":")[1].trim());
                } else if (param.startsWith(DESTINATION)) {
                    message.setDestination(param.split(":")[1].trim());
                } else if (param.startsWith(CONTENT_LENGTH)) {
                    message.setContentLength(Integer.valueOf(param.split(":")[1].trim()));
                }
            }
            String content = queue.poll();
            if (!StringUtils.isEmpty(content)) {
                if (content.startsWith(DESTINATION) && StringUtils.isEmpty(message.getDestination())) {
                    message.setDestination(content.split(":")[1].trim());
                } else {
                    message.setContent(content.trim());
                }
            }
            return message;
        }
        return null;
    }

}