feat: implement WebSocket proxying support (Milestone 1.5)

- Add Connection.hs with ConnectionState type and upgrade detection
- Add Tunnel.hs for bidirectional WebSocket tunneling via responseRaw
- Update Proxy.hs to detect WebSocket upgrades and route to tunnel mode
- Fix request body forwarding (was previously empty)
- Add network, streaming-commons, splice dependencies
- Add WebSocket/SSE test servers in examples/websockets/

WebSocket + regular HTTP now work simultaneously without conflicts.
SSE streaming support pending (Milestone 1.6).

PROJECTS/Aenebris/aenebris.cabal

@@ -23,6 +23,8 @@ library
                      , Aenebris.LoadBalancer
                      , Aenebris.HealthCheck
                      , Aenebris.TLS
+                     , Aenebris.Connection
+                     , Aenebris.Tunnel
                      , Aenebris.Middleware.Security
                      , Aenebris.Middleware.Redirect
   default-language:    Haskell2010
@@ -49,6 +51,9 @@ library
                      , case-insensitive >= 1.2
                      , directory >= 1.3
                      , data-default-class >= 0.1
+                     , network >= 3.1
+                     , streaming-commons >= 0.2
+                     , splice >= 0.6
   ghc-options:         -Wall
                        -Wcompat
                        -Widentities

PROJECTS/Aenebris/examples/test_backend.py

@@ -41,11 +41,17 @@ class TestHandler(BaseHTTPRequestHandler):
         self.send_header('Content-Type', 'application/json')
         self.end_headers()
 
+        try:
+            body_json = json.loads(body.decode())
+        except (json.JSONDecodeError, UnicodeDecodeError):
+            body_json = body.decode('utf-8', errors='replace')
+
         response = {
             'message': 'Received POST',
             'path': self.path,
             'method': 'POST',
-            'body_length': content_length
+            'body_length': content_length,
+            'body_received': body_json
         }
         self.wfile.write(json.dumps(response, indent = 2).encode())
 

PROJECTS/Aenebris/examples/websockets/config_ws_sse.yaml

@@ -0,0 +1,64 @@
+version: 1
+
+listen:
+  - port: 8081
+
+upstreams:
+  - name: http-backend
+    servers:
+      - host: "127.0.0.1:8000"
+        weight: 1
+
+  - name: ws-backend
+    servers:
+      - host: "127.0.0.1:8002"
+        weight: 1
+
+  - name: sse-backend
+    servers:
+      - host: "127.0.0.1:8003"
+        weight: 1
+
+routes:
+  - host: "localhost"
+    paths:
+      - path: /ws
+        upstream: ws-backend
+      - path: /events
+        upstream: sse-backend
+      - path: /stream
+        upstream: sse-backend
+      - path: /
+        upstream: http-backend
+
+  - host: "localhost:8081"
+    paths:
+      - path: /ws
+        upstream: ws-backend
+      - path: /events
+        upstream: sse-backend
+      - path: /stream
+        upstream: sse-backend
+      - path: /
+        upstream: http-backend
+
+  - host: "127.0.0.1:8081"
+    paths:
+      - path: /ws
+        upstream: ws-backend
+      - path: /events
+        upstream: sse-backend
+      - path: /stream
+        upstream: sse-backend
+      - path: /
+        upstream: http-backend
+
+  - host: "ws.localhost"
+    paths:
+      - path: /
+        upstream: ws-backend
+
+  - host: "sse.localhost"
+    paths:
+      - path: /
+        upstream: sse-backend

PROJECTS/Aenebris/examples/websockets/sse_server.py

@@ -0,0 +1,117 @@
+#!/usr/bin/env python3
+"""
+Server-Sent Events (SSE) Streaming Server for Aenebris testing
+"""
+
+import time
+import json
+from http.server import HTTPServer, BaseHTTPRequestHandler
+from datetime import datetime
+
+
+class SSEHandler(BaseHTTPRequestHandler):
+    def do_GET(self):
+        if self.path == "/health":
+            self.send_response(200)
+            self.send_header("Content-Type", "application/json")
+            self.end_headers()
+            self.wfile.write(json.dumps({"status": "healthy"}).encode())
+            return
+
+        if self.path == "/events":
+            self.send_response(200)
+            self.send_header("Content-Type", "text/event-stream")
+            self.send_header("Cache-Control", "no-cache")
+            self.send_header("Connection", "keep-alive")
+            self.send_header("X-Accel-Buffering", "no")
+            self.end_headers()
+
+            print(f"[SSE] Client connected: {self.client_address}")
+
+            try:
+                event_id = 0
+                while True:
+                    event_id += 1
+                    data = {
+                        "id": event_id,
+                        "timestamp": datetime.now().isoformat(),
+                        "message": f"Event #{event_id}"
+                    }
+
+                    event = f"id: {event_id}\nevent: tick\ndata: {json.dumps(data)}\n\n"
+                    self.wfile.write(event.encode())
+                    self.wfile.flush()
+
+                    print(f"[SSE] Sent event #{event_id}")
+                    time.sleep(1)
+
+            except (BrokenPipeError, ConnectionResetError):
+                print(f"[SSE] Client {self.client_address} disconnected")
+            return
+
+        if self.path == "/stream/fast":
+            self.send_response(200)
+            self.send_header("Content-Type", "text/event-stream")
+            self.send_header("Cache-Control", "no-cache")
+            self.send_header("Connection", "keep-alive")
+            self.end_headers()
+
+            print(f"[SSE-FAST] Client connected: {self.client_address}")
+
+            try:
+                for i in range(100):
+                    data = {"seq": i, "ts": datetime.now().isoformat()}
+                    event = f"data: {json.dumps(data)}\n\n"
+                    self.wfile.write(event.encode())
+                    self.wfile.flush()
+                    time.sleep(0.05)
+
+                self.wfile.write(b"event: done\ndata: complete\n\n")
+                self.wfile.flush()
+                print(f"[SSE-FAST] Stream complete")
+
+            except (BrokenPipeError, ConnectionResetError):
+                print(f"[SSE-FAST] Client disconnected early")
+            return
+
+        self.send_response(200)
+        self.send_header("Content-Type", "text/html")
+        self.end_headers()
+        html = """<!DOCTYPE html>
+<html>
+<head><title>SSE Test</title></head>
+<body>
+<h1>SSE Test Endpoints</h1>
+<ul>
+  <li><a href="/events">/events</a> - Continuous 1-second ticks</li>
+  <li><a href="/stream/fast">/stream/fast</a> - Fast burst (100 events)</li>
+  <li><a href="/health">/health</a> - Health check</li>
+</ul>
+<h2>Live Events:</h2>
+<pre id="output"></pre>
+<script>
+const es = new EventSource('/events');
+es.onmessage = (e) => {
+  document.getElementById('output').textContent += e.data + '\\n';
+};
+es.addEventListener('tick', (e) => {
+  document.getElementById('output').textContent += e.data + '\\n';
+});
+</script>
+</body>
+</html>"""
+        self.wfile.write(html.encode())
+
+    def log_message(self, format, *args):
+        print(f"[HTTP] {format % args}")
+
+
+if __name__ == "__main__":
+    port = 8003
+    server = HTTPServer(("localhost", port), SSEHandler)
+    print(f"SSE server running on http://localhost:{port}")
+    print("Endpoints:")
+    print("  /events      - Continuous SSE stream (1 event/sec)")
+    print("  /stream/fast - Fast burst stream (100 events)")
+    print("  /health      - Health check")
+    server.serve_forever()

PROJECTS/Aenebris/examples/websockets/test_client.py

@@ -0,0 +1,124 @@
+#!/usr/bin/env python3
+"""
+Test client for WebSocket and SSE through Aenebris proxy
+"""
+
+import asyncio
+import sys
+import json
+
+try:
+    import websockets
+except ImportError:
+    print("Install websockets: pip install websockets")
+    websockets = None
+
+try:
+    import httpx
+except ImportError:
+    print("Install httpx: pip install httpx")
+    httpx = None
+
+
+async def test_websocket(host="localhost", port=8081):
+    if websockets is None:
+        print("SKIP: websockets not installed")
+        return False
+
+    uri = f"ws://{host}:{port}/ws"
+    print(f"\n[TEST] WebSocket: {uri}")
+
+    try:
+        async with websockets.connect(uri, extra_headers={"Host": "ws.localhost"}) as ws:
+            print("[OK] WebSocket connected")
+
+            await ws.send("ping")
+            response = await ws.recv()
+            assert response == "pong", f"Expected 'pong', got '{response}'"
+            print(f"[OK] ping -> {response}")
+
+            await ws.send("Hello Aenebris!")
+            response = await ws.recv()
+            data = json.loads(response)
+            assert data["echo"] == "Hello Aenebris!"
+            print(f"[OK] echo -> {data}")
+
+            await ws.send("time")
+            response = await ws.recv()
+            print(f"[OK] time -> {response}")
+
+            print("[PASS] WebSocket test passed!")
+            return True
+
+    except Exception as e:
+        print(f"[FAIL] WebSocket error: {e}")
+        return False
+
+
+async def test_sse(host="localhost", port=8081):
+    if httpx is None:
+        print("SKIP: httpx not installed")
+        return False
+
+    url = f"http://{host}:{port}/events"
+    print(f"\n[TEST] SSE: {url}")
+
+    try:
+        async with httpx.AsyncClient() as client:
+            async with client.stream("GET", url, headers={"Host": "sse.localhost"}) as response:
+                print(f"[OK] SSE connected, status: {response.status_code}")
+                print(f"[OK] Content-Type: {response.headers.get('content-type')}")
+
+                count = 0
+                async for line in response.aiter_lines():
+                    if line.startswith("data:"):
+                        data = json.loads(line[5:].strip())
+                        print(f"[OK] Event: {data}")
+                        count += 1
+                        if count >= 3:
+                            break
+
+                print("[PASS] SSE test passed!")
+                return True
+
+    except Exception as e:
+        print(f"[FAIL] SSE error: {e}")
+        return False
+
+
+async def test_both_simultaneously(host="localhost", port=8081):
+    print("\n[TEST] Running WebSocket + SSE simultaneously...")
+
+    ws_task = asyncio.create_task(test_websocket(host, port))
+    sse_task = asyncio.create_task(test_sse(host, port))
+
+    ws_ok, sse_ok = await asyncio.gather(ws_task, sse_task, return_exceptions=True)
+
+    if isinstance(ws_ok, Exception):
+        print(f"[FAIL] WebSocket: {ws_ok}")
+        ws_ok = False
+    if isinstance(sse_ok, Exception):
+        print(f"[FAIL] SSE: {sse_ok}")
+        sse_ok = False
+
+    if ws_ok and sse_ok:
+        print("\n" + "=" * 50)
+        print("[SUCCESS] Both WebSocket AND SSE work simultaneously!")
+        print("          nginx's conflict has been SOLVED!")
+        print("=" * 50)
+        return True
+    else:
+        print("\n[PARTIAL] Some tests failed")
+        return False
+
+
+if __name__ == "__main__":
+    host = sys.argv[1] if len(sys.argv) > 1 else "localhost"
+    port = int(sys.argv[2]) if len(sys.argv) > 2 else 8081
+
+    print("=" * 50)
+    print("Aenebris WebSocket + SSE Test Suite")
+    print("=" * 50)
+    print(f"Target: {host}:{port}")
+
+    asyncio.run(test_both_simultaneously(host, port))

PROJECTS/Aenebris/examples/websockets/ws_echo_server.py

@@ -0,0 +1,59 @@
+#!/usr/bin/env python3
+"""
+WebSocket Echo Server for Aenebris testing
+"""
+
+import asyncio
+import json
+from datetime import datetime
+
+try:
+    import websockets
+except ImportError:
+    print("Install websockets: pip install websockets")
+    exit(1)
+
+
+async def echo_handler(websocket):
+    client_addr = websocket.remote_address
+    print(f"[WS] Client connected: {client_addr}")
+
+    try:
+        async for message in websocket:
+            print(f"[WS] Received from {client_addr}: {message[:100]}...")
+
+            if message == "ping":
+                await websocket.send("pong")
+            elif message == "time":
+                await websocket.send(datetime.now().isoformat())
+            elif message == "info":
+                info = {
+                    "server": "Aenebris WebSocket Test Server",
+                    "client": str(client_addr),
+                    "protocol": "WebSocket",
+                    "timestamp": datetime.now().isoformat()
+                }
+                await websocket.send(json.dumps(info))
+            else:
+                response = {
+                    "echo": message,
+                    "length": len(message),
+                    "timestamp": datetime.now().isoformat()
+                }
+                await websocket.send(json.dumps(response))
+
+    except websockets.exceptions.ConnectionClosed as e:
+        print(f"[WS] Client {client_addr} disconnected: {e.code} {e.reason}")
+
+
+async def main():
+    port = 8002
+    async with websockets.serve(echo_handler, "localhost", port):
+        print(f"WebSocket echo server running on ws://localhost:{port}")
+        print("Commands: 'ping' -> 'pong', 'time' -> timestamp, 'info' -> server info")
+        print("Other messages are echoed back as JSON")
+        await asyncio.Future()
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

PROJECTS/Aenebris/src/Aenebris/Connection.hs

@@ -0,0 +1,99 @@
+{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE RecordWildCards #-}
+
+module Aenebris.Connection
+  ( ConnectionState(..)
+  , ConnectionType(..)
+  , TimeoutConfig(..)
+  , defaultTimeoutConfig
+  , detectConnectionType
+  , isWebSocketUpgrade
+  , isStreamingResponse
+  , getTimeout
+  ) where
+
+import Data.ByteString (ByteString)
+import qualified Data.ByteString as BS
+import qualified Data.ByteString.Char8 as BS8
+import Data.CaseInsensitive (CI)
+import qualified Data.CaseInsensitive as CI
+import Data.Maybe (isJust, fromMaybe)
+import Network.HTTP.Types (HeaderName, Status, statusCode)
+
+data ConnectionState
+  = HttpRequest
+  | HttpResponse
+  | ProtocolUpgrade
+  | TunnelMode
+  | StreamingResponse
+  deriving (Eq, Show)
+
+data ConnectionType
+  = RegularHttp
+  | WebSocket
+  | ServerSentEvents
+  | ChunkedStream
+  deriving (Eq, Show)
+
+data TimeoutConfig = TimeoutConfig
+  { tcHttpIdle :: Int
+  , tcWebSocketTunnel :: Int
+  , tcStreamingResponse :: Int
+  , tcProxyPingInterval :: Int
+  , tcPongTimeout :: Int
+  , tcConnectTimeout :: Int
+  }
+
+defaultTimeoutConfig :: TimeoutConfig
+defaultTimeoutConfig = TimeoutConfig
+  { tcHttpIdle = 60
+  , tcWebSocketTunnel = 3600
+  , tcStreamingResponse = 3600
+  , tcProxyPingInterval = 30
+  , tcPongTimeout = 10
+  , tcConnectTimeout = 5
+  }
+
+getTimeout :: TimeoutConfig -> ConnectionState -> Int
+getTimeout TimeoutConfig{..} state = case state of
+  HttpRequest -> tcHttpIdle
+  HttpResponse -> tcHttpIdle
+  ProtocolUpgrade -> tcHttpIdle
+  TunnelMode -> tcWebSocketTunnel
+  StreamingResponse -> tcStreamingResponse
+
+detectConnectionType :: [(HeaderName, ByteString)] -> ConnectionType
+detectConnectionType headers
+  | isWebSocketUpgrade headers = WebSocket
+  | otherwise = RegularHttp
+
+isWebSocketUpgrade :: [(HeaderName, ByteString)] -> Bool
+isWebSocketUpgrade headers =
+  hasUpgradeWebsocket && hasConnectionUpgrade
+  where
+    hasUpgradeWebsocket = case lookup "Upgrade" headers of
+      Just val -> CI.mk val == CI.mk ("websocket" :: ByteString)
+      Nothing -> False
+
+    hasConnectionUpgrade = case lookup "Connection" headers of
+      Just val -> "upgrade" `BS.isInfixOf` CI.foldedCase (CI.mk val)
+      Nothing -> False
+
+isStreamingResponse :: Status -> [(HeaderName, ByteString)] -> Bool
+isStreamingResponse status headers =
+  isSSE || isChunkedWithoutLength || isUnknownLength
+  where
+    isSSE = case lookup "Content-Type" headers of
+      Just ct -> "text/event-stream" `BS.isInfixOf` ct
+      Nothing -> False
+
+    isChunkedWithoutLength =
+      hasTransferEncodingChunked && not hasContentLength
+
+    hasTransferEncodingChunked = case lookup "Transfer-Encoding" headers of
+      Just te -> "chunked" `BS.isInfixOf` CI.foldedCase (CI.mk te)
+      Nothing -> False
+
+    hasContentLength = isJust (lookup "Content-Length" headers)
+
+    isUnknownLength = statusCode status == 200 && not hasContentLength && not hasTransferEncodingChunked

PROJECTS/Aenebris/src/Aenebris/Proxy.hs

@@ -12,12 +12,14 @@ module Aenebris.Proxy
 
 import Aenebris.Backend
 import Aenebris.Config
+import Aenebris.Connection
 import Aenebris.HealthCheck
 import Aenebris.LoadBalancer
 import Aenebris.TLS
+import Aenebris.Tunnel
 import Aenebris.Middleware.Security
 import Aenebris.Middleware.Redirect
-import Control.Concurrent.Async (Async, async, waitAnyCancel, mapConcurrently_)
+import Control.Concurrent.Async (Async, async, waitAnyCancel)
 import Control.Exception (try, SomeException)
 import Data.Function ((&))
 import Data.List (sortBy)
@@ -35,9 +37,9 @@ import Network.Wai
 import Network.Wai.Handler.Warp (run, defaultSettings, setPort)
 import Network.Wai.Handler.WarpTLS (runTLS)
 import System.IO (hPutStrLn, stderr)
+import Data.ByteString (ByteString)
 import qualified Data.ByteString as BS
 import qualified Data.ByteString.Char8 as BS8
-import qualified Data.ByteString.Lazy as LBS
 
 -- | Proxy runtime state
 data ProxyState = ProxyState
@@ -207,16 +209,15 @@ launchHTTPSWithSNI port tlsConfig app = do
 -- | Main proxy application (WAI)
 proxyApp :: Config -> Map Text LoadBalancer -> Manager -> Application
 proxyApp config loadBalancers manager req respond = do
-  -- Log incoming request
   logRequest req
 
-  -- Find matching route based on Host header and path
   let hostHeader = lookup "Host" (requestHeaders req)
       requestPath = rawPathInfo req
+      headers = requestHeaders req
+      connType = detectConnectionType headers
 
   case selectRoute config hostHeader requestPath of
     Nothing -> do
-      -- No matching route found - return 404
       hPutStrLn stderr $ "ERROR: No route found for request"
       respond $ responseLBS
         status404
@@ -224,7 +225,6 @@ proxyApp config loadBalancers manager req respond = do
         "Not Found: No route configured for this host/path"
 
     Just (upstreamName, _pathRoute) -> do
-      -- Find the load balancer for this upstream
       case Map.lookup upstreamName loadBalancers of
         Nothing -> do
           hPutStrLn stderr $ "ERROR: Load balancer not found: " ++ T.unpack upstreamName
@@ -234,7 +234,6 @@ proxyApp config loadBalancers manager req respond = do
             "Internal Server Error: Upstream configuration error"
 
         Just loadBalancer -> do
-          -- Select a backend using load balancing
           mBackend <- selectBackend loadBalancer
 
           case mBackend of
@@ -246,23 +245,57 @@ proxyApp config loadBalancers manager req respond = do
                 "Service Unavailable: No healthy backends available"
 
             Just backend -> do
-              -- Track this connection and forward request
+              case connType of
-              result <- try $ trackConnection backend $
+                WebSocket -> do
-                forwardRequest manager req (rbHost backend)
+                  hPutStrLn stderr $ "[WS] WebSocket upgrade detected"
+                  handleWebSocketUpgrade req respond backend
 
-              case result of
+                RegularHttp -> do
-                Left (err :: SomeException) -> do
+                  result <- try $ trackConnection backend $
-                  -- Handle errors gracefully
+                    forwardRequest manager req (rbHost backend)
-                  hPutStrLn stderr $ "ERROR: " ++ show err
-                  respond $ responseLBS
-                    status502
-                    [("Content-Type", "text/plain")]
-                    "Bad Gateway: Could not connect to backend server"
 
-                Right response -> do
+                  case result of
-                  -- Log response status
+                    Left (err :: SomeException) -> do
-                  logResponse response
+                      hPutStrLn stderr $ "ERROR: " ++ show err
-                  respond response
+                      respond $ responseLBS
+                        status502
+                        [("Content-Type", "text/plain")]
+                        "Bad Gateway: Could not connect to backend server"
+
+                    Right response -> do
+                      logResponse response
+                      respond response
+
+                _ -> do
+                  result <- try $ trackConnection backend $
+                    forwardRequest manager req (rbHost backend)
+
+                  case result of
+                    Left (err :: SomeException) -> do
+                      hPutStrLn stderr $ "ERROR: " ++ show err
+                      respond $ responseLBS
+                        status502
+                        [("Content-Type", "text/plain")]
+                        "Bad Gateway: Could not connect to backend server"
+
+                    Right response -> do
+                      logResponse response
+                      respond response
+
+handleWebSocketUpgrade :: Request -> (Response -> IO ResponseReceived) -> RuntimeBackend -> IO ResponseReceived
+handleWebSocketUpgrade req respond backend = do
+  let backendHost = rbHost backend
+      backupResponse = responseLBS
+        status502
+        [("Content-Type", "text/plain")]
+        "WebSocket upgrade failed"
+
+  respond $ responseRaw (wsHandler req backendHost) backupResponse
+
+wsHandler :: Request -> Text -> IO ByteString -> (ByteString -> IO ()) -> IO ()
+wsHandler req backendHost recv send = do
+  hPutStrLn stderr $ "[WS] Starting WebSocket tunnel to " ++ T.unpack backendHost
+  tunnelWebSocket req backendHost send recv
 
 -- | Select a route based on Host header and path
 selectRoute :: Config -> Maybe BS.ByteString -> BS.ByteString -> Maybe (Text, PathRoute)
@@ -297,33 +330,31 @@ selectUpstream config hostHeader requestPath =
 -- | Forward request to backend server
 forwardRequest :: Manager -> Request -> Text -> IO Response
 forwardRequest manager clientReq backendHost = do
-  -- Parse backend host:port
+  requestBody <- strictRequestBody clientReq
+
   let backendUrl = "http://" ++ T.unpack backendHost ++
                    BS8.unpack (rawPathInfo clientReq) ++
                    BS8.unpack (rawQueryString clientReq)
 
-  -- Parse and build backend request
   initReq <- parseRequest backendUrl
 
   let backendReq = initReq
         { HTTP.method = requestMethod clientReq
         , HTTP.requestHeaders = filterHeaders (requestHeaders clientReq)
-        , HTTP.requestBody = RequestBodyLBS LBS.empty  -- TODO: Forward request body
+        , HTTP.requestBody = RequestBodyLBS requestBody
         }
 
-  -- Make request to backend
   backendResponse <- httpLbs backendReq manager
 
-  -- Convert backend response to WAI response
   let status = HTTP.responseStatus backendResponse
       headers = HTTP.responseHeaders backendResponse
       body = HTTP.responseBody backendResponse
 
   return $ responseLBS status headers body
 
--- | Filter headers (remove hop-by-hop headers)
+-- | Filter headers for regular HTTP (remove hop-by-hop headers)
 filterHeaders :: [(HeaderName, BS.ByteString)] -> [(HeaderName, BS.ByteString)]
-filterHeaders = filter (\(name, _) -> name `notElem` hopByHopHeaders)
+filterHeaders headers = filter (\(name, _) -> name `notElem` hopByHopHeaders) headers
   where
     hopByHopHeaders =
       [ "Connection"
@@ -336,6 +367,18 @@ filterHeaders = filter (\(name, _) -> name `notElem` hopByHopHeaders)
       , "Upgrade"
       ]
 
+-- | Filter headers for WebSocket upgrade (preserve Upgrade and Connection)
+filterHeadersForUpgrade :: [(HeaderName, BS.ByteString)] -> [(HeaderName, BS.ByteString)]
+filterHeadersForUpgrade headers = filter (\(name, _) -> name `notElem` hopByHopHeaders) headers
+  where
+    hopByHopHeaders =
+      [ "Keep-Alive"
+      , "Proxy-Authenticate"
+      , "Proxy-Authorization"
+      , "TE"
+      , "Trailers"
+      ]
+
 -- | Log incoming request
 logRequest :: Request -> IO ()
 logRequest req = do

PROJECTS/Aenebris/src/Aenebris/Tunnel.hs

@@ -0,0 +1,167 @@
+{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE RecordWildCards #-}
+
+module Aenebris.Tunnel
+  ( tunnelWebSocket
+  , streamResponse
+  , bidirectionalCopy
+  ) where
+
+import Control.Concurrent.Async (race_)
+import Control.Exception (SomeException, try, bracket)
+import Data.ByteString (ByteString)
+import qualified Data.ByteString as BS
+import qualified Data.ByteString.Char8 as BS8
+import Data.CaseInsensitive (original)
+import Data.Text (Text)
+import qualified Data.Text as T
+import Network.HTTP.Types (HeaderName)
+import Network.Socket (Socket)
+import qualified Network.Socket as Socket
+import qualified Network.Socket.ByteString as SocketBS
+import Network.Wai
+import System.IO (hPutStrLn, stderr)
+
+tunnelWebSocket
+  :: Request
+  -> Text
+  -> (ByteString -> IO ())
+  -> IO ByteString
+  -> IO ()
+tunnelWebSocket clientReq backendHost clientSend clientRecv = do
+  hPutStrLn stderr $ "[WS] Initiating tunnel to " ++ T.unpack backendHost
+
+  result <- try $ do
+    let (host, port) = parseHostPort backendHost
+
+    bracket
+      (connectToBackend host port)
+      Socket.close
+      $ \backendSocket -> do
+          sendUpgradeRequest backendSocket clientReq
+          upgradeResponse <- receiveUpgradeResponse backendSocket
+
+          case parseUpgradeStatus upgradeResponse of
+            Just 101 -> do
+              hPutStrLn stderr "[WS] Backend accepted upgrade (101)"
+              clientSend upgradeResponse
+              bidirectionalCopy clientRecv clientSend
+                (SocketBS.recv backendSocket 65536)
+                (SocketBS.sendAll backendSocket)
+
+            Just code -> do
+              hPutStrLn stderr $ "[WS] Backend rejected upgrade: " ++ show code
+              clientSend upgradeResponse
+
+            Nothing -> do
+              hPutStrLn stderr "[WS] Invalid upgrade response"
+              clientSend "HTTP/1.1 502 Bad Gateway\r\n\r\n"
+
+  case result of
+    Left (e :: SomeException) ->
+      hPutStrLn stderr $ "[WS] Tunnel error: " ++ show e
+    Right () ->
+      hPutStrLn stderr "[WS] Tunnel closed"
+
+bidirectionalCopy
+  :: IO ByteString
+  -> (ByteString -> IO ())
+  -> IO ByteString
+  -> (ByteString -> IO ())
+  -> IO ()
+bidirectionalCopy clientRecv clientSend backendRecv backendSend = do
+  hPutStrLn stderr "[TUNNEL] Starting bidirectional copy"
+
+  race_
+    (copyLoop "client->backend" clientRecv backendSend)
+    (copyLoop "backend->client" backendRecv clientSend)
+
+  hPutStrLn stderr "[TUNNEL] Bidirectional copy ended"
+
+copyLoop :: String -> IO ByteString -> (ByteString -> IO ()) -> IO ()
+copyLoop name recv send = go
+  where
+    go = do
+      chunk <- recv
+      if BS.null chunk
+        then hPutStrLn stderr $ "[TUNNEL] " ++ name ++ ": connection closed"
+        else do
+          send chunk
+          go
+
+streamResponse
+  :: (ByteString -> IO ())
+  -> IO ByteString
+  -> IO ()
+streamResponse clientSend backendRecv = do
+  hPutStrLn stderr "[STREAM] Starting streaming response"
+  go
+  where
+    go = do
+      chunk <- backendRecv
+      if BS.null chunk
+        then hPutStrLn stderr "[STREAM] Backend closed"
+        else do
+          clientSend chunk
+          go
+
+parseHostPort :: Text -> (String, Int)
+parseHostPort hostPort =
+  case T.splitOn ":" hostPort of
+    [host, portStr] ->
+      case reads (T.unpack portStr) of
+        [(port, "")] -> (T.unpack host, port)
+        _ -> (T.unpack host, 80)
+    [host] -> (T.unpack host, 80)
+    _ -> (T.unpack hostPort, 80)
+
+connectToBackend :: String -> Int -> IO Socket
+connectToBackend host port = do
+  addrInfos <- Socket.getAddrInfo
+    (Just Socket.defaultHints { Socket.addrSocketType = Socket.Stream })
+    (Just host)
+    (Just $ show port)
+
+  case addrInfos of
+    [] -> error $ "Cannot resolve: " ++ host ++ ":" ++ show port
+    (addr:_) -> do
+      sock <- Socket.socket
+        (Socket.addrFamily addr)
+        Socket.Stream
+        Socket.defaultProtocol
+      Socket.connect sock (Socket.addrAddress addr)
+      return sock
+
+sendUpgradeRequest :: Socket -> Request -> IO ()
+sendUpgradeRequest sock req = do
+  let method = requestMethod req
+      path = rawPathInfo req <> rawQueryString req
+      headers = requestHeaders req
+
+      requestLine = method <> " " <> path <> " HTTP/1.1\r\n"
+      headerLines = BS.concat
+        [ original name <> ": " <> value <> "\r\n"
+        | (name, value) <- headers
+        ]
+      fullRequest = requestLine <> headerLines <> "\r\n"
+
+  SocketBS.sendAll sock fullRequest
+
+receiveUpgradeResponse :: Socket -> IO ByteString
+receiveUpgradeResponse sock = do
+  chunk <- SocketBS.recv sock 4096
+  if "\r\n\r\n" `BS.isInfixOf` chunk
+    then return chunk
+    else do
+      rest <- receiveUpgradeResponse sock
+      return $ chunk <> rest
+
+parseUpgradeStatus :: ByteString -> Maybe Int
+parseUpgradeStatus response =
+  case BS8.words (head $ BS8.lines response) of
+    (_:codeBS:_) ->
+      case reads (BS8.unpack codeBS) of
+        [(code, "")] -> Just code
+        _ -> Nothing
+    _ -> Nothing