tweaked writing bytes data to socket stream

zNitche · May 15, 2024 · 14ef374 · 14ef374
1 parent aea31d5
commit 14ef374
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 5 deletions.
diff --git a/lightberry/core/sockets_servers/http/app_server.py b/lightberry/core/sockets_servers/http/app_server.py
@@ -1,6 +1,6 @@
 from lightberry.config import ServerConfig as Config
 from lightberry.core.sockets_servers.http import HttpSocketServer
-from lightberry.utils import files_utils
+from lightberry.utils import files_utils, requests_utils
 import ssl
 import time
 import asyncio
@@ -32,14 +32,14 @@ async def __requests_handler(self, client_r: StreamReader, client_w: StreamWrite
                 response = await asyncio.wait_for(self.__app.requests_handler(request), self.__app.config.TIMEOUT)
 
                 if response.is_payload_streamed:
-                    client_w.write(bytes(f"{response.get_headers()}\r\n\r\n", "utf-8"))
+                    requests_utils.write_to_stream(client_w, f"{response.get_headers()}\r\n\r\n")
                     await client_w.drain()
 
                     for chunk in response.get_body():
-                        client_w.write(bytes(chunk, "utf-8"))
+                        requests_utils.write_to_stream(client_w, chunk)
                         await client_w.drain()
                 else:
-                    client_w.write(bytes(response.get_response_string(), "utf-8"))
+                    requests_utils.write_to_stream(client_w, response.get_response_string())
                     await client_w.drain()
 
         except Exception as e:

diff --git a/lightberry/core/sockets_servers/http/ssl_proxy_server.py b/lightberry/core/sockets_servers/http/ssl_proxy_server.py
@@ -1,6 +1,7 @@
 from lightberry.config import ServerConfig as Config
 from lightberry.core.communication.response import Response
 from lightberry.core.sockets_servers.http import HttpSocketServer
+from lightberry.utils import requests_utils
 import asyncio
 
 from lightberry.typing import TYPE_CHECKING
@@ -27,7 +28,7 @@ async def __requests_handler(self, client_r: StreamReader, client_w: StreamWrite
                 response = Response(301)
                 response.add_header("LOCATION", f"https://{self.hostname}{request.url}")
 
-                client_w.write(bytes(response.get_response_string(), "utf-8"))
+                requests_utils.write_to_stream(client_w, response.get_response_string())
                 await client_w.drain()
 
         except Exception as e:

diff --git a/lightberry/utils/requests_utils.py b/lightberry/utils/requests_utils.py
@@ -39,3 +39,7 @@ async def load_request_header_from_stream(stream: StreamReader) -> str:
         request_header_string += request_line
 
     return request_header_string
+
+
+def write_to_stream(stream, data, encoding="utf-8"):
+    stream.write(bytes(data, encoding))