From cdc664e5568cdc3af05902e72bffb2430712fb9f Mon Sep 17 00:00:00 2001
From: noO0ob <wujiasheng03@meituan.com>
Date: Wed, 18 Dec 2024 20:00:13 +0800
Subject: [PATCH 1/2] web not proxy websocket

---
 lyrebird/mitm/mitm_script.py | 13 +++++++++++--
 1 file changed, 11 insertions(+), 2 deletions(-)

diff --git a/lyrebird/mitm/mitm_script.py b/lyrebird/mitm/mitm_script.py
index 5892a660..fd1d735b 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
 
@@ -14,6 +13,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 to_mock_server(flow: http.HTTPFlow):
     raw_url = urlparse(flow.request.url)
     raw_host = raw_url.hostname
@@ -43,6 +50,8 @@ def request(flow: http.HTTPFlow):
     if 'mitm.it' in flow.request.url:
         # Support mitm.it
         return
+    if is_websocket_request(flow):
+        return
     to_mock_server(flow)
 
 

From 57bef919ff9cedb3fc58809533692d30df29b9ac Mon Sep 17 00:00:00 2001
From: noO0ob <wujiasheng03@meituan.com>
Date: Wed, 18 Dec 2024 20:43:27 +0800
Subject: [PATCH 2/2] add upgrade request report log

---
 lyrebird/mock/extra_mock_server/__init__.py |  1 -
 lyrebird/mock/extra_mock_server/server.py   | 33 +++++++++++++++++++--
 2 files changed, 30 insertions(+), 4 deletions(-)

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 20962c63..64ae73eb 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)
+    })
+
+
 async def make_response_header(proxy_resp_headers: dict, context: LyrebirdProxyContext, data=None):
     response_headers = CIMultiDict()
     for k, v in proxy_resp_headers.items():
@@ -116,6 +126,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:
@@ -198,8 +209,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": {
@@ -211,10 +222,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))
 
@@ -222,7 +249,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)