Use sa-Token to resolve WebSocket handshake authentication

preface

Compared with the single communication mode of Http, WebSocket can actively push messages from the server to the browser, which can help us complete some specific services such as order message push, IM real-time chat and so on.

However, WebSocket itself does not provide direct support for “identity authentication”, the default connection to the client is “whatever comes”, so authentication and authorization of this matter, we have to start by ourselves.

Sa-token is a Java permission authentication framework, which mainly solves a series of permission related problems such as login authentication, permission authentication, single sign-on (SSO), OAuth2, and micro service network authentication. GitHub open Source: github.com/dromara/sa-…

Next, we introduce how to integrate SA-Token authentication into WebSocket to ensure connection security.

Two ways to integrate

We’ll take a look at the two most common ways to integrate WebSockets today:

• Java native: javax.websocket.session
• Spring package: WebSocketSession

Without further ado, let’s get straight to work:

Method 1: Java native javax.websocket.session

The first is the introduction of the POM.xml dependency

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

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

<! -- Sa-Token permission authentication: http://sa-token.dev33.cn/ -->
<dependency>
	<groupId>cn.dev33</groupId>
	<artifactId>sa-token-spring-boot-starter</artifactId>
	<version>1.29.0</version>
</dependency>
Copy the code

2. Login interface for obtaining the session token

/** * Login test */
@RestController
@RequestMapping("/acc/")
public class LoginController {

	/ / test login - http://localhost:8081/acc/doLogin? name=zhang&pwd=123456
	@RequestMapping("doLogin")
	public SaResult doLogin(String name, String pwd) {
		// This is only a simulation example. Real projects need to query data from the database for comparison
		if("zhang".equals(name) && "123456".equals(pwd)) {
			StpUtil.login(10001);
			return SaResult.ok("Login successful").set("token", StpUtil.getTokenValue());
		}
		return SaResult.error("Login failed");
	}

	// ... 
	
}
Copy the code

3. WebSocket connection processing

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

    /** ** Fixed prefix */
    private static final String USER_ID = "user_id_";
	
	 Javax.websocket.session) */
    private static ConcurrentHashMap<String, Session> sessionMap = new ConcurrentHashMap<>();
    
	// Listening: Connection succeeded
	@OnOpen
	public void onOpen(Session session, @PathParam("satoken") String satoken) throws IOException {
		
		// Obtain the corresponding userId based on the token
		Object loginId = StpUtil.getLoginIdByToken(satoken);
		if(loginId == null) {
			session.close();
			throw new SaTokenException("Connection failed, invalid Token:" + satoken);
		}
		
		// put to collection for subsequent operations
		long userId = SaFoxUtil.getValueByType(loginId, long.class);
		sessionMap.put(USER_ID + userId, session);
		
		// Give a hint
		String tips = "Web-socket connection successful, sid=" + session.getId() + "，userId=" + userId;
		System.out.println(tips);
		sendMessage(session, tips);
	}

	// Listen: connection closed
	@OnClose
	public void onClose(Session session) {
		System.out.println("Connection closed, sid=" + session.getId());
		for (String key : sessionMap.keySet()) {
			if(sessionMap.get(key).getId().equals(session.getId())) { sessionMap.remove(key); }}}// Listening: receives the message sent by the client
	@OnMessage
	public void onMessage(Session session, String message) {
		System.out.println(Sid: "" + session.getId() + ", sent:" + message);
	}
	
	// Listen: an exception occurred
	@OnError
	public void onError(Session session, Throwable error) {
		System.out.println(Sid: "" + session.getId() + ", error occurred");
		error.printStackTrace();
	}
	
	// ---------
	
	// Push a message to the specified client
	public static void sendMessage(Session session, String message) {
		try {
			System.out.println("To sid: + session.getId() + ", send:" + message);
			session.getBasicRemote().sendText(message);
		} catch (IOException e) {
			throw newRuntimeException(e); }}// Push a message to the specified user
	public static void sendMessage(long userId, String message) {
		Session session = sessionMap.get(USER_ID + userId);
		if(session ! =null) { sendMessage(session, message); }}}Copy the code

4. WebSocket configuration

/** * Enable WebSocket support */
@Configuration  
public class WebSocketConfig { 
	
	@Bean  
	public ServerEndpointExporter serverEndpointExporter(a) {  
		return newServerEndpointExporter(); }}Copy the code

5. Start classes

@SpringBootApplication
public class SaTokenWebSocketApplication {

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

Set up, start the project

6, test,

1. First we access the login interface and get the session token

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

As shown in the figure:

2, then we can grab a WebSocket to connect online test page, for example: www.bejson.com/httputil/we…

Connection Address:

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

As shown in the figure:

3. What happens if we enter a wrong token?

As you can see, the connection is immediately disconnected!

Method 2: Spring package: WebSocketSession

1. Same as above: First, introduce the POM.xml dependency

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

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

<! -- Sa-Token permission authentication: http://sa-token.dev33.cn/ -->
<dependency>
	<groupId>cn.dev33</groupId>
	<artifactId>sa-token-spring-boot-starter</artifactId>
	<version>1.29.0</version>
</dependency>
Copy the code

2. Login interface for obtaining the session token

/** * Login test */
@RestController
@RequestMapping("/acc/")
public class LoginController {

	/ / test login - http://localhost:8081/acc/doLogin? name=zhang&pwd=123456
	@RequestMapping("doLogin")
	public SaResult doLogin(String name, String pwd) {
		// This is only a simulation example. Real projects need to query data from the database for comparison
		if("zhang".equals(name) && "123456".equals(pwd)) {
			StpUtil.login(10001);
			return SaResult.ok("Login successful").set("token", StpUtil.getTokenValue());
		}
		return SaResult.error("Login failed");
	}

	// ... 
	
}
Copy the code

3. WebSocket connection processing

/** * Handles WebSocket connections */
public class MyWebSocketHandler extends TextWebSocketHandler {

    /** ** Fixed prefix */
    private static final String USER_ID = "user_id_";
    
    /** * hold Session set, convenient push message */
    private static ConcurrentHashMap<String, WebSocketSession> webSocketSessionMaps = new ConcurrentHashMap<>();

    // Listen: the connection is open
    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {

    	// put to collection for subsequent operations
        String userId = session.getAttributes().get("userId").toString();
        webSocketSessionMaps.put(USER_ID + userId, session);
        

		// Give a hint
		String tips = "Web-socket connection successful, sid=" + session.getId() + "，userId=" + userId;
		System.out.println(tips);
		sendMessage(session, tips);
    }
    
    // Listen: connection closed
    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
    	// Remove from collection
        String userId = session.getAttributes().get("userId").toString();
        webSocketSessionMaps.remove(USER_ID + userId);
        
        // Give a hint
        String tips = "Web-socket connection closed, sid=" + session.getId() + "，userId=" + userId;
    	System.out.println(tips);
    }

    // Received the message
    @Override
    public void handleTextMessage(WebSocketSession session, TextMessage message) throws IOException {
    	System.out.println(Sid: "" + session.getId() + ", sent:" + message);
    }

    // ----------- 
    
    // Push a message to the specified client
 	public static void sendMessage(WebSocketSession session, String message) {
 		try {
 			System.out.println("To sid: + session.getId() + ", send:" + message);
 			session.sendMessage(new TextMessage(message));
 		} catch (IOException e) {
 			throw newRuntimeException(e); }}// Push a message to the specified user
 	public static void sendMessage(long userId, String message) {
 		WebSocketSession session = webSocketSessionMaps.get(USER_ID + userId);
		if(session ! =null) { sendMessage(session, message); }}}Copy the code

4. WebSocket front interceptor

/** * WebSocket handshake preblocker */
public class WebSocketInterceptor implements HandshakeInterceptor {

	// Trigger before handshake (return true only if handshake succeeds)
	@Override
	public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler handler, Map
       
         attr)
       ,> {
		
		System.out.println("---- triggered before handshake" + StpUtil.getTokenValue());
		
		// Refuse to shake hands without logging in
		if(StpUtil.isLogin() == false) {
			System.out.println("---- unauthorized client, connection failed");
			return false;
		}
		
		// mark userId and the handshake is successful
		attr.put("userId", StpUtil.getLoginIdAsLong());
		return true;
	}

	// Triggered after the handshake
	@Override
	public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Exception exception) {
		System.out.println("---- triggered after handshake"); }}Copy the code

5. WebSocket configuration

/** * WebSocket configuration */
@Configuration
@EnableWebSocket
public class WebSocketConfig implements WebSocketConfigurer {
	
	// Register the WebSocket handler
    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry webSocketHandlerRegistry) {
        webSocketHandlerRegistry
        		// WebSocket connects to the handler
                .addHandler(new MyWebSocketHandler(), "/ws-connect")
                // WebSocket interceptor
                .addInterceptors(new WebSocketInterceptor())
                // Allow cross-domain
                .setAllowedOrigins("*"); }}Copy the code

6. Start classes

/** * Sa-token integration WebSocket authentication example */
@SpringBootApplication
public class SaTokenWebSocketSpringApplication {

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

Start the project and start testing

7, test,

1. Access the login interface first and get the session token

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

As shown in the figure:

2, then open the WebSocket connect online test page, for example: www.bejson.com/httputil/we…

Connection Address:

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

As shown in the figure:

Note: Url Token transfer is used here because it is convenient for third-party test pages. In real projects, session tokens can be passed through any of the three methods, Cookie, Header parameter or URL parameter, with the same effect

3. If an incorrect Token is entered

Connection failed!

The sample address

The above code has been uploaded to Git, sample address: code cloud: sa-token-demo-websocket

The resources

• Gitee address: gitee.com/dromara/sa-…
• GitHub address: github.com/dromara/sa-…
• Sa-token official website: sa-token.dev33.cn/