From 427cc7260a8e92f32b8a8d68b72a1265d07be45b Mon Sep 17 00:00:00 2001
From: helloplhm-qwq <helloplhm-qwq@outlook.com>
Date: Sat, 13 Jan 2024 14:59:30 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20=E4=B8=80=E4=B8=AAplus=E7=89=88?=
 =?UTF-8?q?=E7=9A=84=E5=8F=8D=E4=BB=A3=E8=87=AA=E5=AE=9A=E4=B9=89=E9=85=8D?=
 =?UTF-8?q?=E7=BD=AE=E5=AE=9E=E7=8E=B0=20#22?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 common/config.py | 11 +++++++++++
 main.py          | 18 ++++++++++++------
 2 files changed, 23 insertions(+), 6 deletions(-)

diff --git a/common/config.py b/common/config.py
index 19203e3..ac214a9 100644
--- a/common/config.py
+++ b/common/config.py
@@ -62,6 +62,17 @@ class ConfigReadException(Exception):
                 "privkey": "/path/to/your/private/key",
             },
         },
+        "reverse_proxy": {
+            "desc": "针对类似于nginx一类的反代的配置",
+            "allow_proxy": True,
+            "_allow_proxy-desc": "是否允许反代",
+            "proxy_whitelist_remote": [
+                "反代时允许的ip来源列表，通常为127.0.0.1",
+                "127.0.0.1"
+            ],
+            "real_ip_header": 'X-Real-IP',
+            "_real_ip_header-desc": "反代来源ip的来源头，不懂请保持默认",
+        },
         "debug_mode": False,
         "_debug_mode-desc": "是否开启调试模式",
         "log_length_limit": 500,
diff --git a/main.py b/main.py
index d40ca69..d311d14 100644
--- a/main.py
+++ b/main.py
@@ -32,7 +32,7 @@
 import time
 import os
 
-def handleResult(dic, status = 200):
+def handleResult(dic, status = 200) -> Response:
     return Response(body = json.dumps(dic, indent=2, ensure_ascii=False), content_type='application/json', status = status)
 
 logger = log.log("main")
@@ -46,16 +46,22 @@ def handleResult(dic, status = 200):
 else:
     stopEvent = asyncio.exceptions.CancelledError
 
-def start_checkcn_thread():
+def start_checkcn_thread() -> None:
     threading.Thread(target=Httpx.checkcn).start()
 
 # check request info before start
 async def handle_before_request(app, handler):
     async def handle_request(request):
         try:
-            # nginx proxy header
-            if (request.headers.get("X-Real-IP")):
-                request.remote_addr = request.headers.get("X-Real-IP")
+            if (config.read_config('common.reverse_proxy.enable')):
+                if (request.headers.get(config.read_config('common.reverse_proxy.real_ip_header'))):
+                    # proxy header
+                    if (request.remote in config.read_config('common.reverse_proxy.proxy_whitelist_remote')):
+                        request.remote_addr = request.headers.get(config.read_config('common.reverse_proxy.real_ip_header'))
+                    else:
+                        return handleResult({"code": 1, "msg": "反代客户端远程地址不在反代ip白名单中", "data": None}, 403)
+                else:
+                    request.remote_addr = request.remote
             else:
                 request.remote_addr = request.remote
             # check ip
@@ -91,7 +97,7 @@ async def handle_request(request):
                 resp = handleResult(resp)
             elif (not isinstance(resp, Response)):
                 resp = Response(body = str(resp), content_type='text/plain', status = 200)
-            aiologger.info(f'{request.remote_addr} - {request.method} "{request.path}", {resp.status}')
+            aiologger.info(f'{request.remote_addr + "" if (request.remote == request.remote_addr) else f"|proxy@{request.remote}"} - {request.method} "{request.path}", {resp.status}')
             return resp
         except: 
             logger.error(traceback.format_exc())