Integrate WebSocket

Tags:

• Developer
• Domain Service
• Java

🎯context

Set up a WebSocket server in your Domain Java service using Java Stack 2.0 to enable real-time bidirectional communication.

Description

WebSockets enable real-time, two-way communication between clients and servers. Unlike traditional HTTP requests, WebSockets maintain a persistent connection, allowing servers to push data to clients without clients needing to poll for updates.

This How-To demonstrates two approaches for implementing WebSockets in your Java Domain service:

• A basic implementation without authentication
• A secure implementation with JWT authentication

The guide covers managing multiple client connections, handling reconnections, keeping connections alive, and sending targeted messages to specific clients.

Preconditions

• Java Stack 2.0: Your project uses Java Stack 2.0
• Spring Boot: Your service is built with Spring Boot
• Maven/Gradle: You use Maven or Gradle for dependency management
• OpenShift: Your service will be deployed to an OpenShift environment

Step-by-Step Guide

1. Add WebSocket Dependencies

1.1. Update Maven Dependencies:

• Add the Spring Boot WebSocket starter to your pom.xml:

<dependency>
  <groupId>org.springframework.boot</groupId>
  <artifactId>spring-boot-starter-websocket</artifactId>
</dependency>

2. Configure WebSocket Endpoint

2.1. Create WebSocket Configuration:

• Create a configuration class to register WebSocket handlers:

import org.springframework.context.annotation.Configuration;
import org.springframework.web.socket.config.annotation.EnableWebSocket;
import org.springframework.web.socket.config.annotation.WebSocketConfigurer;
import org.springframework.web.socket.config.annotation.WebSocketHandlerRegistry;

@Configuration
@EnableWebSocket
public class WebSocketConfig implements WebSocketConfigurer {

  @Override
  public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
    registry
      .addHandler(new CustomWebSocketHandler(), "/websocket")
      .setAllowedOrigins("*");
  }
}

2.2. Expose WebSocket Endpoint:

• For Kubernetes/OpenShift deployment, expose the WebSocket endpoint using one of these approaches:
• Add the WebSocket endpoint to your API gateway (if it supports WebSockets)
• Create an OpenShift route with path "/websocket", termination type "Reencrypt", and port 8443

💡tip

When connecting to the WebSocket endpoint, clients should use the "wss://" protocol (WebSocket Secure) instead of "https://".

2.3. Disable Security for WebSocket Endpoint:

• By default, your service will block unauthenticated requests to all endpoints
• Follow the How-To: Disable Security for REST endpoint guide to allow connections to the WebSocket endpoint

3. Implement WebSocket Handler

Without Authentication

3.1. Create Basic WebSocket Handler:

• Implement a handler to manage sessions and messages:

        import org.slf4j.Logger;
        import org.slf4j.LoggerFactory;
        import org.springframework.stereotype.Service;
        import org.springframework.web.socket.CloseStatus;
        import org.springframework.web.socket.TextMessage;
        import org.springframework.web.socket.WebSocketSession;
        import org.springframework.web.socket.handler.TextWebSocketHandler;

        import jakarta.annotation.PostConstruct;
        import jakarta.annotation.PreDestroy;
        import java.io.IOException;
        import java.util.ArrayList;
        import java.util.List;
        import java.util.Map;
        import java.util.UUID;
        import java.util.concurrent.ConcurrentHashMap;
        import java.util.concurrent.Executors;
        import java.util.concurrent.ScheduledExecutorService;
        import java.util.concurrent.TimeUnit;

        @Service
        public class CustomWebSocketHandler extends TextWebSocketHandler {

            private static final Logger log = LoggerFactory.getLogger(CustomWebSocketHandler.class);

            private static final Map<String, List<WebSocketSession>> sessions = new ConcurrentHashMap<>();
            private final ScheduledExecutorService scheduler = Executors.newScheduledThreadPool(1);

            @PostConstruct
            public void startKeepAlive() {
                this.scheduler.scheduleAtFixedRate(this::sendKeepAliveMessages, 30, 30, TimeUnit.SECONDS);
            }

            @PreDestroy
            public void stopKeepAlive() {
                this.scheduler.shutdown();
            }

            @Override
            public void handleTextMessage(WebSocketSession session, TextMessage message) throws IOException {
                log.info("Received new message");

                String clientId = message.getPayload();
                if (clientId.isEmpty() || !sessions.containsKey(clientId)) {
                    clientId = UUID.randomUUID().toString();
                    sessions.put(clientId, new ArrayList<>());
                    session.sendMessage(new TextMessage(clientId));
                }
                sessions.get(clientId).add(session);
            }

            @Override
            public void afterConnectionEstablished(WebSocketSession session) {
                log.info("Connection established for session with ID {}", session.getId());
            }

            @Override
            public void afterConnectionClosed(WebSocketSession session, CloseStatus status) {
                log.info("Session with ID {} has been closed with status code {}", session.getId(), status.getCode());

                sessions.entrySet().stream()
                .filter(entry -> entry.getValue().contains(session))
                .findFirst()
                .ifPresent(entry -> entry.getValue().remove(session));
            }

            public void sendMessageToClient(String clientId, String content) {
                log.info("Sending message to client with ID {}", clientId);

                List<WebSocketSession> clientSessions = sessions.get(clientId);
                if (clientSessions == null || clientSessions.isEmpty()) {
                    log.error("No sessions found for client ID {}", clientId);
                    return;
                }

                TextMessage message = new TextMessage(content);
                for (WebSocketSession session : clientSessions) {
                    try {
                        session.sendMessage(message);
                    } catch (IOException e) {
                        log.error("Exception occurred when sending message to client with ID {}. Error message: {}",
                        clientId, e.getMessage());
                    }
                }
            }

            private void sendKeepAliveMessages() {
                for (List<WebSocketSession> clientSessions : sessions.values()) {
                    for (WebSocketSession session : clientSessions) {
                        try {
                            if(session.isOpen()) {
                                session.sendMessage(new TextMessage(WebsocketMessageType.KEEP_ALIVE.toString()));
                                log.debug("Sent KEEP_ALIVE message to session with ID {}", session.getId());
                            }
                        } catch (IOException e) {
                            log.error("Exception occurred when sending KEEP_ALIVE message to session with ID {}. Error message: {}",
                            session.getId(), e.getMessage());
                        }
                    }
                }
            }
        }

With JWT Authentication

3.1. Create JWT-Authenticated WebSocket Handler:

• Implement a handler that validates JWT tokens:

    @Override
    public void handleTextMessage(WebSocketSession session, TextMessage message) throws IOException {
        log.info("Received new message");

        if (this.isValidAuthToken(message.getPayload())) {
            String clientId = this.getClientIdFromToken(message.getPayload());
            if (clientId == null) {
                log.error("Client ID could not be extracted from JWT");
                return;
            }

            if (!sessions.containsKey(clientId)) {
                sessions.put(clientId, new ArrayList<>());
                session.sendMessage(new TextMessage(WebsocketMessageType.AUTHENTICATION_SUCCEEDED.toString()));
            }
            sessions.get(clientId).add(session);
        } else {
            session.sendMessage(new TextMessage(WebsocketMessageType.AUTHENTICATION_FAILED.toString()));
        }
    }

    private boolean isValidAuthToken(String token) {
        try {
            HttpHeaders httpHeaders = new HttpHeaders();
            httpHeaders.add("Authorization", "Bearer " + token);
            this.apiValidatetoken.validate(httpHeaders);
            return true;
        } catch (Exception e) {
            log.error("Token validation not successful: {}", e.getMessage());
            return false;
        }
    }

    private String getClientIdFromToken(String token) {
        try {
            return JWT.decode(token).getClaim(ConfigValues.JWT_KEY).as(String.class);
        } catch (JWTDecodeException e) {
            log.error("Failed to decode JWT token: {}", e.getMessage());
        }
        return null;
    }

3.2. Update Constructor and Configuration:

• Modify the handler to accept the token validation service:

    private final ValidateApiValidatetoken apiValidatetoken;

    public CustomWebSocketHandler(ValidateApiValidatetoken apiValidatetoken) {
        this.apiValidatetoken = apiValidatetoken;
    }

3.3. Update WebSocket Configuration:

• Update the configuration class to use the authenticated handler:

    import org.springframework.beans.factory.annotation.Autowired;
    import org.springframework.context.annotation.Configuration;
    import org.springframework.web.socket.config.annotation.EnableWebSocket;
    import org.springframework.web.socket.config.annotation.WebSocketConfigurer;
    import org.springframework.web.socket.config.annotation.WebSocketHandlerRegistry;

    @Configuration
    @EnableWebSocket
    public class WebSocketConfig implements WebSocketConfigurer {
        private final ValidateApiValidatetoken apiValidatetoken;

        @Autowired
        public WebSocketConfig(ValidateApiValidatetoken apiValidatetoken) {
            this.apiValidatetoken = apiValidatetoken;
        }

        @Override
        public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
        registry
        .addHandler(new CustomWebSocketHandler(this.apiValidatetoken), "/websocket")
        .setAllowedOrigins("*");
        }
    }

💡tip

If your service is deployed with multiple replicas, consider storing client sessions in an external cache (like Redis) to enable cross-pod communication.

4. Implement Client-Side Connection

4.1. Client Connection Example (React):

• For a React frontend, you can use the "react-use-websocket" library:

  • Install with: npm i react-use-websocket

        import useWebSocket from 'react-use-websocket';
    
        // For basic connection
        const socketUrl = 'wss://YOUR-HOST/websocket';
        let clientId = localStorage.getItem('websocket-client-id') || '';
    
        // For authenticated connection
        const jwtToken = "YOUR-JWT-TOKEN";
    
        const { sendMessage } = useWebSocket(socketUrl, {
          onOpen: () => {
            // For basic connection:
            sendMessage(clientId);
    
            // OR for authenticated connection:
            // sendMessage(jwtToken);
          },
          onMessage: (event) => {
            switch (event.data) {
              case 'KEEP_ALIVE':
                break;
              case 'AUTHENTICATION_SUCCEEDED':
                console.log('Authentication successful');
                break;
              case 'AUTHENTICATION_FAILED':
                console.log('Authentication failed');
                break;
              default:
                // For basic connection, store client ID if new
                if (!clientId) {
                  clientId = event.data;
                  localStorage.setItem('websocket-client-id', clientId);
                }
                console.log('Message received: ' + event.data);
                break;
            }
          },
          onError: (event) => console.error(event),
          shouldReconnect: (closeEvent) => true,
          onClose: (event) => console.log(event),
          share: true,
          retryOnError: true,
          reconnectInterval: 10,
          reconnectAttempts: 5
        });

Conclusion

🌟result

Congratulations! You have successfully integrated WebSockets into your Java Domain service. Your service can now:

• Accept real-time WebSocket connections from clients
• Handle multiple concurrent connections from the same client
• Manage client reconnections
• Keep connections alive with heartbeat messages
• Send targeted messages to specific clients

Further Reading

• Spring WebSocket Documentation
• WebSocket API Specification