!577 发布 5.2.2 正式版 安全性提升 Merge pull request !577 from 疯狂的狮子Li/dev

疯狂的狮子Li

2024-08-26 098d3347a0df808908aab8c554cd7c4febc5e6d9

 ruoyi-common/ruoyi-common-sse/src/main/java/org/dromara/common/sse/core/SseEmitterManager.java

¶Ô±ÈÐÂÎļþ
@@ -0,0 +1,160 @@
package org.dromara.common.sse.core;

import cn.hutool.core.collection.CollUtil;
import lombok.extern.slf4j.Slf4j;
import org.dromara.common.redis.utils.RedisUtils;
import org.dromara.common.sse.dto.SseMessageDto;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.Consumer;

/**
 * ç®¡ç† Server-Sent Events (SSE) è¿žæŽ¥
 *
 * @author Lion Li
 */
@Slf4j
public class SseEmitterManager {

    /**
     * è®¢é˜…的频道
     */
    private final static String SSE_TOPIC = "global:sse";

    private final static Map<Long, Map<String, SseEmitter>> USER_TOKEN_EMITTERS = new ConcurrentHashMap<>();

    /**
     * å»ºç«‹ä¸ŽæŒ‡å®šç”¨æˆ·çš„ SSE è¿žæŽ¥
     *
     * @param userId ç”¨æˆ·çš„唯一标识符，用于区分不同用户的连接
     * @param token  ç”¨æˆ·çš„唯一令牌，用于识别具体的连接
     * @return è¿”回一个 SseEmitter å®žä¾‹ï¼Œå®¢æˆ·ç«¯å¯ä»¥é€šè¿‡è¯¥å®žä¾‹æŽ¥æ”¶ SSE äº‹ä»¶
     */
    public SseEmitter connect(Long userId, String token) {
        // ä»Ž USER_TOKEN_EMITTERS ä¸­èŽ·å–或创建当前用户的 SseEmitter æ˜ å°„表（ConcurrentHashMap）
        // æ¯ä¸ªç”¨æˆ·å¯ä»¥æœ‰å¤šä¸ª SSE è¿žæŽ¥ï¼Œé€šè¿‡ token è¿›è¡ŒåŒºåˆ†
        Map<String, SseEmitter> emitters = USER_TOKEN_EMITTERS.computeIfAbsent(userId, k -> new ConcurrentHashMap<>());

        // åˆ›å»ºä¸€ä¸ªæ–°çš„ SseEmitter å®žä¾‹ï¼Œè¶…时时间设置为 0 è¡¨ç¤ºæ— é™åˆ¶
        SseEmitter emitter = new SseEmitter(0L);

        emitters.put(token, emitter);

        // å½“ emitter å®Œæˆã€è¶…时或发生错误时，从映射表中移除对应的 token
        emitter.onCompletion(() -> emitters.remove(token));
        emitter.onTimeout(() -> emitters.remove(token));
        emitter.onError((e) -> emitters.remove(token));

        try {
            // å‘客户端发送一条连接成功的事件
            emitter.send(SseEmitter.event().comment("connected"));
        } catch (IOException e) {
            // å¦‚果发送消息失败，则从映射表中移除 emitter
            emitters.remove(token);
        }
        return emitter;
    }

    /**
     * æ–­å¼€æŒ‡å®šç”¨æˆ·çš„ SSE è¿žæŽ¥
     *
     * @param userId ç”¨æˆ·çš„唯一标识符，用于区分不同用户的连接
     * @param token  ç”¨æˆ·çš„唯一令牌，用于识别具体的连接
     */
    public void disconnect(Long userId, String token) {
        Map<String, SseEmitter> emitters = USER_TOKEN_EMITTERS.get(userId);
        if (emitters != null) {
            try {
                emitters.get(token).send(SseEmitter.event().comment("disconnected"));
            } catch (Exception ignore) {
            }
            emitters.remove(token);
        }
    }

    /**
     * è®¢é˜…SSE消息主题，并提供一个消费者函数来处理接收到的消息
     *
     * @param consumer å¤„理SSE消息的消费者函数
     */
    public void subscribeMessage(Consumer<SseMessageDto> consumer) {
        RedisUtils.subscribe(SSE_TOPIC, SseMessageDto.class, consumer);
    }

    /**
     * å‘指定的用户会话发送消息
     *
     * @param userId  è¦å‘送消息的用户id
     * @param message è¦å‘送的消息内容
     */
    public void sendMessage(Long userId, String message) {
        Map<String, SseEmitter> emitters = USER_TOKEN_EMITTERS.get(userId);
        if (emitters != null) {
            for (Map.Entry<String, SseEmitter> entry : emitters.entrySet()) {
                try {
                    entry.getValue().send(SseEmitter.event()
                        .name("message")
                        .data(message));
                } catch (Exception e) {
                    emitters.remove(entry.getKey());
                }
            }
        }
    }

    /**
     * æœ¬æœºå…¨ç”¨æˆ·ä¼šè¯å‘送消息
     *
     * @param message è¦å‘送的消息内容
     */
    public void sendMessage(String message) {
        for (Long userId : USER_TOKEN_EMITTERS.keySet()) {
            sendMessage(userId, message);
        }
    }

    /**
     * å‘布SSE订阅消息
     *
     * @param sseMessageDto è¦å‘布的SSE消息对象
     */
    public void publishMessage(SseMessageDto sseMessageDto) {
        List<Long> unsentUserIds = new ArrayList<>();
        // å½“前服务内用户,直接发送消息
        for (Long userId : sseMessageDto.getUserIds()) {
            if (USER_TOKEN_EMITTERS.containsKey(userId)) {
                sendMessage(userId, sseMessageDto.getMessage());
                continue;
            }
            unsentUserIds.add(userId);
        }
        // ä¸åœ¨å½“前服务内用户,发布订阅消息
        if (CollUtil.isNotEmpty(unsentUserIds)) {
            SseMessageDto broadcastMessage = new SseMessageDto();
            broadcastMessage.setMessage(sseMessageDto.getMessage());
            broadcastMessage.setUserIds(unsentUserIds);
            RedisUtils.publish(SSE_TOPIC, broadcastMessage, consumer -> {
                log.info("SSE发送主题订阅消息topic:{} session keys:{} message:{}",
                    SSE_TOPIC, unsentUserIds, sseMessageDto.getMessage());
            });
        }
    }

    /**
     * å‘所有的用户发布订阅的消息(群发)
     *
     * @param message è¦å‘布的消息内容
     */
    public void publishAll(String message) {
        SseMessageDto broadcastMessage = new SseMessageDto();
        broadcastMessage.setMessage(message);
        RedisUtils.publish(SSE_TOPIC, broadcastMessage, consumer -> {
            log.info("SSE发送主题订阅消息topic:{} message:{}", SSE_TOPIC, message);
        });
    }
}