From a9ed0f0b42f6ddb93015959217ccfbab548a465c Mon Sep 17 00:00:00 2001
From: Summer-Gu <37869445+gubinjie@users.noreply.github.com>
Date: Wed, 6 Nov 2024 08:50:57 +0800
Subject: [PATCH] feat: The SSRF request timeout configuration item is added
 (#10292)

---
 api/.env.example                |  5 +++++
 api/configs/feature/__init__.py | 20 ++++++++++++++++++++
 api/core/helper/ssrf_proxy.py   | 12 ++++++++++++
 3 files changed, 37 insertions(+)

diff --git a/api/.env.example b/api/.env.example
index f7bcab6d6d..6fc58263c4 100644
--- a/api/.env.example
+++ b/api/.env.example
@@ -320,9 +320,14 @@ ETL_TYPE=dify
 UNSTRUCTURED_API_URL=
 UNSTRUCTURED_API_KEY=
 
+#ssrf
 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=
 
 BATCH_UPLOAD_LIMIT=10
 KEYWORD_DATA_SOURCE_TYPE=database
diff --git a/api/configs/feature/__init__.py b/api/configs/feature/__init__.py
index 533a24dcbd..517b92fda4 100644
--- a/api/configs/feature/__init__.py
+++ b/api/configs/feature/__init__.py
@@ -286,6 +286,26 @@ class HttpConfig(BaseSettings):
         default=None,
     )
 
+    SSRF_DEFAULT_TIME_OUT: PositiveFloat = Field(
+        description="The default timeout period used for network requests (SSRF)",
+        default=5,
+    )
+
+    SSRF_DEFAULT_CONNECT_TIME_OUT: PositiveFloat = Field(
+        description="The default connect timeout period used for network requests (SSRF)",
+        default=5,
+    )
+
+    SSRF_DEFAULT_READ_TIME_OUT: PositiveFloat = Field(
+        description="The default read timeout period used for network requests (SSRF)",
+        default=5,
+    )
+
+    SSRF_DEFAULT_WRITE_TIME_OUT: PositiveFloat = Field(
+        description="The default write timeout period used for network requests (SSRF)",
+        default=5,
+    )
+
     RESPECT_XFORWARD_HEADERS_ENABLED: bool = Field(
         description="Enable or disable the X-Forwarded-For Proxy Fix middleware from Werkzeug"
         " to respect X-* headers to redirect clients",
diff --git a/api/core/helper/ssrf_proxy.py b/api/core/helper/ssrf_proxy.py
index 6793e41978..df812ca83f 100644
--- a/api/core/helper/ssrf_proxy.py
+++ b/api/core/helper/ssrf_proxy.py
@@ -12,6 +12,10 @@ 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"))
 
 proxy_mounts = (
     {
@@ -32,6 +36,14 @@ def make_request(method, url, max_retries=SSRF_DEFAULT_MAX_RETRIES, **kwargs):
         if "follow_redirects" not in kwargs:
             kwargs["follow_redirects"] = allow_redirects
 
+    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,
+        )
+
     retries = 0
     while retries <= max_retries:
         try: