Mastering WebSocket in Spring Boot 2.4: Custom HandlerMapping and Real‑Time Chat

Environment: Spring Boot 2.4.13

WebSocket Introduction

The WebSocket protocol (RFC 6455) provides a standardized way to establish a full‑duplex, bidirectional communication channel over a single TCP connection between client and server. It works on top of HTTP, uses ports 80/443 and can reuse existing firewall rules.

The interaction starts with an HTTP request that includes an Upgrade header to switch to the WebSocket protocol. An example request is shown below:

GET /spring-websocket-portfolio/portfolio HTTP/1.1 Host: localhost:8080 Upgrade: websocket Connection: Upgrade Sec-WebSocket-Key: Uc9l9TMkWGbHFD2qnFHltg== Sec-WebSocket-Protocol: v10.stomp, v11.stomp Sec-WebSocket-Version: 13 Origin: http://localhost:8080

A WebSocket‑enabled server responds with a 101 Switching Protocols status instead of the usual 200, for example:

HTTP/1.1 101 Switching Protocols Upgrade: websocket Connection: Upgrade Sec-WebSocket-Accept: 1qVdfYHU9hPOl4JYYNXF623Gzn0= Sec-WebSocket-Protocol: v10.stomp

After a successful handshake the TCP socket remains open, allowing the client and server to continue sending and receiving messages.

A complete description of WebSocket internals is beyond the scope of this article; refer to RFC 6455, the HTML5 specification, or other tutorials for more details.

If the WebSocket server runs behind a web server such as nginx, you may need to configure the proxy to forward upgrade requests.

Custom HandlerMapping

Custom HandlerMapping is used to automatically register WebSocket connections for multiple request paths.

<code>public class WebSocketHandlerMapping extends SimpleUrlHandlerMapping {
  @Override
  public void initApplicationContext() throws BeansException {
    Map<String, WebSocketHandler> handlers = new HashMap<>();
    ApplicationContext context = getApplicationContext();
    Map<String, WebSocketHandler> beans = context.getBeansOfType(WebSocketHandler.class);
    for (WebSocketHandler handler : beans.values()) {
      WebSocketMapping webSocketMapping = AnnotatedElementUtils.findMergedAnnotation(handler.getClass(), WebSocketMapping.class);
      if (webSocketMapping != null) {
        String value = webSocketMapping.value();
        if (StringUtils.hasLength(value)) {
          handlers.put(value, handler);
        }
      }
    }
    if (handlers.size() > 0) {
      this.setUrlMap(handlers);
      super.initApplicationContext();
    }
  }

  @Override
  public int getOrder() {
    return Ordered.HIGHEST_PRECEDENCE;
  }
}</code>

The mapping uses a custom annotation to mark handler beans:

<code>@Target(ElementType.TYPE)
@Retention(RetentionPolicy.RUNTIME)
@Documented
public @interface WebSocketMapping {
  /** request path */
  String value() default "";
}</code>

Beans annotated with @WebSocketMapping are discovered and registered in the URL map.

<code>@Component
@WebSocketMapping("/chat2/{name}")
public class ChatWebSocketHandler2 implements WebSocketHandler {
  private static final Logger logger = LoggerFactory.getLogger(ChatWebSocketHandler2.class);
  public static final Map<String, WebSocketWrapper> sessions = new ConcurrentHashMap<>();

  @Override
  public Mono<Void> handle(WebSocketSession session) {
    System.out.println(session);
    URI uri = session.getHandshakeInfo().getUri();
    String path = uri.getPath();
    String username = path.split("/")[2];
    logger.info("Client id: {} Connected, Request URI: {}", session.getId(), uri);
    HttpHeaders headers = session.getHandshakeInfo().getHeaders();
    logger.info("Request Headers: {}", headers);
    Mono<Void> receive = session.receive()
        .doOnNext(message -> {
          // If header "to" is missing, return null; need to handle to avoid silent closure.
          List<String> tos = headers.get("to");
          if (tos != null && !tos.isEmpty()) {
            String to = tos.get(0);
            WebSocketWrapper wsw = sessions.get(to);
            if (wsw != null) {
              String msg = message.getPayloadAsText();
              logger.info("Sending message to {}: {}", tos, msg);
              wsw.send(msg);
            }
          } else {
            logger.info("Chat received message: {}", message.getPayloadAsText());
          }
        }).onErrorMap(ex -> {
          ex.printStackTrace();
          return ex;
        }).then();
    Mono<Void> sender = session.send(Flux.create(sink -> sessions.put(username, new WebSocketWrapper(session, sink))));
    return Mono.zip(receive, sender).doFinally(signalType -> {
      logger.info("Client id: {}, disconnected. Signal: {}", session.getId(), signalType.name());
      sessions.remove(username);
      session.close();
    }).then();
  }
}</code>

The wrapper class simplifies sending text messages:

<code>public class WebSocketWrapper {
  private WebSocketSession session;
  private FluxSink<WebSocketMessage> sink;
  public void send(String payload) {
    this.sink.next(session.textMessage(payload));
  }
}</code>

Test the point‑to‑point messaging with the screenshots below:

Finished!