mirror of
https://github.com/langgenius/dify.git
synced 2026-11-10 13:26:58 +08:00
8cbb6c0130
Co-authored-by: zxhlyh <jasonapring2015@outlook.com> Co-authored-by: JzoNg <jzongcode@gmail.com> Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
733 lines
25 KiB
Python
733 lines
25 KiB
Python
import json
|
|
import logging
|
|
import re
|
|
import secrets
|
|
import string
|
|
import struct
|
|
import subprocess
|
|
import time
|
|
import uuid
|
|
from collections.abc import Callable, Generator, Iterable, Mapping
|
|
from contextlib import contextmanager
|
|
from datetime import datetime
|
|
from hashlib import sha256
|
|
from typing import TYPE_CHECKING, Annotated, Any, Protocol, cast, overload, override
|
|
from uuid import UUID
|
|
from zoneinfo import available_timezones
|
|
|
|
from flask import Request, Response, stream_with_context
|
|
from flask_restx import fields
|
|
from pydantic import BaseModel, ConfigDict, TypeAdapter, WithJsonSchema, with_config
|
|
from pydantic.functional_validators import AfterValidator
|
|
from typing_extensions import TypedDict
|
|
|
|
from configs import dify_config
|
|
from core.app.features.rate_limiting.rate_limit import RateLimitGenerator
|
|
from extensions.ext_redis import redis_client
|
|
from extensions.redis_names import serialize_redis_name
|
|
from graphon.file import helpers as file_helpers
|
|
from graphon.model_runtime.utils.encoders import jsonable_encoder
|
|
|
|
if TYPE_CHECKING:
|
|
from models import Account
|
|
from models.model import EndUser
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@with_config(ConfigDict(extra="allow"))
|
|
class _TokenData(TypedDict, total=False):
|
|
"""Shared baseline token payload.
|
|
|
|
`extra='allow'` keeps TokenManager from silently stripping business-
|
|
specific metadata keys while still validating the common auth fields.
|
|
Business flows that need stronger guarantees should validate again at
|
|
their own boundary with a dedicated Pydantic model.
|
|
|
|
For the change-email flow specifically, `email_change_phase` is the
|
|
discriminator used by `services.entities.auth_entities.ChangeEmailTokenData`.
|
|
It is declared here so the shared token adapter can still provide baseline
|
|
validation for the state-machine key without taking over the full business
|
|
model.
|
|
"""
|
|
|
|
account_id: str | None
|
|
email: str
|
|
token_type: str
|
|
code: str
|
|
old_email: str
|
|
phase: str
|
|
email_change_phase: str
|
|
|
|
|
|
_token_data_adapter: TypeAdapter[_TokenData] = TypeAdapter(_TokenData)
|
|
|
|
|
|
def _stream_with_request_context(response: object) -> Any:
|
|
"""Bridge Flask's loosely-typed streaming helper without leaking casts into callers."""
|
|
return cast(Any, stream_with_context)(response)
|
|
|
|
|
|
def escape_like_pattern(pattern: str) -> str:
|
|
"""
|
|
Escape special characters in a string for safe use in SQL LIKE patterns.
|
|
|
|
This function escapes the special characters used in SQL LIKE patterns:
|
|
- Backslash (\\) -> \\
|
|
- Percent (%) -> \\%
|
|
- Underscore (_) -> \\_
|
|
|
|
The escaped pattern can then be safely used in SQL LIKE queries with the
|
|
ESCAPE '\\' clause to prevent SQL injection via LIKE wildcards.
|
|
|
|
Args:
|
|
pattern: The string pattern to escape
|
|
|
|
Returns:
|
|
Escaped string safe for use in SQL LIKE queries
|
|
|
|
Examples:
|
|
>>> escape_like_pattern("50% discount")
|
|
'50\\% discount'
|
|
>>> escape_like_pattern("test_data")
|
|
'test\\_data'
|
|
>>> escape_like_pattern("path\\to\\file")
|
|
'path\\\\to\\\\file'
|
|
"""
|
|
if not pattern:
|
|
return pattern
|
|
# Escape backslash first, then percent and underscore
|
|
return pattern.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_")
|
|
|
|
|
|
def extract_tenant_id(user: "Account | EndUser") -> str | None:
|
|
"""
|
|
Extract tenant_id from Account or EndUser object.
|
|
|
|
Args:
|
|
user: Account or EndUser object
|
|
|
|
Returns:
|
|
tenant_id string if available, None otherwise
|
|
|
|
Raises:
|
|
ValueError: If user is neither Account nor EndUser
|
|
"""
|
|
from models import Account
|
|
from models.model import EndUser
|
|
|
|
match user:
|
|
case Account():
|
|
return user.current_tenant_id
|
|
case EndUser():
|
|
return user.tenant_id
|
|
case _:
|
|
raise ValueError(f"Invalid user type: {type(user)}. Expected Account or EndUser.")
|
|
|
|
|
|
def run(script):
|
|
return subprocess.getstatusoutput("source /root/.bashrc && " + script)
|
|
|
|
|
|
class AppIconUrlField(fields.Raw):
|
|
@override
|
|
def schema(self) -> dict[str, object]:
|
|
return {"type": "string", "nullable": True}
|
|
|
|
@override
|
|
def output(self, key, obj, **kwargs):
|
|
if obj is None:
|
|
return None
|
|
|
|
from models.model import App, IconType, Site
|
|
|
|
if isinstance(obj, dict) and "app" in obj:
|
|
obj = obj["app"]
|
|
|
|
if isinstance(obj, App | Site) and obj.icon_type == IconType.IMAGE:
|
|
return build_icon_url(obj.icon_type, obj.icon)
|
|
return None
|
|
|
|
|
|
def build_icon_url(icon_type: Any, icon: str | None) -> str | None:
|
|
if icon is None or icon_type is None:
|
|
return None
|
|
|
|
from models.model import IconType
|
|
|
|
icon_type_value = icon_type.value if isinstance(icon_type, IconType) else str(icon_type)
|
|
if icon_type_value.lower() != IconType.IMAGE:
|
|
return None
|
|
return file_helpers.get_signed_file_url(icon)
|
|
|
|
|
|
def build_avatar_url(avatar: str | None) -> str | None:
|
|
if avatar is None:
|
|
return None
|
|
if avatar.startswith(("http://", "https://")):
|
|
return avatar
|
|
return file_helpers.get_signed_file_url(avatar)
|
|
|
|
|
|
class TimestampField(fields.Raw):
|
|
@override
|
|
def schema(self) -> dict[str, object]:
|
|
return {"type": "integer", "format": "int64"}
|
|
|
|
@override
|
|
def format(self, value) -> int:
|
|
return int(value.timestamp())
|
|
|
|
|
|
class OptionalTimestampField(fields.Raw):
|
|
@override
|
|
def schema(self) -> dict[str, object]:
|
|
return {"type": "integer", "format": "int64", "nullable": True}
|
|
|
|
@override
|
|
def format(self, value) -> int | None:
|
|
if value is None:
|
|
return None
|
|
return int(value.timestamp())
|
|
|
|
|
|
@overload
|
|
def to_timestamp(value: datetime) -> int: ...
|
|
|
|
|
|
@overload
|
|
def to_timestamp(value: int) -> int: ...
|
|
|
|
|
|
@overload
|
|
def to_timestamp(value: None) -> None: ...
|
|
|
|
|
|
def to_timestamp(value: datetime | int | None) -> int | None:
|
|
"""Normalize API response timestamp values to epoch seconds."""
|
|
if isinstance(value, datetime):
|
|
return int(value.timestamp())
|
|
return value
|
|
|
|
|
|
def dump_response(model: type[BaseModel], data: Any, **dump_kwargs: Any) -> dict[str, Any]:
|
|
"""Serialize a Pydantic response model to JSON-compatible dict output.
|
|
|
|
Extra keyword arguments are forwarded to ``model_dump`` (e.g. ``include=...``,
|
|
``exclude=...``, ``exclude_unset=True``).
|
|
"""
|
|
return model.model_validate(data, from_attributes=True).model_dump(mode="json", **dump_kwargs)
|
|
|
|
|
|
def current_timestamp() -> int:
|
|
"""Return the current Unix timestamp in seconds."""
|
|
return int(time.time())
|
|
|
|
|
|
def email(email):
|
|
# Define a regex pattern for email addresses
|
|
pattern = r"^[\w\.!#$%&'*+\-/=?^_`{|}~]+@([\w-]+\.)+[\w-]{2,}$"
|
|
# Reject trailing newlines as well as other unconsumed input.
|
|
if re.fullmatch(pattern, email) is not None:
|
|
return email
|
|
|
|
error = f"{email} is not a valid email."
|
|
raise ValueError(error)
|
|
|
|
|
|
EmailStr = Annotated[str, AfterValidator(email)]
|
|
|
|
|
|
def uuid_value(value: str | UUID) -> str:
|
|
if isinstance(value, UUID):
|
|
return str(value)
|
|
|
|
if value == "":
|
|
return str(value)
|
|
|
|
try:
|
|
uuid_obj = uuid.UUID(value)
|
|
return str(uuid_obj)
|
|
except ValueError as caught_error:
|
|
error = f"{value} is not a valid uuid."
|
|
raise ValueError(error) from caught_error
|
|
|
|
|
|
def normalize_uuid(value: str | UUID) -> str:
|
|
if not value:
|
|
return ""
|
|
|
|
try:
|
|
return uuid_value(value)
|
|
except ValueError as exc:
|
|
raise ValueError("must be a valid UUID") from exc
|
|
|
|
|
|
def parse_uuid_str_or_none(value: str | None) -> str | None:
|
|
"""
|
|
Return None for missing/empty UUID-like values.
|
|
|
|
Keep non-empty values unchanged to avoid changing behavior in paths that
|
|
currently pass placeholder IDs in tests/mocks.
|
|
"""
|
|
if value is None or not str(value).strip():
|
|
return None
|
|
return str(value)
|
|
|
|
|
|
UUIDStrOrEmpty = Annotated[str, AfterValidator(normalize_uuid)]
|
|
|
|
|
|
def _strict_uuid(value: str | UUID) -> str:
|
|
if not value:
|
|
raise ValueError("must be a non-empty valid UUID")
|
|
try:
|
|
return uuid_value(value)
|
|
except ValueError as exc:
|
|
raise ValueError("must be a valid UUID") from exc
|
|
|
|
|
|
UUIDStr = Annotated[
|
|
str,
|
|
AfterValidator(_strict_uuid),
|
|
WithJsonSchema({"format": "uuid", "type": "string"}),
|
|
]
|
|
|
|
|
|
def alphanumeric(value: str):
|
|
# check if the value is alphanumeric and underlined
|
|
# Use re.fullmatch instead of re.match to reject trailing newlines.
|
|
# In Python, '$' matches at end-of-string OR just before a trailing newline,
|
|
# so re.match accepts "tool_name\n". re.fullmatch requires the entire
|
|
# string to match. Regression for #39666 (sibling of #39234 / #39548).
|
|
if re.fullmatch(r"^[a-zA-Z0-9_]+$", value):
|
|
return value
|
|
|
|
raise ValueError(f"{value} is not a valid alphanumeric value")
|
|
|
|
|
|
def timestamp_value(timestamp):
|
|
try:
|
|
int_timestamp = int(timestamp)
|
|
if int_timestamp < 0:
|
|
raise ValueError
|
|
return int_timestamp
|
|
except ValueError as caught_error:
|
|
error = f"{timestamp} is not a valid timestamp."
|
|
raise ValueError(error) from caught_error
|
|
|
|
|
|
class StrLen:
|
|
"""Restrict input to an integer in a range (inclusive)"""
|
|
|
|
def __init__(self, max_length, argument="argument"):
|
|
self.max_length = max_length
|
|
self.argument = argument
|
|
|
|
def __call__(self, value):
|
|
length = len(value)
|
|
if length > self.max_length:
|
|
error = f"Invalid {self.argument}: {value}. {self.argument} cannot exceed length {self.max_length}"
|
|
raise ValueError(error)
|
|
|
|
return value
|
|
|
|
|
|
class DatetimeString:
|
|
def __init__(self, format, argument="argument"):
|
|
self.format = format
|
|
self.argument = argument
|
|
|
|
def __call__(self, value):
|
|
try:
|
|
datetime.strptime(value, self.format)
|
|
except ValueError as caught_error:
|
|
error = f"Invalid {self.argument}: {value}. {self.argument} must be conform to the format {self.format}"
|
|
raise ValueError(error) from caught_error
|
|
|
|
return value
|
|
|
|
|
|
def timezone(timezone_string):
|
|
if timezone_string and timezone_string in available_timezones():
|
|
return timezone_string
|
|
|
|
error = f"{timezone_string} is not a valid timezone."
|
|
raise ValueError(error)
|
|
|
|
|
|
def convert_datetime_to_date(field, target_timezone: str = ":tz"):
|
|
if dify_config.DB_TYPE == "postgresql":
|
|
return f"DATE(DATE_TRUNC('day', {field} AT TIME ZONE 'UTC' AT TIME ZONE {target_timezone}))"
|
|
elif dify_config.DB_TYPE in ["mysql", "oceanbase", "seekdb"]:
|
|
return f"DATE(CONVERT_TZ({field}, 'UTC', {target_timezone}))"
|
|
else:
|
|
raise NotImplementedError(f"Unsupported database type: {dify_config.DB_TYPE}")
|
|
|
|
|
|
def generate_string(n):
|
|
"""
|
|
Generates a cryptographically secure random string of the specified length.
|
|
|
|
This function uses a cryptographically secure pseudorandom number generator (CSPRNG)
|
|
to create a string composed of ASCII letters (both uppercase and lowercase) and digits.
|
|
|
|
Each character in the generated string provides approximately 5.95 bits of entropy
|
|
(log2(62)). To ensure a minimum of 128 bits of entropy for security purposes, the
|
|
length of the string (`n`) should be at least 22 characters.
|
|
|
|
Args:
|
|
n (int): The length of the random string to generate. For secure usage,
|
|
`n` should be 22 or greater.
|
|
|
|
Returns:
|
|
str: A random string of length `n` composed of ASCII letters and digits.
|
|
|
|
Note:
|
|
This function is suitable for generating credentials or other secure tokens.
|
|
"""
|
|
letters_digits = string.ascii_letters + string.digits
|
|
result = ""
|
|
for _ in range(n):
|
|
result += secrets.choice(letters_digits)
|
|
|
|
return result
|
|
|
|
|
|
def extract_remote_ip(request: Request) -> str:
|
|
"""Use the WSGI address, normalized by opt-in trusted proxy middleware."""
|
|
return cast(str, request.remote_addr)
|
|
|
|
|
|
def generate_text_hash(text: str) -> str:
|
|
hash_text = str(text) + "None"
|
|
return sha256(hash_text.encode()).hexdigest()
|
|
|
|
|
|
def compact_generate_response(
|
|
response: Mapping[str, Any] | Iterable[str],
|
|
) -> Response:
|
|
if isinstance(response, Mapping):
|
|
return Response(
|
|
response=json.dumps(jsonable_encoder(response)),
|
|
status=200,
|
|
content_type="application/json; charset=utf-8",
|
|
)
|
|
else:
|
|
stream_response = response
|
|
|
|
def generate() -> Generator[str]:
|
|
yield from stream_response
|
|
|
|
return Response(
|
|
_stream_with_request_context(generate()),
|
|
status=200,
|
|
mimetype="text/event-stream",
|
|
)
|
|
|
|
|
|
def length_prefixed_response(
|
|
magic_number: int,
|
|
response: Mapping[str, Any] | BaseModel | Generator[str | bytes] | RateLimitGenerator,
|
|
) -> Response:
|
|
"""
|
|
This function is used to return a response with a length prefix.
|
|
Magic number is a one byte number that indicates the type of the response.
|
|
|
|
For a compatibility with latest plugin daemon https://github.com/langgenius/dify-plugin-daemon/pull/341
|
|
Avoid using line-based response, it leads a memory issue.
|
|
|
|
We uses following format:
|
|
| Field | Size | Description |
|
|
|---------------|----------|---------------------------------|
|
|
| Magic Number | 1 byte | Magic number identifier |
|
|
| Reserved | 1 byte | Reserved field |
|
|
| Header Length | 2 bytes | Header length (usually 0xa) |
|
|
| Data Length | 4 bytes | Length of the data |
|
|
| Reserved | 6 bytes | Reserved fields |
|
|
| Data | Variable | Actual data content |
|
|
|
|
| Reserved Fields | Header | Data |
|
|
|-----------------|----------|----------|
|
|
| 4 bytes total | Variable | Variable |
|
|
|
|
all data is in little endian
|
|
"""
|
|
|
|
def pack_response_with_length_prefix(response: bytes) -> bytes:
|
|
header_length = 0xA
|
|
data_length = len(response)
|
|
# | Magic Number 1byte | Reserved 1byte | Header Length 2bytes | Data Length 4bytes | Reserved 6bytes | Data
|
|
return struct.pack("<BBHI", magic_number, 0, header_length, data_length) + b"\x00" * 6 + response
|
|
|
|
match response:
|
|
case Mapping():
|
|
return Response(
|
|
response=pack_response_with_length_prefix(json.dumps(jsonable_encoder(response)).encode("utf-8")),
|
|
status=200,
|
|
mimetype="application/json",
|
|
)
|
|
case BaseModel():
|
|
return Response(
|
|
response=pack_response_with_length_prefix(response.model_dump_json().encode("utf-8")),
|
|
status=200,
|
|
mimetype="application/json",
|
|
)
|
|
|
|
stream_response = response
|
|
|
|
def generate() -> Generator[bytes]:
|
|
for chunk in stream_response:
|
|
if isinstance(chunk, str):
|
|
yield pack_response_with_length_prefix(chunk.encode("utf-8"))
|
|
else:
|
|
yield pack_response_with_length_prefix(chunk)
|
|
|
|
return Response(
|
|
_stream_with_request_context(generate()),
|
|
status=200,
|
|
mimetype="text/event-stream",
|
|
)
|
|
|
|
|
|
class TokenManager:
|
|
@classmethod
|
|
def generate_token(
|
|
cls,
|
|
token_type: str,
|
|
account_id: str | None = None,
|
|
email: str | None = None,
|
|
additional_data: dict[str, Any] | None = None,
|
|
) -> str:
|
|
if account_id is None and email is None:
|
|
raise ValueError("Account or email must be provided")
|
|
|
|
if account_id:
|
|
old_token = cls._get_current_token_for_account(account_id, token_type)
|
|
if old_token:
|
|
if isinstance(old_token, bytes):
|
|
old_token = old_token.decode("utf-8")
|
|
cls.revoke_token(old_token, token_type)
|
|
|
|
token = str(uuid.uuid4())
|
|
token_data = {"account_id": account_id, "email": email, "token_type": token_type}
|
|
if additional_data:
|
|
token_data.update(additional_data)
|
|
|
|
expiry_minutes = dify_config.model_dump().get(f"{token_type.upper()}_TOKEN_EXPIRY_MINUTES")
|
|
if expiry_minutes is None:
|
|
raise ValueError(f"Expiry minutes for {token_type} token is not set")
|
|
token_key = cls._get_token_key(token, token_type)
|
|
expiry_seconds = int(expiry_minutes * 60)
|
|
redis_client.setex(token_key, expiry_seconds, json.dumps(token_data))
|
|
|
|
if account_id:
|
|
cls._set_current_token_for_account(account_id, token, token_type, expiry_minutes)
|
|
|
|
return token
|
|
|
|
@classmethod
|
|
def _get_token_key(cls, token: str, token_type: str) -> str:
|
|
return f"{token_type}:token:{token}"
|
|
|
|
@classmethod
|
|
def revoke_token(cls, token: str, token_type: str):
|
|
token_key = cls._get_token_key(token, token_type)
|
|
redis_client.delete(token_key)
|
|
|
|
@classmethod
|
|
def get_token_data(cls, token: str, token_type: str) -> dict[str, Any] | None:
|
|
key = cls._get_token_key(token, token_type)
|
|
token_data_json = redis_client.get(key)
|
|
if token_data_json is None:
|
|
logger.warning("%s token %s not found with key %s", token_type, token, key)
|
|
return None
|
|
return dict(_token_data_adapter.validate_json(token_data_json))
|
|
|
|
@classmethod
|
|
def _get_current_token_for_account(cls, account_id: str, token_type: str) -> str | None:
|
|
key = cls._get_account_token_key(account_id, token_type)
|
|
current_token: str | None = redis_client.get(key)
|
|
return current_token
|
|
|
|
@classmethod
|
|
def _set_current_token_for_account(cls, account_id: str, token: str, token_type: str, expiry_minutes: int | float):
|
|
key = cls._get_account_token_key(account_id, token_type)
|
|
expiry_seconds = int(expiry_minutes * 60)
|
|
redis_client.setex(key, expiry_seconds, token)
|
|
|
|
@classmethod
|
|
def _get_account_token_key(cls, account_id: str, token_type: str) -> str:
|
|
return f"{token_type}:account:{account_id}"
|
|
|
|
|
|
class _RateLimiterRedisClient(Protocol):
|
|
def eval(self, script: str, numkeys: int, *args: str | int) -> object: ...
|
|
def register_script(self, script: str) -> Callable[..., Any]: ...
|
|
|
|
def zadd(self, name: str | bytes, mapping: dict[str | bytes | int | float, float | int | str | bytes]) -> int: ...
|
|
|
|
def zrem(self, name: str | bytes, *values: str | bytes) -> int: ...
|
|
|
|
def zremrangebyscore(self, name: str | bytes, min: str | float, max: str | float) -> int: ...
|
|
|
|
def zcard(self, name: str | bytes) -> int: ...
|
|
|
|
def expire(self, name: str | bytes, time: int) -> bool: ...
|
|
|
|
|
|
def _default_rate_limit_member_factory() -> str:
|
|
current_time = int(time.time())
|
|
return f"{current_time}:{secrets.token_urlsafe(nbytes=8)}"
|
|
|
|
|
|
_ACQUIRE_RATE_LIMIT_LUA = """
|
|
-- EVAL contract: exactly 1 key and 4 arguments.
|
|
-- KEYS[1]: Physical Redis ZSET key for one rate-limit subject. Members are
|
|
-- unique request IDs; scores are Redis Unix timestamps in seconds.
|
|
-- ARGV[1]: Positive rolling-window duration in seconds.
|
|
-- ARGV[2]: Positive maximum number of accepted requests within that window.
|
|
-- ARGV[3]: Minimum seconds between accepted requests, from 0 to ARGV[1].
|
|
-- ARGV[4]: Unique, nonempty request ID to consume a slot if allowed; an empty
|
|
-- string only queries the delay without consuming a slot.
|
|
-- Returns: 0 when allowed; otherwise the positive number of seconds to wait,
|
|
-- rounded up and accounting for both the window limit and cooldown.
|
|
-- Every call removes scores <= now - window, retaining (now - window, now].
|
|
-- Only an allowed, nonempty request ID is inserted and refreshes the key TTL.
|
|
-- Rejected requests and delay queries do not consume quota or refresh TTL.
|
|
-- Redis TIME and atomic execution keep these decisions consistent across nodes.
|
|
local clock = redis.call('TIME')
|
|
local now = tonumber(clock[1]) + tonumber(clock[2]) / 1000000
|
|
local window = tonumber(ARGV[1])
|
|
local maximum = tonumber(ARGV[2])
|
|
local cooldown = tonumber(ARGV[3])
|
|
redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', now - window)
|
|
local count = redis.call('ZCARD', KEYS[1])
|
|
local wait = 0
|
|
if count >= maximum then
|
|
local oldest = redis.call('ZRANGE', KEYS[1], count - maximum, count - maximum, 'WITHSCORES')
|
|
wait = math.max(wait, tonumber(oldest[2]) + window - now)
|
|
end
|
|
local latest = redis.call('ZREVRANGE', KEYS[1], 0, 0, 'WITHSCORES')
|
|
if #latest > 0 then wait = math.max(wait, tonumber(latest[2]) + cooldown - now) end
|
|
if wait > 0 then return math.ceil(wait) end
|
|
if ARGV[4] ~= '' then
|
|
redis.call('ZADD', KEYS[1], now, ARGV[4])
|
|
redis.call('EXPIRE', KEYS[1], math.max(window, cooldown))
|
|
end
|
|
return 0
|
|
"""
|
|
|
|
|
|
class RateLimiter:
|
|
def __init__(
|
|
self,
|
|
prefix: str,
|
|
max_attempts: int,
|
|
time_window: int,
|
|
member_factory: Callable[[], str] = _default_rate_limit_member_factory,
|
|
redis_client: _RateLimiterRedisClient = redis_client,
|
|
):
|
|
self.prefix = prefix
|
|
self.max_attempts = max_attempts
|
|
self.time_window = time_window
|
|
self._member_factory = member_factory
|
|
self._redis_client = redis_client
|
|
|
|
def _get_key(self, email: str) -> str:
|
|
return f"{self.prefix}:{email}"
|
|
|
|
@contextmanager
|
|
def reserve(self, email: str) -> Generator[bool]:
|
|
"""Reserve a slot atomically, releasing only this slot if the operation fails."""
|
|
key = self._get_key(email)
|
|
member = self._member_factory()
|
|
script = self._redis_client.register_script("""
|
|
redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', ARGV[1] - ARGV[2])
|
|
if redis.call('ZSCORE', KEYS[1], ARGV[4]) then return 1 end
|
|
if redis.call('ZCARD', KEYS[1]) >= tonumber(ARGV[3]) then return 0 end
|
|
redis.call('ZADD', KEYS[1], ARGV[1], ARGV[4])
|
|
redis.call('EXPIRE', KEYS[1], ARGV[2] * 2)
|
|
return 1
|
|
""")
|
|
acquired = bool(script(keys=[key], args=[int(time.time()), self.time_window, self.max_attempts, member]))
|
|
try:
|
|
yield acquired
|
|
except BaseException:
|
|
if acquired:
|
|
try:
|
|
self._redis_client.zrem(key, member)
|
|
except Exception:
|
|
logger.exception("Failed to release rate limit reservation")
|
|
raise
|
|
|
|
def is_rate_limited(self, email: str) -> bool:
|
|
"""Legacy non-consuming limit query, retained for existing admission policies.
|
|
|
|
New callers should use acquire() for atomic admission or retry_after()
|
|
for a delay query. Migrating existing callers requires auditing when
|
|
their events should consume quota.
|
|
"""
|
|
key = self._get_key(email)
|
|
current_time = int(time.time())
|
|
window_start_time = current_time - self.time_window
|
|
|
|
self._redis_client.zremrangebyscore(key, "-inf", window_start_time)
|
|
attempts = self._redis_client.zcard(key)
|
|
|
|
return bool(attempts and int(attempts) >= self.max_attempts)
|
|
|
|
def increment_rate_limit(self, email: str):
|
|
"""Record an event unconditionally, without enforcing admission limits.
|
|
|
|
For atomic check-and-consume admission, use acquire() instead of calling
|
|
is_rate_limited() and this method separately.
|
|
"""
|
|
key = self._get_key(email)
|
|
member = self._member_factory()
|
|
current_time = int(time.time())
|
|
|
|
self._redis_client.zadd(key, {member: current_time})
|
|
self._redis_client.expire(key, self.time_window * 2)
|
|
|
|
def acquire(self, subject: str, *, cooldown_seconds: int = 0) -> int:
|
|
"""Atomically consume a slot, or return seconds to wait without consuming it."""
|
|
return self._atomic_wait(subject, cooldown_seconds, self._member_factory())
|
|
|
|
def retry_after(self, subject: str, *, cooldown_seconds: int = 0) -> int:
|
|
"""Read the combined cooldown/window delay without consuming a slot."""
|
|
return self._atomic_wait(subject, cooldown_seconds, "")
|
|
|
|
def _atomic_wait(self, subject: str, cooldown_seconds: int, member: str) -> int:
|
|
if self.max_attempts <= 0 or self.time_window <= 0 or not 0 <= cooldown_seconds <= self.time_window:
|
|
raise ValueError("Atomic rate limits require positive limits and a cooldown within the window")
|
|
# eval bypasses the wrapper's key prefix; Redis time prevents replica clock skew.
|
|
remaining = self._redis_client.eval(
|
|
_ACQUIRE_RATE_LIMIT_LUA,
|
|
1,
|
|
serialize_redis_name(self._get_key(subject)),
|
|
self.time_window,
|
|
self.max_attempts,
|
|
cooldown_seconds,
|
|
member,
|
|
)
|
|
if not isinstance(remaining, int):
|
|
raise ValueError("Invalid Redis rate limit response")
|
|
return remaining
|
|
|
|
def seconds_until_available(self, email: str) -> int:
|
|
"""Seconds until the oldest in-window entry expires, freeing a slot.
|
|
|
|
Defensive floor of 1 second. Caller should only invoke this after
|
|
is_rate_limited() returned True. New callers should use retry_after(),
|
|
which returns 0 when no delay is required.
|
|
"""
|
|
key = self._get_key(email)
|
|
oldest = cast(Any, self._redis_client).zrange(key, 0, 0, withscores=True)
|
|
if not oldest:
|
|
return 1
|
|
_member, score = oldest[0]
|
|
free_at = int(score) + self.time_window
|
|
remaining = free_at - int(time.time())
|
|
return max(remaining, 1)
|