async def build_response( self, message: Any, response: Response, cache_backend: Union[BaseCacheBackend, BaseAsyncCacheBackend], ) -> Any: """ Builds the response object to return to the caller. :param message: The response from the ASGI application :param response: the response object :param cache_backend: The cache backend used to store responses. :return: The message returned from the ASGI application. """ headers = MutableHeaders() headers["Cache-Control"] = str(self.ttl) response.headers.update(dict(headers.items())) try: should_update_cache = self.request.state.update_cache except AttributeError: return message if not should_update_cache: return message cache_key = self.key_func(self.request) if isinstance(cache_backend, BaseAsyncCacheBackend): await cache_backend.set(cache_key, message, self.ttl) else: cache_backend.set(cache_key, message, self.ttl) return message
def test_mutable_headers_merge(): h = MutableHeaders() h = h | MutableHeaders({"a": "1"}) assert isinstance(h, MutableHeaders) assert dict(h) == {"a": "1"} assert h.items() == [("a", "1")] assert h.raw == [(b"a", b"1")]
def test_mutable_headers_update_dict(): h = MutableHeaders() h |= {"a": "1"} assert isinstance(h, MutableHeaders) assert dict(h) == {"a": "1"} assert h.items() == [("a", "1")] assert h.raw == [(b"a", b"1")]
def send_wrapper(message: Message): if message["type"] == "http.response.start": counter = traced_counter.get() if counter and isinstance(counter, Counter): headers = MutableHeaders(scope=message) headers.append("x-dagster-call-counts", json.dumps(counter.counts())) return send(message)
async def send_wrapper(message: Message) -> None: if message['type'] == 'http.response.start': if scope['session']: _data = base64.b64encode( json.dumps(scope['session']).encode('utf-8')) _data = self.signer.sign(_data) headers = MutableHeaders(scope=message) headers.append('session', _data.decode('utf-8')) await send(message)
async def send(self, message, send, request): if message["type"] != "http.response.start": await send(message) return headers = MutableHeaders(scope=message) req_headers = {k.lower(): v for k, v in dict(request.headers).items()} headers.append("x-b3-traceid", req_headers.get("x-b3-traceid", "")) headers.append("x-b3-sampled", req_headers.get("x-b3-sampled", "")) await send(message)
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: if scope["type"] in ("http", "websocket"): connection = HTTPConnection(scope) if "user" in connection.cookies: user_auth = connection.cookies["user"] headers = MutableHeaders(scope=scope) headers.setdefault("Authorization", user_auth) await self.app(scope, receive, send)
def headers(): h = MutableHeaders() h.update({ HeaderKeys.correlation_id: dummy_correlation_id, HeaderKeys.request_id: dummy_request_id, HeaderKeys.date: dummy_date, HeaderKeys.user_agent: dummy_user_agent, HeaderKeys.forwarded_for: dummy_forwarded_for, }) return h
async def send(self, message: Message, send: Send, request_headers: Headers): if message["type"] != "http.response.start": await send(message) return message.setdefault("headers", []) headers = MutableHeaders(scope=message) headers.update(self.__headers) await send(message)
async def enrich_response(self, arg) -> None: value = str(context.get(self.key)) # for ContextMiddleware if isinstance(arg, Response): arg.headers[self.key] = value # for ContextPureMiddleware else: if arg["type"] == "http.response.start": headers = MutableHeaders(scope=arg) headers.append(self.key, value)
async def send_cached(message): # Borrowing from: # https://github.com/encode/starlette/blob/master/starlette/middleware/cors.py if message["type"] != "http.response.start": return await send(message) message.setdefault("headers", []) headers = MutableHeaders(scope=message) headers.update({ 'Cache-Control': f'max-age={max_age}', 'Expires': http_date(time.time() + max_age), }) await send(message)
async def send_wrapper(message: Message) -> None: if message["type"] == "http.response.start": path = scope.get("root_path", "") or "/" if scope["session"]: # We have session data to persist. data = b64encode( json.dumps(scope["session"]).encode("utf-8")) data = self.signer.sign(data) headers = MutableHeaders(scope=message) header_value = "%s=%s; path=%s; Max-Age=%d; %s" % ( self.session_cookie, data.decode("utf-8"), path, self.max_age, self.security_flags, ) headers.append("Set-Cookie", header_value) elif not initial_session_was_empty: # The session has been cleared. headers = MutableHeaders(scope=message) header_value = "{}={}; {}".format( self.session_cookie, f"null; path={path}; expires=Thu, 01 Jan 1970 00:00:00 GMT;", self.security_flags, ) headers.append("Set-Cookie", header_value) await send(message)
def __init__(self) -> None: """Do not use manually.""" self._redirect_to = None self._starlette_resp = StarletteResponse self._body = None self._text = None self._content = None self._json = None self.headers = MutableHeaders() self.cookies = SimpleCookie() self.status_code = HTTPStatus.OK self.streaming: Optional[AsyncGenerator] = None self.reraise = False
async def sender(message: Message) -> None: if message["type"] == "http.response.start": if scope["session"]: # We have session data to persist. data = b64encode( json.dumps(scope["session"]).encode("utf-8")) data = self.signer.sign(data) headers = MutableHeaders(scope=message) header_value = "%s=%s; path=/; Max-Age=%d; %s" % ( self.session_cookie, data.decode("utf-8"), self.max_age, self.security_flags, ) headers.append("Set-Cookie", header_value) elif not was_empty_session: # The session has been cleared. headers = MutableHeaders(scope=message) header_value = "%s=%s; %s" % ( self.session_cookie, "null; path=/; expires=Thu, 01 Jan 1970 00:00:00 GMT;", self.security_flags, ) headers.append("Set-Cookie", header_value) await send(message)
async def send_wrapper(message: Message) -> None: if message["type"] == "http.response.start": if scope["session"]: if "exp" not in scope["session"]: scope["session"]["exp"] = int( time.time()) + self.max_age data = jwt.encode(self.jwt_header, scope["session"], str(self.jwt_secret.encode)) headers = MutableHeaders(scope=message) header_value = "%s=%s; path=/; Max-Age=%d; %s" % ( self.session_cookie, data.decode("utf-8"), self.max_age, self.security_flags, ) if self.domain: # pragma: no cover header_value += f"; domain={self.domain}" headers.append("Set-Cookie", header_value) elif not initial_session_was_empty: # The session has been cleared. headers = MutableHeaders(scope=message) header_value = "%s=%s; %s" % ( self.session_cookie, "null; path=/; expires=Thu, 01 Jan 1970 00:00:00 GMT;", self.security_flags, ) if self.domain: # pragma: no cover header_value += f"; domain={self.domain}" headers.append("Set-Cookie", header_value) await send(message)
async def send_wrapper(message: Message) -> None: if message["type"] == "http.response.start": session: Session = scope["session"] if session.is_modified and not session.is_empty: # We have session data to persist (data was changed, cleared, etc). nonlocal session_id session_id = await scope["session"].persist() headers = MutableHeaders(scope=message) header_value = "%s=%s; path=/; Max-Age=%d; %s" % ( self.session_cookie, session_id, self.max_age, self.security_flags, ) headers.append("Set-Cookie", header_value) elif session.is_loaded and session.is_empty: # no interactions to session were done headers = MutableHeaders(scope=message) header_value = "%s=%s; %s" % ( self.session_cookie, "null; path=/; expires=Thu, 01 Jan 1970 00:00:00 GMT;", self.security_flags, ) headers.append("Set-Cookie", header_value) await send(message)
async def send_wrapper(message: Message) -> None: if message["type"] == "http.response.start": if scope["session"]: # We have session data to persist. data = b64encode( json.dumps(scope["session"]).encode("utf-8")) data = self.signer.sign(data) headers = MutableHeaders(scope=message) header_value = "{session_cookie}={data}; path={path}; {max_age}{security_flags}".format( # noqa E501 session_cookie=self.session_cookie, data=data.decode("utf-8"), path=self.path, max_age=f"Max-Age={self.max_age}; " if self.max_age else "", security_flags=self.security_flags, ) headers.append("Set-Cookie", header_value) elif not initial_session_was_empty: # The session has been cleared. headers = MutableHeaders(scope=message) header_value = "{session_cookie}={data}; path={path}; {expires}{security_flags}".format( # noqa E501 session_cookie=self.session_cookie, data="null", path=self.path, expires="expires=Thu, 01 Jan 1970 00:00:00 GMT; ", security_flags=self.security_flags, ) headers.append("Set-Cookie", header_value) await send(message)
async def send(self, message, send=None, origin=None): if message["type"] != "http.response.start": await send(message) return message.setdefault("headers", []) headers = MutableHeaders(message["headers"]) # If we only allow specific origins, then we have to mirror back # the Origin header in the response. if not self.allow_all_origins and self.is_allowed_origin( origin=origin): headers["Access-Control-Allow-Origin"] = origin headers.update(self.simple_headers) await send(message)
async def wrapped_send_bfyaml(response): nonlocal body, agent_response try: if agent_response: # send response if signature is validated if response.get("type") == "http.response.start": _add_header( response, agent_response[0], agent_response[1] ) # We can add headers as many as we want but it was # not possible to mutate an existing header without # using below approach. # override the Content-Length received from the original # Response. Note: MutableHeaders is present in the minimum # Starlette version used in minimum FastAPI version (0.51.0) from starlette.datastructures import MutableHeaders headers = MutableHeaders( raw=response["headers"] ) headers['Content-Length'] = str(len(body)) elif response.get("type") == "http.response.body": response["body"] = body await send(response) except Exception as e: log.exception(e)
async def send_with_gzip(self, message: Message) -> None: message_type = message["type"] if message_type == "http.response.start": # Don't send the initial message until we've determined how to # modify the ougoging headers correctly. self.initial_message = message elif message_type == "http.response.body" and not self.started: self.started = True body = message.get("body", b"") more_body = message.get("more_body", False) if len(body) < self.minimum_size and not more_body: # Don't apply GZip to small outgoing responses. await self.send(self.initial_message) await self.send(message) elif not more_body: # Standard GZip response. self.gzip_file.write(body) self.gzip_file.close() body = self.gzip_buffer.getvalue() headers = MutableHeaders(raw=self.initial_message["headers"]) headers["Content-Encoding"] = "gzip" headers["Content-Length"] = str(len(body)) headers.add_vary_header("Accept-Encoding") message["body"] = body await self.send(self.initial_message) await self.send(message) else: # Initial body in streaming GZip response. headers = MutableHeaders(raw=self.initial_message["headers"]) headers["Content-Encoding"] = "gzip" headers.add_vary_header("Accept-Encoding") del headers["Content-Length"] self.gzip_file.write(body) message["body"] = self.gzip_buffer.getvalue() self.gzip_buffer.seek(0) self.gzip_buffer.truncate() await self.send(self.initial_message) await self.send(message) elif message_type == "http.response.body": # Remaining body in streaming GZip response. body = message.get("body", b"") more_body = message.get("more_body", False) self.gzip_file.write(body) if not more_body: self.gzip_file.close() message["body"] = self.gzip_buffer.getvalue() self.gzip_buffer.seek(0) self.gzip_buffer.truncate() await self.send(message)
async def send_wrapper(self, msg: Message): if msg["type"] == "http.response.start": result = self.rl_res headers = MutableHeaders(scope=msg) headers.append("X-Rate-Limit-Limit", str(result.consumed_points)) headers.append("X-Rate-Limit-Remaining", str(result.remaining_points)) headers.append("X-Rate-Limit-Reset", str(result.ms_before_next)) await self.send(msg)
async def send_with_brotli(self, message: Message) -> None: """Apply compression using brotli.""" message_type = message["type"] if message_type == "http.response.start": # Don't send the initial message until we've determined how to # modify the outgoing headers correctly. self.initial_message = message elif message_type == "http.response.body" and not self.started: self.started = True body = message.get("body", b"") more_body = message.get("more_body", False) if len(body) < self.minimum_size and not more_body: # Don't apply Brotli to small outgoing responses. await self.send(self.initial_message) await self.send(message) elif not more_body: # Standard Brotli response. body = self.br_file.process(body) + self.br_file.finish() headers = MutableHeaders(raw=self.initial_message["headers"]) headers["Content-Encoding"] = "br" headers["Content-Length"] = str(len(body)) headers.add_vary_header("Accept-Encoding") message["body"] = body await self.send(self.initial_message) await self.send(message) else: # Initial body in streaming Brotli response. headers = MutableHeaders(raw=self.initial_message["headers"]) headers["Content-Encoding"] = "br" headers.add_vary_header("Accept-Encoding") del headers["Content-Length"] self.br_buffer.write( self.br_file.process(body) + self.br_file.flush()) message["body"] = self.br_buffer.getvalue() self.br_buffer.seek(0) self.br_buffer.truncate() await self.send(self.initial_message) await self.send(message) elif message_type == "http.response.body": # Remaining body in streaming Brotli response. body = message.get("body", b"") more_body = message.get("more_body", False) self.br_buffer.write( self.br_file.process(body) + self.br_file.flush()) if not more_body: self.br_buffer.write(self.br_file.finish()) message["body"] = self.br_buffer.getvalue() self.br_buffer.close() await self.send(message) return message["body"] = self.br_buffer.getvalue() self.br_buffer.seek(0) self.br_buffer.truncate() await self.send(message)
async def send_with_caching(self, message: Message) -> None: if message["type"] == "http.response.start": logger.trace(f"patch_cache_control {kvformat(**self.kwargs)}") headers = MutableHeaders(raw=list(message["headers"])) patch_cache_control(headers, **self.kwargs) message["headers"] = headers.raw await self.send(message)
async def body(self) -> bytes: if not hasattr(self, "_body"): body = await super().body() try: self._body = zlib.decompress(body) headers = MutableHeaders(raw=self.scope["headers"]) headers["Content-Length"] = str(len(self._body)) except zlib.error: self._body = body return self._body
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: headers = MutableHeaders(scope=scope) self.should_decode_from_msgpack_to_json = ("application/x-msgpack" in headers.get( "content-type", "")) # Take an initial guess, although we eventually may not # be able to do the conversion. self.should_encode_from_json_to_msgpack = ( "application/x-msgpack" in headers.getlist("accept")) self.receive = receive self.send = send if self.should_decode_from_msgpack_to_json: # We're going to present JSON content to the application, # so rewrite `Content-Type` for consistency and compliance # with possible downstream security checks in some frameworks. # See: https://github.com/florimondmanca/msgpack-asgi/issues/23 headers["content-type"] = "application/json" await self.app(scope, self.receive_with_msgpack, self.send_with_msgpack)
async def send(self, message: Message, send: Send, request_headers: Headers) -> None: if message["type"] != "http.response.start": await send(message) return message.setdefault("headers", []) headers = MutableHeaders(scope=message) headers.update(self.simple_headers) origin = request_headers["Origin"] has_cookie = "cookie" in request_headers # If request includes any cookie headers, then we must respond # with the specific origin instead of '*'. if self.allow_all_origins and has_cookie: self.allow_explicit_origin(headers, origin) # If we only allow specific origins, then we have to mirror back # the Origin header in the response. elif not self.allow_all_origins and self.is_allowed_origin( origin=origin): self.allow_explicit_origin(headers, origin) await send(message)
async def send(self, message, send=None, request_headers=None): if message["type"] != "http.response.start": await send(message) return message.setdefault("headers", []) headers = MutableHeaders(message["headers"]) origin = request_headers["Origin"] has_cookie = "cookie" in request_headers # If request includes any cookie headers, then we must respond # with the specific origin instead of '*'. if self.allow_all_origins and has_cookie: self.simple_headers["Access-Control-Allow-Origin"] = origin # If we only allow specific origins, then we have to mirror back # the Origin header in the response. elif not self.allow_all_origins and self.is_allowed_origin( origin=origin): headers["Access-Control-Allow-Origin"] = origin if "vary" in headers: self.simple_headers["Vary"] = f"{headers.get('vary')}, Origin" headers.update(self.simple_headers) await send(message)
def patch_cache_control(headers: MutableHeaders, **kwargs: typing.Any) -> None: """ Patch headers with an extended version of the initial Cache-Control header by adding all keyword arguments to it. """ cache_control: typing.Dict[str, typing.Any] = {} for field in parse_http_list(headers.get("Cache-Control", "")): try: key, value = field.split("=") except ValueError: cache_control[field] = True else: cache_control[key] = value if "max-age" in cache_control and "max_age" in kwargs: kwargs["max_age"] = min(int(cache_control["max-age"]), kwargs["max_age"]) if "public" in kwargs: raise NotImplementedError( "The 'public' cache control directive isn't supported yet.") if "private" in kwargs: raise NotImplementedError( "The 'private' cache control directive isn't supported yet.") for key, value in kwargs.items(): key = key.replace("_", "-") cache_control[key] = value directives: typing.List[str] = [] for key, value in cache_control.items(): if value is False: continue if value is True: directives.append(key) else: directives.append(f"{key}={value}") patched_cache_control = ", ".join(directives) if patched_cache_control: headers["Cache-Control"] = patched_cache_control else: del headers["Cache-Control"]
def test_mutable_headers(): h = MutableHeaders() assert dict(h) == {} h["a"] = "1" assert dict(h) == {"a": "1"} h["a"] = "2" assert dict(h) == {"a": "2"} h.setdefault("a", "3") assert dict(h) == {"a": "2"} h.setdefault("b", "4") assert dict(h) == {"a": "2", "b": "4"} del h["a"] assert dict(h) == {"b": "4"}
async def send_wrapper(message: Message) -> None: if message['type'] == 'http.response.start': if scope['session']: # We have session data to persist. headers = MutableHeaders(scope=message) header_value = '%s=%s; path=/; Max-Age=%d; %s' % ( self.session_cookie, self.backend.encode(scope['session']), self.max_age, self.security_flags, ) headers.append('Set-Cookie', header_value) elif not initial_session_was_empty: # The session has been cleared. headers = MutableHeaders(scope=message) header_value = '%s=%s; %s' % ( self.session_cookie, 'null; path=/; expires=Thu, 01 Jan 1970 00:00:00 GMT;', self.security_flags, ) headers.append('Set-Cookie', header_value) await send(message)