使用 Sa-Token 解决 WebSocket 握手身份认证

作者: 偗长 | 来源:发表于2022-02-14 00:34 被阅读0次

前言

相比于 Http 的单项通信方式，WebSocket 可以从服务器向浏览器主动推送消息，这一特性可以帮助我们完成诸如订单消息推送、IM实时聊天等一些特定业务。

然而 WebSocket 本身对“身份认证”并没有提供直接的支持，对客户端的连接默认是“来者不拒”，所以认证授权这个事，得我们自己动手。

Sa-Token 是一个 java 权限认证框架，主要解决登录认证、权限认证、单点登录、OAuth2、微服务网关鉴权等一系列权限相关问题。
GitHub 开源地址：https://github.com/dromara/sa-token

下面我们介绍一下如何在 WebSocket 中集成 Sa-Token 身份认证，保证连接的安全性。

两种集成方式

我们将依次介绍目前最常见的两种集成 WebSocket 方式：

Java 原生版：javax.websocket.Session
Spring 封装版：WebSocketSession

废话不多说，直接开搞：

方式一：Java 原生版 javax.websocket.Session

1、首先是引入 pom.xml 依赖

<!-- SpringBoot依赖 -->
<dependency>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-starter-web</artifactId>
</dependency>

<!-- WebScoket 依赖 -->
<dependency>  
    <groupId>org.springframework.boot</groupId>  
    <artifactId>spring-boot-starter-websocket</artifactId>  
</dependency>

<!-- Sa-Token 权限认证, 在线文档：http://sa-token.dev33.cn/ -->
<dependency>
    <groupId>cn.dev33</groupId>
    <artifactId>sa-token-spring-boot-starter</artifactId>
    <version>1.29.0</version>
</dependency>

2、登录接口，用于获取会话token

/**
 * 登录测试 
 */
@RestController
@RequestMapping("/acc/")
public class LoginController {

    // 测试登录  ---- http://localhost:8081/acc/doLogin?name=zhang&pwd=123456
    @RequestMapping("doLogin")
    public SaResult doLogin(String name, String pwd) {
        // 此处仅作模拟示例，真实项目需要从数据库中查询数据进行比对 
        if("zhang".equals(name) && "123456".equals(pwd)) {
            StpUtil.login(10001);
            return SaResult.ok("登录成功").set("token", StpUtil.getTokenValue());
        }
        return SaResult.error("登录失败");
    }

    // ... 
    
}

3、WebSocket连接处理

@Component
@ServerEndpoint("/ws-connect/{satoken}")
public class WebSocketConnect {

    /**
     * 固定前缀 
     */
    private static final String USER_ID = "user_id_";
    
     /** 
      * 存放Session集合，方便推送消息 （javax.websocket.Session）  
      */
    private static ConcurrentHashMap<String, Session> sessionMap = new ConcurrentHashMap<>();
    
    // 监听：连接成功
    @OnOpen
    public void onOpen(Session session, @PathParam("satoken") String satoken) throws IOException {
        
        // 根据 token 获取对应的 userId 
        Object loginId = StpUtil.getLoginIdByToken(satoken);
        if(loginId == null) {
            session.close();
            throw new SaTokenException("连接失败，无效Token：" + satoken);
        }
        
        // put到集合，方便后续操作 
        long userId = SaFoxUtil.getValueByType(loginId, long.class);
        sessionMap.put(USER_ID + userId, session);
        
        // 给个提示 
        String tips = "Web-Socket 连接成功，sid=" + session.getId() + "，userId=" + userId;
        System.out.println(tips);
        sendMessage(session, tips);
    }

    // 监听: 连接关闭
    @OnClose
    public void onClose(Session session) {
        System.out.println("连接关闭，sid=" + session.getId());
        for (String key : sessionMap.keySet()) {
            if(sessionMap.get(key).getId().equals(session.getId())) {
                sessionMap.remove(key);
            }
        }
    }
    
    // 监听：收到客户端发送的消息 
    @OnMessage
    public void onMessage(Session session, String message) {
        System.out.println("sid为：" + session.getId() + "，发来：" + message);
    }
    
    // 监听：发生异常 
    @OnError
    public void onError(Session session, Throwable error) {
        System.out.println("sid为：" + session.getId() + "，发生错误");
        error.printStackTrace();
    }
    
    // ---------
    
    // 向指定客户端推送消息 
    public static void sendMessage(Session session, String message) {
        try {
            System.out.println("向sid为：" + session.getId() + "，发送：" + message);
            session.getBasicRemote().sendText(message);
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }
    
    // 向指定用户推送消息 
    public static void sendMessage(long userId, String message) {
        Session session = sessionMap.get(USER_ID + userId);
        if(session != null) {
            sendMessage(session, message);
        }
    }
    
}

4、WebSocket配置

/**
 * 开启WebSocket支持
 */
@Configuration  
public class WebSocketConfig { 
    
    @Bean  
    public ServerEndpointExporter serverEndpointExporter() {  
        return new ServerEndpointExporter();  
    }
    
}

5、启动类

@SpringBootApplication
public class SaTokenWebSocketApplication {

    public static void main(String[] args) {
        SpringApplication.run(SaTokenWebSocketApplication.class, args); 
    }
    
}

搭建完毕，启动项目

6、测试

1、首先我们访问登录接口，拿到会话token

http://localhost:8081/acc/doLogin?name=zhang&pwd=123456

如图所示：

令牌

2、然后我们随便找一个WebSocket在线测试页面进行连接
，例如：https://www.bejson.com/httputil/websocket/

连接地址：

ws://localhost:8081/ws-connect/302ee2f8-60aa-42aa-8ecb-eeae5ba57015

如图所示：

测试连接

3、如果我们输入一个错误的token，会怎样呢？

连接失败

可以看到，连接会被立即断开！

方式二：Spring 封装版：WebSocketSession

1、同上：首先是引入 pom.xml 依赖

<!-- SpringBoot依赖 -->
<dependency>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-starter-web</artifactId>
</dependency>

<!-- WebScoket 依赖 -->
<dependency>  
    <groupId>org.springframework.boot</groupId>  
    <artifactId>spring-boot-starter-websocket</artifactId>  
</dependency>

<!-- Sa-Token 权限认证, 在线文档：http://sa-token.dev33.cn/ -->
<dependency>
    <groupId>cn.dev33</groupId>
    <artifactId>sa-token-spring-boot-starter</artifactId>
    <version>1.29.0</version>
</dependency>

2、登录接口，用于获取会话token

/**
 * 登录测试 
 */
@RestController
@RequestMapping("/acc/")
public class LoginController {

    // 测试登录  ---- http://localhost:8081/acc/doLogin?name=zhang&pwd=123456
    @RequestMapping("doLogin")
    public SaResult doLogin(String name, String pwd) {
        // 此处仅作模拟示例，真实项目需要从数据库中查询数据进行比对 
        if("zhang".equals(name) && "123456".equals(pwd)) {
            StpUtil.login(10001);
            return SaResult.ok("登录成功").set("token", StpUtil.getTokenValue());
        }
        return SaResult.error("登录失败");
    }

    // ... 
    
}

3、WebSocket 连接处理

/**
 * 处理 WebSocket 连接 
 */
public class MyWebSocketHandler extends TextWebSocketHandler {

    /**
     * 固定前缀 
     */
    private static final String USER_ID = "user_id_";
    
    /**
     * 存放Session集合，方便推送消息
     */
    private static ConcurrentHashMap<String, WebSocketSession> webSocketSessionMaps = new ConcurrentHashMap<>();

    // 监听：连接开启 
    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {

        // put到集合，方便后续操作 
        String userId = session.getAttributes().get("userId").toString();
        webSocketSessionMaps.put(USER_ID + userId, session);
        

        // 给个提示 
        String tips = "Web-Socket 连接成功，sid=" + session.getId() + "，userId=" + userId;
        System.out.println(tips);
        sendMessage(session, tips);
    }
    
    // 监听：连接关闭 
    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        // 从集合移除 
        String userId = session.getAttributes().get("userId").toString();
        webSocketSessionMaps.remove(USER_ID + userId);
        
        // 给个提示 
        String tips = "Web-Socket 连接关闭，sid=" + session.getId() + "，userId=" + userId;
        System.out.println(tips);
    }

    // 收到消息 
    @Override
    public void handleTextMessage(WebSocketSession session, TextMessage message) throws IOException {
        System.out.println("sid为：" + session.getId() + "，发来：" + message);
    }

    // ----------- 
    
    // 向指定客户端推送消息 
    public static void sendMessage(WebSocketSession session, String message) {
        try {
            System.out.println("向sid为：" + session.getId() + "，发送：" + message);
            session.sendMessage(new TextMessage(message));
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }
    
    // 向指定用户推送消息 
    public static void sendMessage(long userId, String message) {
        WebSocketSession session = webSocketSessionMaps.get(USER_ID + userId);
        if(session != null) {
            sendMessage(session, message);
        }
    }
    
}

4、WebSocket 前置拦截器

/**
 * WebSocket 握手的前置拦截器 
 */
public class WebSocketInterceptor implements HandshakeInterceptor {

    // 握手之前触发 (return true 才会握手成功 )
    @Override
    public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler handler,
            Map<String, Object> attr) {
        
        System.out.println("---- 握手之前触发 " + StpUtil.getTokenValue());
        
        // 未登录情况下拒绝握手 
        if(StpUtil.isLogin() == false) {
            System.out.println("---- 未授权客户端，连接失败");
            return false;
        }
        
        // 标记 userId，握手成功 
        attr.put("userId", StpUtil.getLoginIdAsLong());
        return true;
    }

    // 握手之后触发 
    @Override
    public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler,
            Exception exception) {
        System.out.println("---- 握手之后触发 ");
    }
    
}

5、WebSocket 配置

/**
 * WebSocket 相关配置 
 */
@Configuration
@EnableWebSocket
public class WebSocketConfig implements WebSocketConfigurer {
    
    // 注册 WebSocket 处理器 
    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry webSocketHandlerRegistry) {
        webSocketHandlerRegistry
                // WebSocket 连接处理器 
                .addHandler(new MyWebSocketHandler(), "/ws-connect")
                // WebSocket 拦截器 
                .addInterceptors(new WebSocketInterceptor())
                // 允许跨域 
                .setAllowedOrigins("*");
    }

}

6、启动类

/**
 * Sa-Token 整合 WebSocket 鉴权示例 
 */
@SpringBootApplication
public class SaTokenWebSocketSpringApplication {

    public static void main(String[] args) {
        SpringApplication.run(SaTokenWebSocketSpringApplication.class, args); 
    }
    
}

启动项目，开始测试

7、测试

1、首先访问登录接口，拿到会话token

http://localhost:8081/acc/doLogin?name=zhang&pwd=123456

如图所示：

令牌

2、然后打开WebSocket在线测试页面进行连接
，例如：https://www.bejson.com/httputil/websocket/

连接地址：

ws://localhost:8081/ws-connect?satoken=fe6e7dbd-38b8-4de2-ae05-cda7e36bf2f7

如图所示：

测试连接

注：这里采用 url 传递 Token 是因为在第三方测试页面上这样比较方便，真实项目中可以从Cookie、Header参数、url参数三种方式任选其一传递会话令牌，效果同等

3、如果输入一个错误的 Token

连接失败

连接失败！

示例地址

以上代码已经上传git，示例地址：
码云：sa-token-demo-websocket

参考资料

Gitee地址：https://gitee.com/dromara/sa-token
GitHub地址：https://github.com/dromara/sa-token
Sa-Token 官网：https://sa-token.dev33.cn/

使用 Sa-Token 解决 WebSocket 握手身份认证

前言

两种集成方式

方式一：Java 原生版 javax.websocket.Session

1、首先是引入 pom.xml 依赖

2、登录接口，用于获取会话token

3、WebSocket连接处理

4、WebSocket配置

5、启动类

6、测试

方式二：Spring 封装版：WebSocketSession

1、同上：首先是引入 pom.xml 依赖

2、登录接口，用于获取会话token

3、WebSocket 连接处理

4、WebSocket 前置拦截器

5、WebSocket 配置

6、启动类

7、测试

示例地址

参考资料

相关文章

网友评论

延伸阅读

深度阅读

栏目导航

热点阅读

SpringBoot精选

SpringBoot极简教程 · Spring Boot

IT技术篇