Spring boot Websocket 添加授权校验（校验Header）

        最近需要用到websocket,并且确认是在登录成功的情况下才能连接，看了下，websocket 5.8以上已经支持校验了，但是由于我项目是spring boot 2.7.9,为了兼容，只能考虑通过在websocket 的header中带上验证信息，进行校验，详细如下

1.首先往spring boot pom 内加入 websocket的依赖

我这里spring boot 的版本为2.7.9 ，所以对应使用websocket starter 的版本也是2.7.9

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

2.获取并设置连接时Hearder内的信息保存到session中

继承ServerEndpointConfig.Configurator 并重写modifyHandshake方法，在这里拦截header中Authorization的内容，并将其设置到websocket session 中

@Configuration
public class WebSocketConfig extends ServerEndpointConfig.Configurator {
    @Bean
    public ServerEndpointExporter serverEndpointExporter() {
        return new ServerEndpointExporter();
    }

    @Override
    public void modifyHandshake(ServerEndpointConfig sec, HandshakeRequest request, HandshakeResponse response) {
        Map userProperties = sec.getUserProperties();
        Map> headers = request.getHeaders();
        List authorization = headers.get("Authorization");
        if (!CollectionUtils.isEmpty(authorization)) {
            userProperties.put("Authorization", authorization.get(0));
        }
    }
}

3.校验header中的授权信息

获取之前设置入session中的Authorization信息，这里为了方便，仅仅只是校验了有没有传入，没有传入则代表校验失败

@ServerEndpoint(value = "/websocket", configurator = WebSocketConfig.class)
@Service
public class WebSocketTest {
    private static CopyOnWriteArraySet webSocketSet = new CopyOnWriteArraySet();

    private Session session;

    @OnOpen
    public void onOpen(Session session) {
        System.out.println("Authorization:" + session.getUserProperties().get("Authorization"));
        if (StringUtils.isEmpty(session.getUserProperties().get("Authorization"))) {
            throw new RuntimeException("Login Fail");
        }
        this.session = session;
        webSocketSet.add(this);
    }

    @OnClose
    public void onClose() {
        webSocketSet.remove(this); 
    }

    @OnMessage
    public void onMessage(String message, Session session) {

        for (WebSocketTest item : webSocketSet) {
            try {
                item.sendMessage(message);
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }

    @OnError
    public void onError(Session session, Throwable error) {
        System.out.println("发生错误");
        error.printStackTrace();
    }


    public void sendMessage(String message) throws IOException {
        this.session.getBasicRemote().sendText(message);
    }

    public static void sendInfo(String message) throws IOException {
        for (WebSocketTest item : webSocketSet) {
            try {
                item.sendMessage(message);
            } catch (IOException e) {
                continue;
            }
        }
    }
}

4.测试websocket在传入和不传入Authorization内容时的区别

取消送入header内容Authorization时，连接先成功，后失败，则代表已经连接到服务器，但因为校验失败，被断开了连接

日志输出为

添加header内容后，通过校验连接成功

这样一个简易的websocket校验就完成了

2023-04-20 修改

根据前端反馈，在JS 中，websocket使用token貌似并不能自定义header名称

需要使用Sec-WebSocket-Protocol作为header 的名称接收

前端使用代码 new WebSocket(getWebsoctUrl,[token])