Spring WebSocket & Real-Time Communication

@Configuration
@EnableWebSocket
public class WebSocketConfig implements WebSocketConfigurer {

    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
        registry.addHandler(chatHandler(), "/ws/chat")
                .setAllowedOrigins("https://myapp.com")
                .addInterceptors(new HttpSessionHandshakeInterceptor());
    }

    @Bean
    public WebSocketHandler chatHandler() {
        return new ChatWebSocketHandler();
    }
}

public class ChatWebSocketHandler extends TextWebSocketHandler {

    private final Set<WebSocketSession> sessions = 
        ConcurrentHashMap.newKeySet();

    @Override
    public void afterConnectionEstablished(WebSocketSession session) {
        sessions.add(session);
        log.info("Connected: {} (total: {})", 
            session.getId(), sessions.size());
    }

    @Override
    protected void handleTextMessage(WebSocketSession session, 
                                     TextMessage message) throws Exception {
        String payload = message.getPayload();
        ChatMessage msg = objectMapper.readValue(payload, ChatMessage.class);

        // Broadcast to all connected clients
        TextMessage broadcast = new TextMessage(
            objectMapper.writeValueAsString(msg));
        for (WebSocketSession s : sessions) {
            if (s.isOpen()) {
                s.sendMessage(broadcast);
            }
        }
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, 
                                      CloseStatus status) {
        sessions.remove(session);
        log.info("Disconnected: {} (reason: {})", 
            session.getId(), status.getReason());
    }

    @Override
    public void handleTransportError(WebSocketSession session, 
                                     Throwable exception) {
        log.error("Transport error for {}: {}", 
            session.getId(), exception.getMessage());
        sessions.remove(session);
    }
}

@Configuration
@EnableWebSocketMessageBroker
public class WebSocketStompConfig 
        implements WebSocketMessageBrokerConfigurer {

    @Override
    public void configureMessageBroker(MessageBrokerRegistry config) {
        // Enable simple in-memory broker for /topic and /queue prefixes
        config.enableSimpleBroker("/topic", "/queue")
              .setHeartbeatValue(new long[]{10000, 10000})
              .setTaskScheduler(heartbeatScheduler());

        // Prefix for messages bound for @MessageMapping methods
        config.setApplicationDestinationPrefixes("/app");

        // Prefix for user-specific destinations
        config.setUserDestinationPrefix("/user");
    }

    @Override
    public void registerStompEndpoints(StompEndpointRegistry registry) {
        registry.addEndpoint("/ws")
                .setAllowedOriginPatterns("https://*.myapp.com")
                .withSockJS(); // Fallback for non-WebSocket browsers
    }

    @Bean
    public TaskScheduler heartbeatScheduler() {
        ThreadPoolTaskScheduler scheduler = new ThreadPoolTaskScheduler();
        scheduler.setPoolSize(1);
        scheduler.setThreadNamePrefix("ws-heartbeat-");
        return scheduler;
    }
}

@Controller
public class ChatController {

    // Client sends to: /app/chat.send
    // Response goes to: /topic/chat (broadcast)
    @MessageMapping("/chat.send")
    @SendTo("/topic/chat")
    public ChatMessage sendMessage(ChatMessage message) {
        message.setTimestamp(Instant.now());
        return message; // Returned object is sent to /topic/chat
    }

    // With SimpMessageHeaderAccessor for metadata
    @MessageMapping("/chat.send")
    @SendTo("/topic/chat")
    public ChatMessage sendWithHeaders(
            ChatMessage message,
            SimpMessageHeaderAccessor headerAccessor) {

        String sessionId = headerAccessor.getSessionId();
        String user = headerAccessor.getUser().getName();
        message.setSender(user);
        return message;
    }

    // Multiple destinations
    @MessageMapping("/order.create")
    @SendTo({"/topic/orders", "/topic/admin.orders"})
    public OrderEvent createOrder(OrderRequest request) {
        Order order = orderService.create(request);
        return new OrderEvent("CREATED", order);
    }
}

@Controller
public class NotificationController {

    // Client subscribes to: /user/queue/errors
    // Only the sending user receives this response
    @MessageMapping("/trade.execute")
    @SendToUser("/queue/errors")
    public TradeError handleTrade(TradeRequest request) {
        if (!validationService.isValid(request)) {
            return new TradeError("Invalid trade parameters");
        }
        // Normal path — don't return error
        tradeService.execute(request);
        return null; // null = no message sent
    }
}

@Controller
public class DataController {

    // When client subscribes to /app/initial-data,
    // immediately send current state (no broker involved)
    @SubscribeMapping("/initial-data")
    public List<ChatMessage> onSubscribe() {
        return chatService.getRecentMessages(50);
    }
}

@Service
@RequiredArgsConstructor
public class NotificationService {

    private final SimpMessagingTemplate messagingTemplate;

    // Broadcast to all subscribers of a topic
    public void broadcastAlert(SystemAlert alert) {
        messagingTemplate.convertAndSend("/topic/alerts", alert);
    }

    // Send to a specific user (requires user authentication)
    public void notifyUser(String username, Notification notification) {
        messagingTemplate.convertAndSendToUser(
            username,           // Spring Security principal name
            "/queue/notifications",  // Destination
            notification        // Payload
        );
    }

    // Send with custom headers
    public void sendWithHeaders(String destination, Object payload) {
        SimpMessageHeaderAccessor accessor = 
            SimpMessageHeaderAccessor.create(SimpMessageType.MESSAGE);
        accessor.setHeader("priority", "HIGH");
        accessor.setLeaveMutable(true);

        messagingTemplate.convertAndSend(
            destination, payload, accessor.getMessageHeaders());
    }

    // Send from async/scheduled task
    @Scheduled(fixedRate = 5000)
    public void pushLiveMetrics() {
        Metrics metrics = metricsService.getCurrentMetrics();
        messagingTemplate.convertAndSend("/topic/metrics", metrics);
    }
}

// Default — good for development and single-instance deployments
config.enableSimpleBroker("/topic", "/queue");

@Override
public void configureMessageBroker(MessageBrokerRegistry config) {
    config.enableStompBrokerRelay("/topic", "/queue")
          .setRelayHost("rabbitmq.internal")
          .setRelayPort(61613)           // STOMP port
          .setClientLogin("app-user")
          .setClientPasscode("secret")
          .setSystemLogin("system-user")
          .setSystemPasscode("system-secret")
          .setSystemHeartbeatSendInterval(10000)
          .setSystemHeartbeatReceiveInterval(10000);

    config.setApplicationDestinationPrefixes("/app");
    config.setUserDestinationPrefix("/user");
}

@Override
public void registerStompEndpoints(StompEndpointRegistry registry) {
    registry.addEndpoint("/ws")
            .setAllowedOriginPatterns("*")
            .withSockJS()
            .setStreamBytesLimit(512 * 1024)     // 512KB
            .setHttpMessageCacheSize(1000)
            .setDisconnectDelay(30 * 1000)       // 30 seconds
            .setHeartbeatTime(25000);            // 25 seconds
}

@Configuration
@EnableWebSocketMessageBroker
public class WebSocketSecurityConfig 
        implements WebSocketMessageBrokerConfigurer {

    @Override
    public void configureClientInboundChannel(ChannelRegistration reg) {
        reg.interceptors(new ChannelInterceptor() {
            @Override
            public Message<?> preSend(Message<?> message, 
                                      MessageChannel channel) {
                StompHeaderAccessor accessor = 
                    MessageHeaderAccessor.getAccessor(
                        message, StompHeaderAccessor.class);

                if (StompCommand.CONNECT.equals(accessor.getCommand())) {
                    String token = accessor
                        .getFirstNativeHeader("Authorization");

                    // Validate JWT and set principal
                    Authentication auth = 
                        jwtService.validateToken(token);
                    accessor.setUser(auth);
                }
                return message;
            }
        });
    }
}

@Configuration
@EnableWebSocketSecurity
public class WebSocketAuthzConfig {

    @Bean
    AuthorizationManager<Message<?>> messageAuthorizationManager(
            MessageMatcherDelegatingAuthorizationManager.Builder messages) {

        messages
            .simpDestMatchers("/app/**").authenticated()
            .simpSubscribeDestMatchers("/topic/admin/**").hasRole("ADMIN")
            .simpSubscribeDestMatchers("/user/**").authenticated()
            .simpSubscribeDestMatchers("/topic/**").authenticated()
            .anyMessage().denyAll();

        return messages.build();
    }
}

@Override
public void registerStompEndpoints(StompEndpointRegistry registry) {
    registry.addEndpoint("/ws")
            // Restrict origins (NEVER use "*" in production)
            .setAllowedOriginPatterns("https://*.myapp.com")
            .withSockJS();
}

@Configuration
@EnableRedisWebSession
public class SessionConfig {

    @Bean
    public LettuceConnectionFactory connectionFactory() {
        return new LettuceConnectionFactory("redis.internal", 6379);
    }
}

// Ensure WebSocket handshake picks up the HTTP session
@Override
public void registerStompEndpoints(StompEndpointRegistry registry) {
    registry.addEndpoint("/ws")
            .setHandshakeHandler(new DefaultHandshakeHandler())
            .addInterceptors(new HttpSessionHandshakeInterceptor())
            .withSockJS();
}

@Override
public void configureWebSocketTransport(
        WebSocketTransportRegistration registration) {
    registration
        .setMessageSizeLimit(128 * 1024)      // 128KB max message
        .setSendBufferSizeLimit(512 * 1024)   // 512KB send buffer
        .setSendTimeLimit(20 * 1000);         // 20s send timeout
}

@SpringBootTest(webEnvironment = SpringBootTest.WebEnvironment.RANDOM_PORT)
class ChatWebSocketTest {

    @LocalServerPort
    private int port;

    private WebSocketStompClient stompClient;

    @BeforeEach
    void setup() {
        stompClient = new WebSocketStompClient(
            new StandardWebSocketClient());
        stompClient.setMessageConverter(
            new MappingJackson2MessageConverter());
    }

    @Test
    void shouldBroadcastChatMessage() throws Exception {
        // Connect
        StompSession session = stompClient
            .connectAsync(
                "ws://localhost:" + port + "/ws",
                new StompSessionHandlerAdapter() {})
            .get(5, TimeUnit.SECONDS);

        // Subscribe and capture messages
        BlockingQueue<ChatMessage> messages = new LinkedBlockingQueue<>();
        session.subscribe("/topic/chat", new StompFrameHandler() {
            @Override
            public Type getPayloadType(StompHeaders headers) {
                return ChatMessage.class;
            }

            @Override
            public void handleFrame(StompHeaders headers, Object payload) {
                messages.offer((ChatMessage) payload);
            }
        });

        // Send message
        ChatMessage outgoing = new ChatMessage("Hello!", "testUser");
        session.send("/app/chat.send", outgoing);

        // Verify
        ChatMessage received = messages.poll(5, TimeUnit.SECONDS);
        assertThat(received).isNotNull();
        assertThat(received.getContent()).isEqualTo("Hello!");
        assertThat(received.getSender()).isEqualTo("testUser");
    }

    @AfterEach
    void teardown() {
        stompClient.stop();
    }
}

@ExtendWith(MockitoExtension.class)
class ChatControllerTest {

    @InjectMocks
    private ChatController chatController;

    @Test
    void sendMessage_shouldSetTimestamp() {
        ChatMessage input = new ChatMessage("Hi", "user1");

        ChatMessage result = chatController.sendMessage(input);

        assertThat(result.getTimestamp()).isNotNull();
        assertThat(result.getContent()).isEqualTo("Hi");
    }
}

@SpringBootTest
class NotificationServiceTest {

    @Autowired
    private NotificationService notificationService;

    @MockBean
    private SimpMessagingTemplate messagingTemplate;

    @Test
    void shouldSendToSpecificUser() {
        Notification notification = new Notification("New order");

        notificationService.notifyUser("john", notification);

        verify(messagingTemplate).convertAndSendToUser(
            eq("john"),
            eq("/queue/notifications"),
            eq(notification)
        );
    }
}