Springboot中添加原生websocket支持

1、添加配置

@Configuration
@EnableWebSocket
public class WebSocketConfig implements WebSocketConfigurer {@Overridepublic void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {// 注册WebSocket处理器，并允许所有来源的连接（在生产环境中应限制来源）registry.addHandler(new WebSocketHandler(), "/ws/[请求的地址]").setAllowedOrigins("*").addInterceptors(new WebSocketSecurityTokenInterceptor());}
}

2、添加Handler对请求进行处理

@Component
@Slf4j
public class WebSocketHandler extends TextWebSocketHandler {private static final CopyOnWriteArrayList<WebSocketSession> sessions = new CopyOnWriteArrayList<>();private static final ScheduledExecutorService scheduledThreadPool = Executors.newScheduledThreadPool(Runtime.getRuntime().availableProcessors());private static final Map<String,Long> lastPongTimes = new ConcurrentHashMap<>();private static final String PING = "ping";private static final String PONG = "pong";private static final String GET_DATA = "getData";/*** 检测心跳是否正常的周期时间*/private static final Integer heartbeatInterval = 30_000;/*** 检测客户端连接心跳保持时间是否超时的时间*/private static final Integer heartbeatTimeout = 60_000;@Resourceprivate ReportDashboardService reportDashboardService;// 使用Guava 弱引用缓存数据private static final Cache<String, Object> CACHE = CacheBuilder.newBuilder().softValues().expireAfterWrite(3, TimeUnit.SECONDS).build();@PostConstructpublic void init() {scheduledThreadPool.scheduleWithFixedDelay(() -> {try {checkHeartbeats();} catch (IOException e) {throw new RuntimeException(e);}},heartbeatInterval,heartbeatInterval,TimeUnit.SECONDS);}@Overrideprotected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {// 广播消息给所有已连接的客户端String payload = message.getPayload();if (PING.equals(payload)) {session.sendMessage(new TextMessage(PONG));} else if (GET_DATA.equals(payload)) {sendData(session);} else {sendDataByPayload(session,payload);//broadcast(payload);}recordPong(session.getId());}@Overridepublic void afterConnectionEstablished(WebSocketSession session) throws Exception {// 当新连接建立时添加到列表sessions.add(session);//session.sendMessage(new TextMessage(PONG));recordPong(session.getId());sendData(session);}@Overridepublic void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {// 当连接关闭时从列表中移除sessions.remove(session);log.info("Connection closed.sessionId={},status={}",session.getId(),status);}private void sendData(WebSocketSession sess) {try {if (sess.isOpen()) {String query = sess.getUri().getQuery();String[] split = query.split("&");ReqObj req = new ReqObj();String path = "";for (String s : split) {String[] arr = s.split("=");if ("path".equals(arr[0])) {path = arr[1];} else if ("deviceCode".equals(arr[0])) {req.setDeviceCode(arr[1]);} else if ("pointType".equals(arr[0])) {req.setPointType(arr[1]);} else if ("gap".equals(arr[0])) {if (arr[1] != null) {req.setGap(Integer.parseInt(arr[1]));}}}sess.sendMessage(new TextMessage(sendDataByPath(path,req)));}} catch (Exception e) {System.err.println("Failed to send message: " + e.getMessage());}}private boolean closeTimeoutSession(String sessionId) throws IOException {WebSocketSession s = null;for (WebSocketSession sess : sessions) {if (sess.isOpen() && sess.getId().equals(sessionId)) {sess.sendMessage(new TextMessage("当前连接1分钟内未发送心跳消息,即将关闭"));s = sess;}}log.info("关闭心跳超过的连接,sessionId={}",sessionId);return s != null && sessions.remove(s);}private void recordPong(String sessionId) {lastPongTimes.put(sessionId,System.currentTimeMillis());}private boolean isClientAlive(String sessionId) {Long lastPongTime = lastPongTimes.get(sessionId);if (lastPongTime == null){return false;}return System.currentTimeMillis() - lastPongTime <= heartbeatTimeout;}private void checkHeartbeats() throws IOException {log.info("开始检查连接的心跳是否超时......");Set<Map.Entry<String, Long>> entries = lastPongTimes.entrySet();for (Map.Entry<String, Long> entry : entries) {String sessionId = entry.getKey();log.info("sessionId = {}",sessionId);if (!isClientAlive(sessionId)) {closeTimeoutSession(sessionId);}}}private void sendDataByPayload(WebSocketSession sess,String payload){try {if (sess.isOpen()) {ChartDto dto = null;String cacheKey = null;try {cacheKey = MD5Util.encrypt(payload);dto = JSON.parseObject(payload, ChartDto.class);} catch (Exception e) {log.error("将payload转为ChartDto对象失败");}if (dto == null) {JSONObject jsonObject = JSON.parseObject(payload);String path = jsonObject.getString("path");ReqObj req = new ReqObj();req.setGap(jsonObject.getIntValue("gap",0));req.setDeviceCode(jsonObject.getString("deviceCode"));req.setPointType(jsonObject.getString("pointType"));String s = sendDataByPath(path, req);sess.sendMessage(new TextMessage("{\"path\":\""+path+"\",\"data\":"+s+"}"));} else {if (reportDashboardService == null) {reportDashboardService = SpringUtil.getBean(ReportDashboardService.class);}synchronized (Thread.currentThread()) {Object data = CACHE.getIfPresent(cacheKey);if (data == null) {data = reportDashboardService.getChartData(dto);if (data != null) {CACHE.put(cacheKey,data);}}String s = JSON.toJSONString(R.success(data,"success",dto.getId()));sess.sendMessage(new TextMessage(s));}}}} catch (Exception e) {System.err.println("Failed to send message: " + e.getMessage());}}private String sendDataByPath(String path,ReqObj req) {return "{}";}private void broadcast(String message) {for (WebSocketSession sess : sessions) {try {if (sess.isOpen()) {String data = "原始数据:"+message+",翻转后的数据:"+new StringBuilder(message).reverse();sess.sendMessage(new TextMessage(data));}} catch (Exception e) {System.err.println("Failed to send message: " + e.getMessage());}}}
}

3、拦截器握手时进行校验token

@Getter
@Slf4j
@Component
public class WebSocketSecurityTokenInterceptor implements HandshakeInterceptor {private TokenAcquireHandler tokenAcquireHandler;private TokenAnalysisHandler tokenAnalysisHandler;{tokenAcquireHandler = SpringUtil.getOrDefault( TokenAcquireHandler.class, new DefaultTokenAcquireHandler() );tokenAnalysisHandler = SpringUtil.getOrDefault( TokenAnalysisHandler.class, new DefaultTokenAnalysisHandler() );}@Overridepublic boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Map<String, Object> attributes) throws Exception {//  放开的路径直接放行if ( FilterContextHandler.getContext().isExclude() ) {//  如果已经手动setContext 此处不再赋值Emptyif ( SecurityContextHandler.getContext() == null ) {SecurityContextHandler.setContext( SecurityContext.EMPTY );}return true;}String token = getToken(request);if ( !StringUtils.hasText( token ) ) {throw new TokenNotFoundException( "token not found" );}UserDetails userDetails = tokenAnalysisHandler.analysisToken( token );checkUserDetails( token, userDetails );SecurityContextHandler.setContext( new SecurityContext( token, userDetails ) );return true;}/*** 校验用户信息*/private void checkUserDetails( String token, UserDetails userDetails ) {//  解析的UserDetails不能为空if ( userDetails == null ) {throw new TokenAnalysisException( "token analysis userDetails cannot be empty" );}//  判断用户是否启用if ( !userDetails.isEnabled() ) {throw new TokenAnalysisException();}//  判断用户是否过期if ( userDetails.isAccountNonExpired() ) {throw new UserDetailsExpiredException();}//  判断用户是否锁定if ( userDetails.isAccountNonLocked() ) {throw new UserLockException();}//  判断Token是否过期if ( userDetails.isCredentialsNonExpired( token ) ) {throw new TokenExpiredException();}}@Overridepublic void afterHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Exception exception) {// 握手完成后进行一些初始化工作//log.info("握手完成......");}private String getToken( ServerHttpRequest req ) {List< String > headerList = req.getHeaders().get( HttpHeaders.AUTHORIZATION );String token = CollectionUtils.isEmpty( headerList ) ? "" : headerList.get( 0 );if ( StrUtil.isNotBlank( token ) ) {
//            req.setAttribute( HttpHeaders.AUTHORIZATION, token );return token;}List< String > cookies = req.getHeaders().get( HttpHeaders.COOKIE );for (String cookieStr : Optional.ofNullable(cookies).orElse(Collections.emptyList())) {HttpCookie cookie = parseAuthCookie(cookieStr);if ( cookie != null ){return cookie.getValue();}}return null;}private HttpCookie parseAuthCookie(String cookieStr) {if (!StringUtils.hasText(cookieStr)){return null;}List<HttpCookie> cookieList = Arrays.stream(cookieStr.split(";")).map(this::parseCookie).filter(Objects::nonNull).collect(Collectors.toList());for (HttpCookie cookie : cookieList) {if ( HttpHeaders.AUTHORIZATION.equals( cookie.getName() ) ) {return cookie;}}return null;}private HttpCookie parseCookie(String cookieStr) {try {List<HttpCookie> cookies = HttpCookie.parse(cookieStr);return CollectionUtils.isEmpty(cookies) ? null : cookies.get(0);}catch (Exception e){return null;}}
}