mirror of
https://github.com/unshackle-dl/unshackle.git
synced 2026-11-14 13:05:09 +00:00
feat(tunnel): free finished tunnels at once and wait out 429s on --remote
A tunnel closed with code 1000 is freed without the 60s grace, a shared: true API key counts tunnels per X-Client-Id, and --remote waits out Retry-After and a full tunnel cap with a countdown; responses carry X-RateLimit-* headers.
This commit is contained in:
@@ -252,9 +252,13 @@ async def test_rate_limited_request_returns_429(aiohttp_client, monkeypatch) ->
|
||||
setup_routes(app, remote_only=True, dashboard=True)
|
||||
client = await aiohttp_client(app)
|
||||
|
||||
assert (await client.get("/api/services", headers={"X-Secret-Key": "tier-key"})).status == 200
|
||||
resp = await client.get("/api/services", headers={"X-Secret-Key": "tier-key"})
|
||||
assert resp.status == 200
|
||||
assert (resp.headers["X-RateLimit-Limit"], resp.headers["X-RateLimit-Remaining"]) == ("1", "0")
|
||||
assert 0 < int(resp.headers["X-RateLimit-Reset"]) <= 3600
|
||||
resp = await client.get("/api/services", headers={"X-Secret-Key": "tier-key"})
|
||||
assert resp.status == 429 and int(resp.headers["Retry-After"]) > 0
|
||||
assert resp.headers["X-RateLimit-Remaining"] == "0"
|
||||
assert (await client.get("/api/health", headers={"X-Secret-Key": "tier-key"})).status == 200
|
||||
assert (await client.get("/api/health")).status == 401
|
||||
# The dashboard key leaves the middleware before the limit check, so it always answers.
|
||||
|
||||
@@ -4,6 +4,7 @@ from __future__ import annotations
|
||||
|
||||
import json
|
||||
|
||||
import click
|
||||
import pytest
|
||||
import responses
|
||||
|
||||
@@ -244,3 +245,93 @@ def test_proxy_error_page_logs_its_title(client: RemoteClient, caplog: pytest.Lo
|
||||
client.get("/api/session/s/titles")
|
||||
assert "HTTP 524: example.com | 524: A timeout occurred" in caplog.text
|
||||
assert "DOCTYPE" not in caplog.text
|
||||
|
||||
|
||||
@responses.activate
|
||||
def test_a_429_with_retry_after_waits_then_retries(client: RemoteClient, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
waits: list[int] = []
|
||||
monkeypatch.setattr("unshackle.core.remote_service.wait_retry_after", waits.append)
|
||||
url = "http://srv:8786/api/services"
|
||||
responses.add(responses.GET, url, status=429, json={"status": 429}, headers={"Retry-After": "7"})
|
||||
responses.add(responses.GET, url, json={"status": "ok"})
|
||||
|
||||
assert client.get("/api/services") == {"status": "ok"}
|
||||
assert waits == [7]
|
||||
|
||||
|
||||
@responses.activate
|
||||
def test_a_429_without_retry_after_fails_at_once(client: RemoteClient, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr("unshackle.core.remote_service.wait_retry_after", pytest.fail)
|
||||
responses.add(responses.GET, "http://srv:8786/api/services", status=429, json={"error_code": "RATE_LIMITED"})
|
||||
|
||||
with pytest.raises(SystemExit):
|
||||
client.get("/api/services")
|
||||
|
||||
|
||||
def test_wait_retry_after_waits_out_its_own_deadline(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
from unshackle.core import remote_service
|
||||
|
||||
clock = [0.0]
|
||||
monkeypatch.setattr(remote_service.time, "monotonic", lambda: clock[0])
|
||||
monkeypatch.setattr(remote_service.time, "sleep", lambda s: clock.__setitem__(0, clock[0] + s))
|
||||
remote_service.wait_retry_after(3)
|
||||
assert clock[0] == 3.0
|
||||
remote_service.wait_retry_after(99999)
|
||||
assert clock[0] == 3.0 + remote_service.RETRY_AFTER_MAX
|
||||
|
||||
|
||||
def test_a_rate_limited_tunnel_waits_then_opens(client: RemoteClient, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
from unshackle.core import remote_service
|
||||
|
||||
attempts: list[object] = []
|
||||
|
||||
class Refused:
|
||||
def __init__(self, *_: object) -> None:
|
||||
self.retry_after = 5 if not attempts else None
|
||||
self.at_limit = len(attempts) == 1
|
||||
attempts.append(self)
|
||||
|
||||
def start(self) -> str:
|
||||
if self.retry_after or self.at_limit:
|
||||
raise ConnectionError("HTTP 429")
|
||||
return "id"
|
||||
|
||||
waits: list[int] = []
|
||||
monkeypatch.setattr(remote_service, "TunnelClient", Refused)
|
||||
monkeypatch.setattr(remote_service, "wait_retry_after", lambda s, *_: waits.append(s))
|
||||
assert remote_service.start_tunnel(client) is attempts[2]
|
||||
assert waits == [5, remote_service.TUNNEL_SLOT_POLL]
|
||||
|
||||
|
||||
def test_a_tunnel_refused_for_another_reason_fails_at_once(
|
||||
client: RemoteClient, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
from unshackle.core import remote_service
|
||||
|
||||
class Refused:
|
||||
retry_after = None
|
||||
at_limit = False
|
||||
|
||||
def __init__(self, *_: object) -> None:
|
||||
pass
|
||||
|
||||
def start(self) -> str:
|
||||
raise ConnectionError("HTTP 403")
|
||||
|
||||
monkeypatch.setattr(remote_service, "TunnelClient", Refused)
|
||||
monkeypatch.setattr(remote_service, "wait_retry_after", pytest.fail)
|
||||
with pytest.raises(click.ClickException, match="HTTP 403"):
|
||||
remote_service.start_tunnel(client)
|
||||
|
||||
|
||||
@responses.activate
|
||||
def test_the_429_wait_gives_up_after_its_tries(client: RemoteClient, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
from unshackle.core import remote_service
|
||||
|
||||
waits: list[int] = []
|
||||
monkeypatch.setattr(remote_service, "wait_retry_after", waits.append)
|
||||
responses.add(responses.GET, "http://srv:8786/api/services", status=429, json={}, headers={"Retry-After": "0"})
|
||||
|
||||
with pytest.raises(SystemExit):
|
||||
client.get("/api/services")
|
||||
assert waits == [1] * remote_service.RATE_LIMIT_TRIES
|
||||
|
||||
@@ -16,6 +16,7 @@ import pytest
|
||||
import requests
|
||||
import responses
|
||||
from aiohttp import web
|
||||
from multidict import CIMultiDict
|
||||
|
||||
from unshackle.core import tunnel
|
||||
from unshackle.core.api.routes import tunnel as tunnel_route
|
||||
@@ -356,6 +357,7 @@ def test_a_malformed_frame_still_leaves_the_tunnel_to_expire(server: str) -> Non
|
||||
async with http.ws_connect(server + ROUTE) as ws:
|
||||
tunnel_id = (await ws.receive_json())["id"]
|
||||
await ws.send_bytes(b"\x01")
|
||||
await ws.close(code=aiohttp.WSCloseCode.GOING_AWAY)
|
||||
return str(tunnel_id)
|
||||
|
||||
tunnel_id = asyncio.run(send_short_frame())
|
||||
@@ -369,12 +371,112 @@ def test_an_api_key_has_a_tunnel_limit(server: str) -> None:
|
||||
clients = [TunnelClient(server, {"X-Secret-Key": "bob"}) for _ in range(tunnel.MAX_TUNNELS)]
|
||||
for client in clients:
|
||||
client.start()
|
||||
with pytest.raises(ConnectionError, match="429"):
|
||||
with pytest.raises(ConnectionError, match="429, the API key already has the maximum number of tunnels open"):
|
||||
TunnelClient(server, {"X-Secret-Key": "bob"}).start()
|
||||
for client in clients:
|
||||
client.stop()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("shared", [True, False])
|
||||
def test_a_shared_api_key_counts_tunnels_for_each_client(
|
||||
server: str, monkeypatch: pytest.MonkeyPatch, shared: bool
|
||||
) -> None:
|
||||
from unshackle.core.config import config
|
||||
|
||||
monkeypatch.setitem(config.serve, "users", {"tier": {"shared": True} if shared else {}})
|
||||
first = [TunnelClient(server, {"X-Secret-Key": "tier", "X-Client-Id": "u1"}) for _ in range(tunnel.MAX_TUNNELS)]
|
||||
second = TunnelClient(server, {"X-Secret-Key": "tier", "X-Client-Id": "u2"})
|
||||
try:
|
||||
for client in first:
|
||||
client.start()
|
||||
if shared:
|
||||
second.start()
|
||||
else:
|
||||
with pytest.raises(ConnectionError, match="maximum number of tunnels"):
|
||||
second.start()
|
||||
with pytest.raises(ConnectionError, match="maximum number of tunnels"):
|
||||
TunnelClient(server, {"X-Secret-Key": "tier", "X-Client-Id": "u1"}).start()
|
||||
finally:
|
||||
for client in [*first, second]:
|
||||
client.stop()
|
||||
|
||||
|
||||
def test_another_client_of_a_shared_api_key_cannot_reattach(server: str, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
from unshackle.core.config import config
|
||||
|
||||
monkeypatch.setitem(config.serve, "users", {"tier": {"shared": True}})
|
||||
owner = TunnelClient(server, {"X-Secret-Key": "tier", "X-Client-Id": "u1"})
|
||||
tunnel_id = owner.start()
|
||||
try:
|
||||
resp = requests.get(
|
||||
f"{server}{ROUTE}?id={tunnel_id}", headers={"X-Secret-Key": "tier", "X-Client-Id": "u2"}, timeout=10
|
||||
)
|
||||
assert resp.status_code == 404
|
||||
finally:
|
||||
owner.stop()
|
||||
|
||||
|
||||
def test_stopping_one_client_leaves_the_other_tunnels_of_its_api_key_open(
|
||||
server: str, target: int, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setattr(tunnel, "is_public_ip", lambda address: True)
|
||||
monkeypatch.setattr(tunnel, "dialled_itself", lambda writer: False)
|
||||
clients = [TunnelClient(server, {"X-Secret-Key": "erin"}) for _ in range(3)]
|
||||
stopped, *running = [client.start() for client in clients]
|
||||
try:
|
||||
clients[0].stop()
|
||||
assert wait_until(lambda: stopped not in TUNNELS, 5)
|
||||
for tunnel_id in running:
|
||||
proxy = tunnel_proxy(tunnel_id, "erin")
|
||||
assert proxy
|
||||
resp = requests.get(f"http://127.0.0.1:{target}/", proxies={"all": proxy}, timeout=10)
|
||||
assert resp.text == "hello through the tunnel"
|
||||
finally:
|
||||
for client in clients[1:]:
|
||||
client.stop()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("value", "seconds"),
|
||||
[
|
||||
("5", 5),
|
||||
("0", 1),
|
||||
(None, None),
|
||||
("", None),
|
||||
("\u00b2", None),
|
||||
("9" * 5000, None),
|
||||
("Wed, 21 Oct 2026 07:28:00 GMT", None),
|
||||
],
|
||||
ids=["seconds", "zero", "absent", "empty", "superscript", "huge", "http-date"],
|
||||
)
|
||||
def test_parse_retry_after(value: str | None, seconds: int | None) -> None:
|
||||
assert tunnel.parse_retry_after(value) == seconds
|
||||
|
||||
|
||||
def test_a_client_that_exits_frees_its_tunnel_at_once(server: str) -> None:
|
||||
async def open_then_leave() -> str:
|
||||
async with aiohttp.ClientSession(headers={"X-Secret-Key": "dave"}) as http:
|
||||
async with http.ws_connect(server + ROUTE) as ws:
|
||||
return str((await ws.receive_json())["id"])
|
||||
|
||||
tunnel_id = asyncio.run(open_then_leave())
|
||||
assert wait_until(lambda: tunnel_id not in TUNNELS, 5)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("status", "headers", "reason"),
|
||||
[
|
||||
(429, {"Retry-After": "120"}, "HTTP 429, the API key is over its hourly request limit, retry in 120s"),
|
||||
(429, {"retry-after": "120"}, "HTTP 429, the API key is over its hourly request limit, retry in 120s"),
|
||||
(429, {}, "HTTP 429, the API key already has the maximum number of tunnels open"),
|
||||
(404, None, "HTTP 404"),
|
||||
],
|
||||
)
|
||||
def test_a_refused_tunnel_names_the_reason(status: int, headers: dict[str, str] | None, reason: str) -> None:
|
||||
error = aiohttp.WSServerHandshakeError(None, (), status=status, headers=CIMultiDict(headers or {})) # type: ignore[arg-type]
|
||||
assert tunnel.refusal(error) == reason
|
||||
|
||||
|
||||
def test_a_connection_that_stops_reading_is_dropped() -> None:
|
||||
sent: list = []
|
||||
|
||||
|
||||
Reference in New Issue
Block a user