JWT过期时关闭Spring WebSocket连接的最优实现方案

我正在实现Spring WebSocket端点，希望仅对已认证用户开放，当前应用采用纯JWT认证（无OAuth）。目前通过将初始HTTP请求端点置于Spring Security之后，已实现仅允许认证用户建立WebSocket连接，但存在JWT过期后用户仍保持连接的问题。我曾尝试使用OutBoundChannelInterceptor从WebSocket会话属性中验证JWT，验证通过则发送消息，否则抛出异常，但该拦截器无法捕获消息发送事件，且仅在用户触发操作时才会关闭连接。请问解决此问题的最优方案是什么？

核心思路是主动检测会话中JWT的有效性，结合定时检查与即时验证，在JWT过期时主动断开连接，同时收紧WebSocket权限控制。具体实现如下：

1. 完善握手拦截器，存储JWT与过期时间 #

修改WebSocketHandshakeInterceptor，在握手阶段直接验证JWT有效性，并将过期时间存入会话属性，避免后续重复解析：

public class WebSocketHandshakeInterceptor implements HandshakeInterceptor {
    private static final String TOKEN_COOKIE_NAME = "authToken";
    private static final String TOKEN_EXPIRY_ATTR = "tokenExpiry";

    @Autowired
    private TokenHandler tokenHandler;

    @Override
    public boolean beforeHandshake(ServerHttpRequest request,
                                   ServerHttpResponse response,
                                   WebSocketHandler wsHandler,
                                   Map<String, Object> attributes) {
        HttpServletRequest httpRequest = ((ServletServerHttpRequest) request).getServletRequest();
        Cookie[] cookies = httpRequest.getCookies();
        String token = getAuthTokenFromCookie(cookies);

        if (token == null) {
            response.setStatusCode(HttpStatus.UNAUTHORIZED);
            return false;
        }

        try {
            tokenHandler.parseUserFromToken(token);
            Date expiryDate = tokenHandler.getExpiryDateFromToken(token);
            attributes.put(TOKEN_COOKIE_NAME, token);
            attributes.put(TOKEN_EXPIRY_ATTR, expiryDate.getTime());
            return true;
        } catch (Exception e) {
            response.setStatusCode(HttpStatus.UNAUTHORIZED);
            return false;
        }
    }

    private String getAuthTokenFromCookie(Cookie[] cookies) {
        if (cookies == null) return null;
        for (Cookie cookie : cookies) {
            if (TOKEN_COOKIE_NAME.equals(cookie.getName())) {
                return cookie.getValue();
            }
        }
        return null;
    }

    @Override
    public void afterHandshake(ServerHttpRequest serverHttpRequest,
                               ServerHttpResponse serverHttpResponse,
                               WebSocketHandler webSocketHandler,
                               Exception e) {}
}

2. 实现会话管理器，定时检测并断开过期连接 #

创建WebSocket会话管理器，定时遍历活跃会话，检查JWT是否过期，主动断开失效连接：

@Component
public class WebSocketSessionManager {
    private static final String TOKEN_EXPIRY_ATTR = "tokenExpiry";
    private final SimpMessagingTemplate messagingTemplate;
    private final Set<SimpSession> activeSessions = Collections.synchronizedSet(new HashSet<>());

    @Autowired
    public WebSocketSessionManager(SimpMessagingTemplate messagingTemplate) {
        this.messagingTemplate = messagingTemplate;
        // 每30秒检查一次过期会话
        ScheduledExecutorService scheduler = Executors.newSingleThreadScheduledExecutor();
        scheduler.scheduleAtFixedRate(this::checkExpiredSessions, 0, 30, TimeUnit.SECONDS);
    }

    public void addSession(SimpSession session) {
        activeSessions.add(session);
    }

    public void removeSession(SimpSession session) {
        activeSessions.remove(session);
    }

    private void checkExpiredSessions() {
        long currentTime = System.currentTimeMillis();
        Iterator<SimpSession> iterator = activeSessions.iterator();
        while (iterator.hasNext()) {
            SimpSession session = iterator.next();
            Object expiryAttr = session.getAttributes().get(TOKEN_EXPIRY_ATTR);
            if (expiryAttr != null && currentTime > (long) expiryAttr) {
                messagingTemplate.sendToUser(session.getId(), "/queue/expired", "Token expired");
                session.disconnect();
                iterator.remove();
            }
        }
    }
}

在WebSocketConfig中注册会话事件监听器，维护活跃会话列表：

@Configuration
@EnableWebSocketMessageBroker
public class WebSocketConfig implements WebSocketMessageBrokerConfigurer {
    @Autowired
    private WebSocketSessionManager sessionManager;

    @Override
    public void configureMessageBroker(MessageBrokerRegistry config) {
        config.enableSimpleBroker("/topic");
        config.setApplicationDestinationPrefixes("/app");
    }

    @Override
    public void registerStompEndpoints(StompEndpointRegistry registry) {
        registry.addEndpoint("/api/ws-message")
                .setAllowedOrigins("http://localhost:3000")
                .addInterceptors(new WebSocketHandshakeInterceptor())
                .withSockJS();
    }

    @Override
    public void configureClientInboundChannel(ChannelRegistration registration) {
        registration.interceptors(new AuthChannelInterceptor());
    }

    @EventListener
    public void handleSessionConnected(SessionConnectedEvent event) {
        sessionManager.addSession(event.getSession());
    }

    @EventListener
    public void handleSessionDisconnect(SessionDisconnectEvent event) {
        sessionManager.removeSession(event.getSession());
    }
}

3. 优化入站拦截器，即时验证用户操作 #

修改AuthChannelInterceptor，在用户发送消息或订阅时再次验证JWT，避免定时检测间隔内的无效操作：

@Component
public class AuthChannelInterceptor implements ChannelInterceptor {
    private static final String TOKEN_COOKIE_NAME = "authToken";
    @Autowired
    private TokenHandler tokenHandler;

    @Override
    public Message<?> preSend(Message<?> message, MessageChannel channel) {
        StompHeaderAccessor accessor = MessageHeaderAccessor.getAccessor(message, StompHeaderAccessor.class);
        if (accessor == null) return message;

        StompCommand command = accessor.getCommand();
        if (StompCommand.SEND.equals(command) || StompCommand.SUBSCRIBE.equals(command)) {
            Map<String, Object> sessionAttributes = SimpMessageHeaderAccessor.getSessionAttributes(message.getHeaders());
            String authToken = (sessionAttributes != null) ? (String) sessionAttributes.get(TOKEN_COOKIE_NAME) : null;

            if (authToken == null) {
                throw new BadCredentialsException("No auth token found in session attributes");
            }
            try {
                tokenHandler.parseUserFromToken(authToken);
            } catch (Exception e) {
                accessor.setLeaveMutable(true);
                accessor.setError("Token expired or invalid");
                accessor.setCommand(StompCommand.DISCONNECT);
                return MessageBuilder.createMessage(new byte[0], accessor.getMessageHeaders());
            }
        }
        return message;
    }
}

4. 收紧WebSocket安全配置 #

修改WebSocketSecurityConfig，确保仅认证用户能进行消息发送与订阅：

@Configuration
public class WebSocketSecurityConfig extends AbstractSecurityWebSocketMessageBrokerConfigurer {
    @Override
    protected boolean sameOriginDisabled() {
        return true;
    }
    
    @Override
    protected void configureInbound(MessageSecurityMetadataSourceRegistry messages) {
         messages
             .simpTypeMatchers(SimpMessageType.CONNECT, SimpMessageType.DISCONNECT, SimpMessageType.HEARTBEAT).permitAll()
             .simpDestMatchers("/topic/**", "/app/**").authenticated()
             .anyMessage().denyAll();
    }
}

内容的提问来源于stack exchange，提问作者bojackhorseman99