perf: Multiple Domain Validation for Embedded Systems #388

backend/apps/system/api/assistant.py

@@ -17,7 +17,7 @@ from common.core.config import settings
 from common.core.deps import SessionDep, Trans
 from common.core.security import create_access_token
 from common.core.sqlbot_cache import clear_cache
-from common.utils.utils import get_origin_from_referer
+from common.utils.utils import get_origin_from_referer, origin_match_domain
 
 router = APIRouter(tags=["system/assistant"], prefix="/system/assistant")
 
@@ -30,13 +30,15 @@ async def info(request: Request, response: Response, session: SessionDep, trans:
     if not db_model:
         raise RuntimeError(f"assistant application not exist")
     db_model = AssistantModel.model_validate(db_model)
-    response.headers["Access-Control-Allow-Origin"] = db_model.domain
+    
     origin = request.headers.get("origin") or get_origin_from_referer(request)
     if not origin:
         raise RuntimeError(trans('i18n_embedded.invalid_origin', origin=origin or ''))
     origin = origin.rstrip('/')
-    if origin != db_model.domain:
+    if not origin_match_domain(origin, db_model.domain):
         raise RuntimeError(trans('i18n_embedded.invalid_origin', origin=origin or ''))
+    
+    response.headers["Access-Control-Allow-Origin"] = origin
     return db_model
 
 
@@ -48,13 +50,14 @@ async def getApp(request: Request, response: Response, session: SessionDep, tran
     if not db_model:
         raise RuntimeError(f"assistant application not exist")
     db_model = AssistantModel.model_validate(db_model)
-    response.headers["Access-Control-Allow-Origin"] = db_model.domain
     origin = request.headers.get("origin") or get_origin_from_referer(request)
     if not origin:
         raise RuntimeError(trans('i18n_embedded.invalid_origin', origin=origin or ''))
     origin = origin.rstrip('/')
-    if origin != db_model.domain:
+    if not origin_match_domain(origin, db_model.domain):
         raise RuntimeError(trans('i18n_embedded.invalid_origin', origin=origin or ''))
+    
+    response.headers["Access-Control-Allow-Origin"] = origin
     return db_model
 
 

backend/apps/system/crud/assistant.py

@@ -99,17 +99,22 @@ class AssistantOutDs:
     assistant: AssistantHeader
     ds_list: Optional[list[AssistantOutDsSchema]] = None
     certificate: Optional[str] = None
+    request_origin: Optional[str] = None
 
     def __init__(self, assistant: AssistantHeader):
         self.assistant = assistant
         self.ds_list = None
         self.certificate = assistant.certificate
+        self.request_origin = assistant.request_origin
         self.get_ds_from_api()
 
     # @cache(namespace=CacheNamespace.EMBEDDED_INFO, cacheName=CacheName.ASSISTANT_DS, keyExpression="current_user.id")
     def get_ds_from_api(self):
         config: dict[any] = json.loads(self.assistant.configuration)
         endpoint: str = config['endpoint']
+        endpoint = self.get_complete_endpoint(endpoint=endpoint)
+        if not endpoint:
+            raise Exception(f"Failed to get datasource list from {config['endpoint']}, error: [Assistant domain or endpoint miss]")
         certificateList: list[any] = json.loads(self.certificate)
         header = {}
         cookies = {}
@@ -137,6 +142,17 @@ class AssistantOutDs:
         else:
             raise Exception(f"Failed to get datasource list from {endpoint}, status code: {res.status_code}")
 
+    def get_complete_endpoint(self, endpoint: str) -> str | None:
+        if endpoint.startswith("http://") or endpoint.startswith("https://"):
+            return endpoint
+        domain_text = self.assistant.domain
+        if not domain_text:
+            return None
+        if ',' in domain_text:
+            return self.request_origin.strip('/') if self.request_origin else domain_text.split(',')[0].strip('/') + endpoint
+        else:
+            return f"{domain_text}{endpoint}"  
+    
     def get_simple_ds_list(self):
         if self.ds_list:
             return [{'id': ds.id, 'name': ds.name, 'description': ds.comment} for ds in self.ds_list]

backend/apps/system/middleware/auth.py

@@ -16,7 +16,7 @@ from common.core import security
 from common.core.config import settings
 from common.core.schemas import TokenPayload
 from common.utils.locale import I18n
-from common.utils.utils import SQLBotLogUtil
+from common.utils.utils import SQLBotLogUtil, get_origin_from_referer
 from common.utils.whitelist import whiteUtils
 from fastapi.security.utils import get_authorization_scheme_param
 from common.core.deps import get_i18n
@@ -40,6 +40,9 @@ class TokenMiddleware(BaseHTTPMiddleware):
             if validator[0]:
                 request.state.current_user = validator[1]
                 request.state.assistant = validator[2]
+                origin = request.headers.get("origin") or get_origin_from_referer(request)
+                if origin and validator[2]:
+                    request.state.assistant.request_origin = origin
                 return await call_next(request)
             message = trans('i18n_permission.authenticate_invalid', msg = validator[1])
             return JSONResponse(message, status_code=401, headers={"Access-Control-Allow-Origin": "*"})

backend/apps/system/schemas/system_schema.py

@@ -116,6 +116,7 @@ class AssistantHeader(AssistantDTO):
     unique: Optional[str] = None
     certificate: Optional[str] = None
     online: bool = False
+    request_origin: Optional[str] = None
 
 
 class AssistantValidator(BaseModel):

backend/common/utils/utils.py

@@ -263,6 +263,12 @@ def get_origin_from_referer(request: Request):
         SQLBotLogUtil.error(f"解析 Referer 出错: {e}")
         return referer
 
+def origin_match_domain(origin: str, domain: str) -> bool:
+    if not origin or not domain:
+        return False
+    origin_text = origin.rstrip('/')
+    domain_list = domain.replace(" ", "").split(',')
+    return origin_text in [d.rstrip('/') for d in domain_list]
 
 def equals_ignore_case(str1: str, *args: str) -> bool:
     if str1 is None: