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

This commit is contained in:
疯狂的狮子Li 2026-10-19 10:53:11 +08:00
parent 7a88f858ec
commit 7734dd535c
4 changed files with 92 additions and 65 deletions

View File

@ -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);
// 当 emitter 完成、超时或发生错误时,从映射表中移除对应的 token
emitter.onCompletion(() -> {
SseEmitter remove = emitters.remove(token);
if (remove != null) {
remove.complete();
AtomicReference<SseEmitter> oldEmitter = new AtomicReference<>();
USER_TOKEN_EMITTERS.compute(userId, (id, emitters) -> {
if (emitters == null) {
emitters = new ConcurrentHashMap<>();
}
oldEmitter.set(emitters.put(token, emitter));
return emitters;
});
emitter.onTimeout(() -> {
SseEmitter remove = emitters.remove(token);
if (remove != null) {
remove.complete();
if (oldEmitter.get() != null) {
sendKickedMessage(oldEmitter.get());
oldEmitter.get().complete();
}
});
emitter.onError((e) -> {
SseEmitter remove = emitters.remove(token);
if (remove != null) {
remove.complete();
}
});
try {
// 向客户端发送一条连接成功的事件
emitter.send(SseEmitter.event().comment("connected"));
} catch (IOException e) {
// 如果发送消息失败,则从映射表中移除 emitter
emitters.remove(token);
removeEmitter(userId, token, emitter);
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);
if (MapUtil.isNotEmpty(emitters)) {
AtomicReference<SseEmitter> removed = new AtomicReference<>();
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);
} else {
USER_TOKEN_EMITTERS.remove(userId);
sseEmitter.complete();
}
}
@ -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 {
}
});
// 移除空连接用户
if (emitterMap.isEmpty()) {
toRemoveUsers.add(userId);
}
USER_TOKEN_EMITTERS.computeIfPresent(userId, (id, current) -> current.isEmpty() ? null : current);
});
// 循环结束后统一清理空用户,避免并发修改异常
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 (remove != null) {
remove.complete();
if (emitters.remove(entry.getKey(), entry.getValue())) {
entry.getValue().complete();
}
}
}
} else {
USER_TOKEN_EMITTERS.remove(userId);
USER_TOKEN_EMITTERS.computeIfPresent(userId, (id, current) -> current.isEmpty() ? null : current);
}
}

View File

@ -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);
}
}
}

View File

@ -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));
}
}

View File

@ -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("当前账号已在其他设备登录,您已被强制下线", handler.handleNotLoginException(replaced, request).getMsg());
assertEquals("登录已过期,请重新登录", ((R<?>) handler.handleNotLoginException(timeout, request).getBody()).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());
}
/**