fix 修复SSE重连导致连接失管及鉴权异常响应

ruoyi-common/ruoyi-common-push/src/main/java/org/dromara/common/push/core/SseEmitterSessionManager.java

@@ -1,6 +1,5 @@
 package org.dromara.common.push.core;
 
-import cn.hutool.core.collection.CollUtil;
 import cn.hutool.core.map.MapUtil;
 import lombok.extern.slf4j.Slf4j;
 import org.dromara.common.core.utils.ThreadUtils;
@@ -19,6 +18,7 @@ import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Consumer;
 
 /**
@@ -58,52 +58,41 @@ public class SseEmitterSessionManager implements PushSessionManager {
      * @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，防止超过最大连接数
-        SseEmitter oldEmitter = emitters.remove(token);
-        if (oldEmitter != null) {
-            sendKickedMessage(oldEmitter);
-            oldEmitter.complete();
-        }
-
-        // 创建一个新的 SseEmitter 实例，避免连接之后直接关闭浏览器导致连接停滞
         SseEmitter emitter = new SseEmitter(messageProperties.getSseTimeout());
+        emitter.onCompletion(() -> removeEmitter(userId, token, emitter));
+        emitter.onTimeout(() -> removeEmitter(userId, token, emitter));
+        emitter.onError(e -> removeEmitter(userId, token, emitter));
 
-        emitters.put(token, emitter);
+        AtomicReference<SseEmitter> oldEmitter = new AtomicReference<>();
-
+        USER_TOKEN_EMITTERS.compute(userId, (id, emitters) -> {
-        // 当 emitter 完成、超时或发生错误时，从映射表中移除对应的 token
+            if (emitters == null) {
-        emitter.onCompletion(() -> {
+                emitters = new ConcurrentHashMap<>();
-            SseEmitter remove = emitters.remove(token);
-            if (remove != null) {
-                remove.complete();
-            }
-        });
-        emitter.onTimeout(() -> {
-            SseEmitter remove = emitters.remove(token);
-            if (remove != null) {
-                remove.complete();
-            }
-        });
-        emitter.onError((e) -> {
-            SseEmitter remove = emitters.remove(token);
-            if (remove != null) {
-                remove.complete();
             }
+            oldEmitter.set(emitters.put(token, emitter));
+            return emitters;
         });
+        if (oldEmitter.get() != null) {
+            sendKickedMessage(oldEmitter.get());
+            oldEmitter.get().complete();
+        }
 
         try {
             // 向客户端发送一条连接成功的事件
             emitter.send(SseEmitter.event().comment("connected"));
         } catch (IOException e) {
-            // 如果发送消息失败，则从映射表中移除 emitter
+            removeEmitter(userId, token, emitter);
-            emitters.remove(token);
+            emitter.complete();
         }
         return emitter;
     }
 
+    private void removeEmitter(Long userId, String token, SseEmitter emitter) {
+        USER_TOKEN_EMITTERS.computeIfPresent(userId, (id, emitters) -> {
+            emitters.remove(token, emitter);
+            return emitters.isEmpty() ? null : emitters;
+        });
+    }
+
     /**
      * 通知旧连接已被同 token 新连接替换。
      *
@@ -129,17 +118,18 @@ public class SseEmitterSessionManager implements PushSessionManager {
         if (userId == null || token == null) {
             return;
         }
-        Map<String, SseEmitter> emitters = USER_TOKEN_EMITTERS.get(userId);
+        AtomicReference<SseEmitter> removed = new AtomicReference<>();
-        if (MapUtil.isNotEmpty(emitters)) {
+        USER_TOKEN_EMITTERS.computeIfPresent(userId, (id, emitters) -> {
+            removed.set(emitters.remove(token));
+            return emitters.isEmpty() ? null : emitters;
+        });
+        SseEmitter sseEmitter = removed.get();
+        if (sseEmitter != null) {
             try {
-                SseEmitter sseEmitter = emitters.get(token);
                 sseEmitter.send(SseEmitter.event().comment("disconnected"));
-                sseEmitter.complete();
             } catch (Exception ignore) {
             }
-            emitters.remove(token);
+            sseEmitter.complete();
-        } else {
-            USER_TOKEN_EMITTERS.remove(userId);
         }
     }
 
@@ -148,15 +138,7 @@ public class SseEmitterSessionManager implements PushSessionManager {
      */
     public void sseMonitor() {
         final SseEmitter.SseEventBuilder heartbeat = SseEmitter.event().comment("heartbeat");
-        // 记录需要移除的用户ID
-        List<Long> toRemoveUsers = new ArrayList<>();
-
         USER_TOKEN_EMITTERS.forEach((userId, emitterMap) -> {
-            if (CollUtil.isEmpty(emitterMap)) {
-                toRemoveUsers.add(userId);
-                return;
-            }
-
             emitterMap.entrySet().removeIf(entry -> {
                 try {
                     entry.getValue().send(heartbeat);
@@ -171,14 +153,8 @@ public class SseEmitterSessionManager implements PushSessionManager {
                 }
             });
 
-            // 移除空连接用户
+            USER_TOKEN_EMITTERS.computeIfPresent(userId, (id, current) -> current.isEmpty() ? null : current);
-            if (emitterMap.isEmpty()) {
-                toRemoveUsers.add(userId);
-            }
         });
-
-        // 循环结束后统一清理空用户，避免并发修改异常
-        toRemoveUsers.forEach(USER_TOKEN_EMITTERS::remove);
     }
 
     /**
@@ -206,14 +182,13 @@ public class SseEmitterSessionManager implements PushSessionManager {
                         .name("message")
                         .data(message));
                 } catch (Exception e) {
-                    SseEmitter remove = emitters.remove(entry.getKey());
+                    if (emitters.remove(entry.getKey(), entry.getValue())) {
-                    if (remove != null) {
+                        entry.getValue().complete();
-                        remove.complete();
                     }
                 }
             }
         } else {
-            USER_TOKEN_EMITTERS.remove(userId);
+            USER_TOKEN_EMITTERS.computeIfPresent(userId, (id, current) -> current.isEmpty() ? null : current);
         }
     }
 

ruoyi-common/ruoyi-common-push/src/test/java/org/dromara/common/push/PushFunctionTest.java

@@ -3,13 +3,19 @@ package org.dromara.common.push;
 import org.dromara.common.push.annotation.ConditionalOnMessageTransport;
 import org.dromara.common.push.condition.MessageTransportCondition;
 import org.dromara.common.push.enums.MessageTransportEnum;
+import org.dromara.common.push.core.SseEmitterSessionManager;
+import org.dromara.common.push.properties.MessageProperties;
 import org.junit.jupiter.api.DisplayName;
 import org.junit.jupiter.api.Test;
 import org.springframework.context.annotation.ConditionContext;
 import org.springframework.core.type.AnnotatedTypeMetadata;
 import org.springframework.mock.env.MockEnvironment;
+import org.springframework.test.util.ReflectionTestUtils;
+import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
 
 import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ScheduledExecutorService;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
@@ -52,4 +58,28 @@ class PushFunctionTest {
         environment.setProperty("message.enabled", "false");
         assertFalse(condition.matches(context, metadata));
     }
+
+    @Test
+    @DisplayName("旧 SSE 回调不移除重连后的新连接")
+    @SuppressWarnings("unchecked")
+    void shouldKeepReplacementEmitterWhenOldConnectionCompletes() {
+        ScheduledExecutorService scheduler = mock(ScheduledExecutorService.class);
+        SseEmitterSessionManager manager = new SseEmitterSessionManager(scheduler, new MessageProperties());
+        Map<Long, Map<String, SseEmitter>> sessions =
+            (Map<Long, Map<String, SseEmitter>>) ReflectionTestUtils.getField(SseEmitterSessionManager.class, "USER_TOKEN_EMITTERS");
+        long userId = 999_999L;
+        String token = "replacement-test";
+        SseEmitter oldEmitter = new SseEmitter();
+        SseEmitter newEmitter = new SseEmitter();
+        sessions.put(userId, new ConcurrentHashMap<>(Map.of(token, newEmitter)));
+        try {
+            ReflectionTestUtils.invokeMethod(manager, "removeEmitter", userId, token, oldEmitter);
+            assertEquals(newEmitter, sessions.get(userId).get(token));
+
+            ReflectionTestUtils.invokeMethod(manager, "removeEmitter", userId, token, newEmitter);
+            assertFalse(sessions.containsKey(userId));
+        } finally {
+            sessions.remove(userId);
+        }
+    }
 }

ruoyi-common/ruoyi-common-satoken/src/main/java/org/dromara/common/satoken/handler/SaTokenExceptionHandler.java

@@ -7,6 +7,8 @@ import cn.hutool.http.HttpStatus;
 import jakarta.servlet.http.HttpServletRequest;
 import lombok.extern.slf4j.Slf4j;
 import org.dromara.common.core.domain.R;
+import org.springframework.http.ResponseEntity;
+import org.springframework.beans.factory.annotation.Value;
 import org.springframework.web.bind.annotation.ExceptionHandler;
 import org.springframework.web.bind.annotation.RestControllerAdvice;
 
@@ -19,6 +21,9 @@ import org.springframework.web.bind.annotation.RestControllerAdvice;
 @RestControllerAdvice
 public class SaTokenExceptionHandler {
 
+    @Value("${message.path:/resource/message}")
+    private String messagePath = "/resource/message";
+
     /**
      * 处理权限码校验失败异常。
      *
@@ -42,9 +47,12 @@ public class SaTokenExceptionHandler {
      * @return 统一失败响应
      */
     @ExceptionHandler(NotLoginException.class)
-    public R<Void> handleNotLoginException(NotLoginException e, HttpServletRequest request) {
+    public ResponseEntity<?> handleNotLoginException(NotLoginException e, HttpServletRequest request) {
         String requestURI = request.getRequestURI();
-        log.error("请求地址'{}',认证失败'{}',无法访问系统资源", requestURI, e.getMessage());
+        log.warn("请求地址'{}',认证失败,类型'{}'", requestURI, e.getType());
+        if (request.getServletPath().equals(messagePath)) {
+            return ResponseEntity.status(HttpStatus.HTTP_UNAUTHORIZED).build();
+        }
         String msg = switch (e.getType()) {
             case NotLoginException.TOKEN_TIMEOUT,
                  NotLoginException.TOKEN_FREEZE -> "登录已过期，请重新登录";
@@ -52,7 +60,7 @@ public class SaTokenExceptionHandler {
             case NotLoginException.KICK_OUT -> "账号已被管理员强制下线";
             default -> "登录状态异常，请重新登录";
         };
-        return R.fail(HttpStatus.HTTP_UNAUTHORIZED, msg);
+        return ResponseEntity.ok(R.fail(HttpStatus.HTTP_UNAUTHORIZED, msg));
     }
 
 }

ruoyi-common/ruoyi-common-satoken/src/test/java/org/dromara/common/satoken/SaTokenFunctionTest.java

@@ -55,8 +55,22 @@ class SaTokenFunctionTest {
         NotLoginException timeout = NotLoginException.newInstance("login", NotLoginException.TOKEN_TIMEOUT, "timeout", "token");
         NotLoginException replaced = NotLoginException.newInstance("login", NotLoginException.BE_REPLACED, "replaced", "token");
 
-        assertEquals("登录已过期，请重新登录", handler.handleNotLoginException(timeout, request).getMsg());
+        assertEquals("登录已过期，请重新登录", ((R<?>) handler.handleNotLoginException(timeout, request).getBody()).getMsg());
-        assertEquals("当前账号已在其他设备登录，您已被强制下线", handler.handleNotLoginException(replaced, request).getMsg());
+        assertEquals("当前账号已在其他设备登录，您已被强制下线", ((R<?>) handler.handleNotLoginException(replaced, request).getBody()).getMsg());
+    }
+
+    @Test
+    @DisplayName("SSE 登录失效返回无正文 401")
+    void shouldRejectSseWithoutJsonConversion() {
+        SaTokenExceptionHandler handler = new SaTokenExceptionHandler();
+        MockHttpServletRequest request = new MockHttpServletRequest("GET", "/resource/message");
+        request.setServletPath("/resource/message");
+        NotLoginException expired = NotLoginException.newInstance("login", NotLoginException.TOKEN_TIMEOUT, "timeout", "token");
+
+        var response = handler.handleNotLoginException(expired, request);
+
+        assertEquals(HttpStatus.HTTP_UNAUTHORIZED, response.getStatusCode().value());
+        assertEquals(null, response.getBody());
     }
 
     /**