security: add ssrf proxy

api/.env.example

@@ -120,4 +120,7 @@ HOSTED_ANTHROPIC_QUOTA_LIMIT=600000
 HOSTED_ANTHROPIC_PAID_ENABLED=false
 
 ETL_TYPE=dify
-UNSTRUCTURED_API_URL=
+UNSTRUCTURED_API_URL=
+
+SSRF_PROXY_HTTP_URL=
+SSRF_PROXY_HTTPS_URL=

api/core/extension/api_based_extension_requestor.py

@@ -30,10 +30,10 @@ class APIBasedExtensionRequestor:
         try:
             # proxy support for security
             proxies = None
-            if os.environ.get("API_BASED_EXTENSION_HTTP_PROXY") and os.environ.get("API_BASED_EXTENSION_HTTPS_PROXY"):
+            if os.environ.get("SSRF_PROXY_HTTP_URL") and os.environ.get("SSRF_PROXY_HTTPS_URL"):
                 proxies = {
-                    'http': os.environ.get("API_BASED_EXTENSION_HTTP_PROXY"),
-                    'https': os.environ.get("API_BASED_EXTENSION_HTTPS_PROXY"),
+                    'http': os.environ.get("SSRF_PROXY_HTTP_URL"),
+                    'https': os.environ.get("SSRF_PROXY_HTTPS_URL"),
                 }
 
             response = requests.request(

api/core/helper/ssrf_proxy.py

@@ -0,0 +1,42 @@
+"""
+Proxy requests to avoid SSRF
+"""
+
+from httpx import get as _get, post as _post, put as _put, patch as _patch, head as _head, options as _options
+from requests import delete as _delete
+
+import os
+
+SSRF_PROXY_HTTP_URL = os.getenv('SSRF_PROXY_HTTP_URL', '')
+SSRF_PROXY_HTTPS_URL = os.getenv('SSRF_PROXY_HTTPS_URL', '')
+
+requests_proxies = {
+    'http': SSRF_PROXY_HTTP_URL,
+    'https': SSRF_PROXY_HTTPS_URL
+} if SSRF_PROXY_HTTP_URL and SSRF_PROXY_HTTPS_URL else None
+
+httpx_proxies = {
+    'http://': SSRF_PROXY_HTTP_URL,
+    'https://': SSRF_PROXY_HTTPS_URL
+} if SSRF_PROXY_HTTP_URL and SSRF_PROXY_HTTPS_URL else None
+
+def get(url, *args, **kwargs):
+    return _get(url=url, *args, proxies=httpx_proxies, **kwargs)
+
+def post(url, *args, **kwargs):
+    return _post(url=url, *args, proxies=httpx_proxies, **kwargs)
+
+def put(url, *args, **kwargs):
+    return _put(url=url, *args, proxies=httpx_proxies, **kwargs)
+
+def patch(url, *args, **kwargs):
+    return _patch(url=url, *args, proxies=httpx_proxies, **kwargs)
+
+def delete(url, *args, **kwargs):
+    return _delete(url=url, *args, proxies=requests_proxies, **kwargs)
+
+def head(url, *args, **kwargs):
+    return _head(url=url, *args, proxies=httpx_proxies, **kwargs)
+
+def options(url, *args, **kwargs):
+    return _options(url=url, *args, proxies=httpx_proxies, **kwargs)

api/core/tools/tool/api_tool.py

@@ -4,6 +4,7 @@ from typing import Any, Dict, List, Union
 
 import httpx
 import requests
+import core.helper.ssrf_proxy as ssrf_proxy
 from core.tools.entities.tool_bundle import ApiBasedToolBundle
 from core.tools.entities.tool_entities import ToolInvokeMessage
 from core.tools.errors import ToolProviderCredentialValidationError
@@ -201,23 +202,23 @@ class ApiTool(Tool):
         
         # do http request
         if method == 'get':
-            response = httpx.get(url, params=params, headers=headers, cookies=cookies, timeout=10, follow_redirects=True)
+            response = ssrf_proxy.get(url, params=params, headers=headers, cookies=cookies, timeout=10, follow_redirects=True)
         elif method == 'post':
-            response = httpx.post(url, params=params, headers=headers, cookies=cookies, data=body, timeout=10, follow_redirects=True)
+            response = ssrf_proxy.post(url, params=params, headers=headers, cookies=cookies, data=body, timeout=10, follow_redirects=True)
         elif method == 'put':
-            response = httpx.put(url, params=params, headers=headers, cookies=cookies, data=body, timeout=10, follow_redirects=True)
+            response = ssrf_proxy.put(url, params=params, headers=headers, cookies=cookies, data=body, timeout=10, follow_redirects=True)
         elif method == 'delete':
             """
             request body data is unsupported for DELETE method in standard http protocol
             however, OpenAPI 3.0 supports request body data for DELETE method, so we support it here by using requests
             """
-            response = requests.delete(url, params=params, headers=headers, cookies=cookies, data=body, timeout=10, allow_redirects=True)
+            response = ssrf_proxy.delete(url, params=params, headers=headers, cookies=cookies, data=body, timeout=10, allow_redirects=True)
         elif method == 'patch':
-            response = httpx.patch(url, params=params, headers=headers, cookies=cookies, data=body, timeout=10, follow_redirects=True)
+            response = ssrf_proxy.patch(url, params=params, headers=headers, cookies=cookies, data=body, timeout=10, follow_redirects=True)
         elif method == 'head':
-            response = httpx.head(url, params=params, headers=headers, cookies=cookies, timeout=10, follow_redirects=True)
+            response = ssrf_proxy.head(url, params=params, headers=headers, cookies=cookies, timeout=10, follow_redirects=True)
         elif method == 'options':
-            response = httpx.options(url, params=params, headers=headers, cookies=cookies, timeout=10, follow_redirects=True)
+            response = ssrf_proxy.options(url, params=params, headers=headers, cookies=cookies, timeout=10, follow_redirects=True)
         else:
             raise ValueError(f'Invalid http method {method}')