diff --git a/anonstream/__init__.py b/anonstream/__init__.py index ce9dd30..67d5aa3 100644 --- a/anonstream/__init__.py +++ b/anonstream/__init__.py @@ -6,10 +6,10 @@ import secrets import toml from collections import OrderedDict -from quart import Quart from quart_compress import Compress from werkzeug.security import generate_password_hash +from anonstream.quart import Quart from anonstream.utils.captcha import create_captcha_factory, create_captcha_signer from anonstream.utils.colour import color_to_colour from anonstream.utils.user import generate_token diff --git a/anonstream/quart.py b/anonstream/quart.py new file mode 100644 index 0000000..5761e86 --- /dev/null +++ b/anonstream/quart.py @@ -0,0 +1,51 @@ +import asyncio + +from werkzeug.wrappers import Response as WerkzeugResponse +from quart.app import Quart as Quart_ +from quart.asgi import ASGIHTTPConnection as ASGIHTTPConnection_ +from quart.utils import encode_headers + + +RESPONSE_ITERATOR_TIMEOUT = 10 + + +class ASGIHTTPConnection(ASGIHTTPConnection_): + async def _send_response(self, send, response): + await send({ + "type": "http.response.start", + "status": response.status_code, + "headers": encode_headers(response.headers), + }) + + if isinstance(response, WerkzeugResponse): + for data in response.response: + body = data.encode(response.charset) if isinstance(data, str) else data + await asyncio.wait_for( + send({ + "type": "http.response.body", + "body": body, + "more_body": True, + }), + timeout=RESPONSE_ITERATOR_TIMEOUT, + ) + else: + async with response.response as response_body: + async for data in response_body: + body = data.encode(response.charset) if isinstance(data, str) else data + await asyncio.wait_for( + send({ + "type": "http.response.body", + "body": body, + "more_body": True, + }), + timeout=RESPONSE_ITERATOR_TIMEOUT, + ) + await send({ + "type": "http.response.body", + "body": b"", + "more_body": False, + }) + + +class Quart(Quart_): + asgi_http_class = ASGIHTTPConnection