v2.6.1 bugfix: lb proxy protocol parser error (#668)

* update lb proxy protocol

* Add e2e tests for lyrebird proxy protocol parser

* e2etest fix: wait extra mock start

e2e_tests/Readme.md

@@ -1,8 +1,13 @@
 # lyrebird e2e test
+
 ## 测试 lyrebird（POST）request body
 
 ```bash
-nohup python3 serve.py > /dev/null 2>&1 & 
-nohup lyrebird -b > /dev/null 2>&1 & 
+pip install -e .[dev]
+# If you are using zsh you need to escape square brackets or use quotes:
+# pip install -e .\[extra\]
+# or
+# pip install -e ".[extra]"
+cd e2e_tests
 python -m pytest .
 ```

e2e_tests/conftest.py

@@ -68,6 +68,7 @@ def __init__(self):
         self.extra_mock_port = None
         self.api_status = None
         self.uri_mock = None
+        self.uri_extra_mock = None
         self._init_port()
 
     def _init_port(self):
@@ -76,6 +77,7 @@ def _init_port(self):
         self.extra_mock_port = self._find_free_port()
         self.api_status = f'http://127.0.0.1:{self.port}/api/status'
         self.uri_mock = f'http://127.0.0.1:{self.port}/mock/'
+        self.uri_extra_mock = f'http://127.0.0.1:{self.extra_mock_port}/'
 
     def _find_free_port(self):
         with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as s:
@@ -89,6 +91,7 @@ def start(self, checker_path=None):
             cmdline = cmdline + f' --script {checker_path}'
         self.lyrebird_process = subprocess.Popen(cmdline, shell=True, start_new_session=True)
         _wait(requests.get, args=[self.api_status])
+        _wait(requests.get, args=[self.uri_extra_mock])
 
         # Wait for checker to load
         if checker_path:

e2e_tests/test_core.py

@@ -1,5 +1,10 @@
 import pytest
-import os, hashlib, json, gzip, requests
+import os
+import hashlib
+import json
+import gzip
+import requests
+from urllib.parse import quote
 
 
 curPath = os.path.abspath(os.path.dirname(__file__))
@@ -11,34 +16,59 @@ def test_img_data(lyrebird, mock_server):
     r = requests.post(url=lyrebird.uri_mock + mock_server.api_post, data=data)
     assert r.text == hashlib.md5(mock_server.api_post.encode() + data).hexdigest()
 
+
 def test_img_file(lyrebird, mock_server):
     files = {'file': ('1.png', open(f'{curPath}/assets/1.png', 'rb'), 'image/jpg')}
     r = requests.post(url=lyrebird.uri_mock + mock_server.api_post, files=files)
     with open(f'{curPath}/assets/1.png', 'rb') as f:
         data = f.read()
     assert r.text == hashlib.md5(mock_server.api_post.encode() + data).hexdigest()
 
+
 def test_json(lyrebird, mock_server):
     data = json.dumps({"name": {"12": 123}}, ensure_ascii=False)
     headers = {"Content-Type": "application/json"}
     r = requests.post(url=lyrebird.uri_mock + mock_server.api_post, data=data, headers=headers)
     assert r.text == hashlib.md5(mock_server.api_post.encode() + data.encode()).hexdigest()
 
+
 def test_text(lyrebird, mock_server):
     data = "asdasdasd"
     headers = {"Content-Type": "text/plain"}
     r = requests.post(url=lyrebird.uri_mock + mock_server.api_post, data=data, headers=headers)
     assert r.text == hashlib.md5(mock_server.api_post.encode() + data.encode()).hexdigest()
 
+
 def test_form(lyrebird, mock_server):
     data = 'z=9&a=1&a=2&b=1'
     headers = {"Content-Type": "application/x-www-form-urlencoded"}
     r = requests.post(url=lyrebird.uri_mock + mock_server.api_post, data=data, headers=headers)
     assert r.text == hashlib.md5(mock_server.api_post.encode() + data.encode()).hexdigest()
 
+
 def test_json_gzip(lyrebird, mock_server):
     data = {"a": 1}
     ziped_data = gzip.compress(json.dumps(data, ensure_ascii=False).encode())
     headers = {"Content-Type": "application/json", "Content-Encoding": "gzip"}
     r = requests.post(url=lyrebird.uri_mock + mock_server.api_post, data=ziped_data, headers=headers)
     assert r.text == hashlib.md5(mock_server.api_post.encode() + ziped_data).hexdigest()
+
+
+def test_lb_proxy_protocol_target_in_path(lyrebird, mock_server):
+    r = requests.get(url=lyrebird.uri_extra_mock + mock_server.api_status)
+    assert r.text == 'OK'
+
+
+def test_lb_proxy_protocol_target_in_header(lyrebird, mock_server):
+    headers = {
+        'MKScheme': 'http',
+        'MKOriginHost': '127.0.0.1',
+        'MKOriginPort': '5000'
+    }
+    r = requests.get(url=lyrebird.uri_extra_mock + 'status', headers=headers)
+    assert r.text == 'OK'
+
+
+def test_lb_proxy_protocol_target_in_query(lyrebird, mock_server):
+    r = requests.get(url=lyrebird.uri_extra_mock + f'?proxy={quote(mock_server.api_status)}')
+    assert r.text == 'OK'

lyrebird/mock/extra_mock_server.py

@@ -29,7 +29,8 @@ class LyrebirdProxyContext:
     def __init__(self):
         self.request: web.Request = None
         self.netloc = None
-        self.full_url = None
+        self.origin_url = None
+        self.forward_url = None
 
     @classmethod
     def parse(cls, request: web.Request,
@@ -38,6 +39,8 @@ def parse(cls, request: web.Request,
               lb_proxy_port_header_name='LBOriginPort',
               lb_proxy_query_keys: Optional[List] = None):
 
+        global lb_config
+
         # Lyrebird proxy protocol #1
         # Origin target info in path
         # e.g.
@@ -47,7 +50,12 @@ def parse(cls, request: web.Request,
             url = urlparse.urlparse(origin_full_url)
 
             ctx = cls()
-            ctx.full_url = origin_full_url
+            # origin url for proxy
+            ctx.origin_url = origin_full_url
+            # forward url to lyrebird main port 'default 9090'
+            port = lb_config.get('mock.port')
+            ctx.forward_url = f'http://127.0.0.1:{port}/mock/{origin_full_url}'
+
             ctx.netloc = url.netloc
             ctx.request = request
             return ctx
@@ -72,7 +80,11 @@ def parse(cls, request: web.Request,
                 netloc = target_host
 
             ctx = cls()
-            ctx.full_url = origin_full_url
+            # origin url for proxy
+            ctx.origin_url = origin_full_url
+            # forward url to lyrebird main port 'default 9090'
+            port = lb_config.get('mock.port')
+            ctx.forward_url = f'http://127.0.0.1:{port}/mock/{origin_full_url}'
             ctx.netloc = netloc
             ctx.request = request
             return ctx
@@ -81,7 +93,7 @@ def parse(cls, request: web.Request,
         # Origin target info in query
         # default key is "proxy" or "mp_webview_domain_info"
         # e.g.
-        # http://{lyrebird_host}/{origing_path}?{proxy= or custom_query_key=}
+        # http://{lyrebird_host}/{origing_path}?{proxy=encoded-url or custom_query_key=encoded-url}
 
         # set default key=proxy
         if lb_proxy_query_keys is None:
@@ -92,15 +104,20 @@ def parse(cls, request: web.Request,
         def contain_custom_query_key(request: web.Request, lb_proxy_query_keys):
             for query_key in lb_proxy_query_keys:
                 if request.query.get(query_key):
-                    return True
-            return False
+                    return True, query_key
+            return False, query_key
 
-        if contain_custom_query_key(request, lb_proxy_query_keys):
-            origin_full_url = request.path_qs[1:]
-            url = urlparse.urlparse(origin_full_url)
+        has_custom_key, query_key = contain_custom_query_key(request, lb_proxy_query_keys)
+        if has_custom_key:
+            origin_url = request.query.get(query_key)
+            origin_url = urlparse.unquote(origin_url)
+            url = urlparse.urlparse(origin_url)
 
             ctx = cls()
-            ctx.full_url = origin_full_url
+            ctx.origin_url = origin_url
+            # forward url to lyrebird main port 'default 9090'
+            port = lb_config.get('mock.port')
+            ctx.forward_url = f'http://127.0.0.1:{port}/mock/?{query_key}={urlparse.quote(origin_url)}'
             ctx.netloc = url.netloc
             ctx.request = request
             return ctx
@@ -117,7 +134,7 @@ def is_filtered(context: LyrebirdProxyContext):
     global lb_config
     filters = lb_config.get('proxy.filters')
     for _filter in filters:
-        if re.search(_filter, context.full_url):
+        if re.search(_filter, context.origin_url):
             return True
     return False
 
@@ -173,16 +190,13 @@ async def send_request(context: LyrebirdProxyContext, target_url):
 
 
 async def proxy(context: LyrebirdProxyContext):
-    logger.info(f'proxy {context.full_url}')
-    return await send_request(context, context.full_url)
+    logger.info(f'proxy {context.origin_url}')
+    return await send_request(context, context.origin_url)
 
 
 async def forward(context: LyrebirdProxyContext):
-    global lb_config
-    port = lb_config.get('mock.port')
-    url = f'http://127.0.0.1:{port}/mock/{context.full_url}'
-    logger.info(f'forward {url}')
-    return await send_request(context, url)
+    logger.info(f'forward {context.forward_url}')
+    return await send_request(context, context.forward_url)
 
 
 async def req_handler(request: web.Request):

lyrebird/version.py

@@ -1,3 +1,3 @@
-IVERSION = (2, 6, 0)
+IVERSION = (2, 6, 1)
 VERSION = ".".join(str(i) for i in IVERSION)
 LYREBIRD = "Lyrebird " + VERSION