Skip to content

Commit

Permalink
Fix pool poisoning on overloaded systems (encode#550)
Browse files Browse the repository at this point in the history
  • Loading branch information
igor.stulikov authored and igor.stulikov committed Jan 13, 2023
1 parent 7eb2022 commit a406468
Show file tree
Hide file tree
Showing 24 changed files with 112 additions and 0 deletions.
3 changes: 3 additions & 0 deletions httpcore/_async/connection.py
Original file line number Diff line number Diff line change
Expand Up @@ -185,6 +185,9 @@ def is_closed(self) -> bool:
return self._connect_failed
return self._connection.is_closed()

def is_connecting(self) -> bool:
return self._connection is None and not self._connect_failed

def info(self) -> str:
if self._connection is None:
return "CONNECTION FAILED" if self._connect_failed else "CONNECTING"
Expand Down
19 changes: 19 additions & 0 deletions httpcore/_async/connection_pool.py
Original file line number Diff line number Diff line change
Expand Up @@ -140,6 +140,10 @@ def connections(self) -> List[AsyncConnectionInterface]:
"""
return list(self._pool)

@property
def _is_pool_full(self) -> bool:
return len(self._pool) >= self._max_connections

async def _attempt_to_acquire_connection(self, status: RequestStatus) -> bool:
"""
Attempt to provide a connection that can handle the given origin.
Expand Down Expand Up @@ -168,6 +172,21 @@ async def _attempt_to_acquire_connection(self, status: RequestStatus) -> bool:
self._pool.pop(idx)
break

# Attempt to close CONNECTING connections that no one needs
if self._is_pool_full:
for idx, connection in enumerate(self._pool): # Try to check old connections first
if not connection.is_connecting():
continue
for req_status in self._requests:
if req_status is status: # skip current request
continue
if connection.can_handle_request(req_status.request.url.origin):
break
else: # There is no requests that can be handled by this connection
await connection.aclose()
self._pool.pop(idx)
break

# If the pool is still full, then we cannot acquire a connection.
if len(self._pool) >= self._max_connections:
return False
Expand Down
3 changes: 3 additions & 0 deletions httpcore/_async/http11.py
Original file line number Diff line number Diff line change
Expand Up @@ -269,6 +269,9 @@ def is_idle(self) -> bool:
def is_closed(self) -> bool:
return self._state == HTTPConnectionState.CLOSED

def is_connecting(self) -> bool:
return self._state == HTTPConnectionState.NEW

def info(self) -> str:
origin = str(self._origin)
return (
Expand Down
3 changes: 3 additions & 0 deletions httpcore/_async/http2.py
Original file line number Diff line number Diff line change
Expand Up @@ -411,6 +411,9 @@ def is_idle(self) -> bool:
def is_closed(self) -> bool:
return self._state == HTTPConnectionState.CLOSED

def is_connecting(self) -> bool:
return False

def info(self) -> str:
origin = str(self._origin)
return (
Expand Down
6 changes: 6 additions & 0 deletions httpcore/_async/http_proxy.py
Original file line number Diff line number Diff line change
Expand Up @@ -204,6 +204,9 @@ def is_idle(self) -> bool:
def is_closed(self) -> bool:
return self._connection.is_closed()

def is_connecting(self) -> bool:
return self._connection.is_connecting()

def __repr__(self) -> str:
return f"<{self.__class__.__name__} [{self.info()}]>"

Expand Down Expand Up @@ -336,5 +339,8 @@ def is_idle(self) -> bool:
def is_closed(self) -> bool:
return self._connection.is_closed()

def is_connecting(self) -> bool:
return self._connection.is_connecting()

def __repr__(self) -> str:
return f"<{self.__class__.__name__} [{self.info()}]>"
6 changes: 6 additions & 0 deletions httpcore/_async/interfaces.py
Original file line number Diff line number Diff line change
Expand Up @@ -133,3 +133,9 @@ def is_closed(self) -> bool:
returned to the connection pool or not.
"""
raise NotImplementedError() # pragma: nocover

def is_connecting(self) -> bool:
"""
Return `True` if the connection is currently connecting. For HTTP/2 connection always returns `False`
"""
raise NotImplementedError() # pragma: nocover
3 changes: 3 additions & 0 deletions httpcore/_async/socks_proxy.py
Original file line number Diff line number Diff line change
Expand Up @@ -329,6 +329,9 @@ def is_closed(self) -> bool:
return self._connect_failed
return self._connection.is_closed()

def is_connecting(self) -> bool:
return self._connection is None and not self._connect_failed

def info(self) -> str:
if self._connection is None: # pragma: nocover
return "CONNECTION FAILED" if self._connect_failed else "CONNECTING"
Expand Down
3 changes: 3 additions & 0 deletions httpcore/_sync/connection.py
Original file line number Diff line number Diff line change
Expand Up @@ -185,6 +185,9 @@ def is_closed(self) -> bool:
return self._connect_failed
return self._connection.is_closed()

def is_connecting(self) -> bool:
return self._connection is None and not self._connect_failed

def info(self) -> str:
if self._connection is None:
return "CONNECTION FAILED" if self._connect_failed else "CONNECTING"
Expand Down
19 changes: 19 additions & 0 deletions httpcore/_sync/connection_pool.py
Original file line number Diff line number Diff line change
Expand Up @@ -140,6 +140,10 @@ def connections(self) -> List[ConnectionInterface]:
"""
return list(self._pool)

@property
def _is_pool_full(self) -> bool:
return len(self._pool) >= self._max_connections

def _attempt_to_acquire_connection(self, status: RequestStatus) -> bool:
"""
Attempt to provide a connection that can handle the given origin.
Expand Down Expand Up @@ -168,6 +172,21 @@ def _attempt_to_acquire_connection(self, status: RequestStatus) -> bool:
self._pool.pop(idx)
break

# Attempt to close CONNECTING connections that no one needs
if self._is_pool_full:
for idx, connection in enumerate(self._pool): # Try to check old connections first
if not connection.is_connecting():
continue
for req_status in self._requests:
if req_status is status: # skip current request
continue
if connection.can_handle_request(req_status.request.url.origin):
break
else: # There is no requests that can be handled by this connection
connection.close()
self._pool.pop(idx)
break

# If the pool is still full, then we cannot acquire a connection.
if len(self._pool) >= self._max_connections:
return False
Expand Down
3 changes: 3 additions & 0 deletions httpcore/_sync/http11.py
Original file line number Diff line number Diff line change
Expand Up @@ -269,6 +269,9 @@ def is_idle(self) -> bool:
def is_closed(self) -> bool:
return self._state == HTTPConnectionState.CLOSED

def is_connecting(self) -> bool:
return self._state == HTTPConnectionState.NEW

def info(self) -> str:
origin = str(self._origin)
return (
Expand Down
3 changes: 3 additions & 0 deletions httpcore/_sync/http2.py
Original file line number Diff line number Diff line change
Expand Up @@ -411,6 +411,9 @@ def is_idle(self) -> bool:
def is_closed(self) -> bool:
return self._state == HTTPConnectionState.CLOSED

def is_connecting(self) -> bool:
return False

def info(self) -> str:
origin = str(self._origin)
return (
Expand Down
6 changes: 6 additions & 0 deletions httpcore/_sync/http_proxy.py
Original file line number Diff line number Diff line change
Expand Up @@ -204,6 +204,9 @@ def is_idle(self) -> bool:
def is_closed(self) -> bool:
return self._connection.is_closed()

def is_connecting(self) -> bool:
return self._connection.is_connecting()

def __repr__(self) -> str:
return f"<{self.__class__.__name__} [{self.info()}]>"

Expand Down Expand Up @@ -336,5 +339,8 @@ def is_idle(self) -> bool:
def is_closed(self) -> bool:
return self._connection.is_closed()

def is_connecting(self) -> bool:
return self._connection.is_connecting()

def __repr__(self) -> str:
return f"<{self.__class__.__name__} [{self.info()}]>"
6 changes: 6 additions & 0 deletions httpcore/_sync/interfaces.py
Original file line number Diff line number Diff line change
Expand Up @@ -133,3 +133,9 @@ def is_closed(self) -> bool:
returned to the connection pool or not.
"""
raise NotImplementedError() # pragma: nocover

def is_connecting(self) -> bool:
"""
Return `True` if the connection is currently connecting. For HTTP/2 connection always returns `False`
"""
raise NotImplementedError() # pragma: nocover
3 changes: 3 additions & 0 deletions httpcore/_sync/socks_proxy.py
Original file line number Diff line number Diff line change
Expand Up @@ -329,6 +329,9 @@ def is_closed(self) -> bool:
return self._connect_failed
return self._connection.is_closed()

def is_connecting(self) -> bool:
return self._connection is None and not self._connect_failed

def info(self) -> str:
if self._connection is None: # pragma: nocover
return "CONNECTION FAILED" if self._connect_failed else "CONNECTING"
Expand Down
2 changes: 2 additions & 0 deletions tests/_async/test_connection.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@ async def test_http_connection():
assert not conn.is_closed()
assert not conn.is_available()
assert not conn.has_expired()
assert conn.is_connecting()
assert repr(conn) == "<AsyncHTTPConnection [CONNECTING]>"

async with conn.stream("GET", "https://example.com/") as response:
Expand All @@ -45,6 +46,7 @@ async def test_http_connection():
assert not conn.is_closed()
assert conn.is_available()
assert not conn.has_expired()
assert not conn.is_connecting()
assert (
repr(conn)
== "<AsyncHTTPConnection ['https://example.com:443', HTTP/1.1, IDLE, Request Count: 1]>"
Expand Down
5 changes: 5 additions & 0 deletions tests/_async/test_http11.py
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,7 @@ async def test_http11_connection():
assert not conn.is_closed()
assert conn.is_available()
assert not conn.has_expired()
assert not conn.is_connecting()
assert (
repr(conn)
== "<AsyncHTTP11Connection ['https://example.com:443', IDLE, Request Count: 1]>"
Expand Down Expand Up @@ -63,6 +64,7 @@ async def test_http11_connection_unread_response():
assert conn.is_closed()
assert not conn.is_available()
assert not conn.has_expired()
assert not conn.is_connecting()
assert (
repr(conn)
== "<AsyncHTTP11Connection ['https://example.com:443', CLOSED, Request Count: 1]>"
Expand All @@ -85,6 +87,7 @@ async def test_http11_connection_with_remote_protocol_error():
assert conn.is_closed()
assert not conn.is_available()
assert not conn.has_expired()
assert not conn.is_connecting()
assert (
repr(conn)
== "<AsyncHTTP11Connection ['https://example.com:443', CLOSED, Request Count: 1]>"
Expand Down Expand Up @@ -114,6 +117,7 @@ async def test_http11_connection_with_incomplete_response():
assert conn.is_closed()
assert not conn.is_available()
assert not conn.has_expired()
assert not conn.is_connecting()
assert (
repr(conn)
== "<AsyncHTTP11Connection ['https://example.com:443', CLOSED, Request Count: 1]>"
Expand Down Expand Up @@ -146,6 +150,7 @@ async def test_http11_connection_with_local_protocol_error():
assert conn.is_closed()
assert not conn.is_available()
assert not conn.has_expired()
assert not conn.is_connecting()
assert (
repr(conn)
== "<AsyncHTTP11Connection ['https://example.com:443', CLOSED, Request Count: 1]>"
Expand Down
1 change: 1 addition & 0 deletions tests/_async/test_http2.py
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,7 @@ async def test_http2_connection():
assert conn.is_available()
assert not conn.is_closed()
assert not conn.has_expired()
assert not conn.is_connecting()
assert (
conn.info() == "'https://example.com:443', HTTP/2, IDLE, Request Count: 1"
)
Expand Down
3 changes: 3 additions & 0 deletions tests/_async/test_http_proxy.py
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,7 @@ async def test_proxy_forwarding():
assert proxy.connections[0].is_idle()
assert proxy.connections[0].is_available()
assert not proxy.connections[0].is_closed()
assert not proxy.connections[0].is_connecting()

# A connection on a forwarding proxy can only handle HTTP requests to the same origin.
assert proxy.connections[0].can_handle_request(
Expand Down Expand Up @@ -102,6 +103,7 @@ async def test_proxy_tunneling():
assert proxy.connections[0].is_idle()
assert proxy.connections[0].is_available()
assert not proxy.connections[0].is_closed()
assert not proxy.connections[0].is_connecting()

# A connection on a tunneled proxy can only handle HTTPS requests to the same origin.
assert not proxy.connections[0].can_handle_request(
Expand Down Expand Up @@ -193,6 +195,7 @@ async def test_proxy_tunneling_http2():
assert proxy.connections[0].is_idle()
assert proxy.connections[0].is_available()
assert not proxy.connections[0].is_closed()
assert not proxy.connections[0].is_connecting()

# A connection on a tunneled proxy can only handle HTTPS requests to the same origin.
assert not proxy.connections[0].can_handle_request(
Expand Down
2 changes: 2 additions & 0 deletions tests/_async/test_socks_proxy.py
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,7 @@ async def test_socks5_request():
assert proxy.connections[0].is_idle()
assert proxy.connections[0].is_available()
assert not proxy.connections[0].is_closed()
assert not proxy.connections[0].is_connecting()

# A connection on a tunneled proxy can only handle HTTPS requests to the same origin.
assert not proxy.connections[0].can_handle_request(
Expand Down Expand Up @@ -107,6 +108,7 @@ async def test_authenticated_socks5_request():
assert proxy.connections[0].is_idle()
assert proxy.connections[0].is_available()
assert not proxy.connections[0].is_closed()
assert not proxy.connections[0].is_connecting()


@pytest.mark.anyio
Expand Down
2 changes: 2 additions & 0 deletions tests/_sync/test_connection.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@ def test_http_connection():
assert not conn.is_closed()
assert not conn.is_available()
assert not conn.has_expired()
assert conn.is_connecting()
assert repr(conn) == "<HTTPConnection [CONNECTING]>"

with conn.stream("GET", "https://example.com/") as response:
Expand All @@ -45,6 +46,7 @@ def test_http_connection():
assert not conn.is_closed()
assert conn.is_available()
assert not conn.has_expired()
assert not conn.is_connecting()
assert (
repr(conn)
== "<HTTPConnection ['https://example.com:443', HTTP/1.1, IDLE, Request Count: 1]>"
Expand Down
5 changes: 5 additions & 0 deletions tests/_sync/test_http11.py
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,7 @@ def test_http11_connection():
assert not conn.is_closed()
assert conn.is_available()
assert not conn.has_expired()
assert not conn.is_connecting()
assert (
repr(conn)
== "<HTTP11Connection ['https://example.com:443', IDLE, Request Count: 1]>"
Expand Down Expand Up @@ -63,6 +64,7 @@ def test_http11_connection_unread_response():
assert conn.is_closed()
assert not conn.is_available()
assert not conn.has_expired()
assert not conn.is_connecting()
assert (
repr(conn)
== "<HTTP11Connection ['https://example.com:443', CLOSED, Request Count: 1]>"
Expand All @@ -85,6 +87,7 @@ def test_http11_connection_with_remote_protocol_error():
assert conn.is_closed()
assert not conn.is_available()
assert not conn.has_expired()
assert not conn.is_connecting()
assert (
repr(conn)
== "<HTTP11Connection ['https://example.com:443', CLOSED, Request Count: 1]>"
Expand Down Expand Up @@ -114,6 +117,7 @@ def test_http11_connection_with_incomplete_response():
assert conn.is_closed()
assert not conn.is_available()
assert not conn.has_expired()
assert not conn.is_connecting()
assert (
repr(conn)
== "<HTTP11Connection ['https://example.com:443', CLOSED, Request Count: 1]>"
Expand Down Expand Up @@ -146,6 +150,7 @@ def test_http11_connection_with_local_protocol_error():
assert conn.is_closed()
assert not conn.is_available()
assert not conn.has_expired()
assert not conn.is_connecting()
assert (
repr(conn)
== "<HTTP11Connection ['https://example.com:443', CLOSED, Request Count: 1]>"
Expand Down
1 change: 1 addition & 0 deletions tests/_sync/test_http2.py
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,7 @@ def test_http2_connection():
assert conn.is_available()
assert not conn.is_closed()
assert not conn.has_expired()
assert not conn.is_connecting()
assert (
conn.info() == "'https://example.com:443', HTTP/2, IDLE, Request Count: 1"
)
Expand Down
3 changes: 3 additions & 0 deletions tests/_sync/test_http_proxy.py
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,7 @@ def test_proxy_forwarding():
assert proxy.connections[0].is_idle()
assert proxy.connections[0].is_available()
assert not proxy.connections[0].is_closed()
assert not proxy.connections[0].is_connecting()

# A connection on a forwarding proxy can only handle HTTP requests to the same origin.
assert proxy.connections[0].can_handle_request(
Expand Down Expand Up @@ -102,6 +103,7 @@ def test_proxy_tunneling():
assert proxy.connections[0].is_idle()
assert proxy.connections[0].is_available()
assert not proxy.connections[0].is_closed()
assert not proxy.connections[0].is_connecting()

# A connection on a tunneled proxy can only handle HTTPS requests to the same origin.
assert not proxy.connections[0].can_handle_request(
Expand Down Expand Up @@ -193,6 +195,7 @@ def test_proxy_tunneling_http2():
assert proxy.connections[0].is_idle()
assert proxy.connections[0].is_available()
assert not proxy.connections[0].is_closed()
assert not proxy.connections[0].is_connecting()

# A connection on a tunneled proxy can only handle HTTPS requests to the same origin.
assert not proxy.connections[0].can_handle_request(
Expand Down
Loading

0 comments on commit a406468

Please sign in to comment.