108 lines
3.9 KiB
Python
108 lines
3.9 KiB
Python
import asyncio
|
|
import enum
|
|
import logging
|
|
from io import BytesIO
|
|
|
|
from mangum.exceptions import UnexpectedMessage
|
|
from mangum.types import ASGI, Message, Response, Scope
|
|
|
|
|
|
class HTTPCycleState(enum.Enum):
|
|
"""
|
|
The state of the ASGI `http` connection.
|
|
* **REQUEST** - Initial state. The ASGI application instance will be run with the
|
|
connection scope containing the `http` type.
|
|
* **RESPONSE** - The `http.response.start` event has been sent by the application.
|
|
The next expected message is the `http.response.body` event, containing the body
|
|
content. An application may pass the `more_body` argument to send content in chunks,
|
|
however content will always be returned in a single response, never streamed.
|
|
* **COMPLETE** - The body content from the ASGI application has been completely
|
|
read. A disconnect event will be sent to the application, and the response will
|
|
be returned.
|
|
"""
|
|
|
|
REQUEST = enum.auto()
|
|
RESPONSE = enum.auto()
|
|
COMPLETE = enum.auto()
|
|
|
|
|
|
class HTTPCycle:
|
|
def __init__(self, scope: Scope, body: bytes) -> None:
|
|
self.scope = scope
|
|
self.buffer = BytesIO()
|
|
self.state = HTTPCycleState.REQUEST
|
|
self.logger = logging.getLogger("mangum.http")
|
|
self.app_queue: asyncio.Queue[Message] = asyncio.Queue()
|
|
self.app_queue.put_nowait(
|
|
{
|
|
"type": "http.request",
|
|
"body": body,
|
|
"more_body": False,
|
|
}
|
|
)
|
|
|
|
def __call__(self, app: ASGI) -> Response:
|
|
asgi_instance = self.run(app)
|
|
loop = asyncio.get_event_loop()
|
|
asgi_task = loop.create_task(asgi_instance)
|
|
loop.run_until_complete(asgi_task)
|
|
|
|
return {
|
|
"status": self.status,
|
|
"headers": self.headers,
|
|
"body": self.body,
|
|
}
|
|
|
|
async def run(self, app: ASGI) -> None:
|
|
try:
|
|
await app(self.scope, self.receive, self.send)
|
|
except BaseException:
|
|
self.logger.exception("An error occurred running the application.")
|
|
if self.state is HTTPCycleState.REQUEST:
|
|
await self.send(
|
|
{
|
|
"type": "http.response.start",
|
|
"status": 500,
|
|
"headers": [[b"content-type", b"text/plain; charset=utf-8"]],
|
|
}
|
|
)
|
|
await self.send(
|
|
{
|
|
"type": "http.response.body",
|
|
"body": b"Internal Server Error",
|
|
"more_body": False,
|
|
}
|
|
)
|
|
elif self.state is not HTTPCycleState.COMPLETE:
|
|
self.status = 500
|
|
self.body = b"Internal Server Error"
|
|
self.headers = [[b"content-type", b"text/plain; charset=utf-8"]]
|
|
|
|
async def receive(self) -> Message:
|
|
return await self.app_queue.get() # pragma: no cover
|
|
|
|
async def send(self, message: Message) -> None:
|
|
if self.state is HTTPCycleState.REQUEST and message["type"] == "http.response.start":
|
|
self.status = message["status"]
|
|
self.headers = message.get("headers", [])
|
|
self.state = HTTPCycleState.RESPONSE
|
|
elif self.state is HTTPCycleState.RESPONSE and message["type"] == "http.response.body":
|
|
body = message.get("body", b"")
|
|
more_body = message.get("more_body", False)
|
|
self.buffer.write(body)
|
|
if not more_body:
|
|
self.body = self.buffer.getvalue()
|
|
self.buffer.close()
|
|
|
|
self.state = HTTPCycleState.COMPLETE
|
|
await self.app_queue.put({"type": "http.disconnect"})
|
|
|
|
self.logger.info(
|
|
"%s %s %s",
|
|
self.scope["method"],
|
|
self.scope["path"],
|
|
self.status,
|
|
)
|
|
else:
|
|
raise UnexpectedMessage(f"Unexpected {message['type']}")
|