使用 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/