feat(cos): add self domain

MorvanZhou · May 24, 2024 · 8fa25c1 · 8fa25c1
1 parent 7c6adc3
commit 8fa25c1
Show file tree

Hide file tree

Showing 6 changed files with 23 additions and 7 deletions.
diff --git a/requirements.txt b/requirements.txt
@@ -11,7 +11,7 @@ oauthlib>=3.2.2
 httpx>=0.25.0
 captcha>=0.5.0
 python-multipart>=0.0.9
-cos-python-sdk-v5~=1.9.26
+cos-python-sdk-v5~=1.9.29
 whoosh~=2.7.4
 jieba>=0.42.1
 starlette>=0.37.2

diff --git a/setup.cfg b/setup.cfg
@@ -64,4 +64,4 @@ build =
 remote =
     motor>=3.3.2
     elasticsearch[async]~=8.11.0
-    cos-python-sdk-v5~=1.9.26
+    cos-python-sdk-v5~=1.9.29
diff --git a/src/retk/config.py b/src/retk/config.py
@@ -49,8 +49,9 @@ class Settings(BaseSettings):
     OAUTH_CLIENT_SEC_GOOGLE: str = Field(env='OAUTH_CLIENT_SEC_GOOGLE', default="")
     COS_SECRET_ID: str = Field(env='COS_SECRET_ID', default="")
     COS_SECRET_KEY: str = Field(env="COS_SECRET_KEY", default="")
-    COS_REGION: str = Field(env="COS_REGION", default="")
+    COS_REGION: Optional[str] = Field(env="COS_REGION", default=None)
     COS_BUCKET_NAME: str = Field(env="COS_BUCKET_NAME", default="")
+    COS_DOMAIN: Optional[str] = Field(env="COS_DOMAIN", default=None)
     MD_BACKUP_INTERVAL: int = Field(env="MD_BACKUP_INTERVAL", default=60 * 5)  # 5 minutes
 
     model_config = SettingsConfigDict(
@@ -88,6 +89,11 @@ def __init__(self):
         self.REFRESH_TOKEN_EXPIRE_DELTA = datetime.timedelta(days=self.JWT_REFRESH_EXPIRED_DAYS)
         self.ACCESS_TOKEN_EXPIRE_DELTA = datetime.timedelta(minutes=self.JWT_ACCESS_EXPIRED_MINS)
 
+        if self.COS_DOMAIN == "":
+            self.COS_DOMAIN = None
+        if self.COS_REGION == "":
+            self.COS_REGION = None
+
 
 @lru_cache()
 def get_settings() -> Settings:

diff --git a/src/retk/core/files/saver.py b/src/retk/core/files/saver.py
@@ -119,7 +119,7 @@ async def save_remote(self, uid: str, file: File):
         secret_id = settings.COS_SECRET_ID
         secret_key = settings.COS_SECRET_KEY
         region = settings.COS_REGION
-        domain = None
+        domain = settings.COS_DOMAIN
         cos_conf = CosConfig(
             Region=region,
             SecretId=secret_id,
@@ -132,7 +132,8 @@ async def save_remote(self, uid: str, file: File):
 
         key = f"userData/{uid}/{file.hashed_filename}"
 
-        url = f"https://{settings.COS_BUCKET_NAME}.cos.{settings.COS_REGION}.myqcloud.com/{key}"
+        domain = settings.COS_DOMAIN or f"{settings.COS_BUCKET_NAME}.cos.{region}.myqcloud.com"
+        url = f"https://{domain}/{key}"
 
         doc = await client.coll.user_file.find_one({"uid": uid, "fid": file.hashed_filename})
         if doc:

diff --git a/src/retk/core/node/backup.py b/src/retk/core/node/backup.py
@@ -228,7 +228,7 @@ def __get_client_and_key(settings: BaseSettings, uid: str, nid: str, version: st
             SecretId=settings.COS_SECRET_ID,
             SecretKey=settings.COS_SECRET_KEY,
             Token=None,
-            Domain=None,
+            Domain=settings.COS_DOMAIN,
             Scheme='https',
         )
     )

diff --git a/tests/test_utils.py b/tests/test_utils.py
@@ -124,6 +124,7 @@ async def test_get_title_description_from_link(self, mock_get_settings, mock_get
         s.COS_BUCKET_NAME = "rethink-dev-1258395282"
         s.COS_REGION = "ap-hongkong"
         s.DB_HOST = "127.0.0.1"
+        s.COS_DOMAIN = None
         mock_get_settings.return_value = s
         for url, content, res in [
             (
@@ -221,19 +222,27 @@ def test_mask_email(self):
         "retk.config.get_settings",
     )
     def test_ssrf(self, mock_get_settings):
-        for pre, url in [
+        for pre, url, domain in [
             (
                     "rethink-product-1258395282",
                     "https://rethink-product-1258395282.cos.ap-hongkong.myqcloud.com/userData/Rro/as.png",
+                    None,
             ),
             (
                     "rethink-dev-1258395282",
                     "https://rethink-dev-1258395282.cos.ap-hongkong.myqcloud.com/userData/qwqd/qwww.png",
+                    None,
+            ),
+            (
+                    "rethink-dev-1258395282",
+                    "https://files.rethink.run/userData/qwqd/qwww.png",
+                    "files.rethink.run"
             ),
         ]:
             s = config.Settings
             s.COS_BUCKET_NAME = pre
             s.COS_REGION = "ap-hongkong"
             s.DB_HOST = "127.0.0.1"
+            s.COS_DOMAIN = domain
             mock_get_settings.return_value = s
             self.assertTrue(utils.ssrf_check(url), msg=url)