From a60133bfb329ca1fd67634ab4e6099885008fd6c Mon Sep 17 00:00:00 2001
From: Bowen Liang <liangbowen@gf.com.cn>
Date: Fri, 8 Nov 2024 09:33:12 +0800
Subject: [PATCH] fix: config violations when running db migtration ci tests
 (#10428)

---
 .github/workflows/db-migration-test.yml |  1 +
 api/.env.example                        |  8 +++----
 api/configs/feature/__init__.py         | 10 +++++++++
 api/core/helper/ssrf_proxy.py           | 30 ++++++++++---------------
 4 files changed, 27 insertions(+), 22 deletions(-)

diff --git a/.github/workflows/db-migration-test.yml b/.github/workflows/db-migration-test.yml
index b8246aacb3..c6fe87264d 100644
--- a/.github/workflows/db-migration-test.yml
+++ b/.github/workflows/db-migration-test.yml
@@ -6,6 +6,7 @@ on:
       - main
     paths:
       - api/migrations/**
+      - .github/workflows/db-migration-test.yml
 
 concurrency:
   group: db-migration-test-${{ github.ref }}
diff --git a/api/.env.example b/api/.env.example
index a92490608f..1214154389 100644
--- a/api/.env.example
+++ b/api/.env.example
@@ -324,10 +324,10 @@ UNSTRUCTURED_API_KEY=
 SSRF_PROXY_HTTP_URL=
 SSRF_PROXY_HTTPS_URL=
 SSRF_DEFAULT_MAX_RETRIES=3
-SSRF_DEFAULT_TIME_OUT=
-SSRF_DEFAULT_CONNECT_TIME_OUT=
-SSRF_DEFAULT_READ_TIME_OUT=
-SSRF_DEFAULT_WRITE_TIME_OUT=
+SSRF_DEFAULT_TIME_OUT=5
+SSRF_DEFAULT_CONNECT_TIME_OUT=5
+SSRF_DEFAULT_READ_TIME_OUT=5
+SSRF_DEFAULT_WRITE_TIME_OUT=5
 
 BATCH_UPLOAD_LIMIT=10
 KEYWORD_DATA_SOURCE_TYPE=database
diff --git a/api/configs/feature/__init__.py b/api/configs/feature/__init__.py
index 3ac2c28c1f..5babae8810 100644
--- a/api/configs/feature/__init__.py
+++ b/api/configs/feature/__init__.py
@@ -276,6 +276,16 @@ class HttpConfig(BaseSettings):
         default=1 * 1024 * 1024,
     )
 
+    SSRF_DEFAULT_MAX_RETRIES: PositiveInt = Field(
+        description="Maximum number of retries for network requests (SSRF)",
+        default=3,
+    )
+
+    SSRF_PROXY_ALL_URL: Optional[str] = Field(
+        description="Proxy URL for HTTP or HTTPS requests to prevent Server-Side Request Forgery (SSRF)",
+        default=None,
+    )
+
     SSRF_PROXY_HTTP_URL: Optional[str] = Field(
         description="Proxy URL for HTTP requests to prevent Server-Side Request Forgery (SSRF)",
         default=None,
diff --git a/api/core/helper/ssrf_proxy.py b/api/core/helper/ssrf_proxy.py
index df812ca83f..374bd9d57b 100644
--- a/api/core/helper/ssrf_proxy.py
+++ b/api/core/helper/ssrf_proxy.py
@@ -3,26 +3,20 @@ Proxy requests to avoid SSRF
 """
 
 import logging
-import os
 import time
 
 import httpx
 
-SSRF_PROXY_ALL_URL = os.getenv("SSRF_PROXY_ALL_URL", "")
-SSRF_PROXY_HTTP_URL = os.getenv("SSRF_PROXY_HTTP_URL", "")
-SSRF_PROXY_HTTPS_URL = os.getenv("SSRF_PROXY_HTTPS_URL", "")
-SSRF_DEFAULT_MAX_RETRIES = int(os.getenv("SSRF_DEFAULT_MAX_RETRIES", "3"))
-SSRF_DEFAULT_TIME_OUT = float(os.getenv("SSRF_DEFAULT_TIME_OUT", "5"))
-SSRF_DEFAULT_CONNECT_TIME_OUT = float(os.getenv("SSRF_DEFAULT_CONNECT_TIME_OUT", "5"))
-SSRF_DEFAULT_READ_TIME_OUT = float(os.getenv("SSRF_DEFAULT_READ_TIME_OUT", "5"))
-SSRF_DEFAULT_WRITE_TIME_OUT = float(os.getenv("SSRF_DEFAULT_WRITE_TIME_OUT", "5"))
+from configs import dify_config
+
+SSRF_DEFAULT_MAX_RETRIES = dify_config.SSRF_DEFAULT_MAX_RETRIES
 
 proxy_mounts = (
     {
-        "http://": httpx.HTTPTransport(proxy=SSRF_PROXY_HTTP_URL),
-        "https://": httpx.HTTPTransport(proxy=SSRF_PROXY_HTTPS_URL),
+        "http://": httpx.HTTPTransport(proxy=dify_config.SSRF_PROXY_HTTP_URL),
+        "https://": httpx.HTTPTransport(proxy=dify_config.SSRF_PROXY_HTTPS_URL),
     }
-    if SSRF_PROXY_HTTP_URL and SSRF_PROXY_HTTPS_URL
+    if dify_config.SSRF_PROXY_HTTP_URL and dify_config.SSRF_PROXY_HTTPS_URL
     else None
 )
 
@@ -38,17 +32,17 @@ def make_request(method, url, max_retries=SSRF_DEFAULT_MAX_RETRIES, **kwargs):
 
     if "timeout" not in kwargs:
         kwargs["timeout"] = httpx.Timeout(
-            SSRF_DEFAULT_TIME_OUT,
-            connect=SSRF_DEFAULT_CONNECT_TIME_OUT,
-            read=SSRF_DEFAULT_READ_TIME_OUT,
-            write=SSRF_DEFAULT_WRITE_TIME_OUT,
+            timeout=dify_config.SSRF_DEFAULT_TIME_OUT,
+            connect=dify_config.SSRF_DEFAULT_CONNECT_TIME_OUT,
+            read=dify_config.SSRF_DEFAULT_READ_TIME_OUT,
+            write=dify_config.SSRF_DEFAULT_WRITE_TIME_OUT,
         )
 
     retries = 0
     while retries <= max_retries:
         try:
-            if SSRF_PROXY_ALL_URL:
-                with httpx.Client(proxy=SSRF_PROXY_ALL_URL) as client:
+            if dify_config.SSRF_PROXY_ALL_URL:
+                with httpx.Client(proxy=dify_config.SSRF_PROXY_ALL_URL) as client:
                     response = client.request(method=method, url=url, **kwargs)
             elif proxy_mounts:
                 with httpx.Client(mounts=proxy_mounts) as client: