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