webflux websocket 实现简单im聊天

package com.kongjs.im.socket;

import com.alibaba.fastjson2.JSON;
import com.alibaba.fastjson2.JSONObject;
import jakarta.annotation.Resource;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.data.mongodb.core.ReactiveMongoTemplate;
import org.springframework.data.redis.connection.ReactiveSubscription;
import org.springframework.data.redis.core.ReactiveStringRedisTemplate;
import org.springframework.data.redis.listener.ChannelTopic;
import org.springframework.data.redis.listener.ReactiveRedisMessageListenerContainer;
import org.springframework.http.HttpHeaders;
import org.springframework.security.core.context.ReactiveSecurityContextHolder;
import org.springframework.stereotype.Component;
import org.springframework.util.ObjectUtils;
import org.springframework.util.StringUtils;
import org.springframework.web.reactive.socket.HandshakeInfo;
import org.springframework.web.reactive.socket.WebSocketHandler;
import org.springframework.web.reactive.socket.WebSocketMessage;
import org.springframework.web.reactive.socket.WebSocketSession;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

//@Component
public class ImMessageHandler implements WebSocketHandler, InitializingBean {

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

    // 本地内存：userId -> WebSocketSession（半状态架构核心）
    private static final Map<String, WebSocketSession> USER_SESSIONS = new ConcurrentHashMap<>();

    // sessionId -> userId（断开连接清理用）
    private static final Map<String, String> SESSION_USER = new ConcurrentHashMap<>();

    @Resource
    private ReactiveRedisMessageListenerContainer container;
    @Resource
    private ReactiveStringRedisTemplate reactiveStringRedisTemplate;
    @Resource
    private ReactiveMongoTemplate reactiveMongoTemplate;

    private static final String CHAT_MSG = "im:chat:msg:realtime";
    private static final String USER_CHANNELS = "im:user:channels:%s";
    private static final String USER_CHANNEL = "im:channel:user:%s";    // 单聊
    private static final String GROUP_CHANNEL = "im:channel:group:%s";  // 群聊
    private static final String ROOM_CHANNEL = "im:channel:room:%s";    // 房间
    private static final ChannelTopic TOPIC = ChannelTopic.of(CHAT_MSG);


    @Override
    public Mono<Void> handle(WebSocketSession session) {
        HandshakeInfo handshakeInfo = session.getHandshakeInfo();
        HttpHeaders headers = handshakeInfo.getHeaders();
        Mono<String> userIdMono = ReactiveSecurityContextHolder.getContext()
                .filter(a -> !ObjectUtils.isEmpty(a.getAuthentication()) && StringUtils.hasText(a.getAuthentication().getName()))
                .map(a -> a.getAuthentication().getName())
                .doOnNext(userId -> {
                    // 2. 建立映射
                    USER_SESSIONS.put(userId, session);
                    SESSION_USER.put(session.getId(), userId);
                });;

        Mono<Void> input = userIdMono.flatMap(userId ->
                session.receive()
                        .concatMap(message -> sendMsg(userId, session, message))
                        .then()
        );

        Flux<WebSocketMessage> outMessage = userIdMono.flatMapMany(userId ->
                        reactiveStringRedisTemplate.opsForHash()
                                .entries(USER_CHANNELS.formatted(userId))
                                .filter(entry -> !ObjectUtils.isEmpty(entry.getKey()) && !ObjectUtils.isEmpty(entry.getValue()))
                                .collectMap(entry -> entry.getKey().toString(), entry -> entry.getValue().toString())
                                .flatMapMany(map ->
                                        container.receive(TOPIC)
                                                .filter(msg -> filterMsg(userId, map, msg))
                                                .map(ReactiveSubscription.Message::getMessage)
                                                .doOnNext(msg -> reactiveMongoTemplate.save(msg, "message:realtime")
                                                        .doOnError(e -> log.error("消息保存失败", e))
                                                        .subscribe())
                                ))
                .map(session::textMessage);

        Mono<Void> output = session.send(outMessage);
        Mono<Void> cleanup = session.closeStatus()
                .doOnNext(s -> {
                    String userId = SESSION_USER.remove(session.getId());
                    if (userId != null) {
                        USER_SESSIONS.remove(userId);
                    }
                }).then();
        return input.and(output).and(cleanup).then();
    }

    private Mono<?> sendMsg(String userId, WebSocketSession session, WebSocketMessage message) {
        if (message.getType() != WebSocketMessage.Type.TEXT) return Mono.empty();
        try {
            DataBuffer payload = message.getPayload();
            byte[] bytes = new byte[payload.readableByteCount()];
            payload.read(bytes);
            DataBufferUtils.release(payload);
            JSONObject json = JSON.parseObject(bytes);
            json.put("from", userId);
            return reactiveStringRedisTemplate.convertAndSend(CHAT_MSG, message.getPayloadAsText());
        } catch (Exception e) {
            return sendError(session, "400", "消息格式错误").then(Mono.empty());
        }
    }

    private static boolean filterMsg(String userId, Map<String, String> entry, ReactiveSubscription.Message<String, String> msg) {
        try {
            JSONObject json = JSON.parseObject(msg.getMessage());
            String sendType = json.getString("chatType");
            if (!StringUtils.hasText(sendType)) {
                return false;
            }
            String sendTarget = json.getString("to");
            if (sendType.equals("user")) {
                return sendTarget.equals(userId);
            }
            Object target = entry.get(sendType);
            if (ObjectUtils.isEmpty(target)) {
                return false;
            }
            for (String s : target.toString().split(",")) {
                if (s.equals(sendTarget)) {
                    return true;
                }
            }
            return false;
        } catch (Exception e) {
            log.debug("过滤消息异常{}", e.getMessage());
            return false;
        }
    }

    private Mono<Void> sendError(WebSocketSession session, String code, String msg) {
        if (!session.isOpen()) return Mono.empty();
        JSONObject json = new JSONObject();
        json.put("code", code);
        json.put("msg", msg);
        WebSocketMessage webSocketMessage = session.textMessage(json.toString());
        return session
                .send(Mono.just(webSocketMessage))
                .onErrorResume(err -> {
                    log.debug("发生错误{}", err.getMessage());
                    return Mono.empty();
                });
    }

    @Override
    public void afterPropertiesSet() throws Exception {

    }
}