From 0f3f43850bd83a0eeccfb16f24fe41b9bdd369eb Mon Sep 17 00:00:00 2001
From: noO0oOo0ob <38344038+noO0oOo0ob@users.noreply.github.com>
Date: Tue, 24 Dec 2024 14:27:09 +0800
Subject: [PATCH] not redirct websocket request in mitm (#884)

* web not proxy websocket

* add upgrade request report log

* version

---------

Co-authored-by: noO0ob <625171796@qq.com>
---
 lyrebird/mitm/mitm_script.py                | 14 +++++++--
 lyrebird/mock/extra_mock_server/__init__.py |  1 -
 lyrebird/mock/extra_mock_server/server.py   | 33 +++++++++++++++++++--
 lyrebird/version.py                         |  2 +-
 4 files changed, 43 insertions(+), 7 deletions(-)

diff --git a/lyrebird/mitm/mitm_script.py b/lyrebird/mitm/mitm_script.py
index cdf91356..c17afae8 100644
--- a/lyrebird/mitm/mitm_script.py
+++ b/lyrebird/mitm/mitm_script.py
@@ -3,9 +3,8 @@
 
 Redirect request from proxy server to mock server
 """
-import re
-import json
 import os
+import json
 from mitmproxy import http
 from urllib.parse import urlparse
 
@@ -16,6 +15,14 @@
 PROXY_FILTERS = json.loads(os.environ.get('PROXY_FILTERS'))
 
 
+def is_websocket_request(flow: http.HTTPFlow) -> bool:
+    headers = flow.request.headers
+    return (
+        headers.get('Upgrade', '').lower() == 'websocket' and
+        headers.get('Connection', '').lower() == 'upgrade'
+    )
+
+
 def check_lyrebird_request(flow: http.HTTPFlow):
     parsed_url = urlparse(flow.request.url)
     host = parsed_url.hostname
@@ -59,6 +66,9 @@ def request(flow: http.HTTPFlow):
     if check_lyrebird_request(flow):
         # Avoid internal requests
         return
+    if is_websocket_request(flow):
+        # Avoid Websocket connect requests
+        return
     to_mock_server(flow)
 
 
diff --git a/lyrebird/mock/extra_mock_server/__init__.py b/lyrebird/mock/extra_mock_server/__init__.py
index ce111964..147e6f11 100644
--- a/lyrebird/mock/extra_mock_server/__init__.py
+++ b/lyrebird/mock/extra_mock_server/__init__.py
@@ -12,5 +12,4 @@ def run(self, async_obj, config, *args, **kwargs):
         signal.signal(signal.SIGINT, signal.SIG_IGN)
         log_queue = async_obj['logger_queue']
         msg_queue = async_obj['msg_queue']
-        publish_init_status(msg_queue, 'READY')
         serve(msg_queue, config, log_queue, *args, **kwargs)
diff --git a/lyrebird/mock/extra_mock_server/server.py b/lyrebird/mock/extra_mock_server/server.py
index 980c7f01..5be10b4d 100644
--- a/lyrebird/mock/extra_mock_server/server.py
+++ b/lyrebird/mock/extra_mock_server/server.py
@@ -15,6 +15,7 @@
 
 logger = None
 logger_queue = None
+message_queue = None
 lb_config = {}
 semaphore = None
 
@@ -44,6 +45,15 @@ def make_raw_headers_line(request: web.Request):
     return json.dumps(raw_headers, ensure_ascii=False)
 
 
+def upgrade_request_report(context: LyrebirdProxyContext):
+    if not context.request.headers.get('upgrade'):
+        return
+    publish_metrics_msg('upgrade_request', {
+        'url': context.origin_url,
+        'headers': str(context.request.headers)
+    })
+
+
 def make_request_headers(context: LyrebirdProxyContext, is_proxy):
     headers = {k: v for k, v in context.request.headers.items() if k.lower() not in [
             'cache-control', 'host', 'transfer-encoding']}
@@ -127,6 +137,7 @@ async def req_handler(request: web.Request):
     try:
         global lb_config
         proxy_ctx = LyrebirdProxyContext.parse(request, lb_config)
+        upgrade_request_report(proxy_ctx)
         if is_filtered(proxy_ctx):
             # forward to lyrebird
             async with semaphore:
@@ -211,8 +222,8 @@ def _cancel_tasks(
                 }
             )
 
-def publish_init_status(queue, status):
-    queue.put({
+def publish_init_status(status):
+    message_queue.put({
         'type': 'event',
         "channel": "system",
         "content": {
@@ -224,10 +235,26 @@ def publish_init_status(queue, status):
         }
     })
 
+def publish_metrics_msg(action, info):
+    message_queue.put({
+        'type': 'event',
+        "channel": "lyrebird_metrics",
+        "content": {
+            'lyrebird_metrics': {
+                'sender': 'ExtraMockServer',
+                'action': action,
+                'trace_info': str(info)
+                }
+            }
+    })
 
 def serve(msg_queue, config, log_queue, *args, **kwargs):
     global logger_queue
+    global message_queue
     logger_queue = log_queue
+    message_queue = msg_queue
+
+    publish_init_status('READY')
     loop = asyncio.new_event_loop()
     main_task = loop.create_task(_run_app(config))
 
@@ -235,7 +262,7 @@ def serve(msg_queue, config, log_queue, *args, **kwargs):
         asyncio.set_event_loop(loop)
         loop.run_until_complete(main_task)
     except KeyboardInterrupt:
-        publish_init_status(msg_queue, 'ERROR')
+        publish_init_status('ERROR')
     finally:
         _cancel_tasks({main_task}, loop)
         _cancel_tasks(asyncio.all_tasks(loop), loop)
diff --git a/lyrebird/version.py b/lyrebird/version.py
index 4e53c509..b1599561 100644
--- a/lyrebird/version.py
+++ b/lyrebird/version.py
@@ -1,3 +1,3 @@
-IVERSION = (3, 1, 1)
+IVERSION = (3, 1, 2)
 VERSION = ".".join(str(i) for i in IVERSION)
 LYREBIRD = "Lyrebird " + VERSION