import logging from base64 import b64encode, b64decode from uuid import uuid4 from starlette.endpoints import HTTPEndpoint from starlette.requests import Request from starlette.responses import Response from proxy.utils import get_path_with_query_string from ttun_server.proxy_queue import ProxyQueue 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): for header in self._headers: yield header 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()