diff --git a/.gitignore b/.gitignore index 8c79c64263..15f9ee0d16 100755 --- a/.gitignore +++ b/.gitignore @@ -35,3 +35,4 @@ Thumbs.db # Other resources/ tmp/ +auth.db diff --git a/studio/backend/auth/__init__.py b/studio/backend/auth/__init__.py index 7b67b74084..60ad59885c 100644 --- a/studio/backend/auth/__init__.py +++ b/studio/backend/auth/__init__.py @@ -1,24 +1,46 @@ """ Authentication module for JWT-based auth with SQLite storage. """ -from .authentication import create_access_token, get_current_subject, reload_secret +from .authentication import ( + create_access_token, + create_refresh_token, + refresh_access_token, + get_current_subject, + reload_secret, +) from .storage import ( is_initialized, create_initial_user, get_user_and_secret, load_jwt_secret, + save_setup_token, + consume_setup_token, + has_pending_setup_token, + save_refresh_token, + verify_refresh_token, + revoke_user_refresh_tokens, ) from .hashing import hash_password, verify_password __all__ = [ "create_access_token", + "create_refresh_token", + "refresh_access_token", "get_current_subject", "reload_secret", "is_initialized", "create_initial_user", "get_user_and_secret", "load_jwt_secret", + "save_setup_token", + "consume_setup_token", + "has_pending_setup_token", + "save_refresh_token", + "verify_refresh_token", + "revoke_user_refresh_tokens", "hash_password", "verify_password", ] + + diff --git a/studio/backend/auth/auth.db b/studio/backend/auth/auth.db index b525e4b096..8f688f96a5 100644 Binary files a/studio/backend/auth/auth.db and b/studio/backend/auth/auth.db differ diff --git a/studio/backend/auth/authentication.py b/studio/backend/auth/authentication.py index 33b5125ddb..6ea668db3e 100644 --- a/studio/backend/auth/authentication.py +++ b/studio/backend/auth/authentication.py @@ -6,10 +6,11 @@ from fastapi import Depends, HTTPException, status from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer import jwt -from .storage import load_jwt_secret +from .storage import load_jwt_secret, save_refresh_token, verify_refresh_token ALGORITHM = "HS256" ACCESS_TOKEN_EXPIRE_MINUTES = 60 +REFRESH_TOKEN_EXPIRE_DAYS = 7 # Load stable secret from SQLite (set during first-time setup) # This will raise RuntimeError if auth hasn't been initialized yet @@ -40,6 +41,31 @@ def create_access_token( return jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM) +def create_refresh_token(subject: str) -> str: + """ + Create a random refresh token, store its hash in SQLite, and return it. + + Refresh tokens are opaque (not JWTs) and expire after REFRESH_TOKEN_EXPIRE_DAYS. + """ + token = secrets.token_urlsafe(48) + expires_at = datetime.now(UTC) + timedelta(days=REFRESH_TOKEN_EXPIRE_DAYS) + save_refresh_token(token, subject, expires_at.isoformat()) + return token + + +def refresh_access_token(refresh_token: str) -> Optional[str]: + """ + Validate a refresh token and issue a new access token. + + The refresh token itself is NOT consumed — it stays valid until expiry. + Returns a new access_token or None if the refresh token is invalid/expired. + """ + username = verify_refresh_token(refresh_token) + if username is None: + return None + return create_access_token(subject=username) + + def reload_secret() -> None: """ Reload the JWT secret from SQLite. @@ -77,7 +103,3 @@ async def get_current_subject( status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid or expired token", ) -# token = create_access_token("local-user") -# print(token) - - diff --git a/studio/backend/auth/storage.py b/studio/backend/auth/storage.py index 0e2c9d7388..2864cb9852 100644 --- a/studio/backend/auth/storage.py +++ b/studio/backend/auth/storage.py @@ -1,13 +1,20 @@ """ SQLite storage for authentication data (user credentials + JWT secret). """ +import hashlib import sqlite3 +from datetime import UTC, datetime from pathlib import Path from typing import Optional, Tuple DB_PATH = Path(__file__).parent / "auth.db" +def _hash_token(token: str) -> str: + """SHA-256 hash a setup token for safe storage.""" + return hashlib.sha256(token.encode("utf-8")).hexdigest() + + def get_connection() -> sqlite3.Connection: """Get a connection to the auth database, creating tables if needed.""" conn = sqlite3.connect(DB_PATH) @@ -23,6 +30,24 @@ def get_connection() -> sqlite3.Connection: ); """ ) + conn.execute( + """ + CREATE TABLE IF NOT EXISTS setup_tokens ( + id INTEGER PRIMARY KEY, + token_hash TEXT NOT NULL + ); + """ + ) + conn.execute( + """ + CREATE TABLE IF NOT EXISTS refresh_tokens ( + id INTEGER PRIMARY KEY, + token_hash TEXT NOT NULL, + username TEXT NOT NULL, + expires_at TEXT NOT NULL + ); + """ + ) conn.commit() return conn @@ -99,3 +124,119 @@ def load_jwt_secret() -> str: finally: conn.close() + +def save_setup_token(token: str) -> None: + """ + Store a hashed setup token, replacing any existing one. + """ + token_hash = _hash_token(token) + conn = get_connection() + try: + conn.execute("DELETE FROM setup_tokens") + conn.execute("INSERT INTO setup_tokens (token_hash) VALUES (?)", (token_hash,)) + conn.commit() + finally: + conn.close() + + +def consume_setup_token(token: str) -> bool: + """ + Verify a setup token and delete it if valid. + + Returns True if the token was valid (and is now consumed), False otherwise. + """ + token_hash = _hash_token(token) + conn = get_connection() + try: + cur = conn.execute( + "SELECT id FROM setup_tokens WHERE token_hash = ?", (token_hash,) + ) + row = cur.fetchone() + if row is None: + return False + conn.execute("DELETE FROM setup_tokens WHERE id = ?", (row["id"],)) + conn.commit() + return True + finally: + conn.close() + + +def has_pending_setup_token() -> bool: + """Check if a setup token is waiting to be consumed.""" + conn = get_connection() + try: + cur = conn.execute("SELECT COUNT(*) AS c FROM setup_tokens") + row = cur.fetchone() + return bool(row["c"]) + finally: + conn.close() + + +def save_refresh_token(token: str, username: str, expires_at: str) -> None: + """ + Store a hashed refresh token with its associated username and expiry. + """ + token_hash = _hash_token(token) + conn = get_connection() + try: + conn.execute( + """ + INSERT INTO refresh_tokens (token_hash, username, expires_at) + VALUES (?, ?, ?) + """, + (token_hash, username, expires_at), + ) + conn.commit() + finally: + conn.close() + + +def verify_refresh_token(token: str) -> Optional[str]: + """ + Verify a refresh token and return the username. + + Returns the username if valid and not expired, None otherwise. + The token is NOT consumed — it stays valid until it expires. + """ + token_hash = _hash_token(token) + conn = get_connection() + try: + # Clean up any expired tokens while we're here + conn.execute( + "DELETE FROM refresh_tokens WHERE expires_at < ?", + (datetime.now(UTC).isoformat(),), + ) + conn.commit() + + cur = conn.execute( + """ + SELECT id, username, expires_at FROM refresh_tokens + WHERE token_hash = ? + """, + (token_hash,), + ) + row = cur.fetchone() + if row is None: + return None + + # Check expiry + expires_at = datetime.fromisoformat(row["expires_at"]) + if datetime.now(UTC) > expires_at: + conn.execute("DELETE FROM refresh_tokens WHERE id = ?", (row["id"],)) + conn.commit() + return None + + return row["username"] + finally: + conn.close() + + +def revoke_user_refresh_tokens(username: str) -> None: + """Revoke all refresh tokens for a user (e.g. on logout).""" + conn = get_connection() + try: + conn.execute("DELETE FROM refresh_tokens WHERE username = ?", (username,)) + conn.commit() + finally: + conn.close() + diff --git a/studio/backend/main.py b/studio/backend/main.py index f93f6c2820..4d19e7e32a 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -1,6 +1,9 @@ """ Main FastAPI application for Unsloth UI Backend """ +import secrets +from contextlib import asynccontextmanager + from fastapi import FastAPI from fastapi.middleware.cors import CORSMiddleware from fastapi.staticfiles import StaticFiles @@ -10,12 +13,30 @@ from datetime import datetime # Import routers from routes import training_router, models_router, inference_router, datasets_router, auth_router +from auth import storage + + +@asynccontextmanager +async def lifespan(app: FastAPI): + """Generate and print a setup token on startup if auth is not yet initialized.""" + if not storage.is_initialized(): + setup_token = secrets.token_urlsafe(32) + storage.save_setup_token(setup_token) + print("\n" + "=" * 60) + print("FIRST-TIME SETUP REQUIRED") + print("Use this one-time setup token to create your admin account:\n") + print(f" {setup_token}\n") + print("This token can only be used once.") + print("=" * 60 + "\n") + yield + # Create FastAPI app app = FastAPI( title="Unsloth UI Backend", version="1.0.0", - description="Backend API for Unsloth UI - Training and Model Management" + description="Backend API for Unsloth UI - Training and Model Management", + lifespan=lifespan, ) # CORS middleware diff --git a/studio/backend/models/auth.py b/studio/backend/models/auth.py index f71956db58..43c591fb53 100644 --- a/studio/backend/models/auth.py +++ b/studio/backend/models/auth.py @@ -6,6 +6,7 @@ from pydantic import BaseModel, Field class AuthSetupRequest(BaseModel): """First-time setup: create the initial admin user + password.""" + setup_token: str = Field(..., description="One-time setup token printed to the server console") username: str = Field(..., description="Admin username") password: str = Field(..., min_length=8, description="Admin password (minimum 8 characters)") @@ -16,6 +17,11 @@ class AuthLoginRequest(BaseModel): password: str = Field(..., description="Password") +class RefreshTokenRequest(BaseModel): + """Refresh token payload to obtain new access + refresh tokens.""" + refresh_token: str = Field(..., description="Refresh token from a previous login or refresh") + + class AuthStatusResponse(BaseModel): """Indicate whether auth has been initialized.""" initialized: bool = Field(..., description="True if auth setup has been completed") diff --git a/studio/backend/models/users.py b/studio/backend/models/users.py index ba6bb8f8e7..2f404211a7 100644 --- a/studio/backend/models/users.py +++ b/studio/backend/models/users.py @@ -4,7 +4,7 @@ This module defines the data models used for user authentication and management in the FastAPI application. """ -from pydantic import BaseModel +from pydantic import BaseModel, Field class User(BaseModel): @@ -20,9 +20,10 @@ class UserInDB(BaseModel): class Token(BaseModel): - """Authentication token model with access token and type.""" + """Authentication token model with access and refresh tokens.""" access_token: str + refresh_token: str token_type: str @@ -30,3 +31,4 @@ class TokenData(BaseModel): """Token payload model containing username.""" username: str | None = None + diff --git a/studio/backend/routes/auth.py b/studio/backend/routes/auth.py index 4d0ac38742..21ac5ac1eb 100644 --- a/studio/backend/routes/auth.py +++ b/studio/backend/routes/auth.py @@ -7,11 +7,17 @@ import secrets from models.auth import ( AuthSetupRequest, AuthLoginRequest, + RefreshTokenRequest, AuthStatusResponse, ) from models.users import Token from auth import storage, hashing -from auth.authentication import create_access_token, reload_secret +from auth.authentication import ( + create_access_token, + create_refresh_token, + refresh_access_token, + reload_secret, +) router = APIRouter() @@ -32,6 +38,7 @@ async def setup_auth(payload: AuthSetupRequest) -> Token: """ First-time setup: create the admin user and a JWT secret. + Requires a valid setup token (printed to the server console on startup). Can only be called once. Subsequent calls will return 400. """ if storage.is_initialized(): @@ -40,6 +47,13 @@ async def setup_auth(payload: AuthSetupRequest) -> Token: detail="Auth is already initialized.", ) + # Validate the one-time setup token + if not storage.consume_setup_token(payload.setup_token): + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail="Invalid or expired setup token.", + ) + # Generate a strong random JWT secret for this installation jwt_secret = secrets.token_urlsafe(64) @@ -59,15 +73,20 @@ async def setup_auth(payload: AuthSetupRequest) -> Token: # Reload JWT secret from DB (so authentication.py picks it up) reload_secret() - # Issue a token for the new user + # Issue access + refresh tokens for the new user access_token = create_access_token(subject=payload.username) - return Token(access_token=access_token, token_type="bearer") + refresh_token = create_refresh_token(subject=payload.username) + return Token( + access_token=access_token, + refresh_token=refresh_token, + token_type="bearer", + ) @router.post("/login", response_model=Token) async def login(payload: AuthLoginRequest) -> Token: """ - Login with username/password and receive a JWT. + Login with username/password and receive access + refresh tokens. """ record = storage.get_user_and_secret(payload.username) if record is None: @@ -84,5 +103,31 @@ async def login(payload: AuthLoginRequest) -> Token: ) access_token = create_access_token(subject=payload.username) - return Token(access_token=access_token, token_type="bearer") + refresh_token = create_refresh_token(subject=payload.username) + return Token( + access_token=access_token, + refresh_token=refresh_token, + token_type="bearer", + ) + + +@router.post("/refresh", response_model=Token) +async def refresh(payload: RefreshTokenRequest) -> Token: + """ + Exchange a valid refresh token for a new access token. + + The refresh token itself is reusable until it expires (7 days). + """ + new_access_token = refresh_access_token(payload.refresh_token) + if new_access_token is None: + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail="Invalid or expired refresh token", + ) + + return Token( + access_token=new_access_token, + refresh_token=payload.refresh_token, + token_type="bearer", + ) diff --git a/studio/backend/utils/datasets/dataset_utils.py b/studio/backend/utils/datasets/dataset_utils.py index a1137ae1f6..c16ea3fe63 100644 --- a/studio/backend/utils/datasets/dataset_utils.py +++ b/studio/backend/utils/datasets/dataset_utils.py @@ -12,7 +12,7 @@ All internal utilities have been moved to separate modules: - chat_templates: apply_chat_template_to_dataset, get_tokenizer_chat_template, etc. - vlm_processing: generate_smart_vlm_instruction - data_collators: DeepSeekOCRDataCollator, VLMDataCollator -- model_mappings: TEMPLATE_TO_MODEL_MAPPER, RESPONSE_MARKERS +- model_mappings: TEMPLATE_TO_MODEL_MAPPER """ # Import from modular files @@ -37,7 +37,7 @@ from .chat_templates import ( ) from .vlm_processing import generate_smart_vlm_instruction from .data_collators import DeepSeekOCRDataCollator, VLMDataCollator -from .model_mappings import TEMPLATE_TO_MODEL_MAPPER, RESPONSE_MARKERS +from .model_mappings import TEMPLATE_TO_MODEL_MAPPER def check_dataset_format(dataset, is_vlm: bool = False) -> dict: