Skip to content

G
GeoFlyApi
Commit 1b548210 by 真的三个金的鑫
Options

fix: WSS 按用户裁剪设备推送，并支持编辑用户角色 

Co-authored-by: Cursor <cursoragent@cursor.com>
parent 2930692d
Showing with 272 additions and 15 deletions
sample/src/main/java/com/dji/sample/component/websocket/config/AuthPrincipalHandler.java
......@@ -56,6 +56,8 @@ public class AuthPrincipalHandler extends DefaultHandshakeHandler {
* @param attributes
* @return
*/
public static final String CLAIM_ATTR = "CUSTOM_CLAIM";
@Override
protected Principal determineUser(ServerHttpRequest request, WebSocketHandler wsHandler, Map<String, Object> attributes) {
if (request instanceof ServletServerHttpRequest) {
......@@ -68,6 +70,7 @@ public class AuthPrincipalHandler extends DefaultHandshakeHandler {
return () -> null;
}
attributes.put(CLAIM_ATTR, claim);
return () -> claim.getWorkspaceId() + "/" + claim.getUserType() + "/" + claim.getId();
}
return () -> null;
......
sample/src/main/java/com/dji/sample/component/websocket/config/MyWebSocketHandler.java
......@@ -42,11 +42,15 @@ public class MyWebSocketHandler extends WebSocketDefaultHandler {
@Override
public void afterConnectionEstablished(WebSocketSession session) throws Exception {
Principal principal = session.getPrincipal();
String principalName = principal.getName();
String principalName = principal != null ? principal.getName() : null;
if (StringUtils.hasText(principalName) && !principalName.startsWith("temp-")) {
webSocketManageService.put(principalName, new MyConcurrentWebSocketSession(session));
CustomClaim claim = (CustomClaim) session.getAttributes().get(AuthPrincipalHandler.CLAIM_ATTR);
webSocketManageService.put(principalName, new MyConcurrentWebSocketSession(session), claim);
authenticatedSessions.put(session.getId(), true);
if (claim != null) {
sessionClaims.put(session.getId(), claim);
}
// 记录 principalName 用于 afterConnectionClosed 时清理
sessionPrincipalNames.put(session.getId(), principalName);
log.debug("{} is connected (pre-authenticated). ID: {}. WebSocketSession[current count: {}]",
......@@ -107,7 +111,7 @@ public class MyWebSocketHandler extends WebSocketDefaultHandler {
authenticatedSessions.put(sessionId, true);
sessionClaims.put(sessionId, claim);
webSocketManageService.put(key, new MyConcurrentWebSocketSession(session));
webSocketManageService.put(key, new MyConcurrentWebSocketSession(session), claim);
log.debug("Session {} authenticated successfully. User: {}", sessionId, key);
......
sample/src/main/java/com/dji/sample/component/websocket/service/IWebSocketManageService.java
package com.dji.sample.component.websocket.service;
import com.dji.sample.common.model.CustomClaim;
import com.dji.sample.component.websocket.config.MyConcurrentWebSocketSession;
import java.util.Collection;
import java.util.Optional;
/**
* @author sean
......@@ -13,11 +15,15 @@ public interface IWebSocketManageService {
void put(String key, MyConcurrentWebSocketSession val);
void put(String key, MyConcurrentWebSocketSession val, CustomClaim claim);
void remove(String key, String sessionId);
Collection<MyConcurrentWebSocketSession> getValueWithWorkspace(String workspaceId);
Collection<MyConcurrentWebSocketSession> getValueWithWorkspaceAndUserType(String workspaceId, Integer userType);
Optional<CustomClaim> getClaim(String sessionId);
Long getConnectedCount();
}
sample/src/main/java/com/dji/sample/component/websocket/service/impl/WebSocketManageServiceImpl.java
package com.dji.sample.component.websocket.service.impl;
import com.dji.sample.common.model.CustomClaim;
import com.dji.sample.component.redis.RedisConst;
import com.dji.sample.component.redis.RedisOpsUtils;
import com.dji.sample.component.websocket.config.MyConcurrentWebSocketSession;
......@@ -12,6 +13,7 @@ import org.springframework.util.StringUtils;
import java.util.Collection;
import java.util.Collections;
import java.util.Objects;
import java.util.Optional;
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Collectors;
......@@ -26,8 +28,15 @@ public class WebSocketManageServiceImpl implements IWebSocketManageService {
private static final ConcurrentHashMap<String, MyConcurrentWebSocketSession> SESSIONS = new ConcurrentHashMap<>(16);
private static final ConcurrentHashMap<String, CustomClaim> SESSION_CLAIMS = new ConcurrentHashMap<>(16);
@Override
public void put(String key, MyConcurrentWebSocketSession val) {
put(key, val, null);
}
@Override
public void put(String key, MyConcurrentWebSocketSession val, CustomClaim claim) {
String[] name = key.split("/");
if (name.length != 3) {
log.debug("The key is out of format. [{workspaceId}/{userType}/{userId}]");
......@@ -39,6 +48,9 @@ public class WebSocketManageServiceImpl implements IWebSocketManageService {
RedisOpsUtils.hashSet(workspaceKey, sessionId, name[2]);
RedisOpsUtils.hashSet(userTypeKey, sessionId, name[2]);
SESSIONS.put(sessionId, val);
if (claim != null) {
SESSION_CLAIMS.put(sessionId, claim);
}
RedisOpsUtils.expireKey(workspaceKey, RedisConst.WEBSOCKET_ALIVE_SECOND);
RedisOpsUtils.expireKey(userTypeKey, RedisConst.WEBSOCKET_ALIVE_SECOND);
}
......@@ -53,6 +65,15 @@ public class WebSocketManageServiceImpl implements IWebSocketManageService {
RedisOpsUtils.hashDel(RedisConst.WEBSOCKET_PREFIX + name[0], new String[] {sessionId});
RedisOpsUtils.hashDel(RedisConst.WEBSOCKET_PREFIX + UserTypeEnum.find(Integer.parseInt(name[1])).getDesc(), new String[] {sessionId});
SESSIONS.remove(sessionId);
SESSION_CLAIMS.remove(sessionId);
}
@Override
public Optional<CustomClaim> getClaim(String sessionId) {
if (!StringUtils.hasText(sessionId)) {
return Optional.empty();
}
return Optional.ofNullable(SESSION_CLAIMS.get(sessionId));
}
@Override
......
sample/src/main/java/com/dji/sample/component/websocket/service/impl/WebSocketMessageServiceImpl.java
package com.dji.sample.component.websocket.service.impl;
import com.dji.sample.common.model.CustomClaim;
import com.dji.sample.component.websocket.config.MyConcurrentWebSocketSession;
import com.dji.sample.component.websocket.model.BizCodeEnum;
import com.dji.sample.component.websocket.service.IWebSocketManageService;
import com.dji.sample.component.websocket.service.IWebSocketMessageService;
import com.dji.sample.manage.model.dto.TelemetryDTO;
import com.dji.sample.manage.model.dto.TopologyDeviceDTO;
import com.dji.sample.manage.service.IDeviceService;
import com.dji.sdk.websocket.WebSocketMessageResponse;
import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Lazy;
import org.springframework.stereotype.Service;
import org.springframework.util.StringUtils;
import org.springframework.web.socket.TextMessage;
import java.io.IOException;
import java.lang.reflect.Method;
import java.util.Collection;
import java.util.Map;
import java.util.Objects;
import java.util.Set;
import java.util.stream.Collectors;
/**
* @author sean.zhou
......@@ -24,12 +34,27 @@ import java.util.Objects;
@Slf4j
public class WebSocketMessageServiceImpl implements IWebSocketMessageService {
private static final Set<String> DEVICE_SCOPED_BIZ_CODES = Set.of(
BizCodeEnum.DEVICE_ONLINE.getCode(),
BizCodeEnum.DEVICE_OFFLINE.getCode(),
BizCodeEnum.DEVICE_UPDATE_TOPO.getCode(),
BizCodeEnum.DEVICE_OSD.getCode(),
BizCodeEnum.RC_OSD.getCode(),
BizCodeEnum.DOCK_OSD.getCode(),
BizCodeEnum.DEVICE_HMS.getCode(),
BizCodeEnum.FLIGHT_AREAS_DRONE_LOCATION.getCode()
);
@Autowired
private ObjectMapper mapper;
@Autowired
private IWebSocketManageService webSocketManageService;
@Autowired
@Lazy
private IDeviceService deviceService;
@Override
public void sendMessage(MyConcurrentWebSocketSession session, WebSocketMessageResponse message) {
if (session == null) {
......@@ -89,6 +114,17 @@ public class WebSocketMessageServiceImpl implements IWebSocketMessageService {
webSocketManageService.getValueWithWorkspace(workspaceId) :
webSocketManageService.getValueWithWorkspaceAndUserType(workspaceId, userType);
String deviceSn = extractDeviceSn(bizCode, data);
if (DEVICE_SCOPED_BIZ_CODES.contains(bizCode)) {
if (!StringUtils.hasText(deviceSn)) {
log.warn("Drop device-scoped WSS without sn, bizCode={}", bizCode);
return;
}
sessions = sessions.stream()
.filter(session -> canSessionSeeDevice(workspaceId, session, deviceSn))
.collect(Collectors.toList());
}
this.sendBatch(sessions, new WebSocketMessageResponse()
.setData(Objects.requireNonNullElse(data, ""))
.setTimestamp(System.currentTimeMillis())
......@@ -99,4 +135,40 @@ public class WebSocketMessageServiceImpl implements IWebSocketMessageService {
public void sendBatch(String workspaceId, String bizCode, Object data) {
this.sendBatch(workspaceId, null, bizCode, data);
}
}
\ No newline at end of file
private boolean canSessionSeeDevice(String workspaceId, MyConcurrentWebSocketSession session, String deviceSn) {
CustomClaim claim = webSocketManageService.getClaim(session.getId()).orElse(null);
if (claim == null) {
// 无 claim 时无法做用户级裁剪，保守不推送
return false;
}
return deviceService.isDeviceVisibleToUser(
workspaceId, deviceSn, claim.getId(), claim.getRoleType(), claim.getOrgId());
}
private String extractDeviceSn(String bizCode, Object data) {
if (!StringUtils.hasText(bizCode) || data == null) {
return null;
}
if (!DEVICE_SCOPED_BIZ_CODES.contains(bizCode)) {
return null;
}
if (data instanceof TelemetryDTO) {
return ((TelemetryDTO<?>) data).getSn();
}
if (data instanceof TopologyDeviceDTO) {
return ((TopologyDeviceDTO) data).getSn();
}
if (data instanceof Map) {
Object sn = ((Map<?, ?>) data).get("sn");
return sn != null ? String.valueOf(sn) : null;
}
try {
Method getter = data.getClass().getMethod("getSn");
Object sn = getter.invoke(data);
return sn != null ? String.valueOf(sn) : null;
} catch (Exception ignored) {
return null;
}
}
}
sample/src/main/java/com/dji/sample/manage/controller/UserController.java
......@@ -56,7 +56,7 @@ public class UserController {
}
/**
* Modify user information. Only mqtt account information is included, nothing else can be modified.
* 修改用户信息：MQTT 账号、用户类型、角色（同步当前团队 manage_user_org）。
* @param user
* @param workspaceId
* @param userId
......
sample/src/main/java/com/dji/sample/manage/service/IDeviceService.java
......@@ -341,4 +341,9 @@ public interface IDeviceService extends IService<DeviceEntity> {
*/
Optional<Object> getDeviceOsd(String workspaceId, String deviceSn);
/**
* 判断指定用户是否可见某设备（用于 WSS 推送裁剪，不依赖 HTTP 请求上下文）。
*/
boolean isDeviceVisibleToUser(String workspaceId, String deviceSn, String userId, Integer roleType, String orgId);
}
\ No newline at end of file
sample/src/main/java/com/dji/sample/manage/service/impl/DeviceServiceImpl.java
......@@ -62,8 +62,10 @@ import java.time.ZoneId;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Objects;
import java.util.Optional;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Collectors;
import static com.dji.sample.common.constant.DeviceConstant.CUSTOM_DOCK_START;
......@@ -82,6 +84,20 @@ import static com.dji.sample.common.util.SecurityUtils.*;
@Transactional
public class DeviceServiceImpl extends ServiceImpl<IDeviceMapper, DeviceEntity> implements IDeviceService {
private static final long VISIBILITY_CACHE_MS = 30_000L;
private final ConcurrentHashMap<String, VisibilityCacheEntry> visibilityCache = new ConcurrentHashMap<>();
private static final class VisibilityCacheEntry {
private final Set<String> sns;
private final long expireAt;
private VisibilityCacheEntry(Set<String> sns, long expireAt) {
this.sns = sns;
this.expireAt = expireAt;
}
}
@Autowired
private MqttGatewayPublish messageSender;
......@@ -404,9 +420,12 @@ public class DeviceServiceImpl extends ServiceImpl<IDeviceMapper, DeviceEntity>
@Override
public void pushDeviceOnlineTopo(String workspaceId, String gatewaySn, String deviceSn) {
TopologyDeviceDTO topo = getDeviceTopoForPilot(deviceSn)
.orElseGet(TopologyDeviceDTO::new)
.setSn(deviceSn)
.setGatewaySn(gatewaySn);
webSocketMessageService.sendBatch(
workspaceId, null, com.dji.sdk.websocket.BizCodeEnum.DEVICE_ONLINE.getCode(),
getDeviceTopoForPilot(deviceSn).orElseGet(TopologyDeviceDTO::new).setGatewaySn(gatewaySn));
workspaceId, null, com.dji.sdk.websocket.BizCodeEnum.DEVICE_ONLINE.getCode(), topo);
}
@Override
......@@ -1873,6 +1892,7 @@ public class DeviceServiceImpl extends ServiceImpl<IDeviceMapper, DeviceEntity>
delUserDevice.in(UserDeviceEntity::getDeviceSn, toRemove);
userDeviceService.remove(delUserDevice);
}
invalidateVisibilityCache();
}
@Override
......@@ -1893,6 +1913,7 @@ public class DeviceServiceImpl extends ServiceImpl<IDeviceMapper, DeviceEntity>
delUserDevice.eq(UserDeviceEntity::getOrgId, orgId);
delUserDevice.in(UserDeviceEntity::getDeviceSn, deviceSnList);
userDeviceService.remove(delUserDevice);
invalidateVisibilityCache();
}
@Override
......@@ -1987,6 +2008,7 @@ public class DeviceServiceImpl extends ServiceImpl<IDeviceMapper, DeviceEntity>
del.in(UserDeviceEntity::getDeviceSn, toRemove);
userDeviceService.remove(del);
}
invalidateVisibilityCache();
}
/**
......@@ -2011,7 +2033,8 @@ public class DeviceServiceImpl extends ServiceImpl<IDeviceMapper, DeviceEntity>
}
}
List<String> deviceSnList = resolveVisibleDeviceSnsForCurrentUser(workspaceId, filterOrgId);
List<String> deviceSnList = resolveVisibleDeviceSnsForUser(
workspaceId, getUserId(), getRoleType(), filterOrgId);
if (CollectionUtils.isEmpty(deviceSnList)) {
return false;
}
......@@ -2019,20 +2042,66 @@ public class DeviceServiceImpl extends ServiceImpl<IDeviceMapper, DeviceEntity>
return true;
}
@Override
public boolean isDeviceVisibleToUser(String workspaceId, String deviceSn, String userId,
Integer roleType, String orgId) {
if (!StringUtils.hasText(deviceSn)) {
return false;
}
if (!StringUtils.hasText(workspaceId) || !StringUtils.hasText(userId) || roleType == null) {
return false;
}
if (aboveSysAdminRole(roleType)) {
return true;
}
return getVisibleDeviceSnsCached(workspaceId, userId, roleType, orgId).contains(deviceSn);
}
/** 分配变更后清可见性缓存，避免 30s 窗口越权。 */
public void invalidateVisibilityCache() {
visibilityCache.clear();
}
private Set<String> getVisibleDeviceSnsCached(String workspaceId, String userId,
Integer roleType, String orgId) {
String cacheKey = workspaceId + "|" + userId + "|" + roleType + "|" + Objects.toString(orgId, "");
VisibilityCacheEntry cached = visibilityCache.get(cacheKey);
long now = System.currentTimeMillis();
if (cached != null && cached.expireAt > now) {
return cached.sns;
}
Set<String> sns = new HashSet<>(resolveVisibleDeviceSnsForUser(workspaceId, userId, roleType, orgId));
visibilityCache.put(cacheKey, new VisibilityCacheEntry(sns, now + VISIBILITY_CACHE_MS));
return sns;
}
/**
* 管理员看团队全量；飞手/成员仅看分配给自己的设备。
*/
private List<String> resolveVisibleDeviceSnsForCurrentUser(String workspaceId, String orgId) {
return resolveVisibleDeviceSnsForUser(workspaceId, getUserId(), getRoleType(), orgId);
}
private List<String> resolveVisibleDeviceSnsForUser(String workspaceId, String userId,
Integer roleType, String orgId) {
if (aboveSysAdminRole(roleType)) {
// 超管：不在这里收窄；HTTP 列表不走 org 过滤时直接全量
if (!StringUtils.hasText(orgId)) {
return new ArrayList<>();
}
return listVisibleDeviceSns(workspaceId, orgId);
}
if (!StringUtils.hasText(orgId)) {
return new ArrayList<>();
}
List<String> orgSns = listVisibleDeviceSns(workspaceId, orgId);
if (CollectionUtils.isEmpty(orgSns)) {
return new ArrayList<>();
}
// 团队管理员及以上看本团队全部设备
if (aboveAdminRole()) {
if (aboveAdminRole(roleType)) {
return orgSns;
}
// 飞手/成员：与用户分配取交集
List<String> userSns = listUserDeviceSns(workspaceId, orgId, getUserId());
List<String> userSns = listUserDeviceSns(workspaceId, orgId, userId);
if (CollectionUtils.isEmpty(userSns)) {
return new ArrayList<>();
}
......
sample/src/main/java/com/dji/sample/manage/service/impl/UserServiceImpl.java
......@@ -918,6 +918,9 @@ public class UserServiceImpl extends ServiceImpl<IUserMapper, UserEntity> implem
@Override
public Boolean updateUser(String workspaceId, String userId, UserListDTO user) {
// 管理员以上才能改用户
aboveAdminRoleAndThrowError();
UserEntity userEntity = mapper.selectOne(
new LambdaQueryWrapper<UserEntity>()
.eq(UserEntity::getUserId, userId)
......@@ -925,16 +928,90 @@ public class UserServiceImpl extends ServiceImpl<IUserMapper, UserEntity> implem
if (userEntity == null) {
return false;
}
userEntity.setMqttUsername(user.getMqttUsername());
userEntity.setMqttPassword(user.getMqttPassword());
// 不能改自己角色导致锁死；MQTT 仍可改
boolean editingSelf = userId.equals(getUserId());
// 目标若是系统管理员，仅超管可改
if (userEntity.getRoleType() != null
&& userEntity.getRoleType() == RoleTypeEnum.SYS_ADMIN.getVal()) {
aboveSysAdminRoleAndThrowError();
}
// 不能把别人升成系统管理员（除非自己是超管）
if (user.getRoleType() != null
&& user.getRoleType() == RoleTypeEnum.SYS_ADMIN.getVal()) {
aboveSysAdminRoleAndThrowError();
}
if (user.getMqttUsername() != null) {
userEntity.setMqttUsername(user.getMqttUsername());
}
if (user.getMqttPassword() != null) {
userEntity.setMqttPassword(user.getMqttPassword());
}
Integer newUserType = parseUserType(user.getUserType());
if (newUserType != null) {
userEntity.setUserType(newUserType);
}
Integer newRoleType = user.getRoleType();
if (newRoleType != null && !editingSelf) {
// 用户表角色：非超管账号同步；超管账号保持 SYS_ADMIN
if (userEntity.getRoleType() == null
|| userEntity.getRoleType() != RoleTypeEnum.SYS_ADMIN.getVal()) {
userEntity.setRoleType(newRoleType);
}
}
userEntity.setUpdateTime(System.currentTimeMillis());
int id = mapper.update(userEntity, new LambdaUpdateWrapper<UserEntity>()
.eq(UserEntity::getUserId, userId)
.eq(UserEntity::getWorkspaceId, workspaceId));
// 列表展示的是「当前团队」下 manage_user_org.role_type，必须同步
if (newRoleType != null && !editingSelf) {
String currentOrgId = getOrgId();
if (StringUtils.hasText(currentOrgId)) {
LambdaQueryWrapper<UserOrgEntity> orgWrapper = new LambdaQueryWrapper<>();
orgWrapper.eq(UserOrgEntity::getUserId, userId)
.eq(UserOrgEntity::getOrgId, currentOrgId)
.eq(UserOrgEntity::getWorkspaceId, workspaceId);
UserOrgEntity relation = userOrgService.getOne(orgWrapper, false);
if (relation != null) {
relation.setRoleType(newRoleType);
relation.setUpdateTime(System.currentTimeMillis());
relation.setUpdaterId(getUserId());
relation.setUpdaterName(getUsername());
userOrgService.updateById(relation);
} else {
// 无关系时写入当前团队
ensureUserInOrg(userId, workspaceId, currentOrgId, newRoleType);
}
}
}
return id > 0;
}
private Integer parseUserType(String userType) {
if (!StringUtils.hasText(userType)) {
return null;
}
try {
return Integer.valueOf(userType.trim());
} catch (NumberFormatException e) {
// 兼容前端偶发传描述文案
if ("Web".equalsIgnoreCase(userType) || "WEB".equalsIgnoreCase(userType)) {
return UserTypeEnum.WEB.getVal();
}
if ("Pilot".equalsIgnoreCase(userType) || "PILOT".equalsIgnoreCase(userType)) {
return UserTypeEnum.PILOT.getVal();
}
return null;
}
}
@Override
public Boolean deleteUser(String workspaceId, String userId) {
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment