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).
This commit is contained in:
CarterPerez-dev
parent
8f89000277
commit
984d3d5dd6
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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())
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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()
|
||||
|
|
@ -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))
|
||||
|
|
@ -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())
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
result <- try $ trackConnection backend $
|
||||
forwardRequest manager req (rbHost backend)
|
||||
case connType of
|
||||
WebSocket -> do
|
||||
hPutStrLn stderr $ "[WS] WebSocket upgrade detected"
|
||||
handleWebSocketUpgrade req respond backend
|
||||
|
||||
case result of
|
||||
Left (err :: SomeException) -> do
|
||||
-- Handle errors gracefully
|
||||
hPutStrLn stderr $ "ERROR: " ++ show err
|
||||
respond $ responseLBS
|
||||
status502
|
||||
[("Content-Type", "text/plain")]
|
||||
"Bad Gateway: Could not connect to backend server"
|
||||
RegularHttp -> do
|
||||
result <- try $ trackConnection backend $
|
||||
forwardRequest manager req (rbHost backend)
|
||||
|
||||
Right response -> do
|
||||
-- Log response status
|
||||
logResponse response
|
||||
respond response
|
||||
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
|
||||
|
||||
_ -> 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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
Loading…
Reference in New Issue