webSocket+jwt實現方式

背景：

原項目是通過前端定時器獲取消息，存在消息滯后、空刷服務器、浪費帶寬和資源的問題，在springboot項目集成websocket可以實現實時點對點消息推送。

原項目是在header添加jwt令牌實現認證，由於websocket不支持在頭部添加信息（或許是我打開的方式不對？），最終只能采用在url添加令牌參數實現認證，感覺不夠優雅，后續再想辦法重構改進。

ps：至於放行websocket相關url，完全不要去考慮，危害巨大。

 

1、websocket核心依賴

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

2、config

@Configuration
public class WebSocketConfig {
    @Bean
    public ServerEndpointExporter serverEndpointExporter() {
        return new ServerEndpointExporter();
    }
}

3、WebSocketServer

@Slf4j
@ServerEndpoint("/webSocket/{code}")
@Component
public class WebSocketServer {
    /**
     * concurrent包的線程安全Set，用來存放每個客戶端對應的WebSocket對象。
     */
    private static CopyOnWriteArraySet<WebSocketServer> webSocketSet = new CopyOnWriteArraySet<>();

    /**
     * 與客戶端的連接會話，需要通過它來給客戶端發送數據
     */
    private Session session;

    /**
     * 接收識別碼
     */
    private String code = "";

    /**
     * 連接建立成功調用的方法
     */
    @OnOpen
    public void onOpen(Session session, @PathParam("code") String code) {
        this.session = session;
        //如果存在就先刪除一個，防止重復推送消息，實際這里實現了set，不刪除問題也不大
        webSocketSet.removeIf(webSocket -> webSocket.code.equals(code));
        webSocketSet.add(this);
        this.code = code;
        log.info("建立WebSocket連接，code：" + code+"，當前連接數："+webSocketSet.size());
    }

    /**
     * 連接關閉調用的方法
     */
    @OnClose
    public void onClose() {
        webSocketSet.remove(this);
        log.info("關閉WebSocket連接，code：" + this.code+"，當前連接數："+webSocketSet.size());
    }

    /**
     * 收到客戶端消息后調用的方法
     *
     * @param message 客戶端發送過來的消息
     */
    @OnMessage
    public void onMessage(String message, Session session) {
        log.info("收到來[" + code + "]的信息:" + message);

    }

    @OnError
    public void onError(Session session, Throwable error) {
        log.error("websocket發生錯誤");
        error.printStackTrace();
    }

    /**
     * 實現服務器主動推送
     */
    private void sendMessage(String message) throws IOException {
        this.session.getBasicRemote().sendText(message);
    }


    /**
     * 群發自定義消息
     */
    public void sendAll(String message) {
        log.info("推送消息到" + code + "，推送內容:" + message);
        for (WebSocketServer item : webSocketSet) {
            try {
                item.sendMessage(message);
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }

    /**
     * 定點推送
     */
    public void sendTo(String message, @PathParam("code") String code) {
        for (WebSocketServer item : webSocketSet) {
            try {
                if (item.code.equals(code)) {
                    log.info("推送消息到[" + code + "]，推送內容:" + message);
                    item.sendMessage(message);
                }
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }

    @Override
    public boolean equals(Object o) {
        if (this == o) {
            return true;
        }
        if (o == null || getClass() != o.getClass()) {
            return false;
        }
        WebSocketServer that = (WebSocketServer) o;
        return Objects.equals(session, that.session) &&
                Objects.equals(code, that.code);
    }

    @Override
    public int hashCode() {
        return Objects.hash(session, code);
    }
}

4、令牌過濾器

@Slf4j
@Component
public class JwtTokenFilter extends OncePerRequestFilter {

    @Resource
    JwtProperties jwtProperties;

    @Resource
    TokenProvider tokenProvider;

    @Resource
    OnlineUserService onlineUserService;

    @Override
    protected void doFilterInternal(HttpServletRequest request, HttpServletResponse response, FilterChain filterChain) throws ServletException, IOException {
        // http連接時，客戶端應該是在頭信息中攜帶令牌
        String authorizationHeader = request.getHeader(jwtProperties.getHeader());
        if(StringUtils.isBlank(authorizationHeader)) {
            // websocket連接時，令牌放在url參數上，以后重構
            authorizationHeader = request.getParameter(jwtProperties.getHeader());
        }

        String token = null;
        if(!StringUtils.isEmpty(authorizationHeader) && authorizationHeader.startsWith(jwtProperties.getTokenStartWith())){
            token = authorizationHeader.replace(jwtProperties.getTokenStartWith(),"");
        }
        //驗證token
        if(StringUtils.isNotBlank(token) && tokenProvider.validateToken(token)){
            //驗證token是否在緩存中
            OnlineUserDto onlineUserDto = onlineUserService.getOne(jwtProperties.getOnlineKey() + token);
            if(onlineUserDto!=null){
                Authentication authentication = tokenProvider.getAuthentication(token, request);
                SecurityContextHolder.getContext().setAuthentication(authentication);
                log.debug("set Authentication to security context for '{}', uri: {}", authentication.getName(), request.getRequestURI());
            }
        }

        filterChain.doFilter(request, response);
    }
}

5、在業務中調用方式（偽代碼）

    @Resource
    private WebSocketServer webSocketServer;

    // 向客戶端推送實時消息
    webSocketServer.sendTo(content, sysUser.getId());

6、前端，偽代碼

    getMessageCount() {
      getMyMessageCount().then(res => {
        const count = res
        this.messageCount = count > 0 ? count : null
      })
    },
    initWebSocket() {
      const wsUri = process.env.VUE_APP_WS_API + '/webSocket/' + this.user.id + '?Authorization=' + getToken()
      this.websock = new WebSocket(wsUri)
      this.websock.onopen = this.webSocketOnOpen
      this.websock.onerror = this.webSocketOnError
      this.websock.onmessage = this.webSocketOnMessage
    },
    webSocketOnOpen(e) {
      console.log('websocket 已經連接', e)
    },
    webSocketOnError(e) {
      this.$notify({
        title: 'WebSocket連接發生錯誤',
        type: 'error',
        duration: 0
      })
    },
    webSocketOnMessage(e) {
      const data = e.data
      this.$notify({
        title: '',
        message: data,
        type: 'success',
        dangerouslyUseHTMLString: true,
        duration: 5500
      })
      this.getMessageCount()
    }