From 19f8f938fff113df26857529931e1d07950fbf5d Mon Sep 17 00:00:00 2001
From: retoor <retoor@molodetz.nl>
Date: Wed, 18 Dec 2024 03:09:37 +0100
Subject: [PATCH] Repaired websocket support.

---
 src/zamenyat/app.py | 45 +++++++++++++++++++++++----------------------
 1 file changed, 23 insertions(+), 22 deletions(-)

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