From 2e7e886910d5fd9287e93814638dc8d9cc37f8e7 Mon Sep 17 00:00:00 2001 From: Edgars Date: Mon, 17 Aug 2026 10:12:13 +0100 Subject: [PATCH] feat(rate-limit): meter cheap reads separately and expose usage headers MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Rate limiting was method-agnostic: the middleware checked the quota before the JSON-RPC body was parsed, so `gen_getContractCode` — a single indexed read — cost exactly what a write fanning out to LLM validators cost. That forces the limits to be sized for the expensive case, which starves the cheap read traffic that batch tooling generates. Reads that never enter the GenVM now meter into their own bucket at RATE_LIMIT_READ_MULTIPLIER (default 10x) of the tier limits. The standard bucket keeps its key shape and numbers, so limits in flight at deploy time carry over rather than resetting. The allowlist is maintained by hand and deliberately conservative. Two methods that read like lookups are excluded because they build a Node backed by a GenVMManager (gen_getContractSchema, gen_getContractSchemaForCode), and eth_call is excluded for the same reason despite appearing in the DISABLE_INFO_LOGS_ENDPOINTS env list that otherwise looks like the natural source for this. Anything ambiguous — unparseable, oversized, or a batch with one expensive member — is charged to the stricter bucket. Responses now carry X-RateLimit-Limit/Remaining/Reset/Window/Bucket for the window closest to exhaustion, so clients can pace themselves instead of discovering the ceiling by hitting it. Reset is the time until the oldest entry ages out, which is the honest answer for a sliding window. These are listed in the CORS expose_headers, without which browsers hide them from JS. Finally, gen_getContractCode no longer loads the whole contract state to read one slot out of it. ContractSnapshot pulls the entire `data` JSONB — every storage slot the contract owns — which for a contract holding a large vector store is a substantial fetch and deserialize on a call that batch tooling polls hard. The slot is now extracted in SQL. Postgres still detoasts the JSONB server-side, so this narrows transfer and parse cost rather than eliminating the read; legacy and undeployed row shapes defer to the original path to keep their error semantics intact. The Lua script gains a test that actually executes it under a stubbed redis.call. It was previously covered only through mocked evalsha, so an error in the script body would have surfaced first in production, on every /api request. --- .env.example | 1 + backend/database_handler/contract_snapshot.py | 81 +++++- backend/protocol_rpc/endpoints.py | 8 +- backend/protocol_rpc/fastapi_server.py | 11 + backend/protocol_rpc/rate_limit_methods.py | 104 +++++++ backend/protocol_rpc/rate_limit_middleware.py | 56 +++- backend/protocol_rpc/rate_limiter.py | 147 ++++++++-- docker-compose.yml | 1 + tests/unit/test_contract_code_fetch.py | 142 ++++++++++ tests/unit/test_rate_limit_middleware.py | 264 +++++++++++++++++- tests/unit/test_rate_limiter.py | 146 +++++++++- tests/unit/test_rate_limiter_lua.py | 169 +++++++++++ 12 files changed, 1074 insertions(+), 56 deletions(-) create mode 100644 backend/protocol_rpc/rate_limit_methods.py create mode 100644 tests/unit/test_contract_code_fetch.py create mode 100644 tests/unit/test_rate_limiter_lua.py diff --git a/.env.example b/.env.example index f48ffe41d..7c05304fa 100644 --- a/.env.example +++ b/.env.example @@ -96,6 +96,7 @@ RATE_LIMIT_ENABLED='false' # Enable/disable API key rate limiting RATE_LIMIT_ANON_PER_MINUTE='30' # Anonymous (no API key) per-minute limit RATE_LIMIT_ANON_PER_HOUR='500' # Anonymous per-hour limit RATE_LIMIT_ANON_PER_DAY='5000' # Anonymous per-day limit +RATE_LIMIT_READ_MULTIPLIER='10' # Cheap reads (no GenVM) get this multiple of the tier limits # PENDING-tx queue depth caps for eth_sendRawTransaction (admission control). # Empty / unset = no cap (the default for self-hosted). Public shared diff --git a/backend/database_handler/contract_snapshot.py b/backend/database_handler/contract_snapshot.py index f3e6f3e80..840aa5938 100644 --- a/backend/database_handler/contract_snapshot.py +++ b/backend/database_handler/contract_snapshot.py @@ -1,6 +1,7 @@ # database_handler/contract_snapshot.py from .models import CurrentState from .errors import ContractNotFoundError +from sqlalchemy import func, select from sqlalchemy.orm import Session from typing import Optional, Dict import base64 @@ -87,20 +88,80 @@ def extract_deployed_code_b64(self) -> Optional[str]: slices out the code payload, and returns it base64-encoded. Returns None if missing/invalid. """ - # Import here to avoid circular dependencies at module import time - from backend.node.genvm import get_code_slot - accepted = self.states.get("accepted") or {} try: - code_slot_b64 = base64.b64encode(get_code_slot()).decode("ascii") - stored = accepted.get(code_slot_b64) + stored = accepted.get(_code_slot_b64()) if not stored: return None - - raw = base64.b64decode(stored, validate=True) - code_len = int.from_bytes(raw[0:4], byteorder="little", signed=False) - code_bytes = raw[4 : 4 + code_len] - return base64.b64encode(code_bytes).decode("ascii") + return _decode_code_payload(stored) except Exception: return None + + +def _code_slot_b64() -> str: + """Base64 of the deterministic storage slot the deployed code lives in.""" + # Import here to avoid circular dependencies at module import time + from backend.node.genvm import get_code_slot + + return base64.b64encode(get_code_slot()).decode("ascii") + + +def _decode_code_payload(stored: str) -> Optional[str]: + """Slice the code out of a stored slot blob and re-encode it as base64. + + The blob is a 4-byte little-endian length prefix followed by the code. + """ + raw = base64.b64decode(stored, validate=True) + code_len = int.from_bytes(raw[0:4], byteorder="little", signed=False) + code_bytes = raw[4 : 4 + code_len] + return base64.b64encode(code_bytes).decode("ascii") + + +def fetch_deployed_code_b64(session: Session, contract_address: str) -> Optional[str]: + """Read just the deployed code, without loading the contract's whole state. + + ``ContractSnapshot`` pulls the entire ``data`` JSONB — every storage slot + the contract owns — in order to read one deterministic slot out of it. For a + contract holding a large vector store that is a big fetch and deserialize + per call, which matters because ``gen_getContractCode`` is polled heavily by + batch tooling. Extracting the slot in SQL keeps the state off the wire and + out of Python. + + Postgres still has to detoast the JSONB server-side, so this narrows the + transfer and parse cost rather than eliminating the read entirely. + + Raises ContractNotFoundError when the contract is absent or undeployed, and + returns None when the contract exists but holds no code. + """ + slot = _code_slot_b64() + + row = session.execute( + select( + func.jsonb_typeof(CurrentState.data).label("data_kind"), + func.jsonb_typeof(CurrentState.data["state"]).label("state_kind"), + CurrentState.data["state"]["accepted"][slot].astext.label("nested"), + CurrentState.data["state"][slot].astext.label("flat"), + ).where(CurrentState.id == contract_address) + ).one_or_none() + + if row is None: + raise ContractNotFoundError(contract_address) + + if row.data_kind != "object" or row.state_kind is None: + # Legacy rows store `data` as a JSON string scalar, and undeployed ones + # store an empty object with no `state` key. Both are rare and fiddly, + # so hand them to the original path rather than reimplementing its error + # handling in SQL. + return ContractSnapshot(contract_address, session).extract_deployed_code_b64() + + # Current rows nest slots under `state.accepted`; the pre-migration format + # put them directly under `state`. + stored = row.nested if row.nested is not None else row.flat + if not stored: + return None + + try: + return _decode_code_payload(stored) + except Exception: + return None diff --git a/backend/protocol_rpc/endpoints.py b/backend/protocol_rpc/endpoints.py index 83dc0cfa8..863f8abc3 100644 --- a/backend/protocol_rpc/endpoints.py +++ b/backend/protocol_rpc/endpoints.py @@ -16,7 +16,10 @@ from sqlalchemy.orm import Session import backend.validators as validators -from backend.database_handler.contract_snapshot import ContractSnapshot +from backend.database_handler.contract_snapshot import ( + ContractSnapshot, + fetch_deployed_code_b64, +) from backend.database_handler.llm_providers import LLMProviderRegistry from backend.rollup.consensus_service import ConsensusService from backend.database_handler.models import Base, TransactionStatus @@ -1108,13 +1111,12 @@ async def get_contract_schema_for_code( def get_contract_code(session: Session, contract_address: str) -> str: try: - contract_snapshot = ContractSnapshot(contract_address, session) + code_b64 = fetch_deployed_code_b64(session, contract_address) except ContractNotFoundError: raise NotFoundError( message=f"Contract {contract_address} not found", data={"contract_address": contract_address}, ) - code_b64 = contract_snapshot.extract_deployed_code_b64() if not code_b64: raise InvalidAddressError( contract_address, diff --git a/backend/protocol_rpc/fastapi_server.py b/backend/protocol_rpc/fastapi_server.py index 20e0b67c7..7dcffa86f 100644 --- a/backend/protocol_rpc/fastapi_server.py +++ b/backend/protocol_rpc/fastapi_server.py @@ -68,6 +68,17 @@ async def lifespan(app: FastAPI): allow_credentials=True, allow_methods=["*"], allow_headers=["*"], + # Browsers hide non-simple response headers from JS unless they are listed + # here, so without this the rate limit headers are readable by curl but not + # by genlayer-js in the browser — the client that most needs to self-pace. + expose_headers=[ + "Retry-After", + "X-RateLimit-Bucket", + "X-RateLimit-Window", + "X-RateLimit-Limit", + "X-RateLimit-Remaining", + "X-RateLimit-Reset", + ], ) # Add rate limiting middleware (executes after CORS, before route handler) diff --git a/backend/protocol_rpc/rate_limit_methods.py b/backend/protocol_rpc/rate_limit_methods.py new file mode 100644 index 000000000..effc8a85f --- /dev/null +++ b/backend/protocol_rpc/rate_limit_methods.py @@ -0,0 +1,104 @@ +"""Classification of JSON-RPC methods for rate limiting. + +Requests are split into two buckets: + +- *cheap reads* — methods that never enter the GenVM and never touch an LLM. + These are served from Postgres (or are outright constants), so the cost of + serving one is orders of magnitude below a consensus round. They get their + own, much larger bucket. +- *everything else* — the default. Keeps the pre-existing limits untouched. + +The allowlist below is deliberately conservative and maintained by hand rather +than derived from any other list in the codebase. Two traps make that +necessary: + +- ``DISABLE_INFO_LOGS_ENDPOINTS`` (set in the deployment env) looks like the + natural source, but it contains ``eth_call`` — which runs contract code in + the GenVM and can fan out to LLM validators. Reusing that list would make the + single most expensive call in the system effectively free. +- ``gen_getContractSchema`` and ``gen_getContractSchemaForCode`` read like + metadata lookups, but both build a ``Node`` backed by a ``GenVMManager`` to + derive the schema from bytecode. + +Being too conservative is cheap: an omitted method simply keeps today's limits. +Being too liberal hands out free capacity on a path that costs real money. When +in doubt, leave a method out. +""" + +from __future__ import annotations + +import json +from typing import Any + +# Bodies larger than this are not parsed for classification — they are charged +# to the standard bucket. A cheap read is a handful of bytes; anything this +# large is a contract deployment or a batch, neither of which is cheap. +MAX_CLASSIFY_BODY_BYTES = 64 * 1024 + +CHEAP_READ_METHODS = frozenset( + { + # Constants / trivial responses + "ping", + "net_version", + "eth_chainId", + "eth_syncing", + "eth_gasPrice", + "eth_maxPriorityFeePerGas", + "eth_blockNumber", + "eth_feeHistory", + "eth_getCode", # returns a literal "0x" + "eth_estimateGas", # returns a constant + "sim_getFinalityWindowTime", + "sim_getConsensusContract", + # Indexed database reads + "eth_getBalance", + "eth_getTransactionCount", + "eth_getTransactionByHash", + "eth_getTransactionReceipt", + "eth_getBlockByHash", + "eth_getBlockByNumber", + "gen_getContractCode", + "gen_getContractNonce", + "gen_getTransactionStatus", + "gen_getStudioTransactionByHash", + "sim_getTransactionsForAddress", + } +) + +# Explicitly *not* cheap, recorded here so the reasoning survives future edits: +# eth_call, gen_call, sim_call -> execute contract code in the GenVM +# gen_getContractSchema -> builds a Node + GenVMManager +# gen_getContractSchemaForCode -> builds a Node + GenVMManager +# sim_lintContract -> runs the GenVM linter +# eth_getLogs -> unbounded range scan +# eth_sendRawTransaction, sim_*, admin_*, dev_* -> writes / privileged + + +def is_cheap_read_payload(raw_body: bytes) -> bool: + """Return True if every call in this JSON-RPC body is a cheap read. + + Anything ambiguous — unparseable, oversized, empty, or a batch containing a + single expensive call — is reported as *not* cheap, so uncertainty charges + the stricter bucket rather than the looser one. + """ + if not raw_body or len(raw_body) > MAX_CLASSIFY_BODY_BYTES: + return False + + try: + payload = json.loads(raw_body) + except (ValueError, UnicodeDecodeError): + return False + + if isinstance(payload, list): + # A batch is only cheap if every member is. An empty batch is invalid + # JSON-RPC and is charged normally. + return bool(payload) and all(_is_cheap_call(call) for call in payload) + + return _is_cheap_call(payload) + + +def _is_cheap_call(call: Any) -> bool: + if not isinstance(call, dict): + return False + method = call.get("method") + return isinstance(method, str) and method in CHEAP_READ_METHODS diff --git a/backend/protocol_rpc/rate_limit_middleware.py b/backend/protocol_rpc/rate_limit_middleware.py index e9651ab91..569dae4f5 100644 --- a/backend/protocol_rpc/rate_limit_middleware.py +++ b/backend/protocol_rpc/rate_limit_middleware.py @@ -12,7 +12,8 @@ from starlette.responses import JSONResponse, Response from backend.protocol_rpc.exceptions import RateLimitExceeded -from backend.protocol_rpc.rate_limiter import RateLimiterService +from backend.protocol_rpc.rate_limit_methods import is_cheap_read_payload +from backend.protocol_rpc.rate_limiter import RateLimiterService, RateLimitUsage logger = logging.getLogger(__name__) @@ -63,13 +64,14 @@ async def dispatch(self, request: Request, call_next) -> Response: api_key = request.headers.get("X-API-Key") client_ip = self._client_ip(request) + is_cheap_read = await self._is_cheap_read(request) + usage: Optional[RateLimitUsage] = None try: - await rate_limiter.check_rate_limit(api_key, client_ip) + usage = await rate_limiter.check_rate_limit( + api_key, client_ip, is_cheap_read=is_cheap_read + ) except RateLimitExceeded as exc: - retry_after = "60" - if exc.data and isinstance(exc.data, dict): - retry_after = str(exc.data.get("retry_after_seconds", 60)) return JSONResponse( status_code=429, content={ @@ -77,7 +79,7 @@ async def dispatch(self, request: Request, call_next) -> Response: "error": exc.to_dict(), "id": None, }, - headers={"Retry-After": retry_after}, + headers=self._denial_headers(exc), ) except Exception: logger.warning( @@ -85,7 +87,47 @@ async def dispatch(self, request: Request, call_next) -> Response: exc_info=True, ) - return await call_next(request) + response = await call_next(request) + if usage is not None: + response.headers.update(usage.as_headers()) + return response + + async def _is_cheap_read(self, request: Request) -> bool: + """Classify the request body, charging the stricter bucket on any doubt. + + Reading the body here is safe because Starlette's BaseHTTPMiddleware + wraps the request in a _CachedRequest, which replays the buffered body + downstream. Calling request.stream() instead would starve the route + handler. + """ + try: + raw_body = await request.body() + except Exception: + logger.warning( + "Could not read request body to classify rate limit bucket", + exc_info=True, + ) + return False + return is_cheap_read_payload(raw_body) + + @staticmethod + def _denial_headers(exc: RateLimitExceeded) -> dict: + data = exc.data if isinstance(exc.data, dict) else {} + retry_after = data.get("retry_after_seconds", 60) + headers = {"Retry-After": str(retry_after)} + # An invalid API key is raised before any window is evaluated, so there + # is no usage to report — only the windowed denials carry limits. + if "limit" in data: + headers.update( + { + "X-RateLimit-Bucket": str(data.get("bucket", "standard")), + "X-RateLimit-Window": str(data.get("window", "")), + "X-RateLimit-Limit": str(data["limit"]), + "X-RateLimit-Remaining": "0", + "X-RateLimit-Reset": str(retry_after), + } + ) + return headers def _client_ip(self, request: Request) -> str: peer_host = request.client.host if request.client else "unknown" diff --git a/backend/protocol_rpc/rate_limiter.py b/backend/protocol_rpc/rate_limiter.py index 629aa044d..c33e7aec7 100644 --- a/backend/protocol_rpc/rate_limiter.py +++ b/backend/protocol_rpc/rate_limiter.py @@ -24,6 +24,15 @@ DEFAULT_ANON_PER_HOUR = 500 DEFAULT_ANON_PER_DAY = 5000 +# Cheap reads (see rate_limit_methods) are metered in their own bucket at this +# multiple of the tier's limits. Reads never reach the GenVM, so the tier +# numbers — which have to be sized for consensus rounds — are far stricter than +# a database read warrants. +DEFAULT_READ_MULTIPLIER = 10 + +STANDARD_BUCKET = "standard" +READ_BUCKET = "read" + # Lua script that atomically prunes, checks, and records in one round-trip. # This eliminates the TOCTOU race where concurrent requests could all read the # same stale count before any of them recorded, bypassing the limit. @@ -32,7 +41,14 @@ # ARGV: [now, member, minute_window, minute_limit, hour_window, hour_limit, # day_window, day_limit] # -# Returns: [0] on success, or [1, window_name, limit, count, retry_after] on denial. +# Returns [allowed, window_name, limit, count, reset_seconds] in both the +# allowed (allowed=0) and denied (allowed=1) case. The reported window is the +# one closest to exhaustion, which is what a client needs in order to pace +# itself — reporting all three would just make the caller compute this anyway. +# +# `reset_seconds` is the time until the oldest entry in that window ages out, +# i.e. when capacity actually frees up. For a sliding window that is the honest +# answer; the window length alone would overstate the wait. _CHECK_AND_RECORD_LUA = """ local now = tonumber(ARGV[1]) local member = ARGV[2] @@ -43,12 +59,24 @@ {key = KEYS[3], seconds = tonumber(ARGV[7]), limit = tonumber(ARGV[8]), name = "day"}, } +local function reset_seconds(w) + local oldest = redis.call('ZRANGE', w.key, 0, 0, 'WITHSCORES') + if not oldest[2] then + return w.seconds + end + local reset = math.ceil(tonumber(oldest[2]) + w.seconds - now) + if reset < 1 then + return 1 + end + return reset +end + -- Phase 1: Prune expired entries and check counts for _, w in ipairs(windows) do redis.call('ZREMRANGEBYSCORE', w.key, 0, now - w.seconds) - local count = redis.call('ZCARD', w.key) - if count >= w.limit then - return {1, w.name, w.limit, count, w.seconds} + w.count = redis.call('ZCARD', w.key) + if w.count >= w.limit then + return {1, w.name, w.limit, w.count, reset_seconds(w)} end end @@ -56,9 +84,18 @@ for _, w in ipairs(windows) do redis.call('ZADD', w.key, now, member) redis.call('EXPIRE', w.key, w.seconds + 60) + w.count = w.count + 1 end -return {0} +-- Phase 3: Report whichever window has the least headroom left +local tightest = windows[1] +for _, w in ipairs(windows) do + if (w.limit - w.count) < (tightest.limit - tightest.count) then + tightest = w + end +end + +return {0, tightest.name, tightest.limit, tightest.count, reset_seconds(tightest)} """ @@ -69,6 +106,34 @@ class TierLimits: rate_limit_hour: int rate_limit_day: int + def scaled(self, factor: int) -> "TierLimits": + return TierLimits( + name=self.name, + rate_limit_minute=self.rate_limit_minute * factor, + rate_limit_hour=self.rate_limit_hour * factor, + rate_limit_day=self.rate_limit_day * factor, + ) + + +@dataclass(frozen=True) +class RateLimitUsage: + """Headroom in the window closest to exhaustion, for X-RateLimit-* headers.""" + + bucket: str + window: str + limit: int + remaining: int + reset_seconds: int + + def as_headers(self) -> dict[str, str]: + return { + "X-RateLimit-Bucket": self.bucket, + "X-RateLimit-Window": self.window, + "X-RateLimit-Limit": str(self.limit), + "X-RateLimit-Remaining": str(self.remaining), + "X-RateLimit-Reset": str(self.reset_seconds), + } + class RateLimiterService: """Sliding-window rate limiter backed by Redis sorted sets.""" @@ -81,6 +146,7 @@ def __init__( anon_per_minute: int = DEFAULT_ANON_PER_MINUTE, anon_per_hour: int = DEFAULT_ANON_PER_HOUR, anon_per_day: int = DEFAULT_ANON_PER_DAY, + read_multiplier: int = DEFAULT_READ_MULTIPLIER, ): self._redis = redis_client self._get_session = get_session @@ -91,6 +157,9 @@ def __init__( rate_limit_hour=anon_per_hour, rate_limit_day=anon_per_day, ) + # A multiplier below 1 would make reads *stricter* than writes, which is + # never intended; clamp rather than trust the environment. + self._read_multiplier = max(1, read_multiplier) self._lua_sha: Optional[str] = None @classmethod @@ -112,26 +181,42 @@ def from_environment( anon_per_day=int( os.environ.get("RATE_LIMIT_ANON_PER_DAY", DEFAULT_ANON_PER_DAY) ), + read_multiplier=int( + os.environ.get("RATE_LIMIT_READ_MULTIPLIER", DEFAULT_READ_MULTIPLIER) + ), ) @property def enabled(self) -> bool: return self._enabled - async def check_rate_limit(self, api_key: Optional[str], client_ip: str) -> None: - """Check rate limits. Raises RateLimitExceeded if over limit.""" + async def check_rate_limit( + self, + api_key: Optional[str], + client_ip: str, + is_cheap_read: bool = False, + ) -> Optional[RateLimitUsage]: + """Check rate limits. Raises RateLimitExceeded if over limit. + + Returns the usage of whichever window is closest to exhaustion, or None + when limiting is disabled. + """ if not self._enabled: - return + return None if api_key: identity, limits = await self._resolve_api_key(api_key) - if identity is None: + if identity is None or limits is None: raise RateLimitExceeded(message="Invalid API key") else: identity = f"ip:{client_ip}" limits = self._anon_limits - await self._check_windows(identity, limits) + if is_cheap_read: + return await self._check_windows( + identity, limits.scaled(self._read_multiplier), READ_BUCKET + ) + return await self._check_windows(identity, limits, STANDARD_BUCKET) async def _resolve_api_key( self, raw_key: str @@ -192,15 +277,27 @@ async def _ensure_lua_loaded(self) -> str: self._lua_sha = await self._redis.script_load(_CHECK_AND_RECORD_LUA) return self._lua_sha - async def _check_windows(self, identity: str, limits: TierLimits) -> None: + async def _check_windows( + self, + identity: str, + limits: TierLimits, + bucket: str, + ) -> RateLimitUsage: """Atomically prune, check, and record using a Lua script.""" now = time.time() member = f"{now}:{uuid.uuid4().hex[:8]}" + # The standard bucket keeps its original key shape so that limits in + # flight at deploy time carry over instead of silently resetting. + prefix = ( + f"ratelimit:{identity}" + if bucket == STANDARD_BUCKET + else f"ratelimit:{identity}:{bucket}" + ) keys = [ - f"ratelimit:{identity}:minute", - f"ratelimit:{identity}:hour", - f"ratelimit:{identity}:day", + f"{prefix}:minute", + f"{prefix}:hour", + f"{prefix}:day", ] args = [ str(now), @@ -222,23 +319,31 @@ async def _check_windows(self, identity: str, limits: TierLimits) -> None: sha = await self._ensure_lua_loaded() result = await self._redis.evalsha(sha, len(keys), *keys, *args) + window_name = result[1].decode() if isinstance(result[1], bytes) else result[1] + max_requests = int(result[2]) + count = int(result[3]) + reset_after = int(result[4]) + if result[0] == 1: - window_name = ( - result[1].decode() if isinstance(result[1], bytes) else result[1] - ) - max_requests = int(result[2]) - count = int(result[3]) - retry_after = int(result[4]) raise RateLimitExceeded( message=f"Rate limit exceeded: {max_requests} requests per {window_name}", data={ + "bucket": bucket, "window": window_name, "limit": max_requests, "current": count, - "retry_after_seconds": retry_after, + "retry_after_seconds": reset_after, }, ) + return RateLimitUsage( + bucket=bucket, + window=window_name, + limit=max_requests, + remaining=max(0, max_requests - count), + reset_seconds=reset_after, + ) + async def invalidate_key_cache(self, key_hash: str) -> None: """Invalidate cached tier for an API key (call after deactivation).""" cache_key = f"ratelimit:tier:{key_hash}" diff --git a/docker-compose.yml b/docker-compose.yml index 3e84651eb..780a3a73d 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -125,6 +125,7 @@ services: - RATE_LIMIT_ANON_PER_MINUTE=${RATE_LIMIT_ANON_PER_MINUTE:-30} - RATE_LIMIT_ANON_PER_HOUR=${RATE_LIMIT_ANON_PER_HOUR:-500} - RATE_LIMIT_ANON_PER_DAY=${RATE_LIMIT_ANON_PER_DAY:-5000} + - RATE_LIMIT_READ_MULTIPLIER=${RATE_LIMIT_READ_MULTIPLIER:-10} # Per-contract / per-sender PENDING tx caps (admission control on # eth_sendRawTransaction). Empty/unset = no cap. Set in shared # deployments to keep one heavy user from filling the queue. diff --git a/tests/unit/test_contract_code_fetch.py b/tests/unit/test_contract_code_fetch.py new file mode 100644 index 000000000..64f2061a5 --- /dev/null +++ b/tests/unit/test_contract_code_fetch.py @@ -0,0 +1,142 @@ +"""Unit tests for the narrowed gen_getContractCode read path.""" + +import base64 +from unittest.mock import MagicMock, patch + +import pytest + +from backend.database_handler.contract_snapshot import ( + _code_slot_b64, + _decode_code_payload, + fetch_deployed_code_b64, +) +from backend.database_handler.errors import ContractNotFoundError + +ADDRESS = "0xabc" + + +def _stored_blob(code: bytes) -> str: + """Build a slot blob: 4-byte little-endian length prefix, then the code.""" + return base64.b64encode(len(code).to_bytes(4, "little") + code).decode("ascii") + + +def _make_session(row): + session = MagicMock() + session.execute.return_value.one_or_none.return_value = row + return session + + +def _make_row(data_kind="object", state_kind="object", nested=None, flat=None): + row = MagicMock() + row.data_kind = data_kind + row.state_kind = state_kind + row.nested = nested + row.flat = flat + return row + + +class TestFetchDeployedCode: + def test_returns_code_from_nested_state(self): + code = b"class Contract: pass" + session = _make_session(_make_row(nested=_stored_blob(code))) + + result = fetch_deployed_code_b64(session, ADDRESS) + + assert base64.b64decode(result) == code + + def test_falls_back_to_flat_state_for_old_format(self): + """Pre-migration rows put slots directly under `state`.""" + code = b"legacy contract" + session = _make_session(_make_row(nested=None, flat=_stored_blob(code))) + + result = fetch_deployed_code_b64(session, ADDRESS) + + assert base64.b64decode(result) == code + + def test_prefers_nested_over_flat(self): + session = _make_session( + _make_row(nested=_stored_blob(b"new"), flat=_stored_blob(b"old")) + ) + + assert base64.b64decode(fetch_deployed_code_b64(session, ADDRESS)) == b"new" + + def test_missing_row_raises_not_found(self): + session = _make_session(None) + + with pytest.raises(ContractNotFoundError): + fetch_deployed_code_b64(session, ADDRESS) + + def test_missing_slot_returns_none(self): + """Contract exists but holds no code — distinct from not existing.""" + session = _make_session(_make_row(nested=None, flat=None)) + + assert fetch_deployed_code_b64(session, ADDRESS) is None + + def test_corrupt_blob_returns_none_rather_than_raising(self): + session = _make_session(_make_row(nested="!!! not base64 !!!")) + + assert fetch_deployed_code_b64(session, ADDRESS) is None + + @pytest.mark.parametrize( + "data_kind,state_kind", + [ + ("string", "object"), # legacy rows store `data` as a JSON string + ("null", None), + ("object", None), # `{}` — present but never deployed + ], + ) + def test_unusual_shapes_defer_to_the_snapshot_path(self, data_kind, state_kind): + """Legacy and undeployed rows keep the original error semantics. + + Rather than reimplement those in SQL, the fast path declines and hands + the row to ContractSnapshot. + """ + session = _make_session(_make_row(data_kind=data_kind, state_kind=state_kind)) + + with patch( + "backend.database_handler.contract_snapshot.ContractSnapshot" + ) as snapshot_cls: + snapshot_cls.return_value.extract_deployed_code_b64.return_value = ( + "FALLBACK" + ) + + assert fetch_deployed_code_b64(session, ADDRESS) == "FALLBACK" + + snapshot_cls.assert_called_once_with(ADDRESS, session) + + +class TestCodeSlotHelpers: + def test_code_slot_is_stable(self): + """The slot address is a protocol constant; drift silently breaks reads.""" + assert _code_slot_b64() == _code_slot_b64() + assert len(base64.b64decode(_code_slot_b64())) == 32 + + def test_decode_respects_length_prefix(self): + """Trailing bytes past the declared length must not leak into the code.""" + blob = base64.b64encode((3).to_bytes(4, "little") + b"abc" + b"PADDING") + assert _decode_code_payload(blob.decode()) == base64.b64encode(b"abc").decode() + + +class TestGeneratedSQL: + def test_statement_compiles_against_postgres(self): + """Guards the JSONB path expression, which mocks cannot validate.""" + from sqlalchemy import func, select + from sqlalchemy.dialects import postgresql + + from backend.database_handler.models import CurrentState + + slot = _code_slot_b64() + stmt = select( + func.jsonb_typeof(CurrentState.data).label("data_kind"), + func.jsonb_typeof(CurrentState.data["state"]).label("state_kind"), + CurrentState.data["state"]["accepted"][slot].astext.label("nested"), + CurrentState.data["state"][slot].astext.label("flat"), + ).where(CurrentState.id == ADDRESS) + + sql = str(stmt.compile(dialect=postgresql.dialect())) + + assert "jsonb_typeof" in sql + assert "->>" in sql + assert "current_state" in sql + # The whole point is not selecting the full state blob. + assert "current_state.data \n" not in sql diff --git a/tests/unit/test_rate_limit_middleware.py b/tests/unit/test_rate_limit_middleware.py index 5d5afc7a3..8c7c2fba7 100644 --- a/tests/unit/test_rate_limit_middleware.py +++ b/tests/unit/test_rate_limit_middleware.py @@ -1,5 +1,7 @@ """Unit tests for RateLimitMiddleware.""" +import json + import pytest from unittest.mock import AsyncMock, MagicMock @@ -7,6 +9,7 @@ from backend.protocol_rpc.rate_limit_middleware import RateLimitMiddleware from backend.protocol_rpc.exceptions import RateLimitExceeded +from backend.protocol_rpc.rate_limiter import RateLimitUsage def _make_request( @@ -15,12 +18,14 @@ def _make_request( api_key=None, client_host="127.0.0.1", headers=None, + body=b'{"jsonrpc":"2.0","method":"eth_sendRawTransaction","params":[],"id":1}', ): """Create a mock Starlette Request.""" headers = headers or {} request = MagicMock() request.url.path = path request.method = method + request.body = AsyncMock(return_value=body) request.headers = MagicMock() def get_header(key, default=None): @@ -160,7 +165,9 @@ async def test_passes_api_key_header(self): middleware = RateLimitMiddleware(app=MagicMock()) await middleware.dispatch(request, call_next) - limiter.check_rate_limit.assert_called_once_with("glk_testkey123", "127.0.0.1") + limiter.check_rate_limit.assert_called_once_with( + "glk_testkey123", "127.0.0.1", is_cheap_read=False + ) @pytest.mark.asyncio async def test_uses_forwarded_client_ip_from_trusted_proxy(self): @@ -177,7 +184,9 @@ async def test_uses_forwarded_client_ip_from_trusted_proxy(self): middleware = RateLimitMiddleware(app=MagicMock()) await middleware.dispatch(request, call_next) - limiter.check_rate_limit.assert_called_once_with(None, "198.51.100.7") + limiter.check_rate_limit.assert_called_once_with( + None, "198.51.100.7", is_cheap_read=False + ) @pytest.mark.asyncio async def test_ignores_forwarded_client_ip_from_untrusted_peer(self): @@ -194,7 +203,9 @@ async def test_ignores_forwarded_client_ip_from_untrusted_peer(self): middleware = RateLimitMiddleware(app=MagicMock()) await middleware.dispatch(request, call_next) - limiter.check_rate_limit.assert_called_once_with(None, "198.51.100.9") + limiter.check_rate_limit.assert_called_once_with( + None, "198.51.100.9", is_cheap_read=False + ) @pytest.mark.asyncio async def test_uses_first_forwarded_ip_when_all_hops_are_trusted(self): @@ -211,7 +222,9 @@ async def test_uses_first_forwarded_ip_when_all_hops_are_trusted(self): middleware = RateLimitMiddleware(app=MagicMock()) await middleware.dispatch(request, call_next) - limiter.check_rate_limit.assert_called_once_with(None, "10.0.12.7") + limiter.check_rate_limit.assert_called_once_with( + None, "10.0.12.7", is_cheap_read=False + ) @pytest.mark.asyncio async def test_uses_real_ip_from_trusted_proxy_when_forwarded_for_missing(self): @@ -228,7 +241,9 @@ async def test_uses_real_ip_from_trusted_proxy_when_forwarded_for_missing(self): middleware = RateLimitMiddleware(app=MagicMock()) await middleware.dispatch(request, call_next) - limiter.check_rate_limit.assert_called_once_with(None, "203.0.113.12") + limiter.check_rate_limit.assert_called_once_with( + None, "203.0.113.12", is_cheap_read=False + ) @pytest.mark.asyncio async def test_falls_back_to_peer_when_forwarded_headers_are_invalid(self): @@ -248,7 +263,9 @@ async def test_falls_back_to_peer_when_forwarded_headers_are_invalid(self): middleware = RateLimitMiddleware(app=MagicMock()) await middleware.dispatch(request, call_next) - limiter.check_rate_limit.assert_called_once_with(None, "127.0.0.1") + limiter.check_rate_limit.assert_called_once_with( + None, "127.0.0.1", is_cheap_read=False + ) @pytest.mark.asyncio async def test_invalid_trusted_proxy_config_is_ignored(self, monkeypatch, caplog): @@ -270,7 +287,9 @@ async def test_invalid_trusted_proxy_config_is_ignored(self, monkeypatch, caplog await middleware.dispatch(request, call_next) assert "Ignoring invalid RATE_LIMIT_TRUSTED_PROXIES entry" in caplog.text - limiter.check_rate_limit.assert_called_once_with(None, "198.51.100.7") + limiter.check_rate_limit.assert_called_once_with( + None, "198.51.100.7", is_cheap_read=False + ) @pytest.mark.asyncio async def test_retry_after_header_with_no_data(self): @@ -306,3 +325,234 @@ async def test_fails_open_when_rate_limiter_throws_unexpected_error(self): assert response.status_code == 200 call_next.assert_called_once() + + +def _rpc_body(*methods): + if len(methods) == 1: + return json.dumps( + {"jsonrpc": "2.0", "method": methods[0], "params": [], "id": 1} + ).encode() + return json.dumps( + [ + {"jsonrpc": "2.0", "method": m, "params": [], "id": i} + for i, m in enumerate(methods) + ] + ).encode() + + +class TestBucketClassification: + """The bucket a request lands in is decided from its JSON-RPC method.""" + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "method", + ["gen_getContractCode", "eth_getBalance", "ping", "eth_chainId"], + ) + async def test_cheap_reads_use_read_bucket(self, method): + limiter = AsyncMock() + limiter.enabled = True + limiter.check_rate_limit = AsyncMock(return_value=None) + request = _make_request(body=_rpc_body(method)) + request.app.state.rate_limiter = limiter + + middleware = RateLimitMiddleware(app=MagicMock()) + await middleware.dispatch(request, _make_call_next()) + + assert limiter.check_rate_limit.call_args.kwargs["is_cheap_read"] is True + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "method", + [ + # Each of these looks like a read but reaches the GenVM, so putting + # any of them in the read bucket would hand out free LLM capacity. + "eth_call", + "gen_call", + "gen_getContractSchema", + "gen_getContractSchemaForCode", + "sim_lintContract", + "eth_sendRawTransaction", + ], + ) + async def test_genvm_methods_use_standard_bucket(self, method): + limiter = AsyncMock() + limiter.enabled = True + limiter.check_rate_limit = AsyncMock(return_value=None) + request = _make_request(body=_rpc_body(method)) + request.app.state.rate_limiter = limiter + + middleware = RateLimitMiddleware(app=MagicMock()) + await middleware.dispatch(request, _make_call_next()) + + assert limiter.check_rate_limit.call_args.kwargs["is_cheap_read"] is False + + @pytest.mark.asyncio + async def test_batch_of_reads_is_cheap(self): + limiter = AsyncMock() + limiter.enabled = True + limiter.check_rate_limit = AsyncMock(return_value=None) + request = _make_request(body=_rpc_body("ping", "eth_chainId", "eth_getBalance")) + request.app.state.rate_limiter = limiter + + middleware = RateLimitMiddleware(app=MagicMock()) + await middleware.dispatch(request, _make_call_next()) + + assert limiter.check_rate_limit.call_args.kwargs["is_cheap_read"] is True + + @pytest.mark.asyncio + async def test_batch_with_one_expensive_call_is_not_cheap(self): + """One costly member must taint the whole batch, or it is a free ride.""" + limiter = AsyncMock() + limiter.enabled = True + limiter.check_rate_limit = AsyncMock(return_value=None) + request = _make_request(body=_rpc_body("ping", "eth_call", "ping")) + request.app.state.rate_limiter = limiter + + middleware = RateLimitMiddleware(app=MagicMock()) + await middleware.dispatch(request, _make_call_next()) + + assert limiter.check_rate_limit.call_args.kwargs["is_cheap_read"] is False + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "body", + [ + b"", + b"not json at all", + b"{", + b'{"jsonrpc":"2.0","id":1}', # no method + b'{"jsonrpc":"2.0","method":123,"id":1}', # method not a string + b"[]", # empty batch + b'"just a string"', + b'{"jsonrpc":"2.0","method":"unknown_future_method","id":1}', + ], + ) + async def test_ambiguous_bodies_charge_the_stricter_bucket(self, body): + limiter = AsyncMock() + limiter.enabled = True + limiter.check_rate_limit = AsyncMock(return_value=None) + request = _make_request(body=body) + request.app.state.rate_limiter = limiter + + middleware = RateLimitMiddleware(app=MagicMock()) + await middleware.dispatch(request, _make_call_next()) + + assert limiter.check_rate_limit.call_args.kwargs["is_cheap_read"] is False + + @pytest.mark.asyncio + async def test_oversized_body_is_not_parsed(self): + oversized = b'{"jsonrpc":"2.0","method":"ping","params":["' + (b"x" * 70_000) + limiter = AsyncMock() + limiter.enabled = True + limiter.check_rate_limit = AsyncMock(return_value=None) + request = _make_request(body=oversized) + request.app.state.rate_limiter = limiter + + middleware = RateLimitMiddleware(app=MagicMock()) + await middleware.dispatch(request, _make_call_next()) + + assert limiter.check_rate_limit.call_args.kwargs["is_cheap_read"] is False + + @pytest.mark.asyncio + async def test_unreadable_body_does_not_break_the_request(self): + limiter = AsyncMock() + limiter.enabled = True + limiter.check_rate_limit = AsyncMock(return_value=None) + request = _make_request() + request.body = AsyncMock(side_effect=RuntimeError("stream consumed")) + request.app.state.rate_limiter = limiter + + middleware = RateLimitMiddleware(app=MagicMock()) + response = await middleware.dispatch(request, _make_call_next()) + + assert response.status_code == 200 + assert limiter.check_rate_limit.call_args.kwargs["is_cheap_read"] is False + + +class TestRateLimitHeaders: + @pytest.mark.asyncio + async def test_usage_headers_on_success(self): + limiter = AsyncMock() + limiter.enabled = True + limiter.check_rate_limit = AsyncMock( + return_value=RateLimitUsage( + bucket="read", + window="minute", + limit=6000, + remaining=5987, + reset_seconds=42, + ) + ) + request = _make_request(body=_rpc_body("ping")) + request.app.state.rate_limiter = limiter + + middleware = RateLimitMiddleware(app=MagicMock()) + response = await middleware.dispatch(request, _make_call_next()) + + assert response.status_code == 200 + assert response.headers["X-RateLimit-Bucket"] == "read" + assert response.headers["X-RateLimit-Window"] == "minute" + assert response.headers["X-RateLimit-Limit"] == "6000" + assert response.headers["X-RateLimit-Remaining"] == "5987" + assert response.headers["X-RateLimit-Reset"] == "42" + + @pytest.mark.asyncio + async def test_no_headers_when_limiter_disabled(self): + limiter = AsyncMock() + limiter.enabled = True + limiter.check_rate_limit = AsyncMock(return_value=None) + request = _make_request() + request.app.state.rate_limiter = limiter + + middleware = RateLimitMiddleware(app=MagicMock()) + response = await middleware.dispatch(request, _make_call_next()) + + assert "X-RateLimit-Limit" not in response.headers + + @pytest.mark.asyncio + async def test_usage_headers_on_429(self): + limiter = AsyncMock() + limiter.enabled = True + limiter.check_rate_limit = AsyncMock( + side_effect=RateLimitExceeded( + message="Rate limit exceeded: 600 requests per minute", + data={ + "bucket": "standard", + "window": "minute", + "limit": 600, + "current": 600, + "retry_after_seconds": 17, + }, + ) + ) + request = _make_request() + request.app.state.rate_limiter = limiter + + middleware = RateLimitMiddleware(app=MagicMock()) + response = await middleware.dispatch(request, _make_call_next()) + + assert response.status_code == 429 + assert response.headers["Retry-After"] == "17" + assert response.headers["X-RateLimit-Bucket"] == "standard" + assert response.headers["X-RateLimit-Window"] == "minute" + assert response.headers["X-RateLimit-Limit"] == "600" + assert response.headers["X-RateLimit-Remaining"] == "0" + assert response.headers["X-RateLimit-Reset"] == "17" + + @pytest.mark.asyncio + async def test_invalid_key_429_has_no_usage_headers(self): + """No window was evaluated, so there is no headroom to report.""" + limiter = AsyncMock() + limiter.enabled = True + limiter.check_rate_limit = AsyncMock( + side_effect=RateLimitExceeded(message="Invalid API key", data=None) + ) + request = _make_request() + request.app.state.rate_limiter = limiter + + middleware = RateLimitMiddleware(app=MagicMock()) + response = await middleware.dispatch(request, _make_call_next()) + + assert response.status_code == 429 + assert response.headers["Retry-After"] == "60" + assert "X-RateLimit-Limit" not in response.headers diff --git a/tests/unit/test_rate_limiter.py b/tests/unit/test_rate_limiter.py index ddca9a40c..70a918daf 100644 --- a/tests/unit/test_rate_limiter.py +++ b/tests/unit/test_rate_limiter.py @@ -6,12 +6,15 @@ from backend.protocol_rpc.rate_limiter import RateLimiterService from backend.protocol_rpc.exceptions import RateLimitExceeded +# [allowed, window, limit, count, reset] — the shape the Lua script returns. +_ALLOWED = [0, b"minute", 5, 1, 60] + def _make_redis_mock(): """Create a mock Redis client with Lua script support.""" redis = AsyncMock() redis.script_load = AsyncMock(return_value="fake_sha") - redis.evalsha = AsyncMock(return_value=[0]) # default: allow + redis.evalsha = AsyncMock(return_value=_ALLOWED) # default: allow redis.hgetall = AsyncMock(return_value={}) redis.hset = AsyncMock() redis.expire = AsyncMock() @@ -55,7 +58,7 @@ class TestAnonymousRateLimiting: @pytest.mark.asyncio async def test_allowed_when_under_limit(self): redis = _make_redis_mock() - redis.evalsha = AsyncMock(return_value=[0]) + redis.evalsha = AsyncMock(return_value=_ALLOWED) service, _ = _make_service(redis=redis) # Should not raise await service.check_rate_limit(None, "1.2.3.4") @@ -93,7 +96,7 @@ async def test_raises_when_day_limit_exceeded(self): @pytest.mark.asyncio async def test_identity_uses_ip(self): redis = _make_redis_mock() - redis.evalsha = AsyncMock(return_value=[0]) + redis.evalsha = AsyncMock(return_value=_ALLOWED) service, _ = _make_service(redis=redis) await service.check_rate_limit(None, "10.0.0.1") # Verify keys passed to evalsha contain the IP @@ -104,7 +107,7 @@ async def test_identity_uses_ip(self): @pytest.mark.asyncio async def test_lua_script_loaded_once(self): redis = _make_redis_mock() - redis.evalsha = AsyncMock(return_value=[0]) + redis.evalsha = AsyncMock(return_value=_ALLOWED) service, _ = _make_service(redis=redis) await service.check_rate_limit(None, "1.2.3.4") await service.check_rate_limit(None, "1.2.3.4") @@ -116,7 +119,9 @@ async def test_lua_script_reloaded_on_noscript_error(self): from redis.exceptions import NoScriptError redis_mock = _make_redis_mock() - redis_mock.evalsha = AsyncMock(side_effect=[NoScriptError("NOSCRIPT"), [0]]) + redis_mock.evalsha = AsyncMock( + side_effect=[NoScriptError("NOSCRIPT"), _ALLOWED] + ) redis_mock.script_load = AsyncMock(return_value="new_sha") service, _ = _make_service(redis=redis_mock) # Should not raise — recovers by reloading the script @@ -126,7 +131,7 @@ async def test_lua_script_reloaded_on_noscript_error(self): @pytest.mark.asyncio async def test_passes_correct_limits_as_args(self): redis = _make_redis_mock() - redis.evalsha = AsyncMock(return_value=[0]) + redis.evalsha = AsyncMock(return_value=_ALLOWED) service, _ = _make_service(redis=redis) await service.check_rate_limit(None, "1.2.3.4") call_args = redis.evalsha.call_args[0] @@ -164,7 +169,7 @@ async def test_cached_key_uses_cached_limits(self): "rpd": "50000", } ) - redis.evalsha = AsyncMock(return_value=[0]) + redis.evalsha = AsyncMock(return_value=_ALLOWED) service, _ = _make_service(redis=redis) await service.check_rate_limit("glk_test1234", "1.2.3.4") # Verify it used the cached limits (didn't query DB) @@ -175,7 +180,7 @@ async def test_cache_miss_queries_db(self): redis = _make_redis_mock() # Cache miss redis.hgetall = AsyncMock(return_value={}) - redis.evalsha = AsyncMock(return_value=[0]) + redis.evalsha = AsyncMock(return_value=_ALLOWED) # Mock DB session mock_session = MagicMock() @@ -278,3 +283,128 @@ def test_defaults_to_disabled(self): assert service._anon_limits.rate_limit_minute == 30 assert service._anon_limits.rate_limit_hour == 500 assert service._anon_limits.rate_limit_day == 5000 + + +class TestReadBucket: + """Cheap reads are metered separately from GenVM-bound calls.""" + + @pytest.mark.asyncio + async def test_read_bucket_uses_separate_keys(self): + service, redis = _make_service() + await service.check_rate_limit(None, "1.2.3.4", is_cheap_read=True) + + keys = redis.evalsha.call_args[0][2:5] + assert keys == ( + "ratelimit:ip:1.2.3.4:read:minute", + "ratelimit:ip:1.2.3.4:read:hour", + "ratelimit:ip:1.2.3.4:read:day", + ) + + @pytest.mark.asyncio + async def test_standard_bucket_keeps_original_keys(self): + """Key shape must not change, or limits in flight silently reset.""" + service, redis = _make_service() + await service.check_rate_limit(None, "1.2.3.4") + + keys = redis.evalsha.call_args[0][2:5] + assert keys == ( + "ratelimit:ip:1.2.3.4:minute", + "ratelimit:ip:1.2.3.4:hour", + "ratelimit:ip:1.2.3.4:day", + ) + + @pytest.mark.asyncio + async def test_read_limits_are_scaled_by_multiplier(self): + redis = _make_redis_mock() + service = RateLimiterService( + redis_client=redis, + get_session=MagicMock(), + enabled=True, + anon_per_minute=5, + anon_per_hour=50, + anon_per_day=500, + read_multiplier=10, + ) + await service.check_rate_limit(None, "1.2.3.4", is_cheap_read=True) + + args = redis.evalsha.call_args[0][5:] + # [now, member, 60, rpm, 3600, rph, 86400, rpd] + assert args[3] == "50" + assert args[5] == "500" + assert args[7] == "5000" + + @pytest.mark.asyncio + async def test_standard_limits_are_not_scaled(self): + service, redis = _make_service() + await service.check_rate_limit(None, "1.2.3.4") + + args = redis.evalsha.call_args[0][5:] + assert args[3] == "5" + assert args[5] == "50" + assert args[7] == "500" + + @pytest.mark.asyncio + async def test_multiplier_below_one_is_clamped(self): + """A misconfigured multiplier must never make reads stricter.""" + redis = _make_redis_mock() + service = RateLimiterService( + redis_client=redis, + get_session=MagicMock(), + enabled=True, + anon_per_minute=5, + anon_per_hour=50, + anon_per_day=500, + read_multiplier=0, + ) + await service.check_rate_limit(None, "1.2.3.4", is_cheap_read=True) + + assert redis.evalsha.call_args[0][8] == "5" + + +class TestUsageReporting: + @pytest.mark.asyncio + async def test_returns_usage_for_tightest_window(self): + redis = _make_redis_mock() + redis.evalsha = AsyncMock(return_value=[0, b"day", 500, 498, 3600]) + service, _ = _make_service(redis=redis) + + usage = await service.check_rate_limit(None, "1.2.3.4") + + assert usage.bucket == "standard" + assert usage.window == "day" + assert usage.limit == 500 + assert usage.remaining == 2 + assert usage.reset_seconds == 3600 + + @pytest.mark.asyncio + async def test_usage_reports_read_bucket(self): + service, _ = _make_service() + usage = await service.check_rate_limit(None, "1.2.3.4", is_cheap_read=True) + assert usage.bucket == "read" + + @pytest.mark.asyncio + async def test_remaining_never_goes_negative(self): + redis = _make_redis_mock() + redis.evalsha = AsyncMock(return_value=[0, b"minute", 5, 9, 60]) + service, _ = _make_service(redis=redis) + + usage = await service.check_rate_limit(None, "1.2.3.4") + + assert usage.remaining == 0 + + @pytest.mark.asyncio + async def test_returns_none_when_disabled(self): + service, _ = _make_service(enabled=False) + assert await service.check_rate_limit(None, "1.2.3.4") is None + + @pytest.mark.asyncio + async def test_denial_carries_bucket_for_headers(self): + redis = _make_redis_mock() + redis.evalsha = AsyncMock(return_value=[1, b"minute", 50, 50, 12]) + service, _ = _make_service(redis=redis) + + with pytest.raises(RateLimitExceeded) as exc_info: + await service.check_rate_limit(None, "1.2.3.4", is_cheap_read=True) + + assert exc_info.value.data["bucket"] == "read" + assert exc_info.value.data["retry_after_seconds"] == 12 diff --git a/tests/unit/test_rate_limiter_lua.py b/tests/unit/test_rate_limiter_lua.py new file mode 100644 index 000000000..deb766bf2 --- /dev/null +++ b/tests/unit/test_rate_limiter_lua.py @@ -0,0 +1,169 @@ +"""Executes the rate limiter's Lua script for real. + +Every other test mocks `evalsha`, so the script body itself is otherwise +unexercised — a syntax error or a bad redis call in it would first surface in +production, on every single /api request. Here the script runs under a stubbed +`redis.call` that implements enough sorted-set semantics to be meaningful. + +Skipped when no `lua` interpreter is present. +""" + +import shutil +import subprocess +import tempfile + +import pytest + +from backend.protocol_rpc.rate_limiter import _CHECK_AND_RECORD_LUA + +pytestmark = pytest.mark.skipif( + shutil.which("lua") is None, reason="requires a lua interpreter" +) + +REDIS_STUB = """ +local store = {} +local function zset(k) store[k] = store[k] or {}; return store[k] end + +redis = {} +function redis.call(cmd, key, a, b, c) + local z = zset(key) + if cmd == 'ZREMRANGEBYSCORE' then + local kept = {} + for _, e in ipairs(z) do + if not (e.score >= tonumber(a) and e.score <= tonumber(b)) then + kept[#kept + 1] = e + end + end + store[key] = kept + return 0 + elseif cmd == 'ZCARD' then + return #z + elseif cmd == 'ZADD' then + z[#z + 1] = {score = tonumber(a), member = b} + table.sort(z, function(x, y) return x.score < y.score end) + return 1 + elseif cmd == 'EXPIRE' then + return 1 + elseif cmd == 'ZRANGE' then + if #z == 0 then return {} end + return {z[1].member, tostring(z[1].score)} + end + error('unstubbed redis command: ' .. cmd) +end + +local function run(now, keys, argv) + KEYS = keys + ARGV = argv + return SCRIPT(now) +end + +local function argv(now, member, m, h, d) + return {tostring(now), member, "60", tostring(m), + "3600", tostring(h), "86400", tostring(d)} +end + +local K = {"k:minute", "k:hour", "k:day"} +""" + + +def _run_lua(scenario: str) -> str: + script = ( + REDIS_STUB + + "function SCRIPT(now)\n" + + _CHECK_AND_RECORD_LUA + + "\nend\n" + + scenario + ) + with tempfile.NamedTemporaryFile("w", suffix=".lua") as fh: + fh.write(script) + fh.flush() + result = subprocess.run( + ["lua", fh.name], capture_output=True, text=True, timeout=30 + ) + assert result.returncode == 0, result.stderr + return result.stdout + + +def test_allows_and_reports_tightest_window(): + out = _run_lua( + """ + local r = run(1000, K, argv(1000, "a", 3, 100, 1000)) + assert(r[1] == 0) + assert(r[2] == "minute") + assert(r[3] == 3 and r[4] == 1) + print("ok") + """ + ) + assert "ok" in out + + +def test_denies_once_window_is_full(): + out = _run_lua( + """ + run(1000, K, argv(1000, "a", 3, 100, 1000)) + run(1001, K, argv(1001, "b", 3, 100, 1000)) + run(1002, K, argv(1002, "c", 3, 100, 1000)) + local d = run(1003, K, argv(1003, "d", 3, 100, 1000)) + assert(d[1] == 1) + assert(d[2] == "minute") + -- oldest entry is at t=1000, so capacity returns at t=1060 + assert(d[5] == 57, "reset was " .. tostring(d[5])) + print("ok") + """ + ) + assert "ok" in out + + +def test_capacity_returns_as_window_slides(): + out = _run_lua( + """ + run(1000, K, argv(1000, "a", 1, 100, 1000)) + local denied = run(1030, K, argv(1030, "b", 1, 100, 1000)) + assert(denied[1] == 1) + local allowed = run(1061, K, argv(1061, "c", 1, 100, 1000)) + assert(allowed[1] == 0) + print("ok") + """ + ) + assert "ok" in out + + +def test_day_window_can_be_the_reported_one(): + out = _run_lua( + """ + local r = run(2000, K, argv(2000, "a", 1000, 1000, 2)) + assert(r[2] == "day", "got " .. tostring(r[2])) + print("ok") + """ + ) + assert "ok" in out + + +def test_reset_is_never_below_one_second(): + """A zero would tell clients to retry immediately into another denial.""" + out = _run_lua( + """ + run(3000, K, argv(3000, "a", 1, 100, 1000)) + local edge = run(3059.9, K, argv(3059.9, "b", 1, 100, 1000)) + assert(edge[1] == 1) + assert(edge[5] >= 1, "reset was " .. tostring(edge[5])) + print("ok") + """ + ) + assert "ok" in out + + +def test_denial_does_not_consume_capacity(): + """A rejected request must not push the oldest entry further out.""" + out = _run_lua( + """ + run(1000, K, argv(1000, "a", 2, 100, 1000)) + run(1001, K, argv(1001, "b", 2, 100, 1000)) + local first = run(1002, K, argv(1002, "c", 2, 100, 1000)) + local second = run(1003, K, argv(1003, "d", 2, 100, 1000)) + assert(first[1] == 1 and second[1] == 1) + assert(first[4] == 2 and second[4] == 2, "denied requests were recorded") + print("ok") + """ + ) + assert "ok" in out