diff --git a/app.py b/app.py index a9d27c1..4c9ebf5 100644 --- a/app.py +++ b/app.py @@ -7,7 +7,10 @@ import time from datetime import datetime, timezone from socketserver import ThreadingMixIn import os - +import asyncio +import websockets +import hashlib +import base64 # Define the target server to proxy requests to class ProxyHandler(http.server.BaseHTTPRequestHandler): def __init__(self, configuration, docker_client): @@ -73,6 +76,15 @@ class ProxyHandler(http.server.BaseHTTPRequestHandler): return activity[proxy_host_configuration['domain']] = datetime.now(timezone.utc) + + # Check if this is a WebSocket request + if self.headers.get("Upgrade", "").lower() == "websocket": + print("Request is WS connecting to {0}".format(container['container_name'])) + print("Request is WS connecting to {0}:{1}".format(proxy_host_configuration['proxy_host'],proxy_host_configuration['proxy_port'])) + activity[proxy_host_configuration['domain']] = True + self.upgrade_to_websocket(proxy_host_configuration['proxy_host'], proxy_host_configuration['proxy_port']) + return + # Open a connection to the target server conn = http.client.HTTPConnection(proxy_host_configuration['proxy_host'], proxy_host_configuration['proxy_port']) conn.request(method, self.path, headers=self.headers) @@ -122,6 +134,56 @@ class ProxyHandler(http.server.BaseHTTPRequestHandler): self.wfile.flush() + async def websocket_proxy(self, target_host, target_port): + server_ws = None + try: + client_connection = self.connection + # Establish server connection to backend + server_ws = await websockets.connect(f"ws://{target_host}:{target_port}") + + # Bridge function to handle message forwarding + async def bridge_websockets(): + try: + while True: + # Wait for a message from the client + client_message = await client_connection.recv() + # Send it to the server + await server_ws.send(client_message) + # Wait for a message from the server + server_message = await server_ws.recv() + # Send it to the client + await client_connection.send(server_message) + except websockets.exceptions.ConnectionClosed as e: + print(f"WebSocket connection closed: {e}") + except Exception as e: + print(f"Error during WebSocket communication: {e}") + + # Run the bridge coroutine + await bridge_websockets() + + except Exception as e: + print(f"WebSocket proxy encountered an error: {e}") + finally: + if server_ws: + await server_ws.close() + + def upgrade_to_websocket(self, target_host, target_port): + """ + Handles WebSocket upgrade requests and spawns an asyncio WebSocket proxy. + """ + key = self.headers['Sec-WebSocket-Key'] + accept_val = base64.b64encode(hashlib.sha1((key + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11").encode('utf-8')).digest()).decode('utf-8') + + + self.send_response(101) # Switching Protocols + self.send_header("Upgrade", "websocket") + self.send_header("Connection", "Upgrade") + self.send_header("Sec-WebSocket-Accept", accept_val) + self.end_headers() + + loop = asyncio.new_event_loop() + threading.Thread(target=loop.run_until_complete, args=(self.websocket_proxy(target_host, target_port),)).start() + class ThreadedHTTPServer(ThreadingMixIn, http.server.HTTPServer): """Handle requests in a separate thread.""" @@ -148,6 +210,9 @@ class BackgroundTasks(threading.Thread): if (apps['domain'] in activity): dt = activity[apps['domain']] + if (dt == True): + continue + diff_seconds = (datetime.now(timezone.utc) - dt).total_seconds() if(diff_seconds > apps['proxy_timeout_seconds']): print("stopping container: {0} ({1}) after {2}s".format(container['container_name'], container_object.id, diff_seconds)) @@ -157,6 +222,30 @@ class BackgroundTasks(threading.Thread): time.sleep(sleep_time) +async def websocket_proxy(client_ws, target_host, target_port): + """ + Forwards WebSocket messages between the client and the target container. + """ + try: + async with websockets.connect(f"ws://{target_host}:{target_port}") as server_ws: + # Create tasks to read from both directions + async def forward_client_to_server(): + async for message in client_ws: + await server_ws.send(message) + + async def forward_server_to_client(): + async for message in server_ws: + await client_ws.send(message) + + await asyncio.gather(forward_client_to_server(), forward_server_to_client()) + + except Exception as e: + print(f"WebSocket error: {e}") + finally: + await client_ws.close() + +# MAIN # + if __name__ == '__main__': activity = {} @@ -173,4 +262,5 @@ if __name__ == '__main__': proxy_handler = ProxyHandler(configuration, docker_client) httpd = ThreadedHTTPServer(server_address, proxy_handler) print('Reverse proxy server running on port {0}...'.format(configuration['proxy_port'])) - httpd.serve_forever() \ No newline at end of file + httpd.serve_forever() + diff --git a/compose.yaml b/compose.yaml new file mode 100644 index 0000000..6ad23ed --- /dev/null +++ b/compose.yaml @@ -0,0 +1,12 @@ +services: + py_proxy: + build: + dockerfile: Dockerfile + context: . + restart: unless-stopped + volumes: + - ./config.yaml:/app/config.yaml + - /var/run/docker.sock:/var/run/docker.sock + ports: + - 80:80 +networks: {} \ No newline at end of file diff --git a/config copy.yml b/config copy.yml new file mode 100644 index 0000000..4b13ead --- /dev/null +++ b/config copy.yml @@ -0,0 +1,10 @@ +proxy_port: 80 +proxy_hosts: + - domain: wp.local + containers: + - container_name: wp-dev-db-1 + - container_name: wp-dev-wordpress-1 + proxy_host: localhost + proxy_port: 8888 + proxy_timeout_seconds: 10 + proxy_load_seconds: 5 \ No newline at end of file diff --git a/config.yml b/config.yml index 4b13ead..71d9ee9 100644 --- a/config.yml +++ b/config.yml @@ -1,10 +1,11 @@ proxy_port: 80 proxy_hosts: - - domain: wp.local + - domain: ws.local containers: - - container_name: wp-dev-db-1 - - container_name: wp-dev-wordpress-1 + - container_name: web-socket-test + - container_name: web-socket-test + proxy_host: localhost - proxy_port: 8888 + proxy_port: 8010 proxy_timeout_seconds: 10 proxy_load_seconds: 5 \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index d0174d3..d89b982 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1 +1,3 @@ -docker-py \ No newline at end of file +docker +PyYAML +websockets \ No newline at end of file