Using SA token to solve WebSocket handshake authentication

Posted by dhiren22 on Sun, 13 Feb 2022 17:40:39 +0100

preface

Compared with the single communication mode of Http, WebSocket can actively push messages from the server to the browser. This feature 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", and the default connection to the client is "whoever comes is not refused", so we have to do it 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, OAuth2, micro service gateway authentication and so on.
GitHub open source address: https://github.com/dromara/sa-token

Next, let's introduce how to integrate SA token authentication in WebSocket to ensure the security of the connection.

Two integration modes

We will introduce the two most common ways of integrating WebSocket s in turn:

Java Native: javax websocket. Session
Spring encapsulated version: WebSocketSession

Don't talk too much nonsense, just do it:

Method 1: Java Native javax websocket. Session

1. The first is the introduction of POM XML dependency

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

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

<!-- Sa-Token Authority authentication, Online documentation: http://sa-token.dev33.cn/ -->
<dependency>
	<groupId>cn.dev33</groupId>
	<artifactId>sa-token-spring-boot-starter</artifactId>
	<version>1.29.0</version>
</dependency>

2. Login interface, which is used to obtain 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");
	}

	// ... 
	
}

3. WebSocket connection processing

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

    /**
     * Fixed prefix 
     */
    private static final String USER_ID = "user_id_";
	
	 /** 
	  * Store the session set to facilitate pushing messages (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 {
		
		// Get the corresponding userId according to the token 
		Object loginId = StpUtil.getLoginIdByToken(satoken);
		if(loginId == null) {
			session.close();
			throw new SaTokenException("Connection failed, invalid Token: " + satoken);
		}
		
		// put to the set to facilitate subsequent operations 
		long userId = SaFoxUtil.getValueByType(loginId, long.class);
		sessionMap.put(USER_ID + userId, session);
		
		// Give me a hint 
		String tips = "Web-Socket Successfully connected, sid=" + session.getId() + "，userId=" + userId;
		System.out.println(tips);
		sendMessage(session, tips);
	}

	// Listening: 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);
			}
		}
	}
	
	// Listen: receive the message sent by the client 
	@OnMessage
	public void onMessage(Session session, String message) {
		System.out.println("sid For:" + session.getId() + "，From:" + message);
	}
	
	// Listening: exception occurred 
	@OnError
	public void onError(Session session, Throwable error) {
		System.out.println("sid For:" + session.getId() + "，An error occurred");
		error.printStackTrace();
	}
	
	// ---------
	
	// Push the specified message to the client 
	public static void sendMessage(Session session, String message) {
		try {
			System.out.println("towards sid For:" + session.getId() + "，send out:" + message);
			session.getBasicRemote().sendText(message);
		} catch (IOException e) {
			throw new RuntimeException(e);
		}
	}
	
	// Push message to specified user 
	public static void sendMessage(long userId, String message) {
		Session session = sessionMap.get(USER_ID + userId);
		if(session != null) {
			sendMessage(session, message);
		}
	}
	
}

4. WebSocket configuration

/**
 * Enable WebSocket support
 */
@Configuration  
public class WebSocketConfig { 
	
	@Bean  
	public ServerEndpointExporter serverEndpointExporter() {  
		return new ServerEndpointExporter();  
	}
	
}

5. Startup class

@SpringBootApplication
public class SaTokenWebSocketApplication {

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

After construction, start the project

6. Testing

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

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

As shown in the figure:

2. Then we connect to a WebSocket online test page
For example: https://www.bejson.com/httputil/websocket/

Connection address:

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

As shown in the figure:

3. What happens if we enter a wrong token?

As you can see, the connection will be disconnected immediately!

Method 2: Spring encapsulated version: WebSocketSession

1. Ditto: the first is the introduction of POM XML dependency

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

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

<!-- Sa-Token Authority authentication, Online documentation: http://sa-token.dev33.cn/ -->
<dependency>
	<groupId>cn.dev33</groupId>
	<artifactId>sa-token-spring-boot-starter</artifactId>
	<version>1.29.0</version>
</dependency>

2. Login interface, which is used to obtain 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");
	}

	// ... 
	
}

3. WebSocket connection processing

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

    /**
     * Fixed prefix 
     */
    private static final String USER_ID = "user_id_";
    
    /**
     * Store the Session set to facilitate pushing messages
     */
    private static ConcurrentHashMap<String, WebSocketSession> webSocketSessionMaps = new ConcurrentHashMap<>();

    // Listening: connection on 
    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {

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

		// Give me a hint 
		String tips = "Web-Socket Successfully connected, sid=" + session.getId() + "，userId=" + userId;
		System.out.println(tips);
		sendMessage(session, tips);
    }
    
    // Listening: 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 me a hint 
        String tips = "Web-Socket Connection closed, sid=" + session.getId() + "，userId=" + userId;
    	System.out.println(tips);
    }

    // Received message 
    @Override
    public void handleTextMessage(WebSocketSession session, TextMessage message) throws IOException {
    	System.out.println("sid For:" + session.getId() + "，From:" + message);
    }

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

4. WebSocket pre interceptor

/**
 * WebSocket Handshake front interceptor 
 */
public class WebSocketInterceptor implements HandshakeInterceptor {

	// Triggered before handshake (handshake will succeed only if return true)
	@Override
	public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler handler,
			Map<String, Object> attr) {
		
		System.out.println("---- Triggered before handshake " + StpUtil.getTokenValue());
		
		// Refuse handshake without login 
		if(StpUtil.isLogin() == false) {
			System.out.println("---- Unauthorized client, connection failed");
			return false;
		}
		
		// Mark userId, handshake succeeded 
		attr.put("userId", StpUtil.getLoginIdAsLong());
		return true;
	}

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

5. WebSocket configuration

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

}

6. Startup class

/**
 * Sa-Token Integrating WebSocket authentication example 
 */
@SpringBootApplication
public class SaTokenWebSocketSpringApplication {

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

Start the project and start the test

7. Testing

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

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

As shown in the figure:

2. Then open the WebSocket online test page to connect
For example: https://www.bejson.com/httputil/websocket/

Connection address:

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

As shown in the figure:

Note: the reason for using url to pass Token here is that it is more convenient on the third-party test page. In real projects, you can choose one of Cookie, Header parameter and url parameter to pass session Token, and the effect is the same

3. If you enter an incorrect Token

Connection failed!

Example address

The above code has been uploaded to git. Example address:
Code cloud: SA token demo websocket

reference material

Gitee address: https://gitee.com/dromara/sa-token
GitHub address: https://github.com/dromara/sa-token
SA token official website: https://sa-token.dev33.cn/

Topics: Java websocket sa-token

Programmer Think

Using SA token to solve WebSocket handshake authentication

preface

Two integration modes

Method 1: Java Native javax websocket. Session

1. The first is the introduction of POM XML dependency

2. Login interface, which is used to obtain the session token

3. WebSocket connection processing

4. WebSocket configuration

5. Startup class

6. Testing

Method 2: Spring encapsulated version: WebSocketSession

1. Ditto: the first is the introduction of POM XML dependency

2. Login interface, which is used to obtain the session token

3. WebSocket connection processing

4. WebSocket pre interceptor

5. WebSocket configuration

6. Startup class

7. Testing

Example address

reference material

Hot Topics