SpringBoot整合WebSocket带案例

WebSocket通信过程

客户端构建一个websocket实例，并且为它绑定一个需要连接到的服务器地址，当客户端连接服务端的候，会向服务端发送一个http get报文，告诉服务端需要将通信协议切换到websocket，服务端收到http请求后将通信协议切换到websocket,同时发给客户端一个响应报文，返回的状态码为101，表示同意客户端协议转请求，并转换为websocket协议。以上过程都是利用http通信完成的，称之为websocket协议握手(websocket Protocol handshake)，经过握手之后，客户端和服务端就建立了websocket连接，以后的通信走的都是websocket协议了。
 

1.pom文件

<dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-starter-websocket</artifactId>
        </dependency>

2.启动类上面

@SpringBootApplication
@EnableWebSocket
public class ServiceApplication {
    public static void main(String[] args) {
        SpringApplication.run(ServiceApplication.class);
    }
}

3.配置WebSocketConfig

@Configuration
@SuppressWarnings("all")
public class WebSocketConfig implements WebSocketConfigurer {
    @Autowired
    private TerminalInfoMapper terminalInfoMapper;
    @Autowired
    private TerminalGroupMapper terminalGroupMapper;
    @Autowired
    private FileEncryptionTypeMapper fileEncryptionTypeMapper;
    @Autowired
    private ProcessWhiteInfoMapper processWhiteInfoMapper;
    @Autowired
    private FeatureInfoMapper featureInfoMapper;
    @Autowired
    private CustomDirectoryMapper customDirectoryMapper;

    @Bean
    public TextWebSocketHandler textWebSocketHandler(){
        // 试过在MyWebSocketHandler类上加入@Component注解，然后Autoware注入mapper，但是不能成功，提示mapper为null
        return new MyWebSocketHandler(terminalInfoMapper, terminalGroupMapper, fileEncryptionTypeMapper, processWhiteInfoMapper, featureInfoMapper, customDirectoryMapper);
    }
    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
        registry.addHandler(textWebSocketHandler(), "/api/control/")
                //.addInterceptors(new WebSocketHandshakeInterceptor())
                .setAllowedOrigins("*");
    }
}

4.WebSocket处理类

因为WebSocket是类似客户端服务端的形式(采用ws协议)，那么这里的WebSocketHandler其实就相当于一个ws协议的Controller将此类注册进WebSocketConfig中去，注意要采用@Bean的方式注入，不然Autoware引入其他bean的时候可能会出现空指针异常
下面是具体业务代码：

@Slf4j
@SuppressWarnings("all")
public class MyWebSocketHandler extends TextWebSocketHandler {
    // 统计在线人数
    public static AtomicLong atomicLong = new AtomicLong(0);
    // 分组
    public static final ConcurrentHashMap<WebSocketSession, String> CLIENTS = new ConcurrentHashMap<>();
    // VIP组
    public static final ConcurrentHashMap<WebSocketSession, String> VIP_CLIENTS = new ConcurrentHashMap<>();
    private TerminalInfoMapper terminalInfoMapper;
    private TerminalGroupMapper terminalGroupMapper;
    private FileEncryptionTypeMapper fileEncryptionTypeMapper;
    private ProcessWhiteInfoMapper processWhiteInfoMapper;
    private FeatureInfoMapper featureInfoMapper;
    private CustomDirectoryMapper customDirectoryMapper;

    @Autowired
    private WebsocketUtil websocketUtil;

    public MyWebSocketHandler() {
    }

    public MyWebSocketHandler(TerminalInfoMapper terminalInfoMapper,
                              TerminalGroupMapper terminalGroupMapper,
                              FileEncryptionTypeMapper fileEncryptionTypeMapper,
                              ProcessWhiteInfoMapper processWhiteInfoMapper,
                              FeatureInfoMapper featureInfoMapper,
                              CustomDirectoryMapper customDirectoryMapper) {
        this.terminalInfoMapper = terminalInfoMapper;
        this.terminalGroupMapper = terminalGroupMapper;
        this.fileEncryptionTypeMapper = fileEncryptionTypeMapper;
        this.processWhiteInfoMapper = processWhiteInfoMapper;
        this.featureInfoMapper = featureInfoMapper;
        this.customDirectoryMapper = customDirectoryMapper;
    }

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
//        WebsocketUtil websocketUtil = new WebsocketUtil();
        URI uri = session.getUri();
        String query = uri.getQuery();
        System.out.println("query = " + query);
        Map<String, String> queryMap = websocketUtil.queryValue(query);
        System.out.println("queryMap = " + queryMap);
        if (queryMap == null || queryMap.get("terminal_id") == null){
            session.close();
        }
        websocketUtil.dealTerminal(session, queryMap.get("terminal_id"), terminalInfoMapper, terminalGroupMapper, featureInfoMapper, fileEncryptionTypeMapper, processWhiteInfoMapper, customDirectoryMapper);
        System.out.println("query = " + query);
        System.out.println("uri = " + uri);
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        CLIENTS.remove(session);
        if (VIP_CLIENTS.containsKey(session)) VIP_CLIENTS.remove(session);
        log.info("终端{}\t退出连接", CLIENTS.get(session));
    }

    @Override
    public void handleMessage(WebSocketSession session, WebSocketMessage<?> message) throws Exception {
        log.info("收到的消息是!{}", message.getPayload().toString());
//        WebsocketUtil websocketUtil = new WebsocketUtil();
        websocketUtil.send(session, message);
    }

    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
        CLIENTS.remove(session);
        if (VIP_CLIENTS.containsKey(session)) VIP_CLIENTS.remove(session);
        log.error("终端{}\t发生错误{}", CLIENTS.get(session), exception.getMessage());
    }
}

5.WebSocketUtil

@Slf4j
@Component
@SuppressWarnings("all")
public class WebsocketUtil {
    public Map<String, String> queryValue(String uri){
        if (uri != null){
            System.out.println("uri = " + uri);
            String[] paramsArray = uri.split("&");
            System.out.println("paramsArray = " + Arrays.toString(paramsArray));
            if (paramsArray.length > 0) {
                for (String s : paramsArray) {
                    String[] strings = s.split("=");
                    return dealParams(strings);
                }
            }
        }
        return null;
    }

    private Map<String, String> dealParams(String[] paramsArray) {
        Map<String, String> map = new HashMap<>();
        for (int i = 0; i < paramsArray.length; i = i + 2) {
            switch (paramsArray[i]){
                case "terminal_id" -> map.put("terminal_id", paramsArray[i + 1].replace("&", ""));
                case "virus_version" -> map.put("virus_version", paramsArray[i + 1].replace("&", ""));
                case "edr_version" -> map.put("edr_version", paramsArray[i + 1].replace("&", ""));
                case "scanner_version" -> map.put("scanner_version", paramsArray[i + 1].replace("&", ""));
            }
        }
        return map;
    }

    public void dealTerminal(WebSocketSession session,
                             String terminal_id,
                             TerminalInfoMapper terminalInfoMapper,
                             TerminalGroupMapper terminalGroupMapper,
                             FeatureInfoMapper featureInfoMapper,
                             FileEncryptionTypeMapper fileEncryptionTypeMapper,
                             ProcessWhiteInfoMapper processWhiteInfoMapper,
                             CustomDirectoryMapper customDirectoryMapper) throws IOException {
        System.out.println("terminalInfoMapper = " + terminalInfoMapper);
        TerminalInfo terminalInfo = terminalInfoMapper.selectByTerminalId(terminal_id);
        if (terminalInfo == null){
            log.error("该客户端的mac地址不存在 {}", terminal_id);
            session.close();
            return;
        }
        List<FeatureInfo> featureInfoList = featureInfoMapper.selectAll();
        send(session, featureInfoList);

        List<FileEncryptionType> fileEncryptionTypes = fileEncryptionTypeMapper.selectAll();
        send(session, fileEncryptionTypes);

        List<String> processWhiteInfos = processWhiteInfoMapper.selectAll();
        send(session, processWhiteInfos);

        String groupId = terminalInfo.getGroup();
        if (groupId == "") groupId = null;
        TerminalGroup terminalGroup = terminalGroupMapper.selectByPrimaryKey(Long.valueOf(groupId));
        List<Map<String, Object>> maps = new ArrayList<>();
        Map<String, Object> vip = new HashMap<>();
        vip.put("Vip", false);
        maps.add(vip);
        if (terminalGroup != null){
            String groupName = terminalGroup.getGroupName();
            if (groupName == "Vip"){
                maps.get(0).put("Vip", true);
                List<CustomDirectory> customDirectories = customDirectoryMapper.selectAll();
                List<Object> list = new ArrayList<>();
                list.addAll(maps);
                list.addAll(customDirectories);
                MyWebSocketHandler.VIP_CLIENTS.put(session, terminal_id);
                System.out.println("MyWebSocketHandler.VIP_CLIENTS = " + MyWebSocketHandler.VIP_CLIENTS);
                send(session, list);
            }
        }
        MyWebSocketHandler.CLIENTS.put(session, terminal_id);
    }

    public synchronized void send(WebSocketSession session, Object feature_data) throws IOException {
        String message = null;
        String customDirector = JsonUtils.toStr(feature_data);
        session.sendMessage(new TextMessage(customDirector));
    }
}

6.测试如下