Studio: require confirmation and report the target machine for llama.cpp updates
A remote authenticated browser could swap the host llama.cpp binary with no confirmation and no visible result. POST /api/llama/update/confirm now mints a single-use, build-bound token, and POST /api/llama/update refuses a bodyless or unconfirmed call instead of running the swap. Update responses include the host name and platform so the result is attributable. Confirmation is required uniformly, so a headless remote server is not gated out; the desktop app, local browser, and CLI update path are unchanged.
This commit is contained in:
danielhanchen
parent
74d1a284eb
commit
3b446b6d24
4 changed files with 648 additions and 4 deletions
|
|
@ -14,19 +14,63 @@ never blocks on a missing marker / offline GitHub.
|
|||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import socket
|
||||
import sys
|
||||
import threading
|
||||
from typing import Optional
|
||||
|
||||
from fastapi import APIRouter, Depends, Query
|
||||
from fastapi import APIRouter, Body, Depends, Query
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from auth.authentication import get_current_subject
|
||||
from loggers import get_logger
|
||||
from utils.llama_cpp_update import get_update_status, start_update
|
||||
from utils.update_confirm import (
|
||||
CONFIRM_TOKEN_TTL_SECONDS,
|
||||
consume_confirm_token,
|
||||
mint_confirm_token,
|
||||
)
|
||||
|
||||
logger = get_logger(__name__)
|
||||
router = APIRouter()
|
||||
|
||||
# Messages for a refused apply, keyed by outcome. A refusal never runs the swap.
|
||||
_REFUSAL_MESSAGES = {
|
||||
"confirmation_required": (
|
||||
"Confirm the llama.cpp update before it runs. It will download and swap "
|
||||
"the binary on the machine running Studio."
|
||||
),
|
||||
"invalid_token": (
|
||||
"The confirmation is missing or unrecognized. Re-check for the update and "
|
||||
"confirm again before it runs."
|
||||
),
|
||||
"expired_token": (
|
||||
"The confirmation expired. Re-check for the update and confirm again."
|
||||
),
|
||||
"stale_target": (
|
||||
"The available build changed since you confirmed. Re-check for the update "
|
||||
"and confirm the new build before it runs."
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
class UpdateMachine(BaseModel):
|
||||
"""Host a swap targets, so a remote operator sees which machine would change."""
|
||||
|
||||
hostname: str = Field("", description = "Hostname of the machine running Studio.")
|
||||
platform: str = Field(
|
||||
"", description = "Host OS tag (sys.platform), e.g. 'linux', 'darwin', 'win32'."
|
||||
)
|
||||
|
||||
|
||||
def _current_machine() -> UpdateMachine:
|
||||
"""Best-effort host identity; cross-OS and never branched on."""
|
||||
try:
|
||||
hostname = socket.gethostname() or "unknown"
|
||||
except Exception: # pragma: no cover
|
||||
hostname = "unknown"
|
||||
return UpdateMachine(hostname = hostname, platform = sys.platform)
|
||||
|
||||
|
||||
class LlamaUpdateJob(BaseModel):
|
||||
state: str = Field("idle", description = "idle | running | success | error")
|
||||
|
|
@ -62,13 +106,59 @@ class LlamaUpdateStatusResponse(BaseModel):
|
|||
update_size_bytes: Optional[int] = Field(
|
||||
None, description = "Download size of the prebuilt Update would fetch, in bytes."
|
||||
)
|
||||
machine: UpdateMachine = Field(
|
||||
default_factory = UpdateMachine,
|
||||
description = "The host this status describes; the machine an Update would change.",
|
||||
)
|
||||
job: LlamaUpdateJob = Field(default_factory = LlamaUpdateJob)
|
||||
|
||||
|
||||
class LlamaUpdateRequest(BaseModel):
|
||||
"""Body for POST /update. Empty body means no confirmation, so the swap is
|
||||
refused (safe default) and a stale banner or replay cannot swap the binary."""
|
||||
|
||||
confirm_token: Optional[str] = Field(
|
||||
None,
|
||||
description = "Single-use token from POST /update/confirm, bound to the exact offered build.",
|
||||
)
|
||||
confirmed: bool = Field(
|
||||
False,
|
||||
description = "Explicit confirmation for non-interactive/CLI callers that do not use a token.",
|
||||
)
|
||||
|
||||
|
||||
class LlamaUpdateConfirmResponse(BaseModel):
|
||||
"""Pending swap details plus, when appliable, a confirm token for the prompt."""
|
||||
|
||||
update_available: bool = False
|
||||
appliable: bool = Field(
|
||||
False, description = "True when an update exists and can actually be applied here."
|
||||
)
|
||||
reason: Optional[str] = Field(
|
||||
None, description = "Why not appliable: up_to_date | local_link | ..."
|
||||
)
|
||||
machine: UpdateMachine = Field(default_factory = UpdateMachine)
|
||||
installed_tag: Optional[str] = None
|
||||
latest_tag: Optional[str] = None
|
||||
update_size_bytes: Optional[int] = None
|
||||
confirm_token: Optional[str] = Field(
|
||||
None, description = "Echo back to POST /update to run the swap. Absent when not appliable."
|
||||
)
|
||||
confirm_expires_at: Optional[str] = None
|
||||
confirm_ttl_seconds: Optional[int] = None
|
||||
|
||||
|
||||
class LlamaUpdateActionResponse(BaseModel):
|
||||
started: bool
|
||||
reason: Optional[str] = None
|
||||
message: Optional[str] = None
|
||||
machine: UpdateMachine = Field(
|
||||
default_factory = UpdateMachine,
|
||||
description = "The host the swap ran (or would run) on; names the result's machine.",
|
||||
)
|
||||
installed_tag: Optional[str] = None
|
||||
latest_tag: Optional[str] = None
|
||||
update_size_bytes: Optional[int] = None
|
||||
job: LlamaUpdateJob = Field(default_factory = LlamaUpdateJob)
|
||||
|
||||
|
||||
|
|
@ -102,14 +192,111 @@ async def llama_update_status(
|
|||
) -> LlamaUpdateStatusResponse:
|
||||
# Off the event loop: detection may probe the host and read GitHub.
|
||||
status = await asyncio.to_thread(get_update_status, force_refresh = force_refresh)
|
||||
resp = LlamaUpdateStatusResponse(**status)
|
||||
resp = LlamaUpdateStatusResponse(machine = _current_machine(), **status)
|
||||
_log_llama_update_progress(resp.job)
|
||||
return resp
|
||||
|
||||
|
||||
@router.post("/update/confirm", response_model = LlamaUpdateConfirmResponse)
|
||||
async def llama_update_confirm(
|
||||
current_subject: str = Depends(get_current_subject),
|
||||
) -> LlamaUpdateConfirmResponse:
|
||||
"""Step one of the two-step apply: describe the pending swap and, when it can
|
||||
be applied, mint a single-use token bound to the offered build. Available to
|
||||
any authenticated operator; confirmation, not location, is the gate."""
|
||||
# Force-refresh so the token binds the current build, not a stale cached tag.
|
||||
status = await asyncio.to_thread(get_update_status, force_refresh = True)
|
||||
machine = _current_machine()
|
||||
installed_tag = status.get("installed_tag")
|
||||
latest_tag = status.get("latest_tag")
|
||||
size = status.get("update_size_bytes")
|
||||
|
||||
if not status.get("update_available"):
|
||||
return LlamaUpdateConfirmResponse(
|
||||
update_available = False,
|
||||
appliable = False,
|
||||
reason = status.get("reason") or "up_to_date",
|
||||
machine = machine,
|
||||
installed_tag = installed_tag,
|
||||
latest_tag = latest_tag,
|
||||
update_size_bytes = size,
|
||||
)
|
||||
# Update exists but the tree is a --with-llama-cpp-dir local link: describe, do not apply.
|
||||
if status.get("local_link"):
|
||||
return LlamaUpdateConfirmResponse(
|
||||
update_available = True,
|
||||
appliable = False,
|
||||
reason = "local_link",
|
||||
machine = machine,
|
||||
installed_tag = installed_tag,
|
||||
latest_tag = latest_tag,
|
||||
update_size_bytes = size,
|
||||
)
|
||||
|
||||
token, expires_at = mint_confirm_token(latest_tag or "")
|
||||
return LlamaUpdateConfirmResponse(
|
||||
update_available = True,
|
||||
appliable = True,
|
||||
reason = None,
|
||||
machine = machine,
|
||||
installed_tag = installed_tag,
|
||||
latest_tag = latest_tag,
|
||||
update_size_bytes = size,
|
||||
confirm_token = token,
|
||||
confirm_expires_at = expires_at,
|
||||
confirm_ttl_seconds = CONFIRM_TOKEN_TTL_SECONDS,
|
||||
)
|
||||
|
||||
|
||||
@router.post("/update", response_model = LlamaUpdateActionResponse)
|
||||
async def llama_update(
|
||||
request: Optional[LlamaUpdateRequest] = Body(default = None),
|
||||
current_subject: str = Depends(get_current_subject),
|
||||
) -> LlamaUpdateActionResponse:
|
||||
"""Apply the swap, but only with an explicit, fresh confirmation.
|
||||
|
||||
The installer replaces the host binary, so a caller must either echo the
|
||||
single-use ``confirm_token`` from POST /update/confirm (preferred: replay-safe,
|
||||
bound to the build) or send ``confirmed=true`` (non-interactive callers). With
|
||||
neither, the swap is refused and the binary is left untouched. The gate is the
|
||||
confirmation, not the caller's location, so a headless SSH server confirms like a local one."""
|
||||
req = request or LlamaUpdateRequest()
|
||||
machine = _current_machine()
|
||||
# Cached read: names the host and build for the prompt and binds the token.
|
||||
status = await asyncio.to_thread(get_update_status)
|
||||
installed_tag = status.get("installed_tag")
|
||||
target_tag = status.get("latest_tag")
|
||||
size = status.get("update_size_bytes")
|
||||
|
||||
confirmed = False
|
||||
refuse_reason = "confirmation_required"
|
||||
if req.confirm_token:
|
||||
ok, why = consume_confirm_token(req.confirm_token, target_tag or "")
|
||||
if ok:
|
||||
confirmed = True
|
||||
else:
|
||||
refuse_reason = why or "invalid_token"
|
||||
elif req.confirmed:
|
||||
confirmed = True
|
||||
|
||||
if not confirmed:
|
||||
# Safe default: no confirmation, no swap; return a visible, actionable refusal.
|
||||
return LlamaUpdateActionResponse(
|
||||
started = False,
|
||||
reason = refuse_reason,
|
||||
message = _REFUSAL_MESSAGES.get(refuse_reason, _REFUSAL_MESSAGES["confirmation_required"]),
|
||||
machine = machine,
|
||||
installed_tag = installed_tag,
|
||||
latest_tag = target_tag,
|
||||
update_size_bytes = size,
|
||||
job = LlamaUpdateJob(**status.get("job", {})),
|
||||
)
|
||||
|
||||
action = await asyncio.to_thread(start_update)
|
||||
return LlamaUpdateActionResponse(**action)
|
||||
return LlamaUpdateActionResponse(
|
||||
machine = machine,
|
||||
installed_tag = installed_tag,
|
||||
latest_tag = target_tag,
|
||||
update_size_bytes = size,
|
||||
**action,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -128,6 +128,9 @@ def test_update_handler_runs_off_event_loop(monkeypatch):
|
|||
return {"started": True, "reason": None, "job": {"state": "running"}}
|
||||
|
||||
monkeypatch.setattr(rl, "start_update", fake_start)
|
||||
out = asyncio.run(rl.llama_update(current_subject = "t"))
|
||||
# The swap now needs explicit confirmation; confirm here so the off-loop path runs.
|
||||
out = asyncio.run(
|
||||
rl.llama_update(request = rl.LlamaUpdateRequest(confirmed = True), current_subject = "t")
|
||||
)
|
||||
assert out.started is True
|
||||
assert seen["thread"] is not threading.main_thread()
|
||||
|
|
|
|||
350
studio/backend/tests/test_update_contract.py
Normal file
350
studio/backend/tests/test_update_contract.py
Normal file
|
|
@ -0,0 +1,350 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Contract tests for the confirmation + visible-result llama.cpp update (#7003).
|
||||
|
||||
Properties covered:
|
||||
- remote authenticated + confirmed -> update proceeds, result names the host
|
||||
- remote authenticated, no confirm -> refused with a clear message, NO swap
|
||||
- local / desktop -> identical behavior (no same-machine axis)
|
||||
- unauthenticated -> refused (401), NO swap
|
||||
- stale / expired / replayed token -> refused, NO swap
|
||||
|
||||
routes/llama.py is loaded standalone with stubbed auth / loggers / llama_cpp_update
|
||||
and the real utils.update_confirm, so no heavy backend deps are needed. Both handler
|
||||
calls and a real FastAPI TestClient (HTTP + auth gate) are exercised.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import importlib.util
|
||||
import sys
|
||||
import types
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
import pytest
|
||||
|
||||
pytest.importorskip("fastapi")
|
||||
|
||||
_HERE = Path(__file__).resolve().parent
|
||||
_BACKEND = _HERE.parent
|
||||
|
||||
|
||||
def _install_stubs():
|
||||
"""Register stub packages so routes/llama.py imports cleanly, plus the real
|
||||
update_confirm module under its production name."""
|
||||
# auth.authentication.get_current_subject -> a real FastAPI dependency that
|
||||
# 401s without a valid bearer, so the HTTP tests can prove the auth gate.
|
||||
from fastapi import Header, HTTPException
|
||||
|
||||
def get_current_subject(authorization: Optional[str] = Header(default = None)) -> str:
|
||||
if authorization == "Bearer good":
|
||||
return "operator"
|
||||
raise HTTPException(status_code = 401, detail = "unauthorized")
|
||||
|
||||
auth_pkg = types.ModuleType("auth")
|
||||
auth_pkg.__path__ = []
|
||||
auth_mod = types.ModuleType("auth.authentication")
|
||||
auth_mod.get_current_subject = get_current_subject
|
||||
|
||||
# loggers.get_logger -> no-op structured logger.
|
||||
loggers_mod = types.ModuleType("loggers")
|
||||
|
||||
class _NopLogger:
|
||||
def info(self, *a, **k):
|
||||
pass
|
||||
|
||||
def debug(self, *a, **k):
|
||||
pass
|
||||
|
||||
def warning(self, *a, **k):
|
||||
pass
|
||||
|
||||
loggers_mod.get_logger = lambda *_a, **_k: _NopLogger()
|
||||
|
||||
# utils.llama_cpp_update -> monkeypatchable get_update_status / start_update.
|
||||
utils_pkg = sys.modules.get("utils")
|
||||
if utils_pkg is None:
|
||||
utils_pkg = types.ModuleType("utils")
|
||||
utils_pkg.__path__ = []
|
||||
sys.modules["utils"] = utils_pkg
|
||||
|
||||
lcu_mod = types.ModuleType("utils.llama_cpp_update")
|
||||
|
||||
def _default_status(force_refresh: bool = False):
|
||||
return {
|
||||
"supported": True,
|
||||
"update_available": True,
|
||||
"stale": False,
|
||||
"installed_tag": "b9860",
|
||||
"latest_tag": "b9909",
|
||||
"published_repo": "unslothai/llama.cpp",
|
||||
"installed_at_utc": None,
|
||||
"age_days": None,
|
||||
"source_build": False,
|
||||
"update_size_bytes": 42_000_000,
|
||||
"job": {"state": "idle"},
|
||||
}
|
||||
|
||||
def _default_start():
|
||||
return {"started": True, "reason": None, "job": {"state": "running"}}
|
||||
|
||||
lcu_mod.get_update_status = _default_status
|
||||
lcu_mod.start_update = _default_start
|
||||
|
||||
# The real confirmation-token module, under its production import name.
|
||||
uc_spec = importlib.util.spec_from_file_location(
|
||||
"utils.update_confirm", str(_BACKEND / "utils" / "update_confirm.py")
|
||||
)
|
||||
uc_mod = importlib.util.module_from_spec(uc_spec)
|
||||
uc_spec.loader.exec_module(uc_mod)
|
||||
|
||||
for name, mod in (
|
||||
("auth", auth_pkg),
|
||||
("auth.authentication", auth_mod),
|
||||
("loggers", loggers_mod),
|
||||
("utils.llama_cpp_update", lcu_mod),
|
||||
("utils.update_confirm", uc_mod),
|
||||
):
|
||||
sys.modules[name] = mod
|
||||
return uc_mod
|
||||
|
||||
|
||||
_uc = _install_stubs()
|
||||
|
||||
|
||||
def _load_route():
|
||||
spec = importlib.util.spec_from_file_location(
|
||||
"llama_route_under_test", str(_BACKEND / "routes" / "llama.py")
|
||||
)
|
||||
mod = importlib.util.module_from_spec(spec)
|
||||
sys.modules["llama_route_under_test"] = mod # pydantic forward-ref resolution
|
||||
spec.loader.exec_module(mod)
|
||||
return mod
|
||||
|
||||
|
||||
rl = _load_route()
|
||||
|
||||
|
||||
@pytest.fixture(autouse = True)
|
||||
def _reset(monkeypatch):
|
||||
_uc.reset_tokens_for_tests()
|
||||
# Restore default status/start each test; individual tests override as needed.
|
||||
monkeypatch.setattr(
|
||||
rl,
|
||||
"get_update_status",
|
||||
lambda force_refresh = False: {
|
||||
"supported": True,
|
||||
"update_available": True,
|
||||
"installed_tag": "b9860",
|
||||
"latest_tag": "b9909",
|
||||
"update_size_bytes": 42_000_000,
|
||||
"job": {"state": "idle"},
|
||||
},
|
||||
)
|
||||
yield
|
||||
|
||||
|
||||
def _track_start(monkeypatch):
|
||||
calls = {"n": 0}
|
||||
|
||||
def _start():
|
||||
calls["n"] += 1
|
||||
return {"started": True, "reason": None, "job": {"state": "running"}}
|
||||
|
||||
monkeypatch.setattr(rl, "start_update", _start)
|
||||
return calls
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# update_confirm token unit tests
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def test_token_roundtrip_single_use():
|
||||
tok, exp = _uc.mint_confirm_token("b9909")
|
||||
assert isinstance(tok, str) and tok
|
||||
ok, reason = _uc.consume_confirm_token(tok, "b9909")
|
||||
assert ok and reason is None
|
||||
# Single use: a replay is refused.
|
||||
ok2, reason2 = _uc.consume_confirm_token(tok, "b9909")
|
||||
assert ok2 is False and reason2 == "invalid_token"
|
||||
|
||||
|
||||
def test_token_stale_target_refused():
|
||||
tok, _ = _uc.mint_confirm_token("b9909")
|
||||
ok, reason = _uc.consume_confirm_token(tok, "b9910") # offered build changed
|
||||
assert ok is False and reason == "stale_target"
|
||||
|
||||
|
||||
def test_token_expired_refused():
|
||||
tok, _ = _uc.mint_confirm_token("b9909", ttl_seconds = 0)
|
||||
ok, reason = _uc.consume_confirm_token(tok, "b9909")
|
||||
assert ok is False and reason == "expired_token"
|
||||
|
||||
|
||||
def test_token_missing_refused():
|
||||
ok, reason = _uc.consume_confirm_token(None, "b9909")
|
||||
assert ok is False and reason == "invalid_token"
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Handler-level: the swap only runs with an explicit confirmation
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def test_apply_without_confirmation_is_refused_and_never_swaps(monkeypatch):
|
||||
calls = _track_start(monkeypatch)
|
||||
out = asyncio.run(rl.llama_update(request = None, current_subject = "operator"))
|
||||
assert out.started is False
|
||||
assert out.reason == "confirmation_required"
|
||||
assert "swap" in out.message.lower() or "confirm" in out.message.lower()
|
||||
# The result still names the host + build so the refusal is actionable.
|
||||
assert out.machine.hostname
|
||||
assert out.latest_tag == "b9909"
|
||||
assert calls["n"] == 0 # installer NEVER started
|
||||
|
||||
|
||||
def test_apply_with_confirmed_true_proceeds(monkeypatch):
|
||||
calls = _track_start(monkeypatch)
|
||||
body = rl.LlamaUpdateRequest(confirmed = True)
|
||||
out = asyncio.run(rl.llama_update(request = body, current_subject = "operator"))
|
||||
assert out.started is True
|
||||
assert out.machine.hostname # which machine
|
||||
assert out.latest_tag == "b9909" # which version
|
||||
assert calls["n"] == 1
|
||||
|
||||
|
||||
def test_apply_with_valid_token_proceeds(monkeypatch):
|
||||
calls = _track_start(monkeypatch)
|
||||
tok, _ = _uc.mint_confirm_token("b9909")
|
||||
body = rl.LlamaUpdateRequest(confirm_token = tok)
|
||||
out = asyncio.run(rl.llama_update(request = body, current_subject = "operator"))
|
||||
assert out.started is True
|
||||
assert calls["n"] == 1
|
||||
# Token is single-use: replaying the same body is refused, no second swap.
|
||||
out2 = asyncio.run(rl.llama_update(request = body, current_subject = "operator"))
|
||||
assert out2.started is False
|
||||
assert out2.reason == "invalid_token"
|
||||
assert calls["n"] == 1
|
||||
|
||||
|
||||
def test_apply_with_stale_token_is_refused(monkeypatch):
|
||||
calls = _track_start(monkeypatch)
|
||||
tok, _ = _uc.mint_confirm_token("b9900") # bound to an older offered build
|
||||
body = rl.LlamaUpdateRequest(confirm_token = tok)
|
||||
out = asyncio.run(rl.llama_update(request = body, current_subject = "operator"))
|
||||
assert out.started is False
|
||||
assert out.reason == "stale_target"
|
||||
assert calls["n"] == 0
|
||||
|
||||
|
||||
def test_confirm_endpoint_describes_and_mints(monkeypatch):
|
||||
out = asyncio.run(rl.llama_update_confirm(current_subject = "operator"))
|
||||
assert out.update_available is True
|
||||
assert out.appliable is True
|
||||
assert out.installed_tag == "b9860"
|
||||
assert out.latest_tag == "b9909"
|
||||
assert out.machine.hostname
|
||||
assert out.confirm_token
|
||||
# The freshly minted token applies the update end to end.
|
||||
calls = _track_start(monkeypatch)
|
||||
body = rl.LlamaUpdateRequest(confirm_token = out.confirm_token)
|
||||
applied = asyncio.run(rl.llama_update(request = body, current_subject = "operator"))
|
||||
assert applied.started is True
|
||||
assert calls["n"] == 1
|
||||
|
||||
|
||||
def test_confirm_endpoint_up_to_date_offers_no_token(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
rl,
|
||||
"get_update_status",
|
||||
lambda force_refresh = False: {
|
||||
"update_available": False,
|
||||
"installed_tag": "b9909",
|
||||
"latest_tag": "b9909",
|
||||
"job": {"state": "idle"},
|
||||
},
|
||||
)
|
||||
out = asyncio.run(rl.llama_update_confirm(current_subject = "operator"))
|
||||
assert out.update_available is False
|
||||
assert out.appliable is False
|
||||
assert out.confirm_token is None
|
||||
|
||||
|
||||
def test_status_reports_machine():
|
||||
out = asyncio.run(
|
||||
rl.llama_update_status(force_refresh = False, current_subject = "operator")
|
||||
)
|
||||
assert out.machine.hostname
|
||||
assert out.machine.platform
|
||||
assert out.latest_tag == "b9909"
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# HTTP-level: auth gate + wiring + backwards-compatible bodyless POST
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
def _client():
|
||||
from fastapi import FastAPI
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(rl.router, prefix = "/api/llama")
|
||||
return TestClient(app, raise_server_exceptions = True)
|
||||
|
||||
|
||||
def test_http_unauthenticated_update_is_refused(monkeypatch):
|
||||
calls = _track_start(monkeypatch)
|
||||
client = _client()
|
||||
# No Authorization header at all -> 401 from the auth dependency.
|
||||
r = client.post("/api/llama/update", json = {"confirmed": True})
|
||||
assert r.status_code == 401
|
||||
assert calls["n"] == 0 # installer never reached
|
||||
|
||||
|
||||
def test_http_authenticated_bodyless_post_refused_no_swap(monkeypatch):
|
||||
# Backwards compat: an OLD frontend posts /update with no body. That must be
|
||||
# safe -- refused (confirmation_required), NEVER a silent swap.
|
||||
calls = _track_start(monkeypatch)
|
||||
client = _client()
|
||||
r = client.post("/api/llama/update", headers = {"Authorization": "Bearer good"})
|
||||
assert r.status_code == 200
|
||||
body = r.json()
|
||||
assert body["started"] is False
|
||||
assert body["reason"] == "confirmation_required"
|
||||
assert body["machine"]["hostname"]
|
||||
assert calls["n"] == 0
|
||||
|
||||
|
||||
def test_http_two_step_confirm_then_apply(monkeypatch):
|
||||
calls = _track_start(monkeypatch)
|
||||
client = _client()
|
||||
h = {"Authorization": "Bearer good"}
|
||||
|
||||
step1 = client.post("/api/llama/update/confirm", headers = h)
|
||||
assert step1.status_code == 200
|
||||
tok = step1.json()["confirm_token"]
|
||||
assert tok
|
||||
|
||||
step2 = client.post("/api/llama/update", headers = h, json = {"confirm_token": tok})
|
||||
assert step2.status_code == 200
|
||||
assert step2.json()["started"] is True
|
||||
assert step2.json()["machine"]["hostname"]
|
||||
assert calls["n"] == 1
|
||||
|
||||
|
||||
def test_http_local_and_remote_behave_identically(monkeypatch):
|
||||
# There is no same-machine axis: the contract depends only on auth + confirm,
|
||||
# so a "local" and a "remote" authenticated caller get identical outcomes.
|
||||
calls = _track_start(monkeypatch)
|
||||
client = _client()
|
||||
h = {"Authorization": "Bearer good"}
|
||||
# Simulate a remote caller by adding proxy/forwarding headers that the closed
|
||||
# PR would have treated as "not host" -- here they change nothing.
|
||||
remote_h = {**h, "X-Forwarded-For": "203.0.113.7", "CF-Connecting-IP": "203.0.113.7"}
|
||||
|
||||
local = client.post("/api/llama/update", headers = h, json = {"confirmed": True})
|
||||
remote = client.post("/api/llama/update", headers = remote_h, json = {"confirmed": True})
|
||||
assert local.json()["started"] == remote.json()["started"] == True
|
||||
assert calls["n"] == 2 # both proceeded; location was never the gate
|
||||
104
studio/backend/utils/update_confirm.py
Normal file
104
studio/backend/utils/update_confirm.py
Normal file
|
|
@ -0,0 +1,104 @@
|
|||
# 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 (local, desktop, or remote over
|
||||
SSH/Cloudflare); we do not gate on "same machine" since a headless SSH server never
|
||||
has a host-local session. A token binds the build it was offered for (``target_tag``)
|
||||
and is single-use with a short TTL. The store is in-process, matching Studio's
|
||||
single-process backend; for a multi-worker deploy swap it for an HMAC stateless token.
|
||||
"""
|
||||
|
||||
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 a burst of confirm calls that are never applied cannot grow
|
||||
# memory without bound; oldest entries are evicted 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()
|
||||
Loading…
Add table
Add a link
Reference in a new issue