mirror of
https://gitee.com/dromara/RuoYi-Vue-Plus.git
synced 2026-11-02 13:31:49 +08:00
fix 修复SSE重连导致连接失管及鉴权异常响应
This commit is contained in:
parent
7a88f858ec
commit
7734dd535c
@ -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();
|
||||
}
|
||||
});
|
||||
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();
|
||||
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;
|
||||
});
|
||||
if (oldEmitter.get() != null) {
|
||||
sendKickedMessage(oldEmitter.get());
|
||||
oldEmitter.get().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);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@ -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);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@ -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));
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
@ -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());
|
||||
}
|
||||
|
||||
/**
|
||||
|
||||
Loading…
Reference in New Issue
Block a user