WebSocket实现鉴权方案

---

WebSocket是一种在单个TCP连接上进行全双工通信的协议。它使得客户端和服务器之间的数据交换变得更加简单，允许服务端主动向客户端推送数据。浏览器只需要和服务端完成一次握手，两端就可以建立持久性链接来进行双向通信。但在实际的应用中为了防止恶意的未经授权的客户端发起链接，此时就需要客户端在与服务端建立链接时携带指定的授权信息，服务端在响应握手时对携带的授权协议信息进行鉴权，校验成功后方可放行从而握手成功建立长链接。那么对于websocket无法携带header 信息和cookie我们该如何向服务端传输授权信息进而鉴权呢（如：自定义token，jwt，用户账号信息等），看下文分解。

---

一、springboot+websocket搭建

1.1.使用依赖

// gradle
implementation 'org.springframework.boot:spring-boot-starter-websocket'

// maven
<dependency>
	<groupId>org.springframework.boot</groupId>
	<artifactId>spring-boot-starter-websocket</artifactId>
</dependency>

1.2.WebSocketInterceptor鉴权拦截器

鉴权逻辑将在这里做，待下文分解

import java.util.Map;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.server.HandshakeInterceptor;

@Component
public class WebSocketInterceptor implements HandshakeInterceptor {
    /**
     * 日志
     */
    private static final Logger log = LoggerFactory.getLogger(WebSocketInterceptor.class);

    /**
     * 握手之前
     *
     * @param request    request
     * @param response   response
     * @param wsHandler  handler
     * @param attributes 属性
     * @return 是否握手成功：true-成功，false-失败
     */
    @Override
    public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler,
            Map<String, Object> attributes) {
        //这里做客户端鉴权业务处理，下文分解
        return true;
    }

    /**
     * 握手后
     *
     * @param request   request
     * @param response  response
     * @param wsHandler wsHandler
     * @param exception exception
     */
    @Override
    public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler,
            Exception exception) {
        log.info("handshake success!");
    }
}

1.3.MyWebSocketHandler处理器

import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

public class MyWebSocketHandler extends TextWebSocketHandler {
    /**
     * 日志
     */
    private static final Logger log = LoggerFactory.getLogger(MyWebSocketHandler.class);

    /**
     * 静态变量，用来记录当前在线连接数
     */
    private static AtomicInteger onlineNum = new AtomicInteger();
    
    /**
     * 存放每个客户端连接对象
     */
    private static ConcurrentHashMap<Integer, WebSocketSession> sessionPools = new ConcurrentHashMap<>();

    /**
     * 在线人数加一
     */
    public static void addOnlineCount() { onlineNum.incrementAndGet(); }

    /**
     * 在线人数减一
     */
    public static void subOnlineCount() {
        onlineNum.decrementAndGet();
    }

    /**
     * 接受客户端消息
     *
     * @param session session
     * @param message message
     * @throws IOException e
     */
    @Override
    public void handleTextMessage(WebSocketSession session, TextMessage message) throws IOException {
        session.sendMessage(new TextMessage(String.format("收到用户：【%s】发来的【%s】",
                session.getAttributes().get("uid"),
                message.getPayload())));
    }

    /**
     * 建立连接后发送消息给客户端
     *
     * @param session session
     * @throws Exception e
     */
    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        String uid = session.getAttributes().get("uid").toString();
        WebSocketSession put = sessionPools.put(Integer.parseInt(uid), session);
        if (put == null) {
            addOnlineCount();
        }
        session.sendMessage(new TextMessage("connection to ws succeeded! online number：" + onlineNum));
    }
    
    /**
     * 连接关闭后
     *
     * @param session session
     * @param status  status
     * @throws Exception e
     */
    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        String uid = session.getAttributes().get("uid").toString();
        sessionPools.remove(Integer.parseInt(uid));
        subOnlineCount();
        log.info("disconnect!");
    }
    
    /**
     * 发送广播消息
     *
     * @param message 消息内容
     */
    public static void sendTopic(String message) {
        if (sessionPools.isEmpty()) {
            return;
        }
        for (Map.Entry<Integer, WebSocketSession> entry : sessionPools.entrySet()) {
            try {
                entry.getValue().sendMessage(new TextMessage(message));
            } catch (IOException e) {
                e.printStackTrace();
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
    }

    /**
     * 点对点发送消息
     *
     * @param uid     用户
     * @param message 消息
     */
    public static void sendToUser(String uid, String message) {
        WebSocketSession socketSession = sessionPools.get(Integer.parseInt(uid));
        if (socketSession == null) {
            return;
        }
        try {
            socketSession.sendMessage(new TextMessage(message));
        } catch (IOException e) {
            log.error("send to user:{}, error! data:{}", uid, message);
        }
    }
}

1.4.WebSocketConfig配置

import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.config.annotation.EnableWebSocket;
import org.springframework.web.socket.config.annotation.WebSocketConfigurer;
import org.springframework.web.socket.config.annotation.WebSocketHandlerRegistry;
import javax.annotation.Resource;

@Configuration
@EnableWebSocket
public class WebSocketConfig implements WebSocketConfigurer {

    /**
     * 注入拦截器
     */
    @Resource
    private WebSocketInterceptor webSocketInterceptor;


    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry webSocketHandlerRegistry) {

        webSocketHandlerRegistry
                //添加myHandler消息处理对象，和websocket访问地址
                .addHandler(myHandler(), "/ws")
                //设置允许跨域访问
                .setAllowedOrigins("*")
                //添加拦截器可实现用户链接前进行权限校验等操作
                .addInterceptors(webSocketInterceptor);
    }

    @Bean
    public WebSocketHandler myHandler() {
        return new MyWebSocketHandler();
    }
}

1.5.前端连接

<!DOCTYPE html>
<html lang="en">
<head>
    <meta http-equiv="Content-Type" content="text/html; charset=utf-8" />
    <title>webSocket客户端</title>
</head>
<body>
        <h4>客户端输入：</h4>
        <textarea id = "message" name="message" style="width: 200px;height: 100px"></textarea>
        <br/>
        <input type="button" value="发送到服务器" onclick="sendMessage()" />
    
 
        <h4>服务器返回消息：</h4>
        <textarea id = "responseText" name="message" style="width: 1100px;height: 100px"></textarea>
        <br/>
        <input type="button" οnclick="javascript:document.getElementById('responseText').value=''" value="clear data">
 <script type="text/javascript">
    function send(){
        alert(2);
    }
    var webSocket;
    if(window.WebSocket){
		//建立长链接，地址为你的服务端地址且以ws开头，以WebSocketConfig配置中的路径 .addHandler(myHandler(), "/ws")结尾
		//如果服务端有配置上下文 context-path，需要加入路径
        webSocket = new WebSocket("ws://127.0.0.1:8080/ws");
       
        webSocket.onmessage = function (ev) {
            var contents = document.getElementById("responseText");
            contents.value = contents.value +"\n"+ ev.data;
        }
 
        webSocket.onopen = function (ev) {
            var contents = document.getElementById("responseText");
            contents.value = "与服务器端的websocket连接建立";
            var data = '{"method":"init","identifier":"11VKF7M0020199"}';
            webSocket.send(data);
        }
        webSocket.onclose = function (ev) {
            var contents = document.getElementById("responseText");
            contents.value =  contents.value +"\n"+ "与服务器端的websocket连接断开";
        }
    }else{
        alert("该环境不支持websocket")
    }
 
    function sendMessage() {
        //alert(document.getElementById("message").value);
        if(window.webSocket){
            if(webSocket.readyState == WebSocket.OPEN){
                var data2 = '{"method":"video","serialNumber":"yjdp"}';
                var data= '{"method":"video","identifier":"11VKF7M0020199","toIdentifier":"yjdp","status":"'+document.getElementById("message").value+'","url":"http://127.0.0.1:8081/UAV_FILES/FILES/0UYKG7K002001F.mp4"}';
                webSocket.send(data);
            }else{
                alert("与服务器连接尚未建立")
            }
        }
    }
    
    
</script>   
</body>
</html>

二、websocket鉴权方案

2.1.url传参方案

页面建立连接方式：

//如带账号密码方式（严谨的话账号密码是要密文方式）或者只带一个token
webSocket = new WebSocket("ws://127.0.0.1:8080/ws?userId=10001&pwd=123456");

服务端接受方式（WebSocketInterceptor中 beforeHandshake方法）：

public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Map<String, Object> attributes) {
        ServletServerHttpRequest serverHttpRequest = (ServletServerHttpRequest) request;
        //获取参数
        String userId = serverHttpRequest.getServletRequest().getParameter("userId");
        String pwd = serverHttpRequest.getServletRequest().getParameter("pwd");
        //todo 这里根据自己业务对userId和pwd进行校验，校验通过则返回true, 失败返回false
        //....
        //这里放入uid是个人需求，因为在MyWebSocketHandler我有对在线人数进行统计需要，还有对每个session存储对应关系。
        attributes.put("uid", userId);
        return true;
    }

2.2.websocket头字段Sec-WebSocket-Protocol传参

为什么需要头字段传参呢？主要有两个原因：
1.头部参数相比url参数会更安全一丢丢。
2.对于需要传输jwt这种非常长的token也不适合在url中进行传参。

页面建立连接方式（token为你的jwt授权码或者自定义的授权码）：

//如带账号密码方式（严谨的话账号密码是要密文方式）或者只带一个token
webSocket = new WebSocket("ws://127.0.0.1:8080/ws", [token]);

服务端接受方式（WebSocketInterceptor中 beforeHandshake方法）：

public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Map<String, Object> attributes) {
        ServletServerHttpRequest serverHttpRequest = (ServletServerHttpRequest) request;
        ServletServerHttpResponse serverHttpResponse = (ServletServerHttpResponse) response;
        String authorization = serverHttpRequest.getServletRequest().getHeader("Sec-WebSocket-Protocol");
        //这里对获取到的 authorization 授权码进行业务校验如 jwt 校验
        //...
        //在后端握手时设置一下请求头（Sec-WebSocket-Protocol），前端发来什么授权值，这里就设置什么值，不设置会报错导致建立连接成功后立即被关闭
        serverHttpResponse.getServletResponse().setHeader("Sec-WebSocket-Protocol", authorization);
        log.info("start shaking hands->>>");
        return true;
    }

Sec-WebSocket-Protocol头字段用于WebSocket打开阶段握手。它从客户端发送到服务器端，并从服务器端发回到客户端来确定连接的子协议。如果前端有设置值使用该方式鉴权，那服务端需要原封不动回传该值，不回传会报以下错误。

Error during WebSocket handshake: Sent non-empty ‘Sec-WebSocket-Protocol’ header but no response was received