Implement Real-Time WebSocket Messaging in Spring Boot 2.7 with WebFlux

Environment: Springboot 2.7.12

1. 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 runs on top of HTTP, uses ports 80 and 443, and can reuse existing firewall rules.

The interaction starts with an HTTP request that is upgraded to the WebSocket protocol via the HTTP Upgrade header. Example request:

<code>GET /ws/demo/chat HTTP/1.1
Host: localhost:8080
Upgrade: websocket // The Upgrade header.
Connection: Upgrade // Using the Upgrade connection.
Sec-WebSocket-Key: Uc9l9TMkWGbHFD2qnFHltg==
Sec-WebSocket-Protocol: v10.stomp, v11.stomp
Sec-WebSocket-Version: 13 Origin: http://localhost:8080</code>

A WebSocket‑supporting server responds with a 101 Switching Protocols status instead of the usual 200:

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

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

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

Next, we detail how to implement a WebSocket real‑time communication service in WebFlux.

2. Custom HandlerMapping

Creating a custom HandlerMapping enables automatic handling of WebSocket connections for different 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 {
  /**请求路径*/
  String value() default "";
}</code>

With this setup, any bean annotated with @WebSocketMapping is automatically 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) {
    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 -> {
          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("给 {} 发送消息: {}", tos, msg);
              wsw.send(msg);
            }
          } else {
            logger.info("Chat 接收到消息: {}", 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: {}, 断开连接. 信号: {}", session.getId(), signalType.name());
      sessions.remove(username);
      session.close();
    }).then();
  }
}
</code>

Helper class for sending messages:

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

Testing the point‑to‑point messaging shows the chat UI as illustrated below.

Finished!!!