def extract_redirect_location(response: Response) -> URL: # if the server returned more than one value, use the first header in order location = response.get_first_header(b'Location') if not location: raise MissingLocationForRedirect(response) # if the location cannot be parsed as URL, let exception happen: this might be a redirect to a URN!! # simply don't follows the redirect, and returns the response to the caller try: return URL(location) except InvalidURL: raise UnsupportedRedirect()
class ClientConnection(asyncio.Protocol): __slots__ = ( "loop", "pool", "transport", "open", "_connection_lost", "writing_paused", "writable", "ready", "response_ready", "request_timeout", "headers", "request", "response", "parser", "expect_100_continue", "_pending_task", "_can_release", "_upgraded", ) def __init__(self, loop, pool) -> None: self.loop = loop self.pool = weakref.ref(pool) self.transport = None self.open = False self.writing_paused = False self.writable = asyncio.Event() self.ready = asyncio.Event() self.response_ready = asyncio.Event() self.expect_100_continue = False self.request = None self.request_timeout = 20 self.headers = [] self.response = None self.parser = httptools.HttpResponseParser(self) # type: ignore self._connection_lost = False self._pending_task = None self._can_release = False self._upgraded = False def reset(self) -> None: self.headers = [] self.request = None self.response = None self.writing_paused = False self.writable.set() self.expect_100_continue = False self.parser = httptools.HttpResponseParser(self) # type: ignore self._connection_lost = False self._pending_task = None self._can_release = False self._upgraded = False def pause_writing(self) -> None: super().pause_writing() self.writing_paused = True self.writable.clear() def resume_writing(self) -> None: super().resume_writing() self.writing_paused = False self.writable.set() def connection_made(self, transport) -> None: self.transport = transport self.open = True self.ready.set() async def _wait_response(self) -> Response: await self.response_ready.wait() self._pending_task = False if self._can_release: self.loop.call_soon(self.release) response = self.response assert response is not None if 99 < response.status < 200: # Handle 1xx informational # https://tools.ietf.org/html/rfc7231#section-6.2 if response.status == 101: # 101 Upgrade is a final response as it's used to switch # protocols with WebSockets handshake. # returns the response object with status 101 and access to the # transport self._upgraded = True raise UpgradeResponse(response, self.transport) if response.status == 100 and self.expect_100_continue: assert self.request is not None await self._send_body(self.request) # ignore; self.response_ready.clear() self.headers = [] # await the final response return await self._wait_response() self.response_ready.clear() return response async def _write_chunks(self, request, method) -> Optional[Response]: async for chunk in method(request): if self._can_release: # the server returned a response before # we ended sending the request: can happen for a bad request or # unauthorized while posting a big enough body return await self._wait_response() if not self.open: raise ConnectionClosedError(False) if self.writing_paused: await self.writable.wait() self.transport.write(chunk) async def _send_body(self, request: Request) -> None: await self._write_chunks(request, write_request_body_only) async def send(self, request: Request) -> Response: if not self.open: # NB: if the connection is closed here, it is always possible to # try again with a new connection # instead, if it happens later; we cannot retry because we started # sending a request raise ConnectionClosedError(True) self.request = request self._pending_task = True if request_has_body(request) and request.expect_100_continue(): # don't send the body immediately; instead, wait for HTTP 100 # Continue interim response from server self.expect_100_continue = True self.transport.write(write_request_without_body(request)) return await self._wait_response() if is_small_request(request): self.transport.write(write_small_request(request)) else: response = await self._write_chunks(request, write_request) if response is not None: # this happens if the server sent a response before we completed # sending a body return response return await self._wait_response() def close(self) -> None: if self.open: self.open = False if self.transport: self.transport.close() def data_received(self, data: bytes) -> None: try: self.parser.feed_data(data) except HttpParserCallbackError: self.close() raise except HttpParserError as pex: self.close() raise InvalidResponseFromServer(pex) def connection_lost(self, exc) -> None: self._connection_lost = True self.ready.clear() self.open = False if self._pending_task: self.response_ready.set() def on_header(self, name, value): self.headers.append((name, value)) def on_headers_complete(self) -> None: status = self.parser.get_status_code() self.response = Response(status, self.headers, None) # NB: check if headers declare a content-length if self._has_content(): self.response.content = IncomingContent( self.response.get_single_header(b"content-type")) self.response_ready.set() def _has_content(self) -> bool: content_length = self.response.get_first_header(b"content-length") if content_length: try: content_length_value = int(content_length) except ValueError as value_error: # server returned an invalid content-length value raise InvalidResponseFromServer( value_error, f"The server returned an invalid value for" f"the Content-Length header; value: {content_length}", ) return content_length_value > 0 transfer_encoding = self.response.get_first_header( b"transfer-encoding") if transfer_encoding and b"chunked" in transfer_encoding: return True return False def on_message_complete(self) -> None: if self.response and self.response.content: self.response.content.complete.set() if self._pending_task: # the server returned a response before we ended sending the # request, the connection can be released now - this can happen # for our Bad Requests self._can_release = True self.response_ready.set() else: # request-response cycle completed now, # the connection can be returned to its pool self.loop.call_soon(self.release) def release(self) -> None: if not self.open or self._upgraded: # if the connection was upgraded, its transport is used for # web sockets, it cannot return to its pool for other cycles return if self.parser.should_keep_alive(): self.reset() pool = self.pool() if pool: pool.try_return_connection(self) else: self.close() def on_body(self, value: bytes) -> None: self.response.content.extend_body(value)