From 10df7fa72af8bd519435780406227386e8f9786f Mon Sep 17 00:00:00 2001 From: cdeler Date: Fri, 16 Oct 2020 20:55:56 +0300 Subject: [PATCH] Fixed the problem with socket, which might leak when tunnel proxy connection couldn't be established --- httpcore/_async/http_proxy.py | 78 ++++++++++++++++++++--------------- httpcore/_sync/http_proxy.py | 78 ++++++++++++++++++++--------------- 2 files changed, 88 insertions(+), 68 deletions(-) diff --git a/httpcore/_async/http_proxy.py b/httpcore/_async/http_proxy.py index 8a9f33c2..aee50cb1 100644 --- a/httpcore/_async/http_proxy.py +++ b/httpcore/_async/http_proxy.py @@ -207,41 +207,51 @@ async def _tunnel_request( connect_url = self.proxy_origin + (target,) connect_headers = [(b"Host", target), (b"Accept", b"*/*")] connect_headers = merge_headers(connect_headers, self.proxy_headers) - (proxy_status_code, _, proxy_stream, _) = await proxy_connection.arequest( - b"CONNECT", connect_url, headers=connect_headers, ext=ext - ) - proxy_reason = get_reason_phrase(proxy_status_code) - logger.trace( - "tunnel_response proxy_status_code=%r proxy_reason=%r ", - proxy_status_code, - proxy_reason, - ) - # Read the response data without closing the socket - async for _ in proxy_stream: - pass - - # See if the tunnel was successfully established. - if proxy_status_code < 200 or proxy_status_code > 299: - msg = "%d %s" % (proxy_status_code, proxy_reason) - raise ProxyError(msg) - - # Upgrade to TLS if required - # We assume the target speaks TLS on the specified port - if scheme == b"https": - await proxy_connection.start_tls(host, timeout) - - # The CONNECT request is successful, so we have now SWITCHED PROTOCOLS. - # This means the proxy connection is now unusable, and we must create - # a new one for regular requests, making sure to use the same socket to - # retain the tunnel. - connection = AsyncHTTPConnection( - origin=origin, - http2=self._http2, - ssl_context=self._ssl_context, - socket=proxy_connection.socket, - ) - await self._add_to_pool(connection, timeout) + try: + ( + proxy_status_code, + _, + proxy_stream, + _, + ) = await proxy_connection.arequest( + b"CONNECT", connect_url, headers=connect_headers, ext=ext + ) + + proxy_reason = get_reason_phrase(proxy_status_code) + logger.trace( + "tunnel_response proxy_status_code=%r proxy_reason=%r ", + proxy_status_code, + proxy_reason, + ) + # Read the response data without closing the socket + async for _ in proxy_stream: + pass + + # See if the tunnel was successfully established. + if proxy_status_code < 200 or proxy_status_code > 299: + msg = "%d %s" % (proxy_status_code, proxy_reason) + raise ProxyError(msg) + + # Upgrade to TLS if required + # We assume the target speaks TLS on the specified port + if scheme == b"https": + await proxy_connection.start_tls(host, timeout) + + # The CONNECT request is successful, so we have now SWITCHED PROTOCOLS. + # This means the proxy connection is now unusable, and we must create + # a new one for regular requests, making sure to use the same socket to + # retain the tunnel. + connection = AsyncHTTPConnection( + origin=origin, + http2=self._http2, + ssl_context=self._ssl_context, + socket=proxy_connection.socket, + ) + await self._add_to_pool(connection, timeout) + except Exception: + await proxy_connection.aclose() + raise # Once the connection has been established we can send requests on # it as normal. diff --git a/httpcore/_sync/http_proxy.py b/httpcore/_sync/http_proxy.py index aa3a1ae5..ae94cc9d 100644 --- a/httpcore/_sync/http_proxy.py +++ b/httpcore/_sync/http_proxy.py @@ -207,41 +207,51 @@ def _tunnel_request( connect_url = self.proxy_origin + (target,) connect_headers = [(b"Host", target), (b"Accept", b"*/*")] connect_headers = merge_headers(connect_headers, self.proxy_headers) - (proxy_status_code, _, proxy_stream, _) = proxy_connection.request( - b"CONNECT", connect_url, headers=connect_headers, ext=ext - ) - proxy_reason = get_reason_phrase(proxy_status_code) - logger.trace( - "tunnel_response proxy_status_code=%r proxy_reason=%r ", - proxy_status_code, - proxy_reason, - ) - # Read the response data without closing the socket - for _ in proxy_stream: - pass - - # See if the tunnel was successfully established. - if proxy_status_code < 200 or proxy_status_code > 299: - msg = "%d %s" % (proxy_status_code, proxy_reason) - raise ProxyError(msg) - - # Upgrade to TLS if required - # We assume the target speaks TLS on the specified port - if scheme == b"https": - proxy_connection.start_tls(host, timeout) - - # The CONNECT request is successful, so we have now SWITCHED PROTOCOLS. - # This means the proxy connection is now unusable, and we must create - # a new one for regular requests, making sure to use the same socket to - # retain the tunnel. - connection = SyncHTTPConnection( - origin=origin, - http2=self._http2, - ssl_context=self._ssl_context, - socket=proxy_connection.socket, - ) - self._add_to_pool(connection, timeout) + try: + ( + proxy_status_code, + _, + proxy_stream, + _, + ) = proxy_connection.request( + b"CONNECT", connect_url, headers=connect_headers, ext=ext + ) + + proxy_reason = get_reason_phrase(proxy_status_code) + logger.trace( + "tunnel_response proxy_status_code=%r proxy_reason=%r ", + proxy_status_code, + proxy_reason, + ) + # Read the response data without closing the socket + for _ in proxy_stream: + pass + + # See if the tunnel was successfully established. + if proxy_status_code < 200 or proxy_status_code > 299: + msg = "%d %s" % (proxy_status_code, proxy_reason) + raise ProxyError(msg) + + # Upgrade to TLS if required + # We assume the target speaks TLS on the specified port + if scheme == b"https": + proxy_connection.start_tls(host, timeout) + + # The CONNECT request is successful, so we have now SWITCHED PROTOCOLS. + # This means the proxy connection is now unusable, and we must create + # a new one for regular requests, making sure to use the same socket to + # retain the tunnel. + connection = SyncHTTPConnection( + origin=origin, + http2=self._http2, + ssl_context=self._ssl_context, + socket=proxy_connection.socket, + ) + self._add_to_pool(connection, timeout) + except Exception: + proxy_connection.close() + raise # Once the connection has been established we can send requests on # it as normal.