diff --git a/docs/dev/rest-api/authentication.md b/docs/dev/rest-api/authentication.md index 09313e53..a96c8c4f 100644 --- a/docs/dev/rest-api/authentication.md +++ b/docs/dev/rest-api/authentication.md @@ -224,6 +224,19 @@ A limit comes from `serve.tiers..rate_limit` or from the API key's own [dashboard rate limits](dashboard.md#rate-limits) for the config and the counters, and [services configuration](../../reference/configuration/services.md) for the keys themselves. +Every response to an API key that has a limit carries these headers, the `429` too. A +WebSocket or event stream response sends its headers before the handler returns, so it has none. +The CORS headers expose these headers and `Retry-After` to a browser client. + +| Header | Value | +|---|---| +| `X-RateLimit-Limit` | Requests per hour that the API key may make. | +| `X-RateLimit-Remaining` | Requests left in the current window. | +| `X-RateLimit-Reset` | Seconds until the window resets. | + +A `--remote` client that gets a `429` with `Retry-After` shows a countdown, waits that time +(one hour at most), and then sends the request again. It stops after 3 waits for one request. + !!! warning "`429` is not the `RATE_LIMITED` error body" `RATE_LIMITED` in [Errors](errors.md) is a *service* rate limit that a handler raises, and it arrives in the standard envelope with an `error_code`. The `429` here comes from the @@ -311,6 +324,7 @@ With `--caddy`, unshackle also starts `caddy run` using the `Caddyfile` located | Missing-header body | `{"status": 401, "message": "Secret Key is Empty."}` | | Invalid-key body | `{"status": 401, "message": "Secret Key is Invalid."}` | | Over-limit body | `{"status": 429, "message": "Rate limit exceeded."}`, plus a `Retry-After` header | +| Rate limit headers | `X-RateLimit-Limit`, `X-RateLimit-Remaining`, `X-RateLimit-Reset` on each response to a limited API key | | Rate limit config | `serve.tiers..rate_limit`, overridden by `serve.users..rate_limit` | | Disable auth | `--no-key` (all requests allowed) | | Config location | `serve.api_secret` and `serve.users` in `unshackle.yaml` | diff --git a/docs/dev/rest-api/remote-sessions.md b/docs/dev/rest-api/remote-sessions.md index 16971ede..98179646 100644 --- a/docs/dev/rest-api/remote-sessions.md +++ b/docs/dev/rest-api/remote-sessions.md @@ -448,7 +448,7 @@ The server mounts all these routes, even in `--remote-only` mode. Paths use the | Method | Path | Purpose | |---|---|---| | `POST` | `/api/session/create` | Build + authenticate a service instance; returns a session ID immediately | -| `GET` | `/api/tunnel` | WebSocket that relays the connections of a `tunnel:` proxy to the client; `?id=` reconnects | +| `GET` | `/api/tunnel` | WebSocket that relays the connections of a `tunnel:` proxy to the client; `?id=` reconnects, close code 1000 frees the tunnel at once | | `GET` | `/api/tunnel/check?id=` | Exit IP address and lookup time through the caller's tunnel, for `unshackle tunnel-test` | | `GET` | `/api/session/{id}/titles` | Fetch and cache the title list | | `POST` | `/api/session/{id}/tracks` | Resolve tracks, manifests, chapters, headers/cookies for a title | diff --git a/docs/guide/proxies-and-vpn.md b/docs/guide/proxies-and-vpn.md index ce23b136..74bd3072 100644 --- a/docs/guide/proxies-and-vpn.md +++ b/docs/guide/proxies-and-vpn.md @@ -336,8 +336,10 @@ A service that uses a server account does not work with `--local-proxy`, because logs in with its own account. If the tunnel drops, the client connects again, and a new request from the server waits up to 10 seconds for it. A request that was in progress when the tunnel dropped fails. The server closes a tunnel when its client has not connected for 60 seconds, and -the client then stops with `The tunnel to the server closed`. An API key can have 4 tunnels -open at once. +the client then stops with `The tunnel to the server closed`. A client that exits closes its +tunnel at once. An API key can have 4 tunnels open at once (4 for each person on a key with +`shared: true`, see [services configuration](../reference/configuration/services.md)). A fifth run waits for a free +tunnel and tries again every 15 seconds, for one hour at most. ## Basic (static proxies) diff --git a/docs/reference/configuration/services.md b/docs/reference/configuration/services.md index 9d6656df..182cd10b 100644 --- a/docs/reference/configuration/services.md +++ b/docs/reference/configuration/services.md @@ -240,10 +240,17 @@ server vault with this grant, are the only source, and a missing content key sto `admin` is a boolean that lets the API key run the maintenance endpoints (clear-cache, clear-temp, refresh-services). It is `false` unless the entry sets it. Keys that have no `users` entry, such as `api_secret`, keep that access. +`shared` is a boolean for an API key that a gateway in front of the server gives to many people. +The gateway puts its own ID for each person in the `X-Client-Id` header, and the server then +counts the `--local-proxy` tunnels of that API key for each person. Without `shared: true` the +server ignores the header, because a client that connects to the server directly can set any +value. Set it only when the gateway always replaces the header that the client sends. + A `tier` names an entry under `serve.tiers`, which holds the settings that several API keys share. Today a tier carries `rate_limit`, the requests per hour that API key may make; an API key can also set its own `rate_limit`, which wins over its tier's. An API key with neither has -no limit. Going over the limit answers `429` with a `Retry-After` header. The server keeps a +no limit. Going over the limit answers `429` with a `Retry-After` header, and a `--remote` +client waits that time before it tries again. The server keeps a fixed window rather than a sliding one: it opens on the first request and resets an hour later. A `tier` that names no entry under `serve.tiers`, or a `rate_limit` that is not a positive whole number, stops the server at startup, because the alternative is an API key that @@ -341,6 +348,8 @@ serve: f3a4b5c6: services: [EXAMPLE1] server_vault: [EXAMPLE1] # the client's device licenses; the server vault still serves keys + d7e8f9a0: + shared: true # a gateway gives this key to many people; tunnels count per X-Client-Id ``` !!! note "`dl` keys inside `serve`" diff --git a/tests/remote/unit/test_dashboard.py b/tests/remote/unit/test_dashboard.py index 548a7129..e4d25d47 100644 --- a/tests/remote/unit/test_dashboard.py +++ b/tests/remote/unit/test_dashboard.py @@ -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. diff --git a/tests/remote/unit/test_remote_client.py b/tests/remote/unit/test_remote_client.py index 80e8fa4d..c54bf2f7 100644 --- a/tests/remote/unit/test_remote_client.py +++ b/tests/remote/unit/test_remote_client.py @@ -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 diff --git a/tests/remote/unit/test_tunnel.py b/tests/remote/unit/test_tunnel.py index db0d2196..ec4acf1f 100644 --- a/tests/remote/unit/test_tunnel.py +++ b/tests/remote/unit/test_tunnel.py @@ -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 = [] diff --git a/unshackle/commands/serve.py b/unshackle/commands/serve.py index cf0b5ccc..a8d1ba3f 100644 --- a/unshackle/commands/serve.py +++ b/unshackle/commands/serve.py @@ -21,6 +21,7 @@ from unshackle.core.api.compression import compression_middleware from unshackle.core.api.events import publish_refresh_events, publish_service_event from unshackle.core.api.handlers import ( DASHBOARD_PREFIX, + add_rate_limit_headers, api_key_authentication, dashboard_authentication, dashboard_key, @@ -559,6 +560,8 @@ def serve( f"https://git.gay/ready-dl/pyplayready serve v{pyplayready_version}" ) + if request.path != "/api/health" and secret_key and known: + add_rate_limit_headers(response, secret_key) return response return serve_authentication diff --git a/unshackle/core/api/handlers.py b/unshackle/core/api/handlers.py index 668dd61e..f99bf4f2 100644 --- a/unshackle/core/api/handlers.py +++ b/unshackle/core/api/handlers.py @@ -1354,6 +1354,7 @@ CORS_HEADERS = { "Access-Control-Allow-Methods": "GET, POST, PUT, DELETE, OPTIONS", "Access-Control-Allow-Headers": "Content-Type, X-Secret-Key, Authorization", "Access-Control-Max-Age": "3600", + "Access-Control-Expose-Headers": "Retry-After, X-RateLimit-Limit, X-RateLimit-Remaining, X-RateLimit-Reset", } @@ -1408,10 +1409,19 @@ def rate_limit_response(secret_key: str) -> Optional[web.Response]: return web.json_response( {"status": 429, "message": "Rate limit exceeded."}, status=429, - headers={"Retry-After": str(retry_after)}, + headers={"Retry-After": str(retry_after), **stats.rate_limit_headers(secret_key)}, ) +def add_rate_limit_headers(response: web.StreamResponse, secret_key: str) -> web.StreamResponse: + """Put *secret_key*'s rate-limit headers on *response*; a streamed response has sent its headers already.""" + from unshackle.core.api.stats import stats + + if not response.prepared: + response.headers.update(stats.rate_limit_headers(secret_key)) + return response + + @web.middleware async def api_key_authentication(request: web.Request, handler: Any) -> web.StreamResponse: """Gate every non-dashboard route behind an API key in ``app["config"]["users"]``, then rate limit it. @@ -1433,7 +1443,7 @@ async def api_key_authentication(request: web.Request, handler: Any) -> web.Stre return web.json_response({"status": 401, "message": "Secret Key is Invalid."}, status=401) if request.path == "/api/health": return await handler(request) - return rate_limit_response(secret_key) or await handler(request) + return rate_limit_response(secret_key) or add_rate_limit_headers(await handler(request), secret_key) def caller_key(request: Optional[web.Request] = None) -> str: @@ -5649,19 +5659,30 @@ async def session_delete_handler(session_id: str, request: Optional[web.Request] return web.json_response(response) +def shared_key_client(request: web.Request) -> Optional[str]: + """The person behind a ``shared: true`` API key, from the ``X-Client-Id`` its gateway sets; None for other keys. + + Only a shared key reads the header, so a client that talks to the server directly cannot use it to pass a limit. + """ + if serve_user_config(request_secret_key(request) or "").get("shared") is not True: + return None + return request.headers.get("X-Client-Id", "")[:64] or None + + async def tunnel_handler(request: web.Request) -> web.StreamResponse: """Open a ``--local-proxy`` tunnel, or reattach to the caller's tunnel named by ``?id=``.""" owner = request_secret_key(request) + client = shared_key_client(request) tunnel_id = request.query.get("id") tunnel = TUNNELS.get(tunnel_id) if tunnel_id else None - if tunnel_id and (not tunnel or tunnel.owner != owner): + if tunnel_id and (not tunnel or tunnel.owner != owner or tunnel.client != client): raise APIError(APIErrorCode.NOT_FOUND, "No open tunnel has this id.") - if not tunnel and sum(t.owner == owner for t in TUNNELS.values()) >= MAX_TUNNELS: + if not tunnel and sum(t.owner == owner and t.client == client for t in TUNNELS.values()) >= MAX_TUNNELS: raise APIError(APIErrorCode.RATE_LIMITED, f"An API key can have {MAX_TUNNELS} tunnels open at once.") ws = web.WebSocketResponse(heartbeat=HEARTBEAT, max_msg_size=MAX_MESSAGE) await ws.prepare(request) # a request that is not a WebSocket upgrade fails here, before a tunnel exists if not tunnel: - tunnel = Tunnel(owner) + tunnel = Tunnel(owner, client) await tunnel.start() await tunnel.serve(ws) return ws diff --git a/unshackle/core/api/routes.py b/unshackle/core/api/routes.py index 52de1a3c..718892b5 100644 --- a/unshackle/core/api/routes.py +++ b/unshackle/core/api/routes.py @@ -2074,7 +2074,9 @@ async def tunnel(request: web.Request) -> web.StreamResponse: description: >- A WebSocket that a --remote client holds open. The server sends the requests of a remote session created with proxy tunnel: through it, and the client makes them from its own connection. - The first message is {"id": }. + The first message is {"id": }. A client that has finished closes with code 1000, + which frees the tunnel at once; after any other close the tunnel waits 60 seconds for a reconnect. + An API key with shared: true counts its tunnels for each X-Client-Id value, which its gateway sets. parameters: - name: id in: query @@ -2088,7 +2090,9 @@ async def tunnel(request: web.Request) -> web.StreamResponse: '404': description: No open tunnel with this id for this API key '429': - description: The API key has the maximum number of tunnels open + description: >- + The API key has the maximum number of tunnels open, or, with a Retry-After header, + the API key is over its hourly request limit """ return await tunnel_handler(request) diff --git a/unshackle/core/api/stats.py b/unshackle/core/api/stats.py index 1d09ce19..53136579 100644 --- a/unshackle/core/api/stats.py +++ b/unshackle/core/api/stats.py @@ -183,6 +183,18 @@ class ServerStats: entry.window_used += 1 return None + def rate_limit_headers(self, key: str) -> Dict[str, str]: + """The ``X-RateLimit-*`` headers for ``key``'s current window; none for an API key with no limit.""" + limit = key_rate_limit(key) + if limit is None: + return {} + entry = self.key_stats(key) + return { + "X-RateLimit-Limit": str(limit), + "X-RateLimit-Remaining": str(max(0, limit - entry.window_used)), + "X-RateLimit-Reset": str(max(0, int(entry.window_start + RATE_LIMIT_WINDOW - time.time()))), + } + def to_dict(self) -> Dict[str, Any]: return { "version": __version__, diff --git a/unshackle/core/remote_service.py b/unshackle/core/remote_service.py index c21910d3..1bf8d984 100644 --- a/unshackle/core/remote_service.py +++ b/unshackle/core/remote_service.py @@ -13,6 +13,7 @@ import contextlib import hashlib import json import logging +import math import random import re import shlex @@ -35,9 +36,10 @@ from langcodes import Language from requests.adapters import HTTPAdapter, Retry from rich.padding import Padding from rich.rule import Rule +from rich.spinner import Spinner from unshackle.core.config import config -from unshackle.core.console import console, prompt_user, show_notice +from unshackle.core.console import SyncLive, console, prompt_user, show_notice from unshackle.core.constants import AnyTrack from unshackle.core.credential import Credential from unshackle.core.drm import BlowfishStripe, cbcs, real_kids @@ -52,7 +54,7 @@ from unshackle.core.titles.music import SONG_FIELDS, Album, Song from unshackle.core.tracks import Audio, Chapter, Chapters, Subtitle, Tracks, Video from unshackle.core.tracks.attachment import Attachment from unshackle.core.tracks.track import Track -from unshackle.core.tunnel import TunnelClient +from unshackle.core.tunnel import TunnelClient, parse_retry_after from unshackle.core.utils import run_proxy from unshackle.core.utils.collections import ci_get from unshackle.core.utils.redact import redact_path, redact_text, safe_display_url @@ -69,6 +71,10 @@ DEFAULT_AUTH_HEADERS = ["X-Secret-Key", "X-Api-Key"] SESSION_LOST = ("SESSION_NOT_FOUND", "CONNECT", "TIMEOUT", "TUNNEL_CLOSED") # what a recovering call can meet on its way, so these log only once the recovery fails RECOVERABLE = (*SESSION_LOST, "SERVER_DRAINING") +RETRY_AFTER_MAX = 3600 +RETRY_AFTER_LOCK = Lock() +RATE_LIMIT_TRIES = 3 +TUNNEL_SLOT_POLL = 15 REOPEN_WAIT = 1800.0 # seconds a lost session waits for its server to come back; covers a deploy's whole drain _recovering = local() @@ -92,6 +98,23 @@ def redact_secrets(text: str, data: Optional[Dict[str, Any]] = None) -> str: return redact_text(text, secrets) or "" +def wait_retry_after(seconds: int, reason: str = "Rate limited by the server") -> None: + """Count down a wait on the console, as ``--slow`` does, so the caller can retry; one hour at most, the rate-limit window. + + Threads refused together share one countdown, each still waits out its own deadline. + """ + deadline = time.monotonic() + min(seconds, RETRY_AFTER_MAX) + with RETRY_AFTER_LOCK: + if deadline <= time.monotonic(): + return + log.info(f"{reason}, retrying in {math.ceil(deadline - time.monotonic())}s") + spinner = Spinner("dots") + with SyncLive(Padding(spinner, (0, 5)), console=console, refresh_per_second=12.5, transient=True): + while (left := deadline - time.monotonic()) > 0: + spinner.update(text=f"{reason}, retrying in {math.ceil(left)} seconds...") + time.sleep(min(1.0, left)) + + class RemoteError(SystemExit): """A remote call that failed: logged, then exits a CLI run like ``SystemExit(1)``. @@ -172,6 +195,7 @@ class RemoteClient: {"json": data} if body is None else {"data": body, "headers": {"Content-Type": "application/octet-stream"}} ) gateway_retried = False + rate_waits = 0 while True: try: with self.in_flight(endpoint): @@ -186,6 +210,14 @@ class RemoteClient: # servers differ on which header carries the key, so retry the rest before giving up if resp.status_code == 401 and self.next_auth_header(): continue + if ( + resp.status_code == 429 + and rate_waits < RATE_LIMIT_TRIES + and (retry_after := parse_retry_after(resp.headers.get("Retry-After"))) + ): + rate_waits += 1 + wait_retry_after(retry_after) + continue # a proxy HTML page, not the server's JSON, means a dropped tunnel stream, which one retry gets past if ( resp.status_code in (502, 503, 504) @@ -831,11 +863,21 @@ def resolve_remote_proxy_arg(proxy_arg: str, providers: Optional[list] = None) - def start_tunnel(client: RemoteClient) -> TunnelClient: """Open a tunnel that sends the server's requests out through this machine's own connection.""" headers = {k: str(v) for k, v in client.session.headers.items() if k.lower() != "connection"} - tunnel = TunnelClient(client.server_url, headers) - try: - tunnel.start() - except ConnectionError as e: - raise click.ClickException(f"--local-proxy could not open a tunnel to the server: {e}") + slot_deadline = time.monotonic() + RETRY_AFTER_MAX + rate_waits = 0 + while True: + tunnel = TunnelClient(client.server_url, headers) + try: + tunnel.start() + break + except ConnectionError as e: + if tunnel.retry_after and rate_waits < RATE_LIMIT_TRIES: + rate_waits += 1 + wait_retry_after(tunnel.retry_after) + elif tunnel.at_limit and time.monotonic() < slot_deadline: + wait_retry_after(TUNNEL_SLOT_POLL, "Every tunnel of this API key is in use") + else: + raise click.ClickException(f"--local-proxy could not open a tunnel to the server: {e}") client.tunnel = tunnel return tunnel diff --git a/unshackle/core/tunnel.py b/unshackle/core/tunnel.py index c8ddb3ce..46fb9e0e 100644 --- a/unshackle/core/tunnel.py +++ b/unshackle/core/tunnel.py @@ -20,7 +20,8 @@ import threading from typing import Any, Optional import aiohttp -from aiohttp import WSMsgType, web +from aiohttp import WSCloseCode, WSMsgType, web +from multidict import CIMultiDict log = logging.getLogger("tunnel") @@ -144,10 +145,11 @@ TUNNELS: dict[str, "Tunnel"] = {} class Tunnel: """The server side of one client's tunnel: a loopback SOCKS5 listener and the client's current link.""" - def __init__(self, owner: Optional[str]) -> None: + def __init__(self, owner: Optional[str], client: Optional[str] = None) -> None: self.id = secrets.token_urlsafe(16) self.secret = secrets.token_urlsafe(24) self.owner = owner + self.client = client self.link: Optional[Link] = None self.pending: dict[int, asyncio.Future[bool]] = {} self.sids = itertools.count(1) @@ -290,7 +292,9 @@ class Tunnel: self.link = None self.linked.clear() link.end_all() - if self.id in TUNNELS: + if ws.close_code == WSCloseCode.OK: + self.expire() + elif self.id in TUNNELS: self.expiry = asyncio.get_running_loop().call_later(GRACE, self.expire) def expire(self) -> None: @@ -377,6 +381,27 @@ async def connect_public(host: str, port: int) -> tuple[asyncio.StreamReader, as raise error +def parse_retry_after(value: Optional[str]) -> Optional[int]: + """Whole seconds, at least 1, from a delay-seconds Retry-After; None for an absent, HTTP-date or malformed one.""" + if not value or len(value) > 9 or not (value.isascii() and value.isdigit()): + return None + return max(1, int(value)) + + +def retry_after(e: aiohttp.WSServerHandshakeError) -> Optional[int]: + """The Retry-After seconds of a rate-limit refusal, or None for any other refusal.""" + return parse_retry_after(CIMultiDict(e.headers or {}).get("Retry-After")) if e.status == 429 else None + + +def refusal(e: aiohttp.WSServerHandshakeError) -> str: + """Why the server refused a tunnel. aiohttp drops the response body, so a 429 is told apart by its headers.""" + if e.status != 429: + return f"HTTP {e.status}" + if (wait := retry_after(e)) is not None: + return f"HTTP 429, the API key is over its hourly request limit, retry in {wait}s" + return "HTTP 429, the API key already has the maximum number of tunnels open" + + class TunnelClient: """The client side: a background thread that keeps the tunnel's WebSocket open and dials for the server.""" @@ -385,6 +410,8 @@ class TunnelClient: self.headers = headers self.id: Optional[str] = None self.error: Optional[BaseException] = None + self.retry_after: Optional[int] = None + self.at_limit = False self.lost = False self.ready = threading.Event() self.relays: set[asyncio.Task[None]] = set() @@ -398,7 +425,7 @@ class TunnelClient: self.thread.start() if not self.ready.wait(timeout) or not self.id: self.stop() # or the thread opens the tunnel later and nobody uses it - raise ConnectionError(f"Could not open a tunnel to the server: {self.error or 'timed out'}") + raise ConnectionError(str(self.error or "timed out")) return self.id def stop(self) -> None: @@ -431,7 +458,9 @@ class TunnelClient: backoff = 1 await self.serve(Link(ws)) except aiohttp.WSServerHandshakeError as e: - self.error = e + self.error = ConnectionError(refusal(e)) + self.retry_after = retry_after(e) + self.at_limit = e.status == 429 and self.retry_after is None # a 404 means the tunnel expired or the server has no tunnel route; a 5xx is a proxy in between if not self.id or e.status < 500: if self.id: