diff --git a/src/zamenyat/app.py b/src/zamenyat/app.py index 0895057..f9193aa 100644 --- a/src/zamenyat/app.py +++ b/src/zamenyat/app.py @@ -126,22 +126,23 @@ class Application: req_resp, headers = await self.get_headers(reader) if not headers: break + if headers: + if 'Content-Length' in headers: + while len(data) != headers['Content-Length']: + chunk_size = headers['Content-Length'] - len(data) if self.buffer_size > headers['Content-Length'] - len(data) else self.buffer_size + print("Bef read") + chunk = await reader.read(chunk_size) + if not chunk: + data = None + break + print("Aff read") + data += chunk + await writer.write(self.header_dict_to_bytes(req_resp, headers)) + await writer.drain() + if data: + await writer.write(data) else: data = await reader.read() - if 'Content-Length' in headers: - while not len(data) == headers['Content-Length']: - chunk_size = headers['Content-Length'] - len(data) if self.buffer_size > headers['Content-Length'] - len(data) else self.buffer_size - print("Bef read") - chunk = await reader.read(chunk_size) - if not chunk: - data = None - break - print("Aff read") - data += chunk - print(self.header_dict_to_bytes(req_resp,headers).decode()) - await writer.write(self.header_dict_to_bytes(req_resp, headers)) - await writer.drain() - if data: await writer.write(data) #if not headers.get('Connection') == 'keep-alive': # and not headers.get('Upgrade-Insecure-Requests'): # break @@ -169,19 +170,19 @@ class Application: print(f"Connected to upstream #{self.total_connection_count} server {self.upstream_host}:{self.upstream_port} #{connection_nr} Time: {get_timestamp()}") request_headers = await self.stream(reader, upstream_writer,is_websocket) + + keep_alive = False + if request_headers: + if request_headers.get('Connection') == 'keep-alive': # and not headers.get('Upgrade-Insecure-Requests'): + keep_alive = True + if request_headers.get("Upgrade") == 'websocket': + is_websocket = True await self.stream(upstream_reader, writer, is_websocket) time_end = time.time() time_duration = time_end - time_start print(f"Disconnected upstream #{self.total_connection_count} server {self.upstream_host}:{self.upstream_port} #{connection_nr} Duration: {time_duration:.5f}s") - keep_alive = False - - if request_headers.get('Connection') == 'keep-alive': # and not headers.get('Upgrade-Insecure-Requests'): - keep_alive = True - - if request_headers.get("Upgrade") == 'websocket': - is_websocket = True - + if not any([keep_alive, is_websocket]): break