若依集成 Websocekt 实现用户认证

首先官网集成

framework 模块中添加依赖

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

添加下面四个类

/**
 * 信号量相关处理
 */
public class SemaphoreUtils {

    /**
     * SemaphoreUtils 日志控制器
     */
    private static final Logger LOGGER = LoggerFactory.getLogger(SemaphoreUtils.class);

    /**
     * 获取信号量
     */
    public static boolean tryAcquire(Semaphore semaphore) {
        boolean flag = false;

        try {
            flag = semaphore.tryAcquire();
        } catch (Exception e) {
            LOGGER.error("获取信号量异常", e);
        }

        return flag;
    }

    /**
     * 释放信号量
     */
    public static void release(Semaphore semaphore) {
        try {
            semaphore.release();
        } catch (Exception e) {
            LOGGER.error("释放信号量异常", e);
        }
    }
}

/**
 * websocket 配置
 */
@Configuration
public class WebSocketConfig {
    @Bean
    public ServerEndpointExporter serverEndpointExporter() {
        return new ServerEndpointExporter();
    }
}

重点是下面这段

通过 session.getRequestParameterMap(); 拿到属性集，这个属性集会将拼接的东西整合成一个 Map 数组。

get("token") 前端在 websocket 地址后拼接 token

通过使用 TokenService 方法将token转换成对象

LoginUser loginUser = tokenService.resolve(token);

最终添加到 Security 中

UsernamePasswordAuthenticationToken authentication = new UsernamePasswordAuthenticationToken( loginUser, null, loginUser.getAuthorities());

SecurityContextHolder.getContext().setAuthentication(authentication);

/**
 * websocket 消息处理
 */
@Component
@ServerEndpoint("/websocket/message")
public class WebSocketServer {

    private final TokenService tokenService;

    /**
     * WebSocketServer 日志控制器
     */
    private static final Logger LOGGER = LoggerFactory.getLogger(WebSocketServer.class);

    public WebSocketServer() {
        this.tokenService = SpringUtils.getBean(TokenService.class);
    }

    /**
     * 默认最多允许同时在线人数100
     */
    public static int socketMaxOnlineCount = 100;

    private static Semaphore socketSemaphore = new Semaphore(socketMaxOnlineCount);

    /**
     * 连接建立成功调用的方法
     */
    @OnOpen
    public void onOpen(Session session) throws Exception {
        // 尝试获取信号量
        boolean semaphoreFlag = SemaphoreUtils.tryAcquire(socketSemaphore);
        if (!semaphoreFlag) {
            // 未获取到信号量
            LOGGER.error("\n 当前在线人数超过限制数- {}", socketMaxOnlineCount);
            WebSocketUsers.sendMessageToUserByText(session, "当前在线人数超过限制数：" + socketMaxOnlineCount);
            session.close();
        } else {
            List<String> strings = session.getRequestParameterMap().get("token");
            if (strings.size() > 0) {
                String token = strings.get(0);
                Assert.isTrue(token != null, "请传递Token");
                LoginUser loginUser = tokenService.resolve(token);
                Assert.isTrue(loginUser != null, "Token 不合法");
                // 创建 Authentication 对象
                UsernamePasswordAuthenticationToken authentication = new UsernamePasswordAuthenticationToken(
                        loginUser, null, loginUser.getAuthorities());
                //  设置用户信息到 SecurityContext 中
                SecurityContextHolder.getContext().setAuthentication(authentication);
                // 添加用户
                WebSocketUsers.put(session.getId(), session);
                LOGGER.info("\n 建立连接 - {}", session);
                LOGGER.info("\n 当前人数 - {}", WebSocketUsers.getUsers().size());
                WebSocketUsers.sendMessageToUserByText(session, "连接成功");
            }
        }
    }

    /**
     * 连接关闭时处理
     */
    @OnClose
    public void onClose(Session session) {
        LOGGER.info("\n 关闭连接 - {}", session);
        // 移除用户
        boolean removeFlag = WebSocketUsers.remove(session.getId());
        if (!removeFlag) {
            // 获取到信号量则需释放
            SemaphoreUtils.release(socketSemaphore);
        }
    }

    /**
     * 抛出异常时处理
     */
    @OnError
    public void onError(Session session, Throwable exception) throws Exception {
        if (session.isOpen()) {
            // 关闭连接
            session.close();
        }
        String sessionId = session.getId();
        LOGGER.info("\n 连接异常 - {}", sessionId);
        LOGGER.info("\n 异常信息 - {}", exception);
        // 移出用户
        WebSocketUsers.remove(sessionId);
        // 获取到信号量则需释放
        SemaphoreUtils.release(socketSemaphore);
    }

    /**
     * 服务器接收到客户端消息时调用的方法
     */
    @OnMessage
    public void onMessage(String message, Session session) {
        LOGGER.info("接收到信息：{}", message);
        WebSocketUsers.sendMessageToUserByText(session, message);
    }
}

/**
 * websocket 客户端用户集
 */
public class WebSocketUsers {

    /**
     * WebSocketUsers 日志控制器
     */
    private static final Logger LOGGER = LoggerFactory.getLogger(WebSocketUsers.class);

    /**
     * 用户集
     */
    private static Map<String, Session> USERS = new ConcurrentHashMap<>();

    /**
     * 存储用户
     */
    public static void put(String key, Session session) {
        USERS.put(key, session);
    }

    /**
     * 移除用户
     */
    public static boolean remove(Session session) {
        String key = null;
        boolean flag = USERS.containsValue(session);
        if (flag) {
            Set<Map.Entry<String, Session>> entries = USERS.entrySet();
            for (Map.Entry<String, Session> entry : entries) {
                Session value = entry.getValue();
                if (value.equals(session)) {
                    key = entry.getKey();
                    break;
                }
            }
        } else {
            return true;
        }
        return remove(key);
    }

    /**
     * 移出用户
     */
    public static boolean remove(String key) {
        LOGGER.info("\n 正在移出用户 - {}", key);
        Session remove = USERS.remove(key);
        if (remove != null) {
            boolean containsValue = USERS.containsValue(remove);
            LOGGER.info("\n 移出结果 - {}", containsValue ? "失败" : "成功");
            return containsValue;
        } else {
            return true;
        }
    }

    /**
     * 获取在线用户列表
     */
    public static Map<String, Session> getUsers() {
        return USERS;
    }

    /**
     * 群发消息文本消息
     */
    public static void sendMessageToUsersByText(String message) {
        Collection<Session> values = USERS.values();
        for (Session value : values) {
            sendMessageToUserByText(value, message);
        }
    }

    /**
     * 发送文本消息
     */
    public static void sendMessageToUserByText(Session session, String message) {
        if (session != null) {
            try {
                session.getBasicRemote().sendText(message);
            } catch (IOException e) {
                LOGGER.error("\n[发送消息异常]", e);
            }
        } else {
            LOGGER.info("\n[你已离线]");
        }
    }
}