codekingpro/portable-devtools

1import sys2from collections import defaultdict3from tempfile import SpooledTemporaryFile4 5from asgiref.sync import AsyncToSync, sync_to_async6 7 8class WsgiToAsgi:9    """10    Wraps a WSGI application to make it into an ASGI application.11    """12 13    def __init__(self, wsgi_application, duplicate_header_limit=100):14        self.wsgi_application = wsgi_application15        self.duplicate_header_limit = duplicate_header_limit16 17    async def __call__(self, scope, receive, send):18        """19        ASGI application instantiation point.20        We return a new WsgiToAsgiInstance here with the WSGI app21        and the scope, ready to respond when it is __call__ed.22        """23        await WsgiToAsgiInstance(self.wsgi_application, self.duplicate_header_limit)(24            scope, receive, send25        )26 27 28class WsgiToAsgiInstance:29    """30    Per-socket instance of a wrapped WSGI application31    """32 33    def __init__(self, wsgi_application, duplicate_header_limit=100):34        self.wsgi_application = wsgi_application35        self.duplicate_header_limit = duplicate_header_limit36        self.response_started = False37        self.response_content_length = None38 39    async def __call__(self, scope, receive, send):40        if scope["type"] != "http":41            raise ValueError("WSGI wrapper received a non-HTTP scope")42        self.scope = scope43        with SpooledTemporaryFile(max_size=65536) as body:44            # Alright, wait for the http.request messages45            while True:46                message = await receive()47                if message["type"] != "http.request":48                    raise ValueError("WSGI wrapper received a non-HTTP-request message")49                body.write(message.get("body", b""))50                if not message.get("more_body"):51                    break52            body.seek(0)53            # Wrap send so it can be called from the subthread54            self.sync_send = AsyncToSync(send)55            # Call the WSGI app56            await self.run_wsgi_app(body)57 58    def build_environ(self, scope, body):59        """60        Builds a scope and request body into a WSGI environ object.61        """62        script_name = scope.get("root_path", "").encode("utf8").decode("latin1")63        path_info = scope["path"].encode("utf8").decode("latin1")64        if path_info.startswith(script_name):65            path_info = path_info[len(script_name) :]66        environ = {67            "REQUEST_METHOD": scope["method"],68            "SCRIPT_NAME": script_name,69            "PATH_INFO": path_info,70            "QUERY_STRING": scope["query_string"].decode("ascii"),71            "SERVER_PROTOCOL": "HTTP/%s" % scope["http_version"],72            "wsgi.version": (1, 0),73            "wsgi.url_scheme": scope.get("scheme", "http"),74            "wsgi.input": body,75            "wsgi.errors": sys.stderr,76            "wsgi.multithread": True,77            "wsgi.multiprocess": True,78            "wsgi.run_once": False,79        }80        # Get server name and port - required in WSGI, not in ASGI81        if "server" in scope:82            environ["SERVER_NAME"] = scope["server"][0]83            environ["SERVER_PORT"] = str(scope["server"][1])84        else:85            environ["SERVER_NAME"] = "localhost"86            environ["SERVER_PORT"] = "80"87 88        if scope.get("client") is not None:89            environ["REMOTE_ADDR"] = scope["client"][0]90 91        # Go through headers and make them into environ entries92        _headers = defaultdict(list)93        for name, value in self.scope.get("headers", []):94            name = name.decode("latin1")95            if name == "content-length":96                corrected_name = "CONTENT_LENGTH"97            elif name == "content-type":98                corrected_name = "CONTENT_TYPE"99            else:100                corrected_name = "HTTP_%s" % name.upper().replace("-", "_")101            # HTTPbis say only ASCII chars are allowed in headers, but we latin1 just in case102            value = value.decode("latin1")103            if (104                self.duplicate_header_limit105                and len(_headers[corrected_name]) >= self.duplicate_header_limit106            ):107                raise ValueError(108                    f"Too many duplicate headers: {corrected_name} exceeds limit of"109                    f"{self.duplicate_header_limit}"110                )111            _headers[corrected_name].append(value)112        for name, values in _headers.items():113            environ[name] = ",".join(values)114        return environ115 116    def start_response(self, status, response_headers, exc_info=None):117        """118        WSGI start_response callable.119        """120        # Don't allow re-calling once response has begun121        if self.response_started:122            raise exc_info[1].with_traceback(exc_info[2])123        # Don't allow re-calling without exc_info124        if hasattr(self, "response_start") and exc_info is None:125            raise ValueError(126                "You cannot call start_response a second time without exc_info"127            )128        # Extract status code129        status_code, _ = status.split(" ", 1)130        status_code = int(status_code)131        # Extract headers132        headers = [133            (name.lower().encode("ascii"), value.encode("ascii"))134            for name, value in response_headers135        ]136        # Extract content-length137        self.response_content_length = None138        for name, value in response_headers:139            if name.lower() == "content-length":140                self.response_content_length = int(value)141        # Build and send response start message.142        self.response_start = {143            "type": "http.response.start",144            "status": status_code,145            "headers": headers,146        }147 148    @sync_to_async149    def run_wsgi_app(self, body):150        """151        Called in a subthread to run the WSGI app. We encapsulate like152        this so that the start_response callable is called in the same thread.153        """154        # Translate the scope and incoming request body into a WSGI environ155        try:156            environ = self.build_environ(self.scope, body)157        except ValueError:158            # Return 400 Bad Request if header limit exceeded159            self.sync_send(160                {161                    "type": "http.response.start",162                    "status": 400,163                    "headers": [(b"content-type", b"text/plain")],164                }165            )166            self.sync_send(167                {168                    "type": "http.response.body",169                    "body": b"Bad Request: Too many duplicate headers",170                }171            )172            return173        # Run the WSGI app174        bytes_sent = 0175        for output in self.wsgi_application(environ, self.start_response):176            # If this is the first response, include the response headers177            if not self.response_started:178                self.response_started = True179                self.sync_send(self.response_start)180            # If the application supplies a Content-Length header181            if self.response_content_length is not None:182                # The server should not transmit more bytes to the client than the header allows183                bytes_allowed = self.response_content_length - bytes_sent184                if len(output) > bytes_allowed:185                    output = output[:bytes_allowed]186            self.sync_send(187                {"type": "http.response.body", "body": output, "more_body": True}188            )189            bytes_sent += len(output)190            # The server should stop iterating over the response when enough data has been sent191            if bytes_sent == self.response_content_length:192                break193        # Close connection194        if not self.response_started:195            self.response_started = True196            self.sync_send(self.response_start)197        self.sync_send({"type": "http.response.body"})198