import logging from base64 import b64decode, b64encode from uuid import uuid4 from starlette.endpoints import HTTPEndpoint from starlette.requests import Request from starlette.responses import Response from proxy.queue import ProxyQueue from proxy.utils import get_path_with_query_string from ttun_server.types import HttpMessage, HttpMessageType, HttpRequestData logger = logging.getLogger(__name__) class HeaderMapping: def __init__(self, headers: list[tuple[str, str]]): self._headers = headers def items(self): yield from self._headers class Proxy(HTTPEndpoint): async def dispatch(self) -> None: request = Request(self.scope, self.receive) subdomain = request.path_params["subdomain"] response = Response(content="Not Found", status_code=404) identifier = str(uuid4()) response_queue = await ProxyQueue.create_for_identifier(identifier) try: request_queue = await ProxyQueue.get_for_identifier(subdomain) logger.debug("PROXY %s%s ", subdomain, request.url) await request_queue.enqueue( HttpMessage( type=HttpMessageType.request.value, identifier=identifier, payload=HttpRequestData( method=request.method, path=get_path_with_query_string(request), headers=list(request.headers.items()), body=b64encode(await request.body()).decode(), ), ) ) _response = await response_queue.dequeue() payload = _response["payload"] response = Response( status_code=payload["status"], headers=HeaderMapping(payload["headers"]), content=b64decode(payload["body"].encode()), ) except AssertionError: pass finally: await response(self.scope, self.receive, self.send) await response_queue.delete()