ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Studio: reset-password rotates the credential in place instead of deleting auth.db (#7573)

Browse source 
* reset-password: rotate the admin credential in place instead of deleting auth.db

* reset-password: fix the CI callers and error handling for the in-place rotation

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* reset-password: narrow the CI change to the jobs that read .bootstrap_password

* reset-password: stop over-claiming what the reset revokes and when it takes effect

* auth: bind token issuance to the credential version that was verified

* auth: bind credential-creating writes to the version the request authenticated with

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* auth: bind the change-password and workflow-key writes to their own credential version

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* auth: read the credential version inside the transaction that validated it

* data-recipe: answer 401 when a reset revokes the credential mid job start

* Fix lint blocker and Windows path assertion for PR #7573

Drop the now-unused validate_api_key import from studio/backend/auth/authentication.py.
Every call site moved to validate_api_key_with_credential, so the Source lint job's
import-hoist gate flagged it as a blocker. The wrapper itself stays in storage.py;
test_api_key_expiry.py still exercises it.

Make test_run_reexec_forwards_resolved_frontend_on_public_launch compare against
str(Path(...)) instead of a POSIX literal. _find_frontend_dist returns a Path, so on
Windows the forwarded value is \fake\studio\frontend\dist and the assertion could
never pass there. Pre-existing, surfaced by running unsloth_cli/tests on Windows.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
This commit is contained in:
Nilay 2026-07-29 14:10:12 +05:30 • committed by GitHub
commit 52609fb890
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
21 changed files with 806 additions and 295 deletions
Download patch file Download diff file Expand all files Collapse all files

67
studio/backend/auth/authentication.py
View file

@ -11,11 +11,12 @@ import jwt
from .storage import (
API_KEY_PREFIX,
credential_generation,
get_jwt_secret,
get_user_and_secret,
load_jwt_secret,
save_refresh_token,
validate_api_key,
validate_api_key_with_credential,
verify_refresh_token,
)
@ -54,11 +55,14 @@ def create_access_token(
expires_delta: Optional[timedelta] = None,
*,
desktop: bool = False,
secret: Optional[str] = None,
) -> str:
"""
Create a signed JWT for the given subject (e.g. username).
Valid across restarts: the signing secret is stored in SQLite.
Valid across restarts: the signing secret is stored in SQLite. Callers that
already verified a credential pass ``secret`` so a rotation landing mid-request
cannot sign the token with the credential that just replaced it.
"""
to_encode = {"sub": subject}
if desktop:
@ -69,7 +73,7 @@ def create_access_token(
to_encode.update({"exp": expire})
return jwt.encode(
to_encode,
_get_secret_for_subject(subject),
secret if secret is not None else _get_secret_for_subject(subject),
algorithm = ALGORITHM,
)
@ -96,15 +100,28 @@ def is_desktop_access_token(token: str) -> bool:
return payload.get("sub") == subject and payload.get("desktop") is True
def create_refresh_token(subject: str, *, desktop: bool = False) -> str:
def create_refresh_token(
subject: str,
*,
desktop: bool = False,
secret: Optional[str] = None,
) -> str:
"""
Create a random refresh token, store its hash in SQLite, and return it.
Refresh tokens are opaque (not JWTs); expire after REFRESH_TOKEN_EXPIRE_DAYS.
``secret`` stamps the token with the credential version the caller verified,
so a rotation cannot leave a token minted from the replaced credential valid.
"""
token = secrets.token_urlsafe(48)
expires_at = datetime.now(timezone.utc) + timedelta(days = REFRESH_TOKEN_EXPIRE_DAYS)
save_refresh_token(token, subject, expires_at.isoformat(), is_desktop = desktop)
save_refresh_token(
token,
subject,
expires_at.isoformat(),
is_desktop = desktop,
secret_gen = credential_generation(secret) if secret is not None else None,
)
return token
@ -137,7 +154,22 @@ def reload_secret() -> None:
async def get_current_subject(credentials: HTTPAuthorizationCredentials = Depends(security)) -> str:
"""Validate JWT and require the password-change flow to be completed."""
return await _get_current_subject(
subject, _generation = await _get_current_credential(
credentials,
allow_password_change = False,
)
return subject
async def get_current_credential(
credentials: HTTPAuthorizationCredentials = Depends(security),
) -> Tuple[str, Optional[str]]:
"""As get_current_subject, but also returns the credential generation.
For routes that persist a new credential and must not do so on behalf of one
a concurrent reset has revoked.
"""
return await _get_current_credential(
credentials,
allow_password_change = False,
)
@ -158,10 +190,11 @@ async def get_current_subject_allow_password_change(
credentials: HTTPAuthorizationCredentials = Depends(security),
) -> str:
"""Validate JWT but allow access to the password-change endpoint."""
return await _get_current_subject(
subject, _generation = await _get_current_credential(
credentials,
allow_password_change = True,
)
return subject
# The literal the examples ship with; pasted unedited more often than a revoked key.
@ -179,21 +212,27 @@ def _invalid_api_key_detail(token: str) -> str:
return "Invalid or expired API key"
async def _get_current_subject(
async def _get_current_credential(
credentials: HTTPAuthorizationCredentials, *, allow_password_change: bool
) -> str:
"""FastAPI dependency: validate the JWT and return the subject. Use on protected routes."""
) -> Tuple[str, Optional[str]]:
"""Validate the bearer and return ``(subject, credential generation)``.
The generation is the credential version this request actually authenticated
against. Routes that persist new credentials must bind their write to it, or
a reset landing mid-request would bless what it just revoked.
"""
token = credentials.credentials
# --- API key path (sk-unsloth-...) ---
if token.startswith(API_KEY_PREFIX):
username = validate_api_key(token)
if username is None:
verified = validate_api_key_with_credential(token)
if verified is None:
raise HTTPException(
status_code = status.HTTP_401_UNAUTHORIZED,
detail = _invalid_api_key_detail(token),
)
return username
username, secret = verified
return username, credential_generation(secret)
# --- JWT path ---
subject = _decode_subject_without_verification(token)
@ -224,7 +263,7 @@ async def _get_current_subject(
status_code = status.HTTP_403_FORBIDDEN,
detail = "Password change required",
)
return subject
return subject, credential_generation(jwt_secret)
except jwt.InvalidTokenError:
raise HTTPException(
status_code = status.HTTP_401_UNAUTHORIZED,

180
studio/backend/auth/storage.py
View file

@ -186,7 +186,7 @@ def clear_bootstrap_password() -> None:
# Removal failed (Windows AV, read-only auth dir). The hash is already
# committed, so don't fail the change -- but truncate the file so its
# stale plaintext can't be re-seeded by generate_bootstrap_password()
# if a later reset-password deletes auth.db and re-validates it.
# if auth.db is ever recreated.
try:
_BOOTSTRAP_PW_PATH.write_text("", encoding = "utf-8")
cleared = True
@ -221,6 +221,31 @@ def _hash_token(token: str) -> str:
return hashlib.sha256(token.encode("utf-8")).hexdigest()
class CredentialRotated(Exception):
"""A password reset revoked the credential this request authenticated with."""
def credential_generation(jwt_secret: str) -> str:
"""Marker for the credential version a refresh token was issued under.
Every password change rotates ``jwt_secret``, so a token stamped with the
previous one is rejected even if it was inserted after the revoking DELETE.
"""
return hashlib.sha256(jwt_secret.encode("utf-8")).hexdigest()
def _current_secret(conn: sqlite3.Connection, username: str) -> Optional[str]:
row = conn.execute(
"SELECT jwt_secret FROM auth_user WHERE username = ?", (username,)
).fetchone()
return row["jwt_secret"] if row else None
def _current_generation(conn: sqlite3.Connection, username: str) -> Optional[str]:
secret = _current_secret(conn, username)
return credential_generation(secret) if secret is not None else None
def get_connection() -> sqlite3.Connection:
"""Get a connection to the auth database, creating tables if needed."""
ensure_dir(DB_PATH.parent)
@ -264,7 +289,8 @@ def get_connection() -> sqlite3.Connection:
token_hash TEXT NOT NULL,
username TEXT NOT NULL,
expires_at TEXT NOT NULL,
is_desktop INTEGER NOT NULL DEFAULT 0
is_desktop INTEGER NOT NULL DEFAULT 0,
secret_gen TEXT
);
"""
)
@ -303,6 +329,8 @@ def get_connection() -> sqlite3.Connection:
refresh_columns = {row["name"] for row in conn.execute("PRAGMA table_info(refresh_tokens)")}
if "is_desktop" not in refresh_columns:
conn.execute("ALTER TABLE refresh_tokens ADD COLUMN is_desktop INTEGER NOT NULL DEFAULT 0")
if "secret_gen" not in refresh_columns:
conn.execute("ALTER TABLE refresh_tokens ADD COLUMN secret_gen TEXT")
conn.commit()
return conn
@ -676,12 +704,22 @@ def update_password(
new_password: str,
*,
revoke_refresh_tokens: bool = False,
) -> bool:
expect_password_hash: Optional[str] = None,
) -> Optional[str]:
"""Update password, clear first-login requirement, rotate JWT secret.
Returns the new JWT secret, or None when nothing was updated. Callers that
mint tokens for the caller must sign with the returned secret: re-reading it
would pick up a reset that landed between this commit and the mint.
``revoke_refresh_tokens`` deletes the user's refresh tokens in the SAME
transaction: a separate delete could fail after the password commit and
leave a pre-change token still able to mint access tokens.
``expect_password_hash`` makes the write conditional on the credential the
caller verified still being current, so a request that checked the old
password cannot overwrite a reset that landed while it was in flight.
Returns False when the credential moved underneath it.
"""
from .hashing import hash_password
@ -689,21 +727,32 @@ def update_password(
jwt_secret = secrets.token_urlsafe(64)
conn = get_connection()
try:
cursor = conn.execute(
"""
UPDATE auth_user
SET password_salt = ?, password_hash = ?, jwt_secret = ?, must_change_password = 0
WHERE username = ?
""",
(salt, pwd_hash, jwt_secret, username),
)
if expect_password_hash is None:
cursor = conn.execute(
"""
UPDATE auth_user
SET password_salt = ?, password_hash = ?, jwt_secret = ?, must_change_password = 0
WHERE username = ?
""",
(salt, pwd_hash, jwt_secret, username),
)
else:
cursor = conn.execute(
"""
UPDATE auth_user
SET password_salt = ?, password_hash = ?, jwt_secret = ?, must_change_password = 0
WHERE username = ? AND password_hash = ?
""",
(salt, pwd_hash, jwt_secret, username, expect_password_hash),
)
if revoke_refresh_tokens and cursor.rowcount > 0:
conn.execute("DELETE FROM refresh_tokens WHERE username = ?", (username,))
conn.commit()
if cursor.rowcount > 0:
clear_bootstrap_password()
clear_desktop_secret()
return cursor.rowcount > 0
return jwt_secret
return None
finally:
conn.close()
@ -714,35 +763,49 @@ def save_refresh_token(
expires_at: str,
*,
is_desktop: bool = False,
secret_gen: Optional[str] = None,
) -> None:
"""
Store a hashed refresh token with its associated username and expiry.
``secret_gen`` binds the token to a credential version; it defaults to the
current one, and callers that already verified a credential must pass the
version they verified rather than let this re-read a rotated one.
"""
token_hash = _hash_token(token)
conn = get_connection()
try:
if secret_gen is None:
secret_gen = _current_generation(conn, username)
conn.execute(
"""
INSERT INTO refresh_tokens (token_hash, username, expires_at, is_desktop)
VALUES (?, ?, ?, ?)
INSERT INTO refresh_tokens (token_hash, username, expires_at, is_desktop, secret_gen)
VALUES (?, ?, ?, ?, ?)
""",
(token_hash, username, expires_at, int(is_desktop)),
(token_hash, username, expires_at, int(is_desktop), secret_gen),
)
conn.commit()
finally:
conn.close()
def consume_refresh_token(token: str) -> Optional[Tuple[str, bool]]:
def consume_refresh_token(token: str) -> Optional[Tuple[str, bool, str]]:
"""Atomically validate-and-delete a refresh token for single-use rotation.
DELETE RETURNING fuses validate and delete into one statement so two
concurrent refresh requests cannot both consume the same token.
concurrent refresh requests cannot both consume the same token. Returns
``(username, is_desktop, jwt_secret)``; the caller must mint the replacement
tokens against that secret so a rotation landing mid-refresh cannot issue a
post-rotation session from a pre-rotation token.
"""
token_hash = _hash_token(token)
now = datetime.now(timezone.utc).isoformat()
conn = get_connection()
try:
# One transaction with the delete: an unstamped legacy row has no
# generation to compare, so reading the credential after committing would
# hand a reset's new secret to a token issued before it.
conn.execute("BEGIN IMMEDIATE")
conn.execute(
"DELETE FROM refresh_tokens WHERE expires_at < ?",
(now,),
@ -751,15 +814,21 @@ def consume_refresh_token(token: str) -> Optional[Tuple[str, bool]]:
"""
DELETE FROM refresh_tokens
WHERE token_hash = ? AND expires_at >= ?
RETURNING username, is_desktop
RETURNING username, is_desktop, secret_gen
""",
(token_hash, now),
)
row = cur.fetchone()
conn.commit()
if row is None:
conn.commit()
return None
return row["username"], bool(row["is_desktop"])
secret = _current_secret(conn, row["username"])
conn.commit()
if secret is None:
return None
if row["secret_gen"] is not None and row["secret_gen"] != credential_generation(secret):
return None
return row["username"], bool(row["is_desktop"]), secret
finally:
conn.close()
@ -783,7 +852,7 @@ def verify_refresh_token(token: str) -> Optional[Tuple[str, bool]]:
cur = conn.execute(
"""
SELECT id, username, expires_at, is_desktop FROM refresh_tokens
SELECT id, username, expires_at, is_desktop, secret_gen FROM refresh_tokens
WHERE token_hash = ?
""",
(token_hash,),
@ -792,6 +861,13 @@ def verify_refresh_token(token: str) -> Optional[Tuple[str, bool]]:
if row is None:
return None
if row["secret_gen"] is not None and row["secret_gen"] != _current_generation(
conn, row["username"]
):
conn.execute("DELETE FROM refresh_tokens WHERE id = ?", (row["id"],))
conn.commit()
return None
# Check expiry
expires_at = datetime.fromisoformat(row["expires_at"])
if datetime.now(timezone.utc) > expires_at:
@ -836,30 +912,41 @@ def create_desktop_secret() -> str:
conn.close()
def validate_desktop_secret(raw_secret: str) -> Optional[str]:
"""Return the real admin username when the desktop secret matches."""
def validate_desktop_secret_with_credential(raw_secret: str) -> Optional[Tuple[str, str]]:
"""Validate the desktop secret and return ``(username, jwt_secret)``.
Both reads share one transaction so the returned secret is the credential
version the desktop secret was checked against; a reset landing mid-request
then invalidates the tokens minted from it rather than blessing them.
"""
if not raw_secret.startswith(DESKTOP_SECRET_PREFIX):
return None
if get_user_and_secret(DEFAULT_ADMIN_USERNAME) is None:
return None
secret_hash = _pbkdf2_desktop_secret(raw_secret)
conn = get_connection()
try:
cur = conn.execute(
conn.execute("BEGIN")
row = conn.execute(
"SELECT value FROM app_secrets WHERE key = ?",
(_DESKTOP_SECRET_HASH_KEY,),
)
row = cur.fetchone()
if row is None:
).fetchone()
if row is None or not secrets.compare_digest(row["value"], secret_hash):
return None
if not secrets.compare_digest(row["value"], secret_hash):
jwt_secret = _current_secret(conn, DEFAULT_ADMIN_USERNAME)
if jwt_secret is None:
return None
return DEFAULT_ADMIN_USERNAME
return DEFAULT_ADMIN_USERNAME, jwt_secret
finally:
conn.rollback()
conn.close()
def validate_desktop_secret(raw_secret: str) -> Optional[str]:
"""Return the real admin username when the desktop secret matches."""
verified = validate_desktop_secret_with_credential(raw_secret)
return verified[0] if verified else None
def clear_desktop_secret() -> None:
"""Remove backend-side desktop auth state."""
conn = get_connection()
@ -885,6 +972,7 @@ def create_api_key(
name: str,
expires_at: Optional[str] = None,
internal: bool = False,
expect_gen: Optional[str] = None,
) -> Tuple[str, dict]:
"""Create a new API key for *username*.
@ -893,6 +981,10 @@ def create_api_key(
Pass ``internal=True`` for keys minted by workflows (e.g. data-recipe
runs) that should not appear in user-facing key listings.
``expect_gen`` ties the insert to the credential generation the request
authenticated under, so a session revoked by a concurrent password reset
cannot mint a key that outlives it. Raises ``CredentialRotated`` if it moved.
"""
raw_key = API_KEY_PREFIX + secrets.token_hex(16)
key_hash = _pbkdf2_api_key(raw_key)
@ -901,6 +993,12 @@ def create_api_key(
conn = get_connection()
try:
if expect_gen is not None:
conn.execute("BEGIN IMMEDIATE")
if _current_generation(conn, username) != expect_gen:
raise CredentialRotated(
"The credential this request authenticated with was revoked."
)
conn.execute(
"""
INSERT INTO api_keys (username, key_prefix, key_hash, name, created_at, expires_at, is_internal)
@ -989,15 +1087,25 @@ def revoke_internal_api_key(key_id: int) -> bool:
def validate_api_key(raw_key: str) -> Optional[str]:
"""Validate *raw_key* and return the owning username, or ``None``.
"""Validate *raw_key* and return the owning username, or ``None``."""
verified = validate_api_key_with_credential(raw_key)
return verified[0] if verified else None
Also updates ``last_used_at`` on success.
def validate_api_key_with_credential(raw_key: str) -> Optional[Tuple[str, str]]:
"""Validate *raw_key* and return ``(username, jwt_secret)``, or ``None``.
Also updates ``last_used_at`` on success. The key check and the credential
read share one write transaction, so the returned version is the one the key
was actually valid under: a reset committing right after cannot have its new
generation handed to a request the key it revoked authenticated.
"""
cache_id = _api_key_cache_id(raw_key)
cached_hash = _api_key_hash_cache.get(cache_id)
key_hash = cached_hash if cached_hash is not None else _pbkdf2_api_key(raw_key)
conn = get_connection()
try:
conn.execute("BEGIN IMMEDIATE")
cur = conn.execute(
"SELECT id, username, is_active, expires_at FROM api_keys WHERE key_hash = ?",
(key_hash,),
@ -1017,11 +1125,15 @@ def validate_api_key(raw_key: str) -> Optional[str]:
expires = datetime.fromisoformat(row["expires_at"])
if datetime.now(timezone.utc) > expires:
return None
secret = _current_secret(conn, row["username"])
if secret is None:
return None
conn.execute(
"UPDATE api_keys SET last_used_at = ? WHERE id = ?",
(datetime.now(timezone.utc).isoformat(), row["id"]),
)
conn.commit()
return row["username"]
return row["username"], secret
finally:
conn.rollback()
conn.close()

64
studio/backend/routes/auth.py
View file

@ -31,6 +31,7 @@ from auth import storage, hashing
from auth.authentication import (
create_access_token,
create_refresh_token,
get_current_credential,
get_current_subject,
get_current_subject_allow_password_change,
refresh_access_token,
@ -399,7 +400,7 @@ async def login(payload: AuthLoginRequest, request: Request) -> Token:
detail = f"Incorrect password. To reset it, run this in your terminal: {_reset_password_command()}",
)
salt, pwd_hash, _jwt_secret, must_change_password = record
salt, pwd_hash, jwt_secret, must_change_password = record
if not hashing.verify_password(payload.password, salt, pwd_hash):
_record_login_failure(key)
raise HTTPException(
@ -409,8 +410,10 @@ async def login(payload: AuthLoginRequest, request: Request) -> Token:
_clear_login_bucket(key)
_clear_login_bucket(unknown_key)
access_token = create_access_token(subject = payload.username)
refresh_token = create_refresh_token(subject = payload.username)
# Issue against the credential version just verified, not whatever is in the DB
# now: a concurrent reset-password must not hand this login a post-reset session.
access_token = create_access_token(subject = payload.username, secret = jwt_secret)
refresh_token = create_refresh_token(subject = payload.username, secret = jwt_secret)
return Token(
access_token = access_token,
refresh_token = refresh_token,
@ -438,16 +441,17 @@ async def logout(
@router.post("/desktop-login", response_model = Token)
async def desktop_login(payload: DesktopLoginRequest) -> Token:
"""Exchange a local desktop secret for normal admin-subject tokens."""
username = storage.validate_desktop_secret(payload.secret)
if username is None:
verified = storage.validate_desktop_secret_with_credential(payload.secret)
if verified is None:
raise HTTPException(
status_code = status.HTTP_401_UNAUTHORIZED,
detail = "Desktop authentication failed",
)
username, jwt_secret = verified
return Token(
access_token = create_access_token(subject = username, desktop = True),
refresh_token = create_refresh_token(subject = username, desktop = True),
access_token = create_access_token(subject = username, desktop = True, secret = jwt_secret),
refresh_token = create_refresh_token(subject = username, desktop = True, secret = jwt_secret),
token_type = "bearer",
must_change_password = False,
)
@ -462,9 +466,11 @@ async def refresh(payload: RefreshTokenRequest) -> Token:
status_code = status.HTTP_401_UNAUTHORIZED,
detail = "Invalid or expired refresh token",
)
username, is_desktop = consumed
new_access_token = create_access_token(subject = username, desktop = is_desktop)
new_refresh_token = create_refresh_token(subject = username, desktop = is_desktop)
username, is_desktop, jwt_secret = consumed
new_access_token = create_access_token(subject = username, desktop = is_desktop, secret = jwt_secret)
new_refresh_token = create_refresh_token(
subject = username, desktop = is_desktop, secret = jwt_secret
)
return Token(
access_token = new_access_token,
@ -507,13 +513,25 @@ async def change_password(
# Single transaction: a separate refresh-token purge could fail after the
# password commit, leaving pre-change tokens able to mint access tokens.
storage.update_password(current_subject, payload.new_password, revoke_refresh_tokens = True)
# Conditional on the hash just verified: a reset-password that landed while
# this request was in flight must not be overwritten by it.
new_secret = storage.update_password(
current_subject,
payload.new_password,
revoke_refresh_tokens = True,
expect_password_hash = pwd_hash,
)
if new_secret is None:
raise HTTPException(
status_code = status.HTTP_409_CONFLICT,
detail = "The password changed while this request was in flight. Sign in again.",
)
try:
request.app.state.bootstrap_password = None
except AttributeError:
pass
access_token = create_access_token(subject = current_subject)
refresh_token = create_refresh_token(subject = current_subject)
access_token = create_access_token(subject = current_subject, secret = new_secret)
refresh_token = create_refresh_token(subject = current_subject, secret = new_secret)
return Token(
access_token = access_token,
refresh_token = refresh_token,
@ -541,20 +559,28 @@ def _row_to_api_key_response(row: dict) -> ApiKeyResponse:
@router.post("/api-keys", response_model = CreateApiKeyResponse)
async def create_api_key(
payload: CreateApiKeyRequest, current_subject: str = Depends(get_current_subject)
payload: CreateApiKeyRequest, credential: tuple = Depends(get_current_credential)
) -> CreateApiKeyResponse:
"""Create a new API key. The raw key is returned once and cannot be retrieved later."""
current_subject, generation = credential
expires_at = None
if payload.expires_in_days is not None:
expires_at = (
datetime.now(timezone.utc) + timedelta(days = payload.expires_in_days)
).isoformat()
raw_key, row = storage.create_api_key(
username = current_subject,
name = payload.name,
expires_at = expires_at,
)
try:
raw_key, row = storage.create_api_key(
username = current_subject,
name = payload.name,
expires_at = expires_at,
expect_gen = generation,
)
except storage.CredentialRotated:
raise HTTPException(
status_code = status.HTTP_401_UNAUTHORIZED,
detail = "Invalid or expired token",
)
return CreateApiKeyResponse(
key = raw_key,
api_key = _row_to_api_key_response(row),

24
studio/backend/routes/data_recipe/jobs.py
View file

@ -10,7 +10,10 @@ from datetime import datetime, timedelta, timezone
from typing import Any, Optional
from urllib.parse import urlparse
from fastapi import APIRouter, HTTPException, Query, Request
from fastapi import APIRouter, Depends, HTTPException, Query, Request
from auth.authentication import get_current_credential
from auth.storage import CredentialRotated
from fastapi.responses import JSONResponse, StreamingResponse
from pydantic import ValidationError
@ -257,7 +260,11 @@ def _inject_local_structured_response_format(
model_configs.extend(new_configs)
def _inject_local_providers(recipe: dict[str, Any], request: Request) -> Optional[int]:
def _inject_local_providers(
recipe: dict[str, Any],
request: Request,
expect_gen: Optional[str] = None,
) -> Optional[int]:
"""Mutate recipe in-place: point is_local providers at this server and mint
a short-lived internal sk-unsloth-* key for workflow auth.
@ -313,6 +320,7 @@ def _inject_local_providers(recipe: dict[str, Any], request: Request) -> Optiona
name = "data-recipe workflow",
expires_at = expires_at,
internal = True,
expect_gen = expect_gen,
)
internal_key_id = int(row["id"])
@ -375,7 +383,11 @@ def _normalize_run_name(value: Any) -> str | None:
@router.post("/jobs", response_class = JSONResponse, response_model = JobCreateResponse)
def create_job(payload: RecipePayload, request: Request):
def create_job(
payload: RecipePayload,
request: Request,
credential: tuple = Depends(get_current_credential),
):
recipe = payload.recipe
if not recipe.get("columns"):
raise HTTPException(status_code = 400, detail = "Recipe must include columns.")
@ -406,7 +418,11 @@ def create_job(payload: RecipePayload, request: Request):
) from exc
try:
internal_api_key_id = _inject_local_providers(recipe, request)
internal_api_key_id = _inject_local_providers(recipe, request, credential[1])
except CredentialRotated as exc:
# A reset-password landed after this request authenticated; the workflow key
# is refused, so answer like any other revoked credential rather than 500.
raise HTTPException(status_code = 401, detail = "Invalid or expired token") from exc
except ValueError as exc:
raise log_and_http_error(
exc,

3
studio/backend/run.py
View file

@ -1328,7 +1328,8 @@ def _apply_supplied_password(password_value: "Optional[str]") -> None:
if not _auth_storage.requires_password_change(_admin):
print(
"Error: an Unsloth admin password is already set; --password only sets "
"the initial password. Run `unsloth studio reset-password` first.",
"the initial password. Change it in the UI, or run `unsloth studio "
"reset-password` for a new one.",
file = sys.stderr,
flush = True,
)

8
studio/backend/tests/test_change_password_policy.py
View file

@ -67,9 +67,11 @@ def test_rejects_password_containing_spaces(_user):
def test_allows_password_without_spaces(_user, monkeypatch):
monkeypatch.setattr(auth_routes.storage, "update_password", lambda *args, **kwargs: True)
monkeypatch.setattr(auth_routes, "create_access_token", lambda subject: "at")
monkeypatch.setattr(auth_routes, "create_refresh_token", lambda subject: "rt")
monkeypatch.setattr(
auth_routes.storage, "update_password", lambda *args, **kwargs: "rotated-secret"
)
monkeypatch.setattr(auth_routes, "create_access_token", lambda subject, **kwargs: "at")
monkeypatch.setattr(auth_routes, "create_refresh_token", lambda subject, **kwargs: "rt")
token = _change("correct-horse-battery")
assert token.access_token == "at"
assert token.must_change_password is False

255
studio/backend/tests/test_credential_rotation_race.py Normal file
View file

@ -0,0 +1,255 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""A password rotation must not leave a session minted from the replaced credential.
`unsloth studio reset-password` rotates in place against a live server, so a login
can verify the old password, have the rotation land, and only then mint its tokens.
Issuance is bound to the credential version that was verified, so such a login gets
tokens that are already dead rather than a session that outlives the reset.
"""
import secrets
from datetime import datetime, timedelta, timezone
import jwt
import pytest
from auth import hashing, storage
from auth.authentication import ALGORITHM, create_access_token, create_refresh_token
@pytest.fixture(autouse = True)
def isolated_auth_db(tmp_path, monkeypatch):
monkeypatch.setattr(storage, "DB_PATH", tmp_path / "auth.db")
monkeypatch.setattr(storage, "_BOOTSTRAP_PW_PATH", tmp_path / ".bootstrap_password")
monkeypatch.setattr(storage, "_bootstrap_password", None)
monkeypatch.setattr(storage, "_api_key_pbkdf2_salt_cache", None)
yield
@pytest.fixture
def admin():
storage.create_initial_user(
username = storage.DEFAULT_ADMIN_USERNAME,
password = "old-password-123",
jwt_secret = secrets.token_urlsafe(64),
)
return storage.DEFAULT_ADMIN_USERNAME
def _verified_secret(username):
return storage.get_user_and_secret(username)[2]
def test_access_token_from_the_replaced_credential_is_rejected(admin):
secret = _verified_secret(admin)
storage.update_password(admin, "new-password-456", revoke_refresh_tokens = True)
token = create_access_token(subject = admin, secret = secret)
with pytest.raises(jwt.InvalidTokenError):
jwt.decode(token, storage.get_jwt_secret(admin), algorithms = [ALGORITHM])
def test_refresh_token_from_the_replaced_credential_is_rejected(admin):
secret = _verified_secret(admin)
# Inserted AFTER the rotation's DELETE, so revocation alone cannot catch it.
storage.update_password(admin, "new-password-456", revoke_refresh_tokens = True)
token = create_refresh_token(subject = admin, secret = secret)
assert storage.verify_refresh_token(token) is None
assert storage.consume_refresh_token(token) is None
def test_a_rejected_refresh_token_is_dropped(admin):
secret = _verified_secret(admin)
storage.update_password(admin, "new-password-456", revoke_refresh_tokens = True)
token = create_refresh_token(subject = admin, secret = secret)
storage.verify_refresh_token(token)
conn = storage.get_connection()
try:
assert conn.execute("SELECT COUNT(*) AS c FROM refresh_tokens").fetchone()["c"] == 0
finally:
conn.close()
def test_tokens_from_the_current_credential_still_work(admin):
secret = _verified_secret(admin)
access = create_access_token(subject = admin, secret = secret)
refresh = create_refresh_token(subject = admin, secret = secret)
jwt.decode(access, storage.get_jwt_secret(admin), algorithms = [ALGORITHM])
assert storage.verify_refresh_token(refresh) == (admin, False)
def test_refresh_cannot_outlive_a_rotation_it_raced(admin):
# /refresh consumes, then mints. A rotation landing in between must not let
# the replacement pair be signed with the credential that just replaced it.
secret = _verified_secret(admin)
token = create_refresh_token(subject = admin, secret = secret)
consumed = storage.consume_refresh_token(token)
assert consumed is not None
_username, _is_desktop, consumed_secret = consumed
storage.update_password(admin, "new-password-456", revoke_refresh_tokens = True)
access = create_access_token(subject = admin, secret = consumed_secret)
refresh = create_refresh_token(subject = admin, secret = consumed_secret)
with pytest.raises(jwt.InvalidTokenError):
jwt.decode(access, storage.get_jwt_secret(admin), algorithms = [ALGORITHM])
assert storage.verify_refresh_token(refresh) is None
def test_desktop_login_cannot_outlive_a_rotation_it_raced(admin):
# The reset deletes the desktop secret, so a desktop-login that validated it
# just beforehand must not mint a session that survives.
raw = storage.create_desktop_secret()
verified = storage.validate_desktop_secret_with_credential(raw)
assert verified is not None
_username, verified_secret = verified
storage.update_password(admin, "new-password-456", revoke_refresh_tokens = True)
access = create_access_token(subject = admin, desktop = True, secret = verified_secret)
refresh = create_refresh_token(subject = admin, desktop = True, secret = verified_secret)
with pytest.raises(jwt.InvalidTokenError):
jwt.decode(access, storage.get_jwt_secret(admin), algorithms = [ALGORITHM])
assert storage.verify_refresh_token(refresh) is None
def test_change_password_cannot_overwrite_a_rotation_it_raced(admin):
# A change-password that verified the old hash must not clobber a reset that
# committed while it was in flight.
_salt, verified_hash, _secret, _must_change = storage.get_user_and_secret(admin)
storage.update_password(admin, "reset-by-the-cli-789", revoke_refresh_tokens = True)
assert not storage.update_password(
admin,
"attacker-chosen-000",
revoke_refresh_tokens = True,
expect_password_hash = verified_hash,
)
salt, pwd_hash, _s, _m = storage.get_user_and_secret(admin)
assert hashing.verify_password("reset-by-the-cli-789", salt, pwd_hash)
def test_api_key_creation_from_a_revoked_credential_is_refused(admin):
generation = storage.credential_generation(_verified_secret(admin))
storage.update_password(admin, "new-password-456", revoke_refresh_tokens = True)
with pytest.raises(storage.CredentialRotated):
storage.create_api_key(username = admin, name = "k", expect_gen = generation)
conn = storage.get_connection()
try:
assert conn.execute("SELECT COUNT(*) AS c FROM api_keys").fetchone()["c"] == 0
finally:
conn.close()
def test_api_key_creation_under_the_current_credential_still_works(admin):
generation = storage.credential_generation(_verified_secret(admin))
raw_key, _row = storage.create_api_key(username = admin, name = "k", expect_gen = generation)
assert storage.validate_api_key(raw_key) == admin
def test_change_password_tokens_are_bound_to_its_own_write(admin):
# The tokens returned to a successful change-password must be signed with the
# secret that write produced, not whatever a later reset put in the DB.
_salt, verified_hash, _secret, _must = storage.get_user_and_secret(admin)
new_secret = storage.update_password(
admin,
"chosen-by-the-user",
revoke_refresh_tokens = True,
expect_password_hash = verified_hash,
)
assert new_secret is not None
storage.update_password(admin, "reset-by-the-cli-789", revoke_refresh_tokens = True)
access = create_access_token(subject = admin, secret = new_secret)
refresh = create_refresh_token(subject = admin, secret = new_secret)
with pytest.raises(jwt.InvalidTokenError):
jwt.decode(access, storage.get_jwt_secret(admin), algorithms = [ALGORITHM])
assert storage.verify_refresh_token(refresh) is None
def test_internal_api_key_minting_honours_the_request_generation(admin):
generation = storage.credential_generation(_verified_secret(admin))
storage.update_password(admin, "new-password-456", revoke_refresh_tokens = True)
with pytest.raises(storage.CredentialRotated):
storage.create_api_key(
username = admin,
name = "data-recipe workflow",
internal = True,
expect_gen = generation,
)
def test_api_key_auth_reports_the_version_the_key_was_valid_under(admin):
# The generation must come from the same transaction as the key check, or a
# revoked key could hand a route the post-reset generation and mint again.
raw, _row = storage.create_api_key(username = admin, name = "agent")
verified = storage.validate_api_key_with_credential(raw)
assert verified is not None
_user, secret = verified
generation = storage.credential_generation(secret)
storage.update_password(admin, "new-password-456", revoke_refresh_tokens = True)
conn = storage.get_connection()
try:
conn.execute("DELETE FROM api_keys")
conn.commit()
finally:
conn.close()
assert storage.validate_api_key(raw) is None
with pytest.raises(storage.CredentialRotated):
storage.create_api_key(username = admin, name = "after", expect_gen = generation)
def test_consuming_a_legacy_token_reports_the_pre_reset_credential(admin):
# An unstamped row has no generation to compare, so consume must read the
# credential inside the delete transaction rather than after committing it.
token = secrets.token_urlsafe(48)
expires_at = (datetime.now(timezone.utc) + timedelta(days = 7)).isoformat()
storage.save_refresh_token(token, admin, expires_at, secret_gen = None)
conn = storage.get_connection()
try:
conn.execute("UPDATE refresh_tokens SET secret_gen = NULL")
conn.commit()
finally:
conn.close()
consumed = storage.consume_refresh_token(token)
assert consumed is not None
_username, _is_desktop, consumed_secret = consumed
storage.update_password(admin, "new-password-456", revoke_refresh_tokens = True)
access = create_access_token(subject = admin, secret = consumed_secret)
with pytest.raises(jwt.InvalidTokenError):
jwt.decode(access, storage.get_jwt_secret(admin), algorithms = [ALGORITHM])
def test_unstamped_legacy_tokens_still_verify(admin):
# Rows written before the secret_gen column existed must not log users out.
token = secrets.token_urlsafe(48)
expires_at = (datetime.now(timezone.utc) + timedelta(days = 7)).isoformat()
storage.save_refresh_token(token, admin, expires_at, secret_gen = None)
conn = storage.get_connection()
try:
conn.execute("UPDATE refresh_tokens SET secret_gen = NULL")
conn.commit()
finally:
conn.close()
assert storage.verify_refresh_token(token) == (admin, False)

59
studio/backend/tests/test_desktop_auth.py
View file

@ -445,7 +445,7 @@ def test_consume_refresh_token_second_call_returns_none():
storage.save_refresh_token(raw, storage.DEFAULT_ADMIN_USERNAME, expires)
first = storage.consume_refresh_token(raw)
assert first == (storage.DEFAULT_ADMIN_USERNAME, False)
assert first[:2] == (storage.DEFAULT_ADMIN_USERNAME, False)
second = storage.consume_refresh_token(raw)
assert second is None
@ -474,7 +474,7 @@ def test_consume_refresh_token_concurrent_only_one_succeeds(tmp_path, monkeypatc
successes = [r for r in results if r is not None]
assert len(successes) == 1, f"expected exactly one consumer to win, got {len(successes)}"
assert successes[0] == (storage.DEFAULT_ADMIN_USERNAME, False)
assert successes[0][:2] == (storage.DEFAULT_ADMIN_USERNAME, False)
def test_consume_refresh_token_expired_returns_none():
@ -548,6 +548,28 @@ def test_local_recipe_token_authenticates_as_admin_for_web_user(loaded_local_mod
assert asyncio.run(get_current_subject(credentials)) == storage.DEFAULT_ADMIN_USERNAME
def test_rotated_credential_job_start_is_401_not_500(loaded_local_model):
# A reset-password landing mid-request makes the workflow-key mint refuse.
# That must reach the client as a revoked credential, not an unhandled error.
from fastapi import HTTPException
seed_user()
jobs_route = data_recipe_jobs_module()
stale_gen = storage.credential_generation(secrets.token_urlsafe(64))
with pytest.raises(storage.CredentialRotated):
jobs_route._inject_local_providers(local_recipe(), local_recipe_request("t"), stale_gen)
def _boom(*_a, **_k):
raise storage.CredentialRotated("revoked")
jobs_route._inject_local_providers = _boom
payload = SimpleNamespace(recipe = local_recipe(), run = {})
with pytest.raises(HTTPException) as excinfo:
jobs_route.create_job(payload, local_recipe_request("t"), ("unsloth", stale_gen))
assert excinfo.value.status_code == 401
def test_desktop_login_rejects_invalid_secret():
seed_user(must_change_password = False)
client = auth_client()
@ -580,18 +602,31 @@ def test_reset_password_removes_desktop_secret_files(tmp_path, monkeypatch):
from unsloth_cli.commands import studio as studio_cli
auth_dir = tmp_path / "auth"
auth_dir.mkdir()
(auth_dir / "auth.db").write_text("db")
(auth_dir / ".bootstrap_password").write_text("boot")
(auth_dir / ".desktop_secret").write_text("new")
monkeypatch.setattr(studio_cli, "STUDIO_HOME", tmp_path)
secret = studio_cli._create_desktop_secret_in_cli()
studio_cli._write_auth_secret(auth_dir / studio_cli.DESKTOP_SECRET_FILE, secret)
(auth_dir / studio_cli.BOOTSTRAP_PASSWORD_FILE).write_text("boot")
result = CliRunner().invoke(studio_cli.studio_app, ["reset-password"])
assert result.exit_code == 0
assert not (auth_dir / "auth.db").exists()
assert not (auth_dir / ".bootstrap_password").exists()
assert not (auth_dir / ".desktop_secret").exists()
assert result.exit_code == 0, result.output
# The DB survives on purpose: a running server keeps serving from its admin row.
assert (auth_dir / "auth.db").exists()
assert not (auth_dir / studio_cli.BOOTSTRAP_PASSWORD_FILE).exists()
assert not (auth_dir / studio_cli.DESKTOP_SECRET_FILE).exists()
conn = studio_cli._connect_auth_db()
try:
surviving = conn.execute(
"SELECT COUNT(*) FROM app_secrets WHERE key IN (?, ?)",
(
studio_cli.DESKTOP_SECRET_HASH_KEY,
studio_cli.DESKTOP_SECRET_CREATED_AT_KEY,
),
).fetchone()[0]
finally:
conn.close()
assert surviving == 0
def test_reset_password_removes_desktop_secret_files_without_db(tmp_path, monkeypatch):
@ -846,7 +881,7 @@ def test_update_password_clears_desktop_secret():
assert storage.validate_desktop_secret(raw) == storage.DEFAULT_ADMIN_USERNAME
changed = storage.update_password(storage.DEFAULT_ADMIN_USERNAME, "new-admin-password")
assert changed is True
assert changed
assert storage.validate_desktop_secret(raw) is None
@ -855,7 +890,7 @@ def test_update_password_on_unknown_user_leaves_desktop_secret_intact():
raw = storage.create_desktop_secret()
changed = storage.update_password("not-a-user", "irrelevant")
assert changed is False
assert not changed
assert storage.validate_desktop_secret(raw) == storage.DEFAULT_ADMIN_USERNAME

4
studio/backend/tests/test_password_prompt_backstop.py
View file

@ -247,8 +247,8 @@ def test_lifespan_honors_bootstrap_suppression_in_source():
def test_clear_bootstrap_password_truncates_when_unlink_fails(monkeypatch, tmp_path):
# If the file cannot be unlinked (Windows AV / read-only auth dir), clear must
# truncate it so its stale plaintext cannot be re-seeded by
# generate_bootstrap_password() after a later reset-password deletes auth.db,
# which would re-validate the revoked bootstrap password.
# generate_bootstrap_password() if auth.db is ever recreated, which would
# re-validate the revoked bootstrap password.
import pathlib
pw_path = tmp_path / ".bootstrap_password"