springBoot集成webSocket,包含鉴权,同一用户多点登录

webSocket集成

pom集成

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

配置类

@Configuration
@EnableWebSocket
public class WebSocketConfig implements WebSocketConfigurer {
    /**
     * 注入拦截器
     */
    @Resource
    private WebSocketInterceptor webSocketInterceptor;


    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry webSocketHandlerRegistry) {

        webSocketHandlerRegistry
                //添加myHandler消息处理对象，和websocket访问地址
                .addHandler(myHandler(), "/ws")
                //设置允许跨域访问
                .setAllowedOrigins("*")
                //添加拦截器可实现用户链接前进行权限校验等操作
                .addInterceptors(webSocketInterceptor);
    }

    @Bean
    public WebSocketHandler myHandler() {
        return new MyWebSocketHandler();
    }


}

拦截器

@Component
@Slf4j
public class WebSocketInterceptor implements HandshakeInterceptor {

    @Autowired
    private CustomSessionManager customSessionManager;

    /**
     * 握手之前
     *
     * @param request    request
     * @param response   response
     * @param wsHandler  handler
     * @param attributes 属性
     * @return 是否握手成功：true-成功，false-失败
     */
    @Override
    public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler,
                                   Map<String, Object> attributes) {
        ServletServerHttpRequest serverHttpRequest = (ServletServerHttpRequest) request;
        String token = serverHttpRequest.getServletRequest().getParameter("token");
        if (StringUtils.isEmpty(token)) {
            return false;
        }
        AdminSession adminSession = null;
        try {
             adminSession = customSessionManager.getSessionById(token);
        } catch (Exception e) {
            log.info("webSocket鉴权失败 token = {}", token);
        }

        if (ObjectUtil.isEmpty(adminSession)) {
            return false;
        }
        attributes.put("uid", adminSession.getId());
        return true;
    }

    /**
     * 握手后
     *
     * @param request   request
     * @param response  response
     * @param wsHandler wsHandler
     * @param exception exception
     */
    @Override
    public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler,
                               Exception exception) {
    }
}

业务实现

@Slf4j
@Component
public class MyWebSocketHandler extends TextWebSocketHandler {
    /**
     * 静态变量，用来记录当前在线连接数
     */
    private static AtomicInteger onlineNum = new AtomicInteger();

    /**
     * 存放每个客户端连接对象
     */
    private static ConcurrentHashMap<Integer, List<WebSocketSession>> sessionPools = new ConcurrentHashMap<>();

    /**
     * 在线人数加一
     */
    public static void addOnlineCount() {
        onlineNum.incrementAndGet();
    }

    /**
     * 在线人数减一
     */
    public static void subOnlineCount() {
        onlineNum.decrementAndGet();
    }

    /**
     * 接受客户端消息
     *
     * @param session session
     * @param message message
     * @throws IOException e
     */
    @Override
    public void handleTextMessage(WebSocketSession session, TextMessage message) throws IOException {
        log.info("接收到用户: {} 发来的消息: {}", session.getAttributes().get("uid"), message.getPayload());
        NewOrderGoods newOrderGoods = new NewOrderGoods();
        switch(message.getPayload()){
            case "1":
                //发送订单创建信息
                newOrderGoods.setId(9109L);
                newOrderGoods.setOrderCode("1017431853084180480");
                sendOrderCreate(newOrderGoods);
                break;
            case "2":
                //发送订单退款消息
                newOrderGoods.setId(9110L);
                newOrderGoods.setOrderCode("1017432921394380800");
                sendOrderRefund(newOrderGoods);
                break;
            default:
                break;
        }
    }

    /**
     * 建立连接后发送消息给客户端
     *
     * @param session session
     * @throws Exception e
     */
    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        String uid = session.getAttributes().get("uid").toString();
        int uidInteger = Integer.parseInt(uid);
        boolean flag = sessionPools.containsKey(uidInteger);
        List<WebSocketSession> put =null;
        if (flag){
            List<WebSocketSession> webSocketSessions = sessionPools.get(uidInteger);
            webSocketSessions.add(session);
            put = sessionPools.put(uidInteger, webSocketSessions);

        }else{
            ArrayList<WebSocketSession> webSocketSessions = new ArrayList<>();
            webSocketSessions.add(session);
            put =  sessionPools.put(uidInteger,webSocketSessions);
        }
        if (put == null) {
            addOnlineCount();
        }
        log.info("webSocket鉴权成功! 用户id = {}", session.getAttributes().get("uid"));
        session.sendMessage(new TextMessage("连接成功 当前在线人数: " + onlineNum));
    }

    /**
     * 连接关闭后
     *
     * @param session session
     * @param status  status
     * @throws Exception e
     */
    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        String uid = session.getAttributes().get("uid").toString();
        sessionPools.remove(Integer.parseInt(uid));
        subOnlineCount();
        log.info("用户: {}断开连接",uid);
    }

    /**
     * 发送广播消息
     *
     * @param message 消息内容
     */
    public  void sendTopic(String message) {
        if (sessionPools.isEmpty()) {
            return;
        }
        TextMessage textMessage = new TextMessage(message);
        for (Map.Entry<Integer, List<WebSocketSession>> entry : sessionPools.entrySet()) {
            try {
                entry.getValue().stream().forEach(item->{
                    try {
                        item.sendMessage(textMessage);
                    } catch (IOException e) {
                        e.printStackTrace();
                    }
                });

            } catch (Exception e) {
                e.printStackTrace();
            }
        }
    }

    /**
     * 点对点发送消息
     *
     * @param uid     用户
     * @param message 消息
     */
    public void sendToUser(String uid, String message) {
        List<WebSocketSession> socketSession = sessionPools.get(Integer.parseInt(uid));
        if (CollectionUtils.isEmpty(socketSession)) {
            return;
        }
        TextMessage textMessage = new TextMessage(message);
        socketSession.stream().forEach(item->{
            try {
                item.sendMessage(textMessage);
            } catch (IOException e) {
                e.printStackTrace();
            }
        });

    }

    /**
     * 群发订单创建消息
     */
    public void sendOrderCreate(NewOrderGoods newOrderGoods) {
        WebSocketResponse response = new WebSocketResponse(WebSocketConstant.ORDER_CREATE,newOrderGoods);
        sendTopic(JSON.toJSONString(response));
    }

    /**
     * 群发订单退款消息
     */
    public void sendOrderRefund(NewOrderGoods newOrderGoods) {
        WebSocketResponse response = new WebSocketResponse(WebSocketConstant.ORDER_REFUND,newOrderGoods);
        sendTopic(JSON.toJSONString(response));
    }

}

@Data
@AllArgsConstructor
@NoArgsConstructor
@Accessors(chain = true)
public class WebSocketResponse {
    String action;
    Object data;
}

public interface WebSocketConstant {

    /**
     * webSocket连接成功
     */
    String CONNECT_SUCCESS = "connect_success";
    /**
     * 订单创建成功
     */
    String ORDER_CREATE = "order_create";

    /**
     * 订单创建成功
     */
    String ORDER_REFUND = "order_refund";
}