如何为已实现的WebSocket聊天系统添加JWT认证？

Got it, let's walk through implementing JWT authentication for your WebSocket chat project. I've built similar setups with Spring WebSocket + Spring Security, so here's a practical, step-by-step approach that covers all your requirements:

WebSocket connections start with an HTTP handshake, so we can hook into that step to validate the JWT before allowing the connection to upgrade. We'll use a HandshakeInterceptor to check the token and attach the authenticated user ID to the WebSocket session.

First, create the handshake interceptor:

@Component
public class JwtHandshakeInterceptor implements HandshakeInterceptor {

    private final JwtTokenProvider jwtTokenProvider;

    public JwtHandshakeInterceptor(JwtTokenProvider jwtTokenProvider) {
        this.jwtTokenProvider = jwtTokenProvider;
    }

    @Override
    public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Map<String, Object> attributes) throws Exception {
        // Extract JWT from Authorization header
        String token = request.getHeaders().getFirst("Authorization");
        if (token != null && token.startsWith("Bearer ")) {
            token = token.substring(7);
            try {
                // Validate token and get user ID
                String userId = jwtTokenProvider.getUserIdFromToken(token);
                if (userId != null) {
                    // Store user ID in session attributes for later use
                    attributes.put("userId", userId);
                    return true;
                }
            } catch (JwtException e) {
                // Invalid token: reject handshake
                response.setStatusCode(HttpStatus.UNAUTHORIZED);
                return false;
            }
        }
        // No valid token: reject handshake
        response.setStatusCode(HttpStatus.UNAUTHORIZED);
        return false;
    }

    @Override
    public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Exception exception) {
        // No action needed here
    }
}

Then register this interceptor in your WebSocket config:

@Configuration
@EnableWebSocketMessageBroker
public class WebSocketConfig implements WebSocketMessageBrokerConfigurer {

    private final JwtHandshakeInterceptor jwtHandshakeInterceptor;
    private final JwtChannelInterceptor jwtChannelInterceptor;

    public WebSocketConfig(JwtHandshakeInterceptor jwtHandshakeInterceptor, JwtChannelInterceptor jwtChannelInterceptor) {
        this.jwtHandshakeInterceptor = jwtHandshakeInterceptor;
        this.jwtChannelInterceptor = jwtChannelInterceptor;
    }

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

    @Override
    public void registerStompEndpoints(StompEndpointRegistry registry) {
        registry.addEndpoint("/chatting")
                .addInterceptors(jwtHandshakeInterceptor)
                .setAllowedOrigins("https://your-frontend-domain.com"); // Replace with your actual frontend origin
    }

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

Note: Replace JwtTokenProvider with your existing JWT utility class that handles token parsing, signature validation, and extracting the user ID.

We need to make sure users can only subscribe to /user/{userId}/queue/chatting where userId matches their authenticated ID. Use a ChannelInterceptor to check subscription destinations before allowing them.

Create the channel interceptor:

@Component
public class JwtChannelInterceptor implements ChannelInterceptor {

    @Override
    public Message<?> preSend(Message<?> message, MessageChannel channel) {
        StompHeaderAccessor accessor = MessageHeaderAccessor.getAccessor(message, StompHeaderAccessor.class);
        
        // Check if the message is a SUBSCRIBE command
        if (StompCommand.SUBSCRIBE.equals(accessor.getCommand())) {
            String destination = accessor.getDestination();
            String currentUserId = (String) accessor.getSessionAttributes().get("userId");
            
            // Parse the target user ID from the subscription destination
            Pattern inboxPattern = Pattern.compile("/user/(.*)/queue/chatting");
            Matcher matcher = inboxPattern.matcher(destination);
            
            if (matcher.matches()) {
                String targetUserId = matcher.group(1);
                if (!currentUserId.equals(targetUserId)) {
                    throw new AccessDeniedException("You cannot subscribe to another user's inbox");
                }
            }
        }
        return message;
    }
}

When a user sends a message via /app/message, we need to ensure the from field in the payload matches their authenticated user ID. Add this check in your message controller:

@Controller
public class ChatController {

    private final SimpMessagingTemplate simpMessagingTemplate;

    public ChatController(SimpMessagingTemplate simpMessagingTemplate) {
        this.simpMessagingTemplate = simpMessagingTemplate;
    }

    @MessageMapping("/message")
    public void handleMessage(@Payload ChatMessage chatMessage, StompHeaderAccessor accessor) {
        String currentUserId = (String) accessor.getSessionAttributes().get("userId");
        
        // Validate that the 'from' field matches the authenticated user
        if (!currentUserId.equals(chatMessage.getFrom())) {
            throw new AccessDeniedException("The 'from' field does not match your authenticated identity");
        }
        
        // Forward the message to the recipient's inbox
        simpMessagingTemplate.convertAndSendToUser(
            chatMessage.getTo(),
            "/queue/chatting",
            chatMessage
        );
    }

    // DTO for chat messages
    public static class ChatMessage {
        private String message;
        private String from;
        private String to;

        // Getters and setters
        public String getMessage() { return message; }
        public void setMessage(String message) { this.message = message; }
        public String getFrom() { return from; }
        public void setFrom(String from) { this.from = from; }
        public String getTo() { return to; }
        public void setTo(String to) { this.to = to; }
    }
}

• Client Side: When connecting via Stomp.js, make sure to include the JWT in the handshake headers:

  const socket = new SockJS('/chatting');
  const stompClient = Stomp.over(socket);
  
  stompClient.connect(
    { Authorization: `Bearer ${yourJwtToken}` },
    (frame) => {
      // Subscribe to your own inbox after successful connection
      stompClient.subscribe(`/user/${yourUserId}/queue/chatting`, (message) => {
        // Handle incoming messages
        const chatMsg = JSON.parse(message.body);
        console.log("Received message:", chatMsg);
      });
    },
    (error) => {
      console.error("Connection failed:", error);
    }
  );
  

• Error Handling: Add a custom exception handler for WebSocket to return user-friendly error messages instead of generic 500s.
• Production Security: Avoid using setAllowedOrigins("*"); specify your exact frontend domain(s) to prevent CSRF risks.

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