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

前言

相比于 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/