46 lines
2.2 KiB
Python
46 lines
2.2 KiB
Python
import httpx
|
|
import asyncio
|
|
from websockets.client import connect
|
|
from fastapi import Request, Response, WebSocket
|
|
|
|
hop_by_hop_headers = ["transfer-encoding", "connection", "keep-alive", "upgrade", "proxy-authenticate", "proxy-authorization", "te", "trailers", "content-encoding", "content-length"]
|
|
|
|
def make_http_proxy(base_url_http: str, headers: dict = {}, remove_prefix_path: bool = False):
|
|
async def http_proxy(request: Request, path: str = "") -> Response:
|
|
url = f"{base_url_http}/{path}" if remove_prefix_path else f"{base_url_http}{request.url.path}"
|
|
merged_headers = dict(request.headers)
|
|
merged_headers.pop("accept-encoding", None)
|
|
merged_headers |= {k.lower(): v for k, v in headers.items()}
|
|
raw_headers = [(k.encode("latin-1"), v.encode("latin-1")) for k, v in merged_headers.items()]
|
|
async with httpx.AsyncClient() as client:
|
|
resp = await client.request(
|
|
method=request.method,
|
|
url=url,
|
|
headers=raw_headers,
|
|
content=await request.body(),
|
|
params=request.query_params
|
|
)
|
|
response = Response(content=resp.content, status_code=resp.status_code)
|
|
for k, v in resp.headers.multi_items():
|
|
if k.lower() not in hop_by_hop_headers:
|
|
response.headers.append(k, v)
|
|
return response
|
|
return http_proxy
|
|
|
|
def make_websocket_proxy(base_url_websocket: str, remove_prefix_path: bool = False):
|
|
async def websocket_proxy(client_ws: WebSocket, path: str = ""):
|
|
url = f"{base_url_websocket}/{path}" if remove_prefix_path else f"{base_url_websocket}{client_ws.url.path}"
|
|
await client_ws.accept()
|
|
async with connect(url) as server_ws:
|
|
async def client_to_server():
|
|
async for message in client_ws.iter_bytes():
|
|
await server_ws.send(message)
|
|
async def server_to_client():
|
|
async for message in server_ws:
|
|
await client_ws.send_bytes(message)
|
|
await asyncio.gather(
|
|
asyncio.create_task(client_to_server()),
|
|
asyncio.create_task(server_to_client())
|
|
)
|
|
return websocket_proxy
|