Skip to content

Commit

Permalink
fix: large request/response #24
Browse files Browse the repository at this point in the history
  • Loading branch information
kamarton committed Jun 11, 2024
1 parent 6906aa2 commit 9e964db
Show file tree
Hide file tree
Showing 4 changed files with 36 additions and 14 deletions.
5 changes: 4 additions & 1 deletion environment.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ def __init__(self):
self.PROXY_DESTINATION: str = f'http://127.0.0.1:{self.TESTER_LISTEN_PORT}'
self.TESTER_DESTINATION: str = f'http://127.0.0.1:{self.PROXY_LISTEN_PORT}'
self.STORAGE_PATH: str = os.path.join('.', 'storage')
self.HTTP_MAX_REQUEST_SIZE_MB: int = 16

@staticmethod
def from_os_env() -> 'Environment':
Expand All @@ -28,7 +29,8 @@ def from_os_env() -> 'Environment':
env.TESTER_MODE_ENABLED = env.TESTER_LISTEN_HOST is not None and env.TESTER_LISTEN_PORT > 0
env.PROXY_DESTINATION = os.getenv('PROXY_DESTINATION',
f'http://{env.TESTER_LISTEN_HOST}:{env.TESTER_LISTEN_PORT}')
env.STORAGE_PATH = str = os.getenv('STORAGE_PATH', os.path.join('.', 'storage'))
env.STORAGE_PATH = os.getenv('STORAGE_PATH', os.path.join('.', 'storage'))
env.HTTP_MAX_REQUEST_SIZE_MB = int(os.getenv('HTTP_MAX_REQUEST_SIZE_MB', 16))
return env

@staticmethod
Expand All @@ -49,6 +51,7 @@ def print(self):
print('PROXY_DESTINATION:', self.PROXY_DESTINATION)
print('PROXY_OVERRIDE_HOST_HEADER:', self.PROXY_OVERRIDE_HOST_HEADER)
print('STORAGE_PATH:', self.STORAGE_PATH)
print('HTTP_MAX_REQUEST_SIZE_MB:', self.HTTP_MAX_REQUEST_SIZE_MB)
if self.TESTER_MODE_ENABLED:
print('TESTER_LISTEN_HOST:', self.TESTER_LISTEN_HOST)
print('TESTER_LISTEN_PORT:', self.TESTER_LISTEN_PORT)
Expand Down
9 changes: 5 additions & 4 deletions proxy_log.py
Original file line number Diff line number Diff line change
Expand Up @@ -196,7 +196,8 @@ def __init__(self, phase: ProxyLogPhase, id: int | None = None, start_time: floa
self._response_headers: HttpHeaders | None = None
self._response_body: ContentItem | None = None
self._response_body_mime_type: str | None = None
self.exception: Exception | None = None
self.exception_message: str | None = None
self.exception_type: str | None = None
self.exception_traceback: str | None = None

def mutate(self, phase: ProxyLogPhase) -> 'RequestEntry':
Expand Down Expand Up @@ -275,10 +276,10 @@ def to_map(self) -> dict:
encoded_content.update({'mimeType': self._response_body_mime_type})
response['body'] = self._clean_dict(encoded_content)
exception = None
if self.exception is not None:
if self.exception_type is not None:
exception = {
'message': str(self.exception),
'type': str(type(self.exception)),
'message': self.exception_message,
'type': self.exception_type,
'traceback': self.exception_traceback,
}
result = {
Expand Down
6 changes: 4 additions & 2 deletions proxy_server.py
Original file line number Diff line number Diff line change
Expand Up @@ -49,9 +49,11 @@ async def handle_proxy(request: aiohttp.web.Request):
) as response:
return await handel_proxy_response(response, log, proxy_log)
except Exception as e:
print(f"Exception: {e}")
if log is not None:
log = log.mutate(ProxyLogPhase.END)
log.exception = e
log.exception_type = str(type(e))
log.exception_message = str(e)
log.exception_traceback = traceback.format_exc()
proxy_log.put(log)
finally:
Expand Down Expand Up @@ -88,7 +90,7 @@ async def storage_garbage_task(app):


def run_proxy_server(environment: Environment, proxy_log: ProxyLog):
app = web.Application()
app = web.Application(client_max_size=environment.HTTP_MAX_REQUEST_SIZE_MB * 1024 ** 2)
app.setdefault('environment', environment)
app.setdefault('proxy_log', proxy_log)
app.router.add_route('*', '/{any:.*}', handle_proxy)
Expand Down
30 changes: 23 additions & 7 deletions tester_server.py
Original file line number Diff line number Diff line change
Expand Up @@ -55,20 +55,34 @@ async def handle(request):
})


def run_rest_api_server(host: str, port: int):
app = web.Application()
def run_rest_api_server(environment: Environment):
app = web.Application(client_max_size=environment.HTTP_MAX_REQUEST_SIZE_MB * 1024 ** 2)
app.router.add_route('*', '/{any:.*}', handle)
try:
web.run_app(app, host=host, port=port)
web.run_app(app, host=environment.TESTER_LISTEN_HOST, port=environment.TESTER_LISTEN_PORT)
except KeyboardInterrupt:
app.shutdown()


def run_rest_api_client(tester_destination: str):
async def fetch(session, method, url, body):
print(f'Tester fetch {method} {url} {body}...')
async with session.request(method, url, json=body) as response:
return await response.text()
body_str = ''
if body is str:
body_str = body
elif body is dict:
body_str = json.dumps(body)
if len(body_str) > 128:
body_str = body_str[:128] + '...'
print(f'Tester fetch {method} {url} {body_str}...')
if body is dict:
async with session.request(method, url, json=body) as response:
return await response.text()
else:
headers = {
'content-type': 'text/plain',
}
async with session.request(method, url, data=body, headers=headers) as response:
return await response.text()

async def main():
async with aiohttp.ClientSession() as session:
Expand All @@ -86,6 +100,8 @@ async def main():
if random_boolean():
query_parameters_string = '?' + '&'.join(
[f'key_{i}={random.randint(0, 100)}' for i in range(random.randint(1, 8))])
if random.randint(0, 15) < 1:
body = ''.join(random.choices(string.ascii_lowercase, k=2 * 1024 * 1024))
response = await fetch(session, method, f'{tester_destination}/{path}{query_parameters_string}',
body=body)
if len(response) > 150:
Expand All @@ -101,7 +117,7 @@ async def main():


def run_rest_api_tester(environment: Environment):
Process(target=run_rest_api_server, args=(environment.TESTER_LISTEN_HOST, environment.TESTER_LISTEN_PORT,)).start()
Process(target=run_rest_api_server, args=(environment,)).start()
time.sleep(2)
for _ in range(1):
Process(target=run_rest_api_client, args=(environment.TESTER_DESTINATION,)).start()

0 comments on commit 9e964db

Please sign in to comment.