99 lines
3.5 KiB
Python
99 lines
3.5 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
|
|
|
"""Single-use, short-lived confirmation tokens for the llama.cpp host-binary swap.
|
|
|
|
The update runs an OS installer that replaces the binary on the machine running
|
|
Studio, so it must not fire from an unconfirmed click, a stale banner, or a replay.
|
|
Confirmation is required uniformly for every caller (not gated on "same machine",
|
|
since a headless SSH server has no host-local session). A token binds its offered
|
|
build (``target_tag``), is single-use with a short TTL, and lives in-process
|
|
(matching the single-process backend; use an HMAC stateless token for multi-worker).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import secrets
|
|
import threading
|
|
import time
|
|
from datetime import datetime, timezone
|
|
from typing import Optional, Tuple
|
|
|
|
# How long a freshly minted confirmation token stays valid.
|
|
CONFIRM_TOKEN_TTL_SECONDS = 300
|
|
# Cap the store so unapplied confirm calls can't grow memory unbounded; oldest first.
|
|
_MAX_TOKENS = 64
|
|
|
|
_lock = threading.Lock()
|
|
# token -> (target_tag, expires_at_monotonic)
|
|
_tokens: "dict[str, Tuple[str, float]]" = {}
|
|
|
|
|
|
def _iso(ts_epoch: float) -> str:
|
|
return (
|
|
datetime.fromtimestamp(ts_epoch, tz = timezone.utc)
|
|
.replace(microsecond = 0)
|
|
.isoformat()
|
|
.replace("+00:00", "Z")
|
|
)
|
|
|
|
|
|
def _purge_expired_locked(now: float) -> None:
|
|
expired = [tok for tok, (_tag, exp) in _tokens.items() if exp <= now]
|
|
for tok in expired:
|
|
_tokens.pop(tok, None)
|
|
|
|
|
|
def mint_confirm_token(
|
|
target_tag: str, *, ttl_seconds: int = CONFIRM_TOKEN_TTL_SECONDS
|
|
) -> Tuple[str, str]:
|
|
"""Mint a single-use token bound to ``target_tag``.
|
|
|
|
Returns ``(token, expires_at_iso)``. ``expires_at_iso`` is wall-clock UTC for
|
|
display; validity itself is tracked on a monotonic clock so a system time
|
|
change cannot extend or shorten it.
|
|
"""
|
|
now_mono = time.monotonic()
|
|
now_wall = time.time()
|
|
token = secrets.token_urlsafe(32)
|
|
with _lock:
|
|
_purge_expired_locked(now_mono)
|
|
if len(_tokens) >= _MAX_TOKENS:
|
|
oldest = min(_tokens, key = lambda t: _tokens[t][1]) # evict nearest expiry
|
|
_tokens.pop(oldest, None)
|
|
_tokens[token] = (target_tag, now_mono + ttl_seconds)
|
|
return token, _iso(now_wall + ttl_seconds)
|
|
|
|
|
|
def consume_confirm_token(
|
|
token: Optional[str], current_target_tag: str
|
|
) -> Tuple[bool, Optional[str]]:
|
|
"""Validate and consume a token for a swap to ``current_target_tag``.
|
|
|
|
Returns ``(ok, reason)``. On success ``(True, None)`` and the token is burned
|
|
(single use). On failure ``ok`` is False and ``reason`` is one of:
|
|
``invalid_token`` (missing/unknown), ``expired_token``, ``stale_target``
|
|
(bound to a different build than the one about to install).
|
|
"""
|
|
if not token:
|
|
return False, "invalid_token"
|
|
now_mono = time.monotonic()
|
|
with _lock:
|
|
# Pop first so any outcome consumes the token, and an expired hit still
|
|
# reports "expired" rather than collapsing to "invalid".
|
|
entry = _tokens.pop(token, None)
|
|
_purge_expired_locked(now_mono)
|
|
if entry is None:
|
|
return False, "invalid_token"
|
|
target_tag, expires_at = entry
|
|
if expires_at <= now_mono:
|
|
return False, "expired_token"
|
|
if target_tag != (current_target_tag or ""):
|
|
return False, "stale_target"
|
|
return True, None
|
|
|
|
|
|
def reset_tokens_for_tests() -> None:
|
|
"""Test-only: clear the token store."""
|
|
with _lock:
|
|
_tokens.clear()
|