From c0b71f8286e90f575b1c1eb8d0cbd081595cab26 Mon Sep 17 00:00:00 2001
From: kurokobo <kuro664@gmail.com>
Date: Wed, 9 Oct 2024 15:42:30 +0900
Subject: [PATCH] feat: respect x-* headers for redirections (#9054)

---
 api/.env.example                |  3 +++
 api/app.py                      |  2 ++
 api/configs/feature/__init__.py |  6 ++++++
 api/extensions/ext_proxy_fix.py | 10 ++++++++++
 4 files changed, 21 insertions(+)
 create mode 100644 api/extensions/ext_proxy_fix.py

diff --git a/api/.env.example b/api/.env.example
index fb0fc045b..71f0e5db8 100644
--- a/api/.env.example
+++ b/api/.env.example
@@ -271,6 +271,9 @@ HTTP_REQUEST_MAX_WRITE_TIMEOUT=600
 HTTP_REQUEST_NODE_MAX_BINARY_SIZE=10485760
 HTTP_REQUEST_NODE_MAX_TEXT_SIZE=1048576
 
+# Respect X-* headers to redirect clients
+RESPECT_XFORWARD_HEADERS_ENABLED=false
+
 # Log file path
 LOG_FILE=
 
diff --git a/api/app.py b/api/app.py
index 0662d94e1..a251ef5f0 100644
--- a/api/app.py
+++ b/api/app.py
@@ -36,6 +36,7 @@ from extensions import (
     ext_login,
     ext_mail,
     ext_migrate,
+    ext_proxy_fix,
     ext_redis,
     ext_sentry,
     ext_storage,
@@ -156,6 +157,7 @@ def initialize_extensions(app):
     ext_mail.init_app(app)
     ext_hosting_provider.init_app(app)
     ext_sentry.init_app(app)
+    ext_proxy_fix.init_app(app)
 
 
 # Flask-Login configuration
diff --git a/api/configs/feature/__init__.py b/api/configs/feature/__init__.py
index 9218d529c..93dbc1367 100644
--- a/api/configs/feature/__init__.py
+++ b/api/configs/feature/__init__.py
@@ -247,6 +247,12 @@ class HttpConfig(BaseSettings):
         default=None,
     )
 
+    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",
+        default=False,
+    )
+
 
 class InnerAPIConfig(BaseSettings):
     """
diff --git a/api/extensions/ext_proxy_fix.py b/api/extensions/ext_proxy_fix.py
new file mode 100644
index 000000000..c106a4384
--- /dev/null
+++ b/api/extensions/ext_proxy_fix.py
@@ -0,0 +1,10 @@
+from flask import Flask
+
+from configs import dify_config
+
+
+def init_app(app: Flask):
+    if dify_config.RESPECT_XFORWARD_HEADERS_ENABLED:
+        from werkzeug.middleware.proxy_fix import ProxyFix
+
+        app.wsgi_app = ProxyFix(app.wsgi_app)