使用 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 依赖


	org.springframework.boot
	spring-boot-starter-web



  
	org.springframework.boot  
	spring-boot-starter-websocket  




	cn.dev33
	sa-token-spring-boot-starter
	1.29.0

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 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

如图所示:

使用 Sa-Token 解决 WebSocket 握手身份认证_第1张图片

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

连接地址:

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

如图所示:

使用 Sa-Token 解决 WebSocket 握手身份认证_第2张图片

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

使用 Sa-Token 解决 WebSocket 握手身份认证_第3张图片

可以看到,连接会被立即断开!

方式二:Spring 封装版:WebSocketSession

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


	org.springframework.boot
	spring-boot-starter-web



  
	org.springframework.boot  
	spring-boot-starter-websocket  




	cn.dev33
	sa-token-spring-boot-starter
	1.29.0

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 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 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

如图所示:

使用 Sa-Token 解决 WebSocket 握手身份认证_第4张图片

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

连接地址:

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

如图所示:

使用 Sa-Token 解决 WebSocket 握手身份认证_第5张图片

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

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

使用 Sa-Token 解决 WebSocket 握手身份认证_第6张图片

连接失败!

示例地址

以上代码已经上传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 握手身份认证)