Fix tunnel proxy and support HTTPS

httpcore/_async/connection.py

@@ -1,7 +1,7 @@
 from ssl import SSLContext
 from typing import List, Optional, Tuple, Union
 
-from .._backends.auto import AsyncLock, AutoBackend
+from .._backends.auto import AsyncLock, AsyncSocketStream, AutoBackend
 from .._types import URL, Headers, Origin, TimeoutDict
 from .base import (
  AsyncByteStream,
@@ -15,7 +15,11 @@
 
 class AsyncHTTPConnection(AsyncHTTPTransport):
  def __init__(
- self, origin: Origin, http2: bool = False, ssl_context: SSLContext = None,
+ self,
+ origin: Origin,
+ http2: bool = False,
+ ssl_context: SSLContext = None,
+ socket: AsyncSocketStream = None,
  ):
  self.origin = origin
  self.http2 = http2
@@ -30,6 +34,7 @@ def __init__(
  self.connect_failed = False
  self.expires_at: Optional[float] = None
  self.backend = AutoBackend()
+ self.socket = socket
 
  @property
  def request_lock(self) -> AsyncLock:
@@ -48,14 +53,16 @@ async def request(
  timeout: TimeoutDict = None,
  ) -> Tuple[bytes, int, bytes, List[Tuple[bytes, bytes]], AsyncByteStream]:
  assert url[:3] == self.origin
-
  async with self.request_lock:
  if self.state == ConnectionState.PENDING:
- try:
- await self._connect(timeout)
- except Exception:
- self.connect_failed = True
- raise
+ if self.socket:
+ self._create_connection(self.socket)
+ else:
+ try:
+ await self._connect(timeout)
+ except Exception:
+ self.connect_failed = True
+ raise
  elif self.state in (ConnectionState.READY, ConnectionState.IDLE):
  pass
  elif self.state == ConnectionState.ACTIVE and self.is_http2:
@@ -70,16 +77,23 @@ async def _connect(self, timeout: TimeoutDict = None) -> None:
  scheme, hostname, port = self.origin
  timeout = {} if timeout is None else timeout
  ssl_context = self.ssl_context if scheme == b"https" else None
- socket = await self.backend.open_tcp_stream(
+self.socket = await self.backend.open_tcp_stream(
  hostname, port, ssl_context, timeout
  )
+ self._create_connection(self.socket)
+
+ def _create_connection(self, socket: AsyncSocketStream) -> None:
  http_version = socket.get_http_version()
  if http_version == "HTTP/2":
  self.is_http2 = True
- self.connection = AsyncHTTP2Connection(socket=socket, backend=self.backend)
+ self.connection = AsyncHTTP2Connection(
+ socket=socket, backend=self.backend, ssl_context=self.ssl_context
+ )
  else:
  self.is_http11 = True
- self.connection = AsyncHTTP11Connection(socket=socket)
+ self.connection = AsyncHTTP11Connection(
+ socket=socket, ssl_context=self.ssl_context
+ )
 
  @property
  def state(self) -> ConnectionState:
@@ -99,3 +113,4 @@ def mark_as_ready(self) -> None:
  async def start_tls(self, hostname: bytes, timeout: TimeoutDict = None) -> None:
  if self.connection is not None:
  await self.connection.start_tls(hostname, timeout)
+ self.socket = self.connection.socket

httpcore/_async/http11.py

@@ -123,7 +123,7 @@ async def _receive_response_data(
  event = await self._receive_event(timeout)
  if isinstance(event, h11.Data):
  yield bytes(event.data)
- elif isinstance(event, h11.EndOfMessage):
+ elif isinstance(event, (h11.EndOfMessage, h11.PAUSED)):
  break
 
  async def _receive_event(self, timeout: TimeoutDict) -> H11Event:

httpcore/_async/http_proxy.py

@@ -140,47 +140,57 @@ async def _tunnel_request(
  connection = await self._get_connection_from_pool(origin)
 
  if connection is None:
- connection = AsyncHTTPConnection(
- origin=origin, http2=False, ssl_context=self._ssl_context,
+ # First, create a connection to the proxy server
+ proxy_connection = AsyncHTTPConnection(
+ origin=self.proxy_origin, http2=False, ssl_context=self._ssl_context,
  )
  async with self._thread_lock:
  self._connections.setdefault(origin, set())
- self._connections[origin].add(connection)
+ self._connections[origin].add(proxy_connection)
 
- # Establish the connection by issuing a CONNECT request...
+ # Issue a CONNECT request...
 
  # CONNECT www.example.org:80 HTTP/1.1
  # [proxy-headers]
  target = b"%b:%d" % (url[1], url[2])
  connect_url = self.proxy_origin + (target,)
- connect_headers = self.proxy_headers
- proxy_response = await connection.request(
+ connect_headers = self.proxy_headers or [(b"host", self.proxy_origin[1])]
+ proxy_response = await proxy_connection.request(
  b"CONNECT", connect_url, headers=connect_headers, timeout=timeout
  )
  proxy_status_code = proxy_response[1]
  proxy_reason_phrase = proxy_response[2]
  proxy_stream = proxy_response[4]
 
- # Ingest any request body.
-await read_body(proxy_stream)
+ # Read the response data without closing the socket
+_ = b"".join([chunk async for chunk in proxy_stream])
 
  # If the proxy responds with an error, then drop the connection
  # from the pool, and raise an exception.
  if proxy_status_code < 200 or proxy_status_code > 299:
- async with self._thread_lock:
- self._connections[connection.origin].remove(connection)
- if not self._connections[connection.origin]:
- del self._connections[connection.origin]
  msg = "%d %s" % (proxy_status_code, proxy_reason_phrase.decode("ascii"))
  raise ProxyError(msg)
 
- # Upgrade to TLS.
- await connection.start_tls(target, timeout)
+ # Upgrade to TLS if required
+ # We assume the target speaks TLS on the specified port
+ if url[0] == b"https":
+ await proxy_connection.start_tls(url[1], timeout)
+
+ # Create a new connection to the target
+ connection = AsyncHTTPConnection(
+ origin=origin,
+ http2=False,
+ ssl_context=self._ssl_context,
+ socket=proxy_connection.socket,
+ )
+ async with self._thread_lock:
+ self._connections[origin].remove(proxy_connection)
+ self._connections[origin].add(connection)
 
  # Once the connection has been established we can send requests on
  # it as normal.
  response = await connection.request(
- method, url, headers=headers, stream=stream, timeout=timeout
+ method, url, headers=headers, stream=stream, timeout=timeout,
  )
  wrapped_stream = ResponseByteStream(
  response[4], connection=connection, callback=self._response_closed

httpcore/_backends/asyncio.py

@@ -107,9 +107,9 @@ async def start_tls(
 
  transport = await asyncio.wait_for(
  loop_start_tls(
- transport=transport,
- protocol=protocol,
-sslcontext=ssl_context,
+ transport,
+ protocol,
+ ssl_context,
  server_hostname=hostname.decode("ascii"),
  ),
  timeout=timeout.get("connect"),

httpcore/_sync/connection.py

@@ -1,7 +1,7 @@
 from ssl import SSLContext
 from typing import List, Optional, Tuple, Union
 
-from .._backends.auto import SyncLock, SyncBackend
+from .._backends.auto import SyncLock, SyncSocketStream, SyncBackend
 from .._types import URL, Headers, Origin, TimeoutDict
 from .base import (
  SyncByteStream,
@@ -15,7 +15,11 @@
 
 class SyncHTTPConnection(SyncHTTPTransport):
  def __init__(
- self, origin: Origin, http2: bool = False, ssl_context: SSLContext = None,
+ self,
+ origin: Origin,
+ http2: bool = False,
+ ssl_context: SSLContext = None,
+ socket: SyncSocketStream = None,
  ):
  self.origin = origin
  self.http2 = http2
@@ -30,6 +34,7 @@ def __init__(
  self.connect_failed = False
  self.expires_at: Optional[float] = None
  self.backend = SyncBackend()
+ self.socket = socket
 
  @property
  def request_lock(self) -> SyncLock:
@@ -48,14 +53,16 @@ def request(
  timeout: TimeoutDict = None,
  ) -> Tuple[bytes, int, bytes, List[Tuple[bytes, bytes]], SyncByteStream]:
  assert url[:3] == self.origin
-
  with self.request_lock:
  if self.state == ConnectionState.PENDING:
- try:
- self._connect(timeout)
- except Exception:
- self.connect_failed = True
- raise
+ if self.socket:
+ self._create_connection(self.socket)
+ else:
+ try:
+ self._connect(timeout)
+ except Exception:
+ self.connect_failed = True
+ raise
  elif self.state in (ConnectionState.READY, ConnectionState.IDLE):
  pass
  elif self.state == ConnectionState.ACTIVE and self.is_http2:
@@ -70,16 +77,23 @@ def _connect(self, timeout: TimeoutDict = None) -> None:
  scheme, hostname, port = self.origin
  timeout = {} if timeout is None else timeout
  ssl_context = self.ssl_context if scheme == b"https" else None
- socket = self.backend.open_tcp_stream(
+self.socket = self.backend.open_tcp_stream(
  hostname, port, ssl_context, timeout
  )
+ self._create_connection(self.socket)
+
+ def _create_connection(self, socket: SyncSocketStream) -> None:
  http_version = socket.get_http_version()
  if http_version == "HTTP/2":
  self.is_http2 = True
- self.connection = SyncHTTP2Connection(socket=socket, backend=self.backend)
+ self.connection = SyncHTTP2Connection(
+ socket=socket, backend=self.backend, ssl_context=self.ssl_context
+ )
  else:
  self.is_http11 = True
- self.connection = SyncHTTP11Connection(socket=socket)
+ self.connection = SyncHTTP11Connection(
+ socket=socket, ssl_context=self.ssl_context
+ )
 
  @property
  def state(self) -> ConnectionState:
@@ -99,3 +113,4 @@ def mark_as_ready(self) -> None:
  def start_tls(self, hostname: bytes, timeout: TimeoutDict = None) -> None:
  if self.connection is not None:
  self.connection.start_tls(hostname, timeout)
+ self.socket = self.connection.socket

httpcore/_sync/http11.py

@@ -123,7 +123,7 @@ def _receive_response_data(
  event = self._receive_event(timeout)
  if isinstance(event, h11.Data):
  yield bytes(event.data)
- elif isinstance(event, h11.EndOfMessage):
+ elif isinstance(event, (h11.EndOfMessage, h11.PAUSED)):
  break
 
  def _receive_event(self, timeout: TimeoutDict) -> H11Event:

httpcore/_sync/http_proxy.py

@@ -140,47 +140,57 @@ def _tunnel_request(
  connection = self._get_connection_from_pool(origin)
 
  if connection is None:
- connection = SyncHTTPConnection(
- origin=origin, http2=False, ssl_context=self._ssl_context,
+ # First, create a connection to the proxy server
+ proxy_connection = SyncHTTPConnection(
+ origin=self.proxy_origin, http2=False, ssl_context=self._ssl_context,
  )
  with self._thread_lock:
  self._connections.setdefault(origin, set())
- self._connections[origin].add(connection)
+ self._connections[origin].add(proxy_connection)
 
- # Establish the connection by issuing a CONNECT request...
+ # Issue a CONNECT request...
 
  # CONNECT www.example.org:80 HTTP/1.1
  # [proxy-headers]
  target = b"%b:%d" % (url[1], url[2])
  connect_url = self.proxy_origin + (target,)
- connect_headers = self.proxy_headers
- proxy_response = connection.request(
+ connect_headers = self.proxy_headers or [(b"host", self.proxy_origin[1])]
+ proxy_response = proxy_connection.request(
  b"CONNECT", connect_url, headers=connect_headers, timeout=timeout
  )
  proxy_status_code = proxy_response[1]
  proxy_reason_phrase = proxy_response[2]
  proxy_stream = proxy_response[4]
 
- # Ingest any request body.
-read_body(proxy_stream)
+ # Read the response data without closing the socket
+_ = b"".join([chunk for chunk in proxy_stream])
 
  # If the proxy responds with an error, then drop the connection
  # from the pool, and raise an exception.
  if proxy_status_code < 200 or proxy_status_code > 299:
- with self._thread_lock:
- self._connections[connection.origin].remove(connection)
- if not self._connections[connection.origin]:
- del self._connections[connection.origin]
  msg = "%d %s" % (proxy_status_code, proxy_reason_phrase.decode("ascii"))
  raise ProxyError(msg)
 
- # Upgrade to TLS.
- connection.start_tls(target, timeout)
+ # Upgrade to TLS if required
+ # We assume the target speaks TLS on the specified port
+ if url[0] == b"https":
+ proxy_connection.start_tls(url[1], timeout)
+
+ # Create a new connection to the target
+ connection = SyncHTTPConnection(
+ origin=origin,
+ http2=False,
+ ssl_context=self._ssl_context,
+ socket=proxy_connection.socket,
+ )
+ with self._thread_lock:
+ self._connections[origin].remove(proxy_connection)
+ self._connections[origin].add(connection)
 
  # Once the connection has been established we can send requests on
  # it as normal.
  response = connection.request(
- method, url, headers=headers, stream=stream, timeout=timeout
+ method, url, headers=headers, stream=stream, timeout=timeout,
  )
  wrapped_stream = ResponseByteStream(
  response[4], connection=connection, callback=self._response_closed

requirements.txt

@@ -18,3 +18,4 @@ autoflake
 mypy
 isort
 mitmproxy
+trustme