Spaces:
Running on Zero
Running on Zero
| import hashlib | |
| import hmac | |
| import html | |
| import json | |
| import logging | |
| import os | |
| import secrets | |
| import threading | |
| import time | |
| from dataclasses import dataclass | |
| from pathlib import Path | |
| from typing import Any, Dict, List, Optional | |
| from uuid import uuid4 | |
| from quota_fallback import ( | |
| GRANT, | |
| REFUND, | |
| embed_parent_from, | |
| embed_parents, | |
| grant_request_nonce, | |
| is_eligible_embed, | |
| issue_ticket, | |
| request_header, | |
| request_host, | |
| safe_message, | |
| signing_secret, | |
| verify_grant, | |
| ) | |
| logger = logging.getLogger(__name__) | |
| FREE_LIMIT_REASON = "LABS_3DAI_FREE_LIMIT" | |
| GRANT_TTL_SECONDS = 180 | |
| SETTLED_GRANT_TTL_SECONDS = 3600 | |
| MAX_GRANT_ENTRIES = 512 | |
| MAX_GRANTS_PER_SESSION = 8 | |
| SIGNUP_URL = "https://app.skyebrowse.com/register" | |
| EMBED_LIMIT_MESSAGE = ( | |
| "You have used your free 3D AI generations. Upgrade to a Commercial plan to keep generating." | |
| ) | |
| EMBED_UNAVAILABLE_MESSAGE = "3D AI could not confirm your SkyeBrowse plan. Please try again." | |
| PUBLIC_LIMIT_MESSAGE = "You have used your {limit} free generations." | |
| PUBLIC_QUOTA_MESSAGE = "Free generations are busy right now." | |
| def env_flag(name: str) -> bool: | |
| return os.environ.get(name, "").strip().lower() in {"1", "true", "yes"} | |
| def env_int(name: str, default: int) -> int: | |
| try: | |
| return max(0, int(os.environ.get(name, "").strip() or default)) | |
| except ValueError: | |
| return default | |
| def grant_parents() -> Optional[frozenset]: | |
| raw = os.environ.get("LABS_EMBED_GRANTS", "").strip() | |
| if env_flag("LABS_EMBED_GRANTS"): | |
| return None | |
| return frozenset(item.strip() for item in raw.split(",") if item.strip()) | |
| def grants_required(request: Any) -> bool: | |
| if not is_eligible_embed(request): | |
| return False | |
| parents = grant_parents() | |
| if parents is None: | |
| return True | |
| return embed_parent_from(request_header(request, "referer"), request_host(request)) in parents | |
| def is_browser_visit(request: Any) -> bool: | |
| referer = request_header(request, "referer") | |
| return referer.startswith(("https://", "http://")) and not is_eligible_embed(request) | |
| def visitor_ip(request: Any) -> str: | |
| return request_header(request, "x-forwarded-for").split(",")[0].strip() | |
| class GrantEntry: | |
| nonce: str | |
| session_hash: str | |
| issued_at: float | |
| settled_at: Optional[float] = None | |
| mode: Optional[str] = None | |
| notice: Optional[str] = None | |
| class Access: | |
| mode: str = "native" | |
| nonce: Optional[str] = None | |
| visitor_key: Optional[str] = None | |
| notice: Optional[str] = None | |
| class GrantBook: | |
| def __init__( | |
| self, | |
| ttl: float = GRANT_TTL_SECONDS, | |
| max_entries: int = MAX_GRANT_ENTRIES, | |
| settled_ttl: float = SETTLED_GRANT_TTL_SECONDS, | |
| max_per_session: int = MAX_GRANTS_PER_SESSION, | |
| ): | |
| self._lock = threading.Lock() | |
| self._entries: Dict[str, GrantEntry] = {} | |
| self._lapsed: Dict[str, List[str]] = {} | |
| self._ttl = ttl | |
| self._settled_ttl = settled_ttl | |
| self._max_entries = max_entries | |
| self._max_per_session = max_per_session | |
| def request(self, session_hash: str) -> Optional[Dict[str, str]]: | |
| if not session_hash: | |
| return None | |
| with self._lock: | |
| self._purge() | |
| if len(self._entries) >= self._max_entries: | |
| return None | |
| if len(self._session_entries(session_hash)) >= self._max_per_session: | |
| return None | |
| entry = GrantEntry(uuid4().hex, session_hash, time.time()) | |
| self._entries[entry.nonce] = entry | |
| ticket = issue_ticket(entry.nonce, GRANT) | |
| return {"ticket": ticket} if ticket else None | |
| def settle(self, session_hash: str, ack: Any, receipt: Any) -> None: | |
| if not isinstance(ack, dict) or not isinstance(receipt, dict) or not receipt.get("ticket"): | |
| return | |
| if ack.get("ticket") != receipt.get("ticket"): | |
| return | |
| nonce = grant_request_nonce(receipt.get("ticket")) | |
| with self._lock: | |
| self._purge() | |
| entry = self._entries.get(nonce) if nonce else None | |
| if entry is None or entry.session_hash != session_hash or entry.settled_at is not None: | |
| return | |
| entry.settled_at = time.time() | |
| grant = verify_grant(ack.get("grant"), entry.nonce) | |
| if grant is not None: | |
| entry.mode = grant["mode"] | |
| elif ack.get("reason") == FREE_LIMIT_REASON: | |
| entry.notice = safe_message(ack.get("error"), EMBED_LIMIT_MESSAGE) | |
| else: | |
| entry.notice = safe_message(ack.get("error"), EMBED_UNAVAILABLE_MESSAGE) | |
| def take(self, session_hash: str) -> Access: | |
| with self._lock: | |
| self._purge() | |
| entry = self._pop_settled(session_hash) | |
| lapsed = None if entry else self._pop_lapsed(session_hash) | |
| if lapsed: | |
| return Access(mode="denied", nonce=lapsed, notice=embed_notice(EMBED_UNAVAILABLE_MESSAGE)) | |
| if entry is None or entry.mode is None: | |
| notice = entry.notice if entry and entry.notice else EMBED_UNAVAILABLE_MESSAGE | |
| return Access(mode="denied", notice=embed_notice(notice)) | |
| return Access(nonce=entry.nonce) | |
| def _pop_settled(self, session_hash: str) -> Optional[GrantEntry]: | |
| settled = [entry for entry in self._session_entries(session_hash) if entry.settled_at is not None] | |
| return self._entries.pop(min(settled, key=lambda entry: entry.settled_at).nonce) if settled else None | |
| def _pop_lapsed(self, session_hash: str) -> Optional[str]: | |
| lapsed = self._lapsed.get(session_hash) | |
| if not lapsed: | |
| return None | |
| nonce = lapsed.pop(0) | |
| if not lapsed: | |
| del self._lapsed[session_hash] | |
| return nonce | |
| def _session_entries(self, session_hash: str) -> List[GrantEntry]: | |
| return [entry for entry in self._entries.values() if entry.session_hash == session_hash] | |
| def _purge(self) -> None: | |
| now = time.time() | |
| for nonce, entry in list(self._entries.items()): | |
| pending = entry.settled_at is None | |
| started = entry.issued_at if pending else entry.settled_at | |
| if started < now - (self._ttl if pending else self._settled_ttl): | |
| del self._entries[nonce] | |
| if entry.mode is not None: | |
| self._lapse(entry) | |
| def _lapse(self, entry: GrantEntry) -> None: | |
| self._lapsed.setdefault(entry.session_hash, []).append(entry.nonce) | |
| while len(self._lapsed) > self._max_entries: | |
| del self._lapsed[next(iter(self._lapsed))] | |
| class PublicAllowance: | |
| def __init__(self, limit: int, window_seconds: float, path: Optional[Path]): | |
| self.limit = limit | |
| self._window = window_seconds | |
| self._path = path | |
| self._lock = threading.Lock() | |
| self._salt = (signing_secret() or secrets.token_hex(32)).encode("utf-8") | |
| self._uses: Dict[str, List[float]] = self._load() | |
| def enabled(self) -> bool: | |
| return self.limit > 0 | |
| def key_for(self, request: Any) -> Optional[str]: | |
| address = visitor_ip(request) | |
| if not address: | |
| logger.warning("Public allowance skipped: request has no X-Forwarded-For") | |
| return None | |
| return hmac.new(self._salt, address.encode("utf-8"), hashlib.sha256).hexdigest()[:32] | |
| def try_consume(self, key: Optional[str]) -> bool: | |
| if key is None: | |
| return True | |
| with self._lock: | |
| uses = self._recent(key) | |
| if len(uses) >= self.limit: | |
| return False | |
| self._uses[key] = uses + [time.time()] | |
| self._save() | |
| return True | |
| def refund(self, key: Optional[str]) -> None: | |
| if key is None: | |
| return | |
| with self._lock: | |
| uses = self._recent(key) | |
| if uses: | |
| self._uses[key] = uses[:-1] | |
| self._save() | |
| def used(self, key: Optional[str]) -> int: | |
| with self._lock: | |
| return len(self._recent(key)) if key else 0 | |
| def _recent(self, key: str) -> List[float]: | |
| cutoff = time.time() - self._window | |
| return [stamp for stamp in self._uses.get(key, []) if stamp > cutoff] | |
| def _load(self) -> Dict[str, List[float]]: | |
| if self._path is None or not self._path.exists(): | |
| return {} | |
| try: | |
| data = json.loads(self._path.read_text()) | |
| except (OSError, ValueError): | |
| return {} | |
| if not isinstance(data, dict): | |
| return {} | |
| return {str(key): [float(stamp) for stamp in value if isinstance(stamp, (int, float))] | |
| for key, value in data.items() if isinstance(value, list)} | |
| def _save(self) -> None: | |
| if self._path is None: | |
| return | |
| cutoff = time.time() - self._window | |
| snapshot = {key: [stamp for stamp in uses if stamp > cutoff] for key, uses in self._uses.items()} | |
| snapshot = {key: uses for key, uses in snapshot.items() if uses} | |
| self._uses = snapshot | |
| try: | |
| temporary = self._path.with_suffix(".tmp") | |
| temporary.write_text(json.dumps(snapshot)) | |
| temporary.replace(self._path) | |
| except OSError as error: | |
| logger.warning("Public allowance not persisted (%s)", type(error).__name__) | |
| def public_usage_path() -> Optional[Path]: | |
| configured = os.environ.get("LABS_PUBLIC_USAGE_PATH", "").strip() | |
| if configured: | |
| return Path(configured) | |
| base = Path("/data") if Path("/data").is_dir() else Path("/tmp") | |
| return base / "labs_3dai_public_usage.json" | |
| def public_allowance_from_env() -> PublicAllowance: | |
| hours = env_int("LABS_PUBLIC_WINDOW_HOURS", 720) or 720 | |
| return PublicAllowance(env_int("LABS_PUBLIC_FREE_LIMIT", 0), hours * 3600, public_usage_path()) | |
| def embed_notice(message: str) -> str: | |
| return f'<div class="labs-access-notice">{html.escape(message)}</div>' | |
| def signup_notice(message: str) -> str: | |
| return ( | |
| '<div class="labs-access-notice">' | |
| f"<strong>{html.escape(message)}</strong> " | |
| "Create a free SkyeBrowse account to keep generating 3D models. " | |
| f'<a href="{SIGNUP_URL}" target="_blank" rel="noopener">Sign up free</a>' | |
| "</div>" | |
| ) | |
| def bridge_script(path: Path) -> str: | |
| return path.read_text().replace("__LABS_EMBED_PARENTS__", json.dumps(list(embed_parents()))) | |
| class AccessPolicy: | |
| def __init__(self, grants: GrantBook, public: PublicAllowance): | |
| self.grants = grants | |
| self.public = public | |
| self._notices: Dict[str, str] = {} | |
| self._refunds: Dict[str, List[str]] = {} | |
| self._lock = threading.Lock() | |
| def authorize(self, request: Any, session_hash: str) -> Access: | |
| if grants_required(request): | |
| access = self.grants.take(session_hash) | |
| if access.mode == "denied": | |
| self.generation_failed(session_hash, access) | |
| return access | |
| if not (self.public.enabled and is_browser_visit(request)): | |
| return Access() | |
| key = self.public.key_for(request) | |
| if not self.public.try_consume(key): | |
| limit_message = PUBLIC_LIMIT_MESSAGE.format(limit=self.public.limit) | |
| return Access(mode="denied", notice=signup_notice(limit_message)) | |
| return Access(visitor_key=key) | |
| def quota_denied(self, access: Access) -> Optional[str]: | |
| self.public.refund(access.visitor_key) | |
| return signup_notice(PUBLIC_QUOTA_MESSAGE) if access.visitor_key else None | |
| def generation_failed(self, session_hash: str, access: Access) -> None: | |
| self.public.refund(access.visitor_key) | |
| if access.nonce: | |
| with self._lock: | |
| self._refunds.setdefault(session_hash, []).append(access.nonce) | |
| def refund_receipt(self, session_hash: str) -> Optional[Dict[str, str]]: | |
| with self._lock: | |
| pending = self._refunds.get(session_hash) or [] | |
| nonce = pending.pop(0) if pending else None | |
| if not pending: | |
| self._refunds.pop(session_hash, None) | |
| ticket = issue_ticket(nonce, REFUND) if nonce else None | |
| return {"ticket": ticket} if ticket else None | |
| def remember(self, session_hash: str, notice: Optional[str]) -> None: | |
| with self._lock: | |
| if notice: | |
| self._notices[session_hash] = notice | |
| else: | |
| self._notices.pop(session_hash, None) | |
| def pop_notice(self, session_hash: str) -> Optional[str]: | |
| with self._lock: | |
| return self._notices.pop(session_hash, None) | |