springboot整合websocket进行鉴权遇到的问题

一、问题背景

之前项目遇到过的问题，就是需要通过websocket给前端和小程序推送数据。因为nginx和wss的问题就不提，终于是通信连上并能发送接收数据了。但是之后有遇到一个需要鉴权的问题，之前用websocket没怎么考虑到鉴权的问题，正常使用起来好像也没有办法带token来过权限，都是直接把security解开来用。

二、自己的想法

一开始首先想到的就是在后面多带几个参数，通过@PathParam取出来只来进行验证，最简单的例如传个明文和一个密文，拿到后对比确认后，才把session放入集合里，再进行发送数据。但是后来想了一下，这样还是不可避免的会被别人连接上也防止不了别人推送数据上来，所以还是不行。

三、网上看到的方法

之前是通过@ServerEndpoint注入的：

@Configuration
public class WebSocketConfig //implements WebSocketConfigurer
{

    /**
     * 使用spring boot时，使用的是spring-boot的内置容器，
     * 如果要使用WebSocket，需要注入ServerEndpointExporter
     *
     * @return
     */
    @Bean
    public ServerEndpointExporter serverEndpointExporter() {
        return new ServerEndpointExporter();
    }
}

@ServerEndpoint(value="/ocwebsocket/{username}")
@Component
@Slf4j
public class WebsocketService {
}

在网上看到了一个方法，是设置拦截器，实现HandshakeInterceptor接口，然后在config里面重写registerWebSocketHandlers方法，把拦截器和文字数据处理都设置进去。拦截器里还是用两个参数一个明文一个密文进行比对判断是否有权限连接。

@Component
public class MyHandshakeInterceptor implements HandshakeInterceptor {
    /**
     * 握手之前，若返回false，则不建立链接 *
     *
     * @param request
     * @param response
     * @param wsHandler
     * @param attributes
     * @return
     */
    @Override
    public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse
            response, WebSocketHandler wsHandler, Map<String, Object> attributes) {
        //将用户id放入socket处理器的会话(WebSocketSession)中
        ServletServerHttpRequest serverHttpRequest = (ServletServerHttpRequest) request;
        //获取参数
        String userId = serverHttpRequest.getServletRequest().getParameter("userId");
        String sign = serverHttpRequest.getServletRequest().getParameter("sign");

        attributes.put("uid", userId);
        attributes.put("sign", sign);
        //可以在此处进行权限验证，当用户权限验证通过后，进行握手成功操作，验证失败返回false
        if (!SecureUtil.md5("54321"+userId+"12345").equals(sign)) {
            System.out.println("握手失败.....");
            return false;
        }
        System.out.println("开始握手。。。。。。。");
        return true;
    }

    @Override
    public void afterHandshake(ServerHttpRequest request, ServerHttpResponse
            response, WebSocketHandler wsHandler, Exception exception) {
        System.out.println("握手成功啦。。。。。。");
    }

}

@Configuration
@EnableWebSocket
public class WebSocketConfig implements WebSocketConfigurer {
    @Resource
    private MyHandshakeInterceptor myHandshakeInterceptor;

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

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

@Component
public class MyWebSocketHandler extends TextWebSocketHandler {

    //静态变量，用来记录当前在线连接数。应该把它设计成线程安全的。
    private static AtomicInteger onlineNum = new AtomicInteger();
    //concurrent包的线程安全Set，用来存放每个客户端对应的WebSocketServer对象。
    private static ConcurrentHashMap<String, WebSocketSession> sessionPools = new ConcurrentHashMap<>();

    @Override
    public void handleTextMessage(WebSocketSession session, TextMessage message)
            throws IOException {
        System.out.println("获取到消息 >> " + message.getPayload());
        session.sendMessage(new TextMessage( message.getPayload()));
    }

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws
            Exception {
        System.out.println("获取到拦截器中用户ID : " + session.getAttributes().get("uid"));
        String uid = session.getAttributes().get("uid").toString();
        //TODO: 重复链接没有进行处理
        sessionPools.put(uid, session);
        addOnlineCount();
        System.out.println(uid + "加入webSocket！当前人数为" + onlineNum);
//        session.sendMessage(new TextMessage("欢迎连接到ws服务! 当前人数为：" + onlineNum));
    }


    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status)
            throws Exception {
        System.out.println("断开连接!");
        String uid = session.getAttributes().get("uid").toString();
        sessionPools.remove(uid);
        subOnlineCount();
    }

    /**
     * 添加链接人数
     */
    public static void addOnlineCount() {
        onlineNum.incrementAndGet();
    }

    /**
     * 移除链接人数
     */
    public static void subOnlineCount() {
        onlineNum.decrementAndGet();
    }

    public void sendMessageToGpsWeb(String message) throws IOException {
        for (WebSocketSession item : sessionPools.values()) {
            item.sendMessage(new TextMessage(message));
        }
    }
}

但是紧接着又遇到了一个问题，在测试服务器http的时候即ws是没有任何问题的，接入socket和推送数据，也确实通过拦截器在握手环节就处理了鉴权连接的问题。但是当放进https域名的服务器里的时候，就一直连不上socket。（后来写这边文章的时候，发现应该是我security没开口，所以连不上。）

四、另一个决定使用的方法

最后也是看到一个网友的想法觉得很妙，也行之有效，就决定采用了这个方法。还是用security进行鉴权，通过把websocket的这个handle或者叫service变成controller注入到spring里面，而不再是Component仅仅注入成一个bean。这样通过了security的鉴权就能正常调用或者说连接这个websocket的接口了。前端呢就通过Sec-WebSocket-Protocol来携带token上来。

• 首先是后端这边一切如常，还是用@ServerEndpoint的方法，但是呢因为通过Sec-WebSocket-Protocol传了token上来，我们就需要对之前校验token的方法进行修改。

  private String getToken(HttpServletRequest request) {
      String token = request.getHeader(header);
      if (StringUtils.isNotEmpty(token) && token.startsWith(Constants.TOKEN_PREFIX)) {
          token = token.replace(Constants.TOKEN_PREFIX, "");
      }else if(StringUtils.isEmpty(token) && ObjectUtil.isNotEmpty(request.getHeader(Constants.WEBSOCKET_PROTOCOL))){
          //如果未从请求头中获取到token,则尝试从sec_websocket_protocol中取出
          token = request.getHeader(Constants.WEBSOCKET_PROTOCOL);
          //如果有值，一定要在response的header中设置，否则还是会断开
          if (StringUtils.isNotEmpty(token)){
              HttpServletResponse response = ServletUtils.getResponse();
              response.setHeader(Constants.WEBSOCKET_PROTOCOL, token);
          }
      }
      return token;
  }
  

  @ServerEndpoint(value="/ocwebsocket/{username}")
  @RestController//用RestController代替了Component之后就可以通过token验证
  @Slf4j
  public class WebsocketService {
  

• 其次就是前端了，也很简单，带上token就完事了。

  socket = new WebSocket("ws://localhost:9999/ocwebsocket/web"+(Math.random()*10000000).toString(16).substr(0,4)+'-'+(new Date()).getTime()+'-'+Math.random().toString().substr(2,5),["eyJhbGciOiJIUzUxMiJ9.eyJsb2dpbl91c2VyX2tleSI6IjA3YmY5YmM3LWIwZDEtNGM3MC1hYjRjLWNhYzdmYzY0YWYzMiJ9.iQV0z8529kovWkRa9LCTLD3tylZF4RSdCXT6odgSUOMvB6XT_1ophys3RljAv7PHaKy8H8Xnqnr3SDXUprVVAA"]);
  

这个方法相比之前那个改动就少了很多。而且也不用开口，不用再设计校验的方法，还是用原本的token就可以了，过期时间什么的都比较方便。

五、结语

有些东西多折腾，会发现很多原本学习的时候没有遇到过实际情况中的小细节问题。