Просмотр исходного кода

1、调整websocket,主库websocket处理

yys 3 месяцев назад
Родитель
Сommit
c65950394d

+ 48 - 3
fs-live-app/src/main/java/com/fs/live/websocket/auth/TenantChannelContext.java

@@ -21,9 +21,10 @@ import org.springframework.security.core.context.SecurityContextHolder;
 import org.springframework.stereotype.Component;
 
 import javax.annotation.Resource;
-import java.util.Collections;
-import java.util.Date;
-import java.util.List;
+import javax.sql.DataSource;
+import java.sql.Connection;
+import java.sql.ResultSet;
+import java.util.*;
 
 /**
  * SaaS WebSocket 租户上下文工具:
@@ -139,6 +140,50 @@ public class TenantChannelContext {
         }
     }
 
+    /**
+     * javax.websocket 握手阶段专用:当未传 tenantCode 时,根据 liveId 遍历所有租户库查找归属租户。
+     * 找到后将 tenantId/dsKey 存入 userProperties,兼容后台/销售端不带租户信息的场景。
+     */
+    public void resolveByLiveId(Long liveId, java.util.Map<String, Object> userProperties) {
+        if (!saasEnabled || liveId == null) return;
+        try {
+            DynamicDataSourceContextHolder.setDataSourceType(DataSourceType.MASTER.name());
+            TenantInfo query = new TenantInfo();
+            query.setStatus(1);
+            List<TenantInfo> tenants = tenantInfoService.selectTenantInfoList(query);
+            if (tenants == null || tenants.isEmpty()) return;
+
+            Date now = new Date();
+            for (TenantInfo tenant : tenants) {
+                if (tenant.getExpireTime() != null && tenant.getExpireTime().before(now)) continue;
+                try {
+                    tenantDataSourceManager.switchTenant(tenant);
+                    DataSource ds = tenantDataSourceManager.getTenantDataSource(tenant.getId());
+                    try (Connection conn = ds.getConnection();
+                         java.sql.PreparedStatement ps = conn.prepareStatement(
+                                 "SELECT 1 FROM live WHERE id = ? LIMIT 1")) {
+                        ps.setLong(1, liveId);
+                        try (ResultSet rs = ps.executeQuery()) {
+                            if (rs.next()) {
+                                String dsKey = "tenant:" + tenant.getId();
+                                userProperties.put("_tenantId", tenant.getId());
+                                userProperties.put("_dsKey", dsKey);
+                                log.info("[SaaS WS] 根据liveId={}反查到租户 tenantId={}, tenantCode={}",
+                                        liveId, tenant.getId(), tenant.getTenantCode());
+                                return;
+                            }
+                        }
+                    }
+                } catch (Exception e) {
+                    log.debug("[SaaS WS] 查询租户 tenantCode={} 的live表失败: {}", tenant.getTenantCode(), e.getMessage());
+                }
+            }
+            log.warn("[SaaS WS] 根据liveId={} 未找到归属租户", liveId);
+        } finally {
+            DynamicDataSourceContextHolder.clearDataSourceType();
+        }
+    }
+
     /**
      * javax.websocket 业务处理前专用:从 userProperties 取 tenantId/dsKey 直接激活,无需查库。
      * 在 onOpen/onMessage/onClose 开头调用。

+ 9 - 2
fs-live-app/src/main/java/com/fs/live/websocket/auth/WebSocketConfigurator.java

@@ -1,6 +1,7 @@
 package com.fs.live.websocket.auth;
 
 import com.fs.common.exception.base.BaseException;
+import com.fs.common.utils.StringUtils;
 import com.fs.live.utils.JwtUtils;
 import com.fs.live.utils.VerifyUtils;
 import com.fs.live.websocket.constant.AttrConstant;
@@ -81,10 +82,16 @@ public class WebSocketConfigurator extends ServerEndpointConfig.Configurator {
         // SaaS:握手阶段查库,将 tenantId/dsKey 存入 userProperties,后续 onOpen/onMessage/onClose 无需再查库
         try {
             TenantChannelContext tenantChannelContext = SpringUtils.getBean(TenantChannelContext.class);
-            tenantChannelContext.resolveAndStore(tenantCode, userProperties);
+            if (StringUtils.isNotBlank(tenantCode)) {
+                // 有 tenantCode:直接根据 tenantCode 查租户
+                tenantChannelContext.resolveAndStore(tenantCode, userProperties);
+            } else {
+                // 无 tenantCode:根据 liveId 遍历租户库反查归属租户(兼容后台/销售端不带租户信息的场景)
+                tenantChannelContext.resolveByLiveId(liveId, userProperties);
+            }
         } catch (Exception e) {
             log.error("[SaaS WS] 握手阶段绑定租户失败, tenantCode={}", tenantCode, e);
-            throw new BaseException("租户信息验证失败: " + tenantCode);
+            throw new BaseException("租户信息验证失败");
         }
 
         // 验证token