From 6622f7a6529e59cef784824c41a79fdba1bd9c34 Mon Sep 17 00:00:00 2001 From: noO0oOo0ob <38344038+noO0oOo0ob@users.noreply.github.com> Date: Thu, 21 Sep 2023 14:22:32 +0800 Subject: [PATCH] change diff mode proxy to asyn (#785) * change diff mode proxy to asyn * update varsion * move proxy handler to core * change proxyhandler init location * remove needless deepcopy * change try_except to in_request_context flag * remove core diff mode code --------- Co-authored-by: wujiasheng03 --- lyrebird/mock/blueprints/apis/flow.py | 3 +-- lyrebird/mock/blueprints/core.py | 7 ------- lyrebird/mock/handlers/handler_context.py | 23 +++++++++++++++++++++-- lyrebird/mock/handlers/proxy_handler.py | 13 +++++++++---- lyrebird/version.py | 2 +- 5 files changed, 32 insertions(+), 16 deletions(-) diff --git a/lyrebird/mock/blueprints/apis/flow.py b/lyrebird/mock/blueprints/apis/flow.py index d4e6b9d39..740c36694 100644 --- a/lyrebird/mock/blueprints/apis/flow.py +++ b/lyrebird/mock/blueprints/apis/flow.py @@ -31,11 +31,10 @@ def get(self, id): # Import decoder for decoding the requested content display_item = {} if is_origin: - display_item.update(item) + display_item = deepcopy(item) else: application.encoders_decoders.decoder_handler(item, output=display_item) if not no_decode: - display_item['request'] = deepcopy(display_item['request']) for key in ('url', 'path', 'query'): url_decode(display_item['request'], key) return application.make_ok_response(data=display_item) diff --git a/lyrebird/mock/blueprints/core.py b/lyrebird/mock/blueprints/core.py index cf69e6d05..1aa9de540 100644 --- a/lyrebird/mock/blueprints/core.py +++ b/lyrebird/mock/blueprints/core.py @@ -65,13 +65,6 @@ def index(path=''): req_context.update_client_resp_time() resp = req_context.response - - if context.application.is_diff_mode == context.MockMode.MULTIPLE and req_context.response_source == 'mock': - proxy_handler.handle(req_context) - if req_context.is_proxiable: - req_context.update_response_headers_code2flow(output_key='proxy_response') - req_context.update_response_data2flow(output_key='proxy_response') - DuplicateHeaderKeyHandler.set_origin_header(resp.headers, req_context.flow['response']['headers']) context.emit('action', 'add flow log') diff --git a/lyrebird/mock/handlers/handler_context.py b/lyrebird/mock/handlers/handler_context.py index d127834c7..d007d3632 100644 --- a/lyrebird/mock/handlers/handler_context.py +++ b/lyrebird/mock/handlers/handler_context.py @@ -11,9 +11,11 @@ from urllib.parse import urlparse, unquote from .http_data_helper import DataHelper from .http_header_helper import HeadersHelper +from .proxy_handler import ProxyHandler logger = get_logger() +proxy_handler = ProxyHandler() class HandlerContext: @@ -28,6 +30,7 @@ def __init__(self, request): self.id = str(uuid.uuid4()) self.request = request self.response = None + self.cookies = request.cookies self.client_req_time = None self.client_resp_time = None self.server_req_time = None @@ -177,14 +180,17 @@ def set_response_source_mock(self): def set_response_source_proxy(self): self.response_source = 'proxy' - def get_request_body(self): + def get_request_body(self, in_request_handler=True): if self.is_request_edited: # TODO Repeated calls, remove it self.flow['request']['headers'] = HeadersHelper.flow2origin(self.flow['request'], chain=self.request_chain) _data = DataHelper.flow2origin(self.flow['request'], chain=self.request_chain) else: - _data = self.request.data or self.request.form or None + if in_request_handler: + _data = self.request.data or self.request.form or None + else: + _data = DataHelper.flow2origin(self.flow['request']) return _data def get_request_headers(self): @@ -200,6 +206,12 @@ def get_request_headers(self): continue headers[name] = value return headers + + def get_request_cookies(self, in_request_handler=True): + if in_request_handler: + return self.request.cookies + else: + return self.cookies def get_response_generator(self): if self.is_response_edited: @@ -299,6 +311,13 @@ def update_client_resp_time(self): parsed_url = self._get_parse_url_dict(url) self.flow['request'].update(parsed_url) + # Diff Mode proxy request + if context.application.is_diff_mode == context.MockMode.MULTIPLE and self.response_source == 'mock': + proxy_handler.handle(self, in_request_handler=False) + if self.is_proxiable: + self.update_response_headers_code2flow(output_key='proxy_response') + self.update_response_data2flow(output_key='proxy_response') + # Import decoder for decoding the requested content decode_flow = {} application.encoders_decoders.decoder_handler(self.flow, output=decode_flow) diff --git a/lyrebird/mock/handlers/proxy_handler.py b/lyrebird/mock/handlers/proxy_handler.py index ef12ced74..153e98fe2 100644 --- a/lyrebird/mock/handlers/proxy_handler.py +++ b/lyrebird/mock/handlers/proxy_handler.py @@ -22,7 +22,7 @@ class ProxyHandler: """ - def handle(self, handler_context): + def handle(self, handler_context, in_request_handler=True): request = handler_context.flow['request'] @@ -42,7 +42,7 @@ def handle(self, handler_context): DuplicateRequest().handle(handler_context) return - data = handler_context.get_request_body() + data = handler_context.get_request_body(in_request_handler) method = request['method'] headers = handler_context.get_request_headers() @@ -53,7 +53,7 @@ def handle(self, handler_context): origin_url, headers=headers, data=data, - cookies=handler_context.request.cookies, + cookies=handler_context.get_request_cookies(in_request_handler), stream=True, verify=False, allow_redirects=False) @@ -88,9 +88,14 @@ def handle(self, handler_context): if r.status_code == 204: handler_context.response = Response(None, status=r.status_code, headers=resp_headers) return + + if in_request_handler: + gen = stream_with_context(r.iter_content(chunk_size=handler_context.response_chunk_size)) + else: + gen = r.iter_content(chunk_size=handler_context.response_chunk_size) # After huangyuanzhen test, we use 2048byte buffer :D handler_context.response = Response( - stream_with_context(r.iter_content(chunk_size=handler_context.response_chunk_size)), + gen, status=r.status_code, headers=resp_headers) diff --git a/lyrebird/version.py b/lyrebird/version.py index b1c9286e4..ae505a180 100644 --- a/lyrebird/version.py +++ b/lyrebird/version.py @@ -1,3 +1,3 @@ -IVERSION = (2, 19, 1) +IVERSION = (2, 19, 2) VERSION = ".".join(str(i) for i in IVERSION) LYREBIRD = "Lyrebird " + VERSION