Spring Boot WebSocket介绍(二)之拦截器

旧城等待， 2022-03-27 08:44 2692阅读 0赞

介绍完[Spring Boot WebSocket介绍(一)][Spring Boot WebSocket]后，我们在看看更高级的WebSocket拦截器功能。 阅读本文前，请先阅读[Spring Boot WebSocket介绍(一)][Spring Boot WebSocket]。

**注意事项**：本文代码是为了介绍拦截过程，示例代码并非可直接用于生产的代码。仅供参看而已。

# 拦截器 #

顾名思义，可以拦截websocket请求，拦截后我们可以做各种行为的改变，本文以收到MESSAGE类型的消息后，对消息的内容进行改变为例。 首先我们假设发送到websocket的消息是json格式，然后我们加上，并获取‘’yqfield1’的内容。同时我们也可以针对websocket连接进行验证，比如不符合一定要求的connect就直接reject，不允许连接websockt

jsonObj.put(logFlag + "ChannelContent2", "add to");
        String value = jsonObj.getString("myfield1");
        jsonObj.put("myfield1", logFlag + " add to " + value);

# 代码 #

完整的代码在[这里][Link 1]，欢迎加星和fork。

package com.yq.config;
    
    import com.alibaba.fastjson.JSON;
    import com.alibaba.fastjson.JSONObject;
    import org.apache.commons.lang3.StringUtils;
    import org.slf4j.Logger;
    import org.slf4j.LoggerFactory;
    import org.springframework.beans.factory.annotation.Autowired;
    import org.springframework.context.annotation.Configuration;
    import org.springframework.lang.Nullable;
    import org.springframework.messaging.Message;
    import org.springframework.messaging.MessageChannel;
    import org.springframework.messaging.MessageHeaders;
    import org.springframework.messaging.converter.MessageConverter;
    import org.springframework.messaging.converter.SimpleMessageConverter;
    import org.springframework.messaging.converter.SmartMessageConverter;
    import org.springframework.messaging.simp.SimpMessageType;
    import org.springframework.messaging.simp.SimpMessagingTemplate;
    import org.springframework.messaging.simp.config.ChannelRegistration;
    import org.springframework.messaging.simp.config.MessageBrokerRegistry;
    import org.springframework.messaging.simp.stomp.StompHeaderAccessor;
    import org.springframework.messaging.support.ChannelInterceptor;
    import org.springframework.messaging.support.ChannelInterceptorAdapter;
    import org.springframework.messaging.support.GenericMessage;
    import org.springframework.stereotype.Service;
    import org.springframework.web.socket.config.annotation.EnableWebSocketMessageBroker;
    import org.springframework.web.socket.config.annotation.StompEndpointRegistry;
    import org.springframework.web.socket.config.annotation.WebSocketMessageBrokerConfigurer;
    
    import java.util.HashMap;
    import java.util.LinkedHashMap;
    import java.util.Map;
    import java.util.concurrent.LinkedBlockingQueue;
    import java.util.concurrent.ThreadPoolExecutor;
    import java.util.concurrent.TimeUnit;
    
    @Configuration
    @EnableWebSocketMessageBroker
    @Service
    public class WebSocketConfig implements WebSocketMessageBrokerConfigurer { 
        private static final Logger log = LoggerFactory.getLogger(WebSocketConfig.class);
    
        @Autowired
        private SimpMessagingTemplate messagingTemplate;
    
        @Override
        public void registerStompEndpoints(StompEndpointRegistry registry) { 
            //开启/myEndPoint端点
            registry.addEndpoint("/myEndPoint")
                    //允许跨域访问
                    .setAllowedOrigins("*")
                    //使用sockJS
                    .withSockJS();
        }
    
        @Override
        public void configureMessageBroker(MessageBrokerRegistry registry) { 
            //订阅topic的前缀，比订阅topic必须是这种格式// /myPrefixes/topic1/xxx
            registry.enableSimpleBroker("/myPrefixes");
        }
    
    
        @Override
        public void configureClientInboundChannel(ChannelRegistration registration) { 
            ChannelInterceptor interceptor = new ChannelInterceptorAdapter() { 
                @Override
                public Message<?> preSend(Message<?> message, MessageChannel channel) { 
                    log.info("Inbound preSend. message={}", message);
                    StompHeaderAccessor accessor = StompHeaderAccessor.wrap(message);
                    MessageHeaders header = message.getHeaders();
                    String sessionId = (String)header.get("simpSessionId");
                    if (accessor != null && accessor.getCommand() !=null && accessor.getCommand().getMessageType() != null) { 
                        SimpMessageType type = accessor.getCommand().getMessageType();
                        if (accessor!= null && SimpMessageType.CONNECT.equals(type)) { 
                            String jwtToken = accessor.getFirstNativeHeader("AuthToken");
    
                            if(StringUtils.isNotBlank(jwtToken)) { 
                                log.info("Inbound preSend: sessionId={}, jwtToken={}", sessionId, jwtToken);
                            }
                            else { 
                                log.error("no token, will be disallowed to connect.");
                                return null;
                            }
                        }else if (type == SimpMessageType.DISCONNECT) { 
                            log.info("Inbound sessionId={} is disconnected", sessionId);
                        }else if (type == SimpMessageType.SUBSCRIBE) { 
                            String topicDest = (String)header.get("simpDestination");
                            log.info("subscribe topicDest={}, message={} SUBSCRIBE", topicDest, message);
                        } else if (type == SimpMessageType.MESSAGE) { 
                            String topicDest = (String)header.get("simpDestination");
                            log.info("之前的消息 topicDest={}, message={} MESSAGE", topicDest, message);
                            message = UpdateMessage(message, "Inbound");
                            log.info("之后的消息e topicDest={}, message={} MESSAGE", topicDest, message);
                        }
                    }
    
                    return message;
                }
                @Override
                public boolean preReceive(MessageChannel channel) { 
                    log.info("Inbound preReceive. channel={}", channel);
                    return true;
                }
    
                @Override
                public Message<?> postReceive(Message<?> message, MessageChannel channel) { 
                    log.info("Inbound postReceive. message={}", message);
                    return message;
                }
    
                @Override
                public void postSend(Message<?> message, MessageChannel channel, boolean sent) { 
                    log.info("Inbound postSend. message={}", message);
                }
    
                @Override
                public void afterSendCompletion(Message<?> message, MessageChannel channel, boolean sent, @Nullable Exception ex) { 
                    log.info("Inbound afterSendCompletion. message={}", message);
                    StompHeaderAccessor accessor = StompHeaderAccessor.wrap(message);
                    MessageHeaders header = message.getHeaders();
                    if (accessor != null && accessor.getCommand() !=null && accessor.getCommand().getMessageType() != null) { 
                        SimpMessageType type = accessor.getCommand().getMessageType();
                        if (type == SimpMessageType.SUBSCRIBE) { 
                            String topicDest = (String)header.get("simpDestination");
                           log.info("afterSenfCompletion. topicDest={}, message={} SUBSCRIBE", topicDest, message);
    
                            String payload = "{\"myfield1\":\"afterSendCompletion初始化消息\"}";
                            messagingTemplate.convertAndSend(topicDest, payload);
                            log.info("send complete. topic={}", topicDest);
                        }
                    }
                }
    
                @Override
                public void afterReceiveCompletion(@Nullable Message<?> message, MessageChannel channel, @Nullable Exception ex) { 
                    log.info("Inbound afterReceiveCompletion. message={}", message);
                }
            };
    
            registration.interceptors(interceptor);
        }
    
        @Override
        public void configureClientOutboundChannel(ChannelRegistration registration) { 
            ChannelInterceptor interceptor = new ChannelInterceptorAdapter() { 
                @Override
                public boolean preReceive(MessageChannel channel) { 
                    log.info("Outbound preReceive: channel={}", channel);
                    return true;
                }
    
                @Override
                public Message<?> preSend(Message<?> message, MessageChannel channel) { 
                    log.info("Outbound preSend: message={}", message);
                    return message;
                }
    
                @Override
                public void postSend(Message<?> message, MessageChannel channel, boolean sent) { 
                    log.info("Outbound postSend. message={}", message);
                }
    
                @Override
                public Message<?> postReceive(Message<?> message, MessageChannel channel) { 
                    log.info("Outbound postReceive. message={}", message);
                    return message;
                }
    
                @Override
                public void afterSendCompletion(Message<?> message, MessageChannel channel, boolean sent, @Nullable Exception ex) { 
                    log.info("Outbound afterSendCompletion. message={}", message);
                }
    
                @Override
                public void afterReceiveCompletion(@Nullable Message<?> message, MessageChannel channel, @Nullable Exception ex) { 
                    log.info("Outbound afterReceiveCompletion. message={}", message);
                }
            };
    
            registration.interceptors(interceptor);
        }
    
        private Message<?> UpdateMessage(Message<?> message, String logFlag) { 
            log.info(logFlag + " preSend: message={}", message);
            MessageHeaders header = message.getHeaders();
            Object obj = message.getPayload();
            //一般都是byte[]
            JSONObject jsonObj = null;
            String strUTF8 = null;
            String strJsonUTF8 = null;
            Message<?> msg = null;
            try { 
                strUTF8 = new String((byte[])obj,"UTF-8");
                jsonObj = JSON.parseObject(strUTF8);
                jsonObj.put(logFlag + "ChannelContent2", "add to");
                String value = jsonObj.getString("myfield1");
                jsonObj.put("myfield1", logFlag + " add to " + value);
                strJsonUTF8 = jsonObj.toJSONString();
                byte[] msgToByte = strJsonUTF8.getBytes("UTF-8");
                msg = new GenericMessage<>(msgToByte, header);
            }
            catch (Exception ex) { 
                log.info("(byte[] to string exception. ex={}", ex.getLocalizedMessage());
            }
    
            if (msg != null) { 
                log.info(logFlag + " preSend Modified: message={}, strUTF8={}, strJsonUTF8={}", msg, strUTF8, strJsonUTF8);
                return msg;
            }
            else { 
                log.info(logFlag + " preSend Original: message={}, strUTF8={}", message, strUTF8);
                return message;
            }
        }
    //"simpMessageType" -> "MESSAGE"
        Message sendInitMsg(Message<?> oldMessage, String dest, Map<String, Object> headers, String payload) { 
            MessageHeaders messageHeaders = null;
            Object conversionHint = headers != null?headers.get("conversionHint"):null;
            Map<String, Object> headersToUse = new HashMap<>();
            headersToUse.put("simpMessageType", SimpMessageType.MESSAGE);
            headersToUse.put("destination", dest);
            headersToUse.put("contentType", "text/plain;charset=UTF-8");
            headersToUse.put("stompCommand", "SEND");
    
            Map<String, Object> nativeHeaders = new LinkedHashMap<>();
            nativeHeaders.put("id", "sub-0");
            nativeHeaders.put("destination", dest);
            headersToUse.put("nativeHeaders", nativeHeaders);
            messageHeaders = new MessageHeaders(headersToUse);
            MessageHeaders oldHeaders = oldMessage.getHeaders()
    
            MessageConverter converter = new SimpleMessageConverter();
            Message<?> message = converter instanceof SmartMessageConverter ?((SmartMessageConverter)converter).toMessage(payload, messageHeaders, conversionHint):converter.toMessage(payload, messageHeaders);
            return message;
        }
    }

# 效果截图 #

![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3J1c3NsZQ_size_16_color_FFFFFF_t_70]

[Spring Boot WebSocket]: https://blog.csdn.net/russle/article/details/85028755
[Link 1]: https://github.com/yqbjtu/websocketTutorials/tree/master/WebSocketInterceptor
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3J1c3NsZQ_size_16_color_FFFFFF_t_70]: /images/20220327/07532f38da25452bb83bab01dcfdf7a0.png