|
@@ -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:
|