Skip to content

Commit

Permalink
change diff mode proxy to asyn (#785)
Browse files Browse the repository at this point in the history
* 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 <noO0oOo0ob>
  • Loading branch information
noO0oOo0ob authored Sep 21, 2023
1 parent 8e8f9d4 commit 6622f7a
Show file tree
Hide file tree
Showing 5 changed files with 32 additions and 16 deletions.
3 changes: 1 addition & 2 deletions lyrebird/mock/blueprints/apis/flow.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
7 changes: 0 additions & 7 deletions lyrebird/mock/blueprints/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -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')
Expand Down
23 changes: 21 additions & 2 deletions lyrebird/mock/handlers/handler_context.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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
Expand Down Expand Up @@ -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):
Expand All @@ -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:
Expand Down Expand Up @@ -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)
Expand Down
13 changes: 9 additions & 4 deletions lyrebird/mock/handlers/proxy_handler.py
Original file line number Diff line number Diff line change
Expand Up @@ -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']

Expand All @@ -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()
Expand All @@ -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)
Expand Down Expand Up @@ -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)
2 changes: 1 addition & 1 deletion lyrebird/version.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,3 @@
IVERSION = (2, 19, 1)
IVERSION = (2, 19, 2)
VERSION = ".".join(str(i) for i in IVERSION)
LYREBIRD = "Lyrebird " + VERSION

0 comments on commit 6622f7a

Please sign in to comment.