WebSocket Containers

/**
 * Abstract base class for WebSocket containers.
 * Manages WebSocket sessions and delegates events to WebSocketListener.
 * Thread-safe for concurrent session access.
 *
 * @since 4.1
 */
public abstract class IntegrationWebSocketContainer implements DisposableBean {
    public static final int DEFAULT_SEND_TIME_LIMIT = 10000; // 10 seconds
    public static final int DEFAULT_SEND_BUFFER_SIZE = 524288; // 512 KB

    /**
     * Set the timeout for sending messages in milliseconds.
     * Default: 10000 (10 seconds)
     * Messages that take longer than this timeout may trigger buffer overflow.
     *
     * @param sendTimeLimit timeout in milliseconds (must be > 0)
     * @throws IllegalArgumentException if sendTimeLimit <= 0
     */
    public void setSendTimeLimit(int sendTimeLimit);

    /**
     * Set the send buffer size limit in bytes.
     * Default: 524288 (512 KB)
     * When buffer exceeds this limit, overflow strategy is applied.
     *
     * @param sendBufferSizeLimit buffer size in bytes (must be > 0)
     * @throws IllegalArgumentException if sendBufferSizeLimit <= 0
     */
    public void setSendBufferSizeLimit(int sendBufferSizeLimit);

    /**
     * Set the send buffer overflow strategy for concurrent sends.
     * Determines behavior when buffer reaches configured limit.
     * Default: TERMINATE
     *
     * @param overflowStrategy the overflow strategy (must not be null)
     * @throws IllegalArgumentException if overflowStrategy is null
     * @since 5.5.19
     */
    public void setSendBufferOverflowStrategy(
        ConcurrentWebSocketSessionDecorator.OverflowStrategy overflowStrategy);

    /**
     * Replace the default WebSocketHandler with a custom one.
     * Useful for applying decorator factories or advanced customization.
     * Must be called before container is started.
     *
     * @param handler the WebSocketHandler to use (must not be null)
     * @throws IllegalStateException if container is already started
     * @since 5.5.18
     */
    protected void setWebSocketHandler(WebSocketHandler handler);

    /**
     * Set the WebSocketListener to receive session events and messages.
     * Typically set to WebSocketInboundChannelAdapter.
     *
     * @param messageListener the listener implementation (must not be null)
     * @throws IllegalArgumentException if messageListener is null
     */
    public void setMessageListener(WebSocketListener messageListener);

    /**
     * Set the supported WebSocket sub-protocols.
     * Replaces any previously configured protocols.
     *
     * @param protocols the sub-protocol names (may be empty, null values ignored)
     */
    public void setSupportedProtocols(String... protocols);

    /**
     * Add supported WebSocket sub-protocols.
     * Adds to existing configured protocols.
     *
     * @param protocols the sub-protocol names to add (null values ignored)
     */
    public void addSupportedProtocols(String... protocols);

    /**
     * Get the internal WebSocket handler.
     *
     * @return the WebSocketHandler instance (may be null if not initialized)
     */
    public WebSocketHandler getWebSocketHandler();

    /**
     * Get all supported sub-protocols from both the listener and container.
     *
     * @return list of supported sub-protocol names (never null, may be empty)
     */
    public List<String> getSubProtocols();

    /**
     * Get all active WebSocket sessions.
     * Thread-safe: returns unmodifiable view of session map.
     *
     * @return unmodifiable map of session ID to WebSocketSession (never null)
     */
    public Map<String, WebSocketSession> getSessions();

    /**
     * Get a specific WebSocket session by ID.
     *
     * @param sessionId the session identifier (must not be null)
     * @return the WebSocketSession (never null)
     * @throws IllegalStateException if session not found
     * @throws IllegalArgumentException if sessionId is null
     */
    public WebSocketSession getSession(String sessionId);

    /**
     * Close a WebSocket session with the specified status.
     *
     * @param session the session to close (must not be null)
     * @param closeStatus the close status (must not be null)
     * @throws Exception if closing fails
     * @throws IllegalArgumentException if session or closeStatus is null
     */
    public void closeSession(WebSocketSession session, CloseStatus closeStatus)
        throws Exception;

    /**
     * Cleanup method that closes all active sessions.
     * Called on container destruction.
     * Thread-safe: closes all sessions safely.
     */
    @Override
    public void destroy();
}

/**
 * Client-side WebSocket container implementation.
 * Manages outbound connection to a WebSocket server.
 * Thread-safe for concurrent operations.
 *
 * @since 4.1
 */
public class ClientWebSocketContainer extends IntegrationWebSocketContainer
        implements SmartLifecycle {

    /**
     * Create client container with URI template and variables.
     *
     * @param client the WebSocketClient implementation (must not be null)
     * @param uriTemplate the URI template (e.g., "ws://localhost:8080/ws/{param}")
     * @param uriVariables values to substitute in template
     * @throws IllegalArgumentException if client or uriTemplate is null
     */
    public ClientWebSocketContainer(
        WebSocketClient client,
        String uriTemplate,
        Object... uriVariables);

    /**
     * Create client container with prepared URI.
     *
     * @param client the WebSocketClient implementation (must not be null)
     * @param uri the complete WebSocket URI (must not be null)
     * @throws IllegalArgumentException if client or uri is null
     * @since 6.1
     */
    public ClientWebSocketContainer(WebSocketClient client, URI uri);

    /**
     * Set the Origin header for the WebSocket handshake.
     *
     * @param origin the origin value (may be null to remove)
     */
    public void setOrigin(String origin);

    /**
     * Set HTTP headers from a map for the WebSocket handshake.
     *
     * @param headers map of header name to comma-delimited values (must not be null)
     * @throws IllegalArgumentException if headers is null
     */
    public void setHeadersMap(Map<String, String> headers);

    /**
     * Set HTTP headers for the WebSocket handshake.
     *
     * @param headers the HttpHeaders instance (must not be null)
     * @throws IllegalArgumentException if headers is null
     */
    public void setHeaders(HttpHeaders headers);

    /**
     * Set the connection timeout in seconds.
     * Default: 10 seconds
     *
     * @param connectionTimeout timeout in seconds (must be > 0)
     * @throws IllegalArgumentException if connectionTimeout <= 0
     * @since 4.2
     */
    public void setConnectionTimeout(int connectionTimeout);

    /**
     * Get the established client session.
     * If not connected and running, attempts to reconnect.
     * Waits up to connectionTimeout for connection.
     *
     * @param sessionId ignored for client container (can be null)
     * @return the established WebSocketSession (never null)
     * @throws IllegalStateException if session cannot be established
     */
    @Override
    public WebSocketSession getSession(String sessionId);

    /**
     * Check if the client session is currently open.
     *
     * @return true if session is open and connected
     * @since 4.2.6
     */
    public boolean isConnected();

    /**
     * Configure auto-startup behavior.
     *
     * @param autoStartup true to start automatically on application context startup
     */
    public void setAutoStartup(boolean autoStartup);

    /**
     * Set the lifecycle phase.
     *
     * @param phase the phase value (lower values start earlier)
     */
    public void setPhase(int phase);

    // SmartLifecycle methods
    @Override
    public boolean isAutoStartup();

    @Override
    public int getPhase();

    @Override
    public boolean isRunning();

    @Override
    public void start();

    @Override
    public void stop();

    @Override
    public void stop(Runnable callback);
}

import org.springframework.integration.websocket.ClientWebSocketContainer;
import org.springframework.web.socket.client.standard.StandardWebSocketClient;
import org.springframework.http.HttpHeaders;

// Create client with URI template
WebSocketClient client = new StandardWebSocketClient();
ClientWebSocketContainer container = new ClientWebSocketContainer(
    client,
    "ws://localhost:8080/websocket/{userId}",
    "user123"
);

// Configure headers
HttpHeaders headers = new HttpHeaders();
headers.add("Authorization", "Bearer token123");
container.setHeaders(headers);

// Set connection timeout
container.setConnectionTimeout(15); // 15 seconds

// Configure send buffer
container.setSendBufferSizeLimit(1024 * 1024); // 1 MB
container.setSendTimeLimit(20000); // 20 seconds
container.setSendBufferOverflowStrategy(
    ConcurrentWebSocketSessionDecorator.OverflowStrategy.DROP_OLDEST
);

// Configure lifecycle
container.setAutoStartup(true);
container.setPhase(100);

// Start connection
container.start();

// Wait for connection
if (!container.isConnected()) {
    // Handle connection failure
    throw new IllegalStateException("Failed to connect");
}

// Get the session
WebSocketSession session = container.getSession(null);

// Check connection status
if (container.isConnected()) {
    System.out.println("Connected: " + session.getId());
}

/**
 * Server-side WebSocket container implementation.
 * Registers WebSocket handlers for specified paths.
 * Thread-safe for concurrent client connections.
 *
 * @since 4.1
 */
public class ServerWebSocketContainer extends IntegrationWebSocketContainer
        implements WebSocketConfigurer, SmartLifecycle {

    /**
     * Create server container for specified WebSocket paths.
     *
     * @param paths the URL paths to handle (e.g., "/ws", "/websocket")
     *              Must not be null or empty
     * @throws IllegalArgumentException if paths is null or empty
     */
    public ServerWebSocketContainer(String... paths);

    /**
     * Set the handshake handler.
     * Required for server operation.
     * Must be called before container is started.
     *
     * @param handshakeHandler the HandshakeHandler implementation (must not be null)
     * @return this container for fluent configuration
     * @throws IllegalArgumentException if handshakeHandler is null
     * @throws IllegalStateException if container is already started
     */
    public ServerWebSocketContainer setHandshakeHandler(HandshakeHandler handshakeHandler);

    /**
     * Set handshake interceptors.
     *
     * @param interceptors the HandshakeInterceptor instances (may be empty)
     * @return this container for fluent configuration
     */
    public ServerWebSocketContainer setInterceptors(HandshakeInterceptor... interceptors);

    /**
     * Configure WebSocket handler decorator factories.
     * Useful for advanced use cases like Spring Security integration.
     *
     * @param factories the WebSocketHandlerDecoratorFactory instances (may be empty)
     * @return this container for fluent configuration
     * @since 4.2
     */
    public ServerWebSocketContainer setDecoratorFactories(
        WebSocketHandlerDecoratorFactory... factories);

    /**
     * Configure allowed CORS origins.
     *
     * @param origins the allowed origin patterns (e.g., "http://localhost:3000", "*")
     *                May be empty to allow all origins
     * @return this container for fluent configuration
     * @since 4.3
     */
    public ServerWebSocketContainer setAllowedOrigins(String... origins);

    /**
     * Enable SockJS fallback support.
     * Can be called without options to use defaults.
     *
     * @param sockJsServiceOptions optional SockJS configuration
     * @return this container for fluent configuration
     */
    public ServerWebSocketContainer withSockJs(SockJsServiceOptions... sockJsServiceOptions);

    /**
     * Set SockJS service options.
     *
     * @param sockJsServiceOptions the SockJS configuration (must not be null)
     * @throws IllegalArgumentException if sockJsServiceOptions is null
     */
    public void setSockJsServiceOptions(SockJsServiceOptions sockJsServiceOptions);

    /**
     * Configure TaskScheduler for SockJS service.
     * Alternative to default SockJS scheduler for runtime registration.
     *
     * @param sockJsTaskScheduler the TaskScheduler instance (must not be null)
     * @throws IllegalArgumentException if sockJsTaskScheduler is null
     * @since 5.5.1
     */
    public void setSockJsTaskScheduler(TaskScheduler sockJsTaskScheduler);

    /**
     * Get the configured SockJS TaskScheduler.
     *
     * @return the TaskScheduler or null if not configured
     */
    public TaskScheduler getSockJsTaskScheduler();

    /**
     * Register WebSocket handlers with the WebSocketHandlerRegistry.
     * Called by Spring WebSocket infrastructure.
     *
     * @param registry the WebSocketHandlerRegistry (must not be null)
     * @throws IllegalArgumentException if registry is null
     */
    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry);

    /**
     * Configure auto-startup behavior.
     *
     * @param autoStartup true to start automatically on application context startup
     */
    public void setAutoStartup(boolean autoStartup);

    /**
     * Set the lifecycle phase.
     *
     * @param phase the phase value (lower values start earlier)
     */
    public void setPhase(int phase);

    // SmartLifecycle methods
    @Override
    public boolean isAutoStartup();

    @Override
    public int getPhase();

    @Override
    public boolean isRunning();

    @Override
    public void start();

    @Override
    public void stop();

    @Override
    public void stop(Runnable callback);
}

import org.springframework.integration.websocket.ServerWebSocketContainer;
import org.springframework.web.socket.server.standard.TomcatRequestUpgradeStrategy;
import org.springframework.web.socket.server.support.DefaultHandshakeHandler;

// Create server container for multiple paths
ServerWebSocketContainer serverContainer = new ServerWebSocketContainer(
    "/websocket",
    "/ws/messages",
    "/ws/notifications"
);

// Configure handshake handler (required)
DefaultHandshakeHandler handshakeHandler = new DefaultHandshakeHandler(
    new TomcatRequestUpgradeStrategy()
);
serverContainer.setHandshakeHandler(handshakeHandler);

// Configure CORS
serverContainer.setAllowedOrigins(
    "http://localhost:3000",
    "https://app.example.com"
);

// Add custom interceptors
serverContainer.setInterceptors(new HttpSessionHandshakeInterceptor());

// Configure send limits
serverContainer.setSendTimeLimit(20000); // 20 seconds
serverContainer.setSendBufferSizeLimit(1024 * 1024); // 1 MB
serverContainer.setSendBufferOverflowStrategy(
    ConcurrentWebSocketSessionDecorator.OverflowStrategy.DROP_OLDEST
);

// Start the container
serverContainer.start();

// Access sessions
Map<String, WebSocketSession> sessions = serverContainer.getSessions();
System.out.println("Active sessions: " + sessions.size());

ClientWebSocketContainer container = new ClientWebSocketContainer(client, uri);

// Configure before starting
container.setConnectionTimeout(10);
container.setAutoStartup(true);
container.setPhase(100);

// Start connection
container.start();

// Verify connection
if (!container.isConnected()) {
    // Handle failure
}

// Stop container
container.stop();

// Or with callback
container.stop(() -> {
    System.out.println("Container stopped");
});

// Container automatically closes session on stop

ServerWebSocketContainer container = new ServerWebSocketContainer("/ws");

// Configure before starting
container.setHandshakeHandler(handshakeHandler);
container.setAllowedOrigins("*");
container.setAutoStartup(true);
container.setPhase(100);

// Start accepting connections
container.start();

// Stop container (closes all sessions)
container.stop();

// Or with callback
container.stop(() -> {
    System.out.println("Container stopped, all sessions closed");
});

try {
    container.start();
} catch (Exception e) {
    // Handle connection failure
    logger.error("Failed to connect", e);
    // Container will not be in running state
}

// Monitor connection status
if (!container.isConnected()) {
    // Attempt reconnection
    try {
        container.start();
    } catch (Exception e) {
        logger.error("Reconnection failed", e);
    }
}

// Configure allowed origins
container.setAllowedOrigins("http://localhost:3000");

// Or allow all (not recommended for production)
container.setAllowedOrigins("*");

try {
    WebSocketSession session = container.getSession(sessionId);
} catch (IllegalStateException e) {
    // Session not found
    logger.warn("Session not found: {}", sessionId);
}

WebSocketSession session = container.getSession(sessionId);
if (session != null && !session.isOpen()) {
    // Session closed
    logger.warn("Session closed: {}", sessionId);
}

// High-frequency, non-critical messages
container.setSendBufferSizeLimit(1024 * 1024); // 1 MB
container.setSendBufferOverflowStrategy(
    ConcurrentWebSocketSessionDecorator.OverflowStrategy.DROP_OLDEST
);

// Critical messages requiring delivery
container.setSendBufferSizeLimit(512 * 1024); // 512 KB
container.setSendBufferOverflowStrategy(
    ConcurrentWebSocketSessionDecorator.OverflowStrategy.TERMINATE
);

// Verify URI and network
URI uri = URI.create("ws://localhost:8080/websocket");
container = new ClientWebSocketContainer(client, uri);

// Check connection status
if (!container.isConnected()) {
    // Implement retry logic
    container.start();
}

// Verify protocols match
container.setSupportedProtocols("v1.protocol");

// Check network connectivity
// Verify firewall rules
// Check server is running

// Ensure handshake handler is set (required)
serverContainer.setHandshakeHandler(new DefaultHandshakeHandler());

// Configure CORS if needed
serverContainer.setAllowedOrigins("*"); // Or specific origins

// Verify paths are correct
ServerWebSocketContainer container = new ServerWebSocketContainer("/websocket");

// Check container is started
if (!container.isRunning()) {
    container.start();
}

// Check buffer overflow
container.setSendBufferOverflowStrategy(
    ConcurrentWebSocketSessionDecorator.OverflowStrategy.DROP_OLDEST
);

// Increase buffer size
container.setSendBufferSizeLimit(1024 * 1024); // 1 MB

// Monitor session state
for (WebSocketSession session : container.getSessions().values()) {
    if (!session.isOpen()) {
        logger.warn("Session closed: {}", session.getId());
    }
}

// Increase buffer size
container.setSendBufferSizeLimit(1024 * 1024); // 1 MB

// Use DROP_OLDEST for non-critical messages
container.setSendBufferOverflowStrategy(
    ConcurrentWebSocketSessionDecorator.OverflowStrategy.DROP_OLDEST
);

// Increase send timeout
container.setSendTimeLimit(30000); // 30 seconds

// Implement message rate limiting

// From Spring Framework
enum ConcurrentWebSocketSessionDecorator.OverflowStrategy {
    /**
     * Throw SessionLimitExceededException when buffer is full.
     * Session is closed with status from exception.
     */
    TERMINATE,

    /**
     * Drop the oldest message when buffer is full.
     * Session remains open.
     */
    DROP_OLDEST,

    /**
     * Drop the current message when buffer is full.
     * Session remains open.
     */
    DROP_CURRENT
}

// From Spring Framework - common implementations
interface WebSocketClient {
    CompletableFuture<WebSocketSession> execute(
        WebSocketHandler webSocketHandler,
        WebSocketHttpHeaders headers,
        URI uri);
}

// Standard Java WebSocket client
class StandardWebSocketClient implements WebSocketClient { }

// Jetty WebSocket client
class JettyWebSocketClient implements WebSocketClient { }

// Undertow WebSocket client
class UndertowWebSocketClient implements WebSocketClient { }

// From Spring Framework
interface HandshakeHandler {
    boolean doHandshake(
        ServerHttpRequest request,
        ServerHttpResponse response,
        WebSocketHandler wsHandler,
        Map<String, Object> attributes);
}

// Default implementation
class DefaultHandshakeHandler implements HandshakeHandler {
    public DefaultHandshakeHandler(RequestUpgradeStrategy upgradeStrategy);
}

// Tomcat upgrade strategy
class TomcatRequestUpgradeStrategy implements RequestUpgradeStrategy { }

// Jetty upgrade strategy
class JettyRequestUpgradeStrategy implements RequestUpgradeStrategy { }

// From Spring Framework
interface HandshakeInterceptor {
    boolean beforeHandshake(
        ServerHttpRequest request,
        ServerHttpResponse response,
        WebSocketHandler wsHandler,
        Map<String, Object> attributes);

    void afterHandshake(
        ServerHttpRequest request,
        ServerHttpResponse response,
        WebSocketHandler wsHandler,
        Exception exception);
}

// HTTP session handshake interceptor
class HttpSessionHandshakeInterceptor implements HandshakeInterceptor { }