diff --git a/studio/backend/main.py b/studio/backend/main.py index b60ad48218..cb277f7007 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -267,7 +267,8 @@ logger = LogConfig.setup_logging( app.add_middleware(LoggingMiddleware) -# Web-search favicons load from *.gstatic.com; everything else is same-origin. +# Citation favicons load from www.google.com/s2/favicons; *.gstatic.com is +# kept for legacy web-search faviconV2 paths. Everything else is same-origin. from starlette.middleware.base import BaseHTTPMiddleware # noqa: E402 from starlette.requests import Request as _StarletteRequest # noqa: E402 @@ -283,7 +284,7 @@ def _build_csp(script_nonce: "str | None" = None) -> str: "default-src 'self'; " "img-src 'self' data: blob: https://t0.gstatic.com " "https://t1.gstatic.com https://t2.gstatic.com " - "https://t3.gstatic.com; " + "https://t3.gstatic.com https://www.google.com; " "connect-src 'self' https://huggingface.co https://datasets-server.huggingface.co; " "style-src 'self' 'unsafe-inline'; " f"{script_src}; " diff --git a/studio/backend/routes/auth.py b/studio/backend/routes/auth.py index 30221c2c93..bb4ce87cd7 100644 --- a/studio/backend/routes/auth.py +++ b/studio/backend/routes/auth.py @@ -7,6 +7,8 @@ Authentication API routes from fastapi import APIRouter, Depends, HTTPException, Request, Response, status +import ipaddress +import os import threading import time from collections import deque @@ -36,47 +38,155 @@ from auth.authentication import ( router = APIRouter() -# In-memory per-IP login rate limiter; multi-process deployment needs a shared store. -_LOGIN_BUCKETS: dict[str, deque] = {} +# Per-(ip, username) bucket + per-IP aggregate. Account bucket stops one user's +# typos from blocking others; the aggregate stops username-rotation spray. +# Single-process only -- multi-worker deployments need a shared store. +_LOGIN_BUCKETS: dict[tuple[str, str], deque] = {} +_LOGIN_IP_BUCKETS: dict[str, deque] = {} _LOGIN_BUCKETS_LOCK = threading.Lock() _LOGIN_WINDOW_SECONDS = 60.0 _LOGIN_MAX_FAILS = 5 +_LOGIN_IP_MAX_FAILS = 30 _LOGIN_LOCKOUT_SECONDS = 60 +# Bucket-dict cap. On overflow we prune stale entries; if still full the +# failure folds into the per-IP aggregate only. +_LOGIN_MAX_BUCKETS = 4096 +# Unrepresentable as a real username (leading NUL); folds unknown-user attempts +# into one slot so attacker cardinality cannot blow the bucket dict. +_UNKNOWN_LOGIN_USER = "\x00unknown-user" -def _client_key(request: Request | None) -> str: - if request is None or request.client is None: +def _trust_forwarded_for() -> bool: + """Honour X-Forwarded-For only when UNSLOTH_STUDIO_TRUST_FORWARDED is set. + + Off by default so a direct caller cannot spoof the header. + """ + return os.environ.get("UNSLOTH_STUDIO_TRUST_FORWARDED", "").lower() in ( + "1", + "true", + "yes", + ) + + +def _normalize_forwarded_addr(value: str) -> str: + """Parse an XFF / Forwarded `for=` value into a bare IP (port-stripped).""" + value = (value or "").strip().strip('"') + if not value or value.lower() == "unknown": + return "" + if value.startswith("["): + # Bracketed IPv6, optionally with port. + end = value.find("]") + if end <= 0: + return "" + host = value[1:end] + elif value.count(":") == 1: + # IPv4:port. Bare IPv6 has multiple colons and takes the else branch. + head, _, tail = value.rpartition(":") + host = head if tail.isdigit() and head else value + else: + host = value + try: + return str(ipaddress.ip_address(host)) + except ValueError: + return "" + + +def _forwarded_for_from_element(element: str) -> str: + """Pick the `for=` token out of a single ``Forwarded`` element.""" + for tok in element.split(";"): + key, sep, val = tok.strip().partition("=") + if sep and key.lower() == "for": + return _normalize_forwarded_addr(val) + return "" + + +def _client_ip(request: Request | None) -> str: + if request is None: return "_unknown" - return request.client.host or "_unknown" + if _trust_forwarded_for(): + xff = request.headers.get("x-forwarded-for", "") + if xff: + # First entry is the originating client. + normalized = _normalize_forwarded_addr(xff.split(",", 1)[0]) + if normalized: + return normalized + fwd = request.headers.get("forwarded", "") + if fwd: + # First element only -- multi-element headers cannot fork buckets. + normalized = _forwarded_for_from_element(fwd.split(",", 1)[0]) + if normalized: + return normalized + return (request.client.host if request.client else None) or "_unknown" -def _record_login_failure(ip: str) -> int: +def _bucket_key(request: Request | None, username: str) -> tuple[str, str]: + return (_client_ip(request), (username or "").casefold()) + + +def _unknown_user_key(request: Request | None) -> tuple[str, str]: + return (_client_ip(request), _UNKNOWN_LOGIN_USER) + + +def _prune_bucket(bucket: deque, now: float) -> None: + while bucket and now - bucket[0] > _LOGIN_WINDOW_SECONDS: + bucket.popleft() + + +def _prune_stale_buckets(now: float) -> None: + """Drop empty / expired account buckets to bound memory under spray.""" + stale: list[tuple[str, str]] = [] + for key, bucket in _LOGIN_BUCKETS.items(): + _prune_bucket(bucket, now) + if not bucket: + stale.append(key) + for key in stale: + _LOGIN_BUCKETS.pop(key, None) + + +def _record_login_failure(key: tuple[str, str]) -> int: now = time.monotonic() + ip, _username = key with _LOGIN_BUCKETS_LOCK: - bucket = _LOGIN_BUCKETS.setdefault(ip, deque()) - while bucket and now - bucket[0] > _LOGIN_WINDOW_SECONDS: - bucket.popleft() - bucket.append(now) - return len(bucket) + ip_bucket = _LOGIN_IP_BUCKETS.setdefault(ip, deque()) + _prune_bucket(ip_bucket, now) + ip_bucket.append(now) + + if key not in _LOGIN_BUCKETS and len(_LOGIN_BUCKETS) >= _LOGIN_MAX_BUCKETS: + _prune_stale_buckets(now) + if key in _LOGIN_BUCKETS or len(_LOGIN_BUCKETS) < _LOGIN_MAX_BUCKETS: + account_bucket = _LOGIN_BUCKETS.setdefault(key, deque()) + _prune_bucket(account_bucket, now) + account_bucket.append(now) + return len(account_bucket) + # Bucket dict is at its cap; per-IP cap still applies via ip_bucket. + return len(ip_bucket) -def _login_blocked(ip: str) -> int: +def _blocked_for(bucket: deque | None, now: float, max_fails: int) -> int: + if not bucket: + return 0 + _prune_bucket(bucket, now) + if len(bucket) >= max_fails: + return max(1, int(_LOGIN_WINDOW_SECONDS - (now - bucket[0]))) + return 0 + + +def _login_blocked(key: tuple[str, str]) -> int: """Return seconds until the next attempt is allowed, or 0.""" now = time.monotonic() + ip, _username = key with _LOGIN_BUCKETS_LOCK: - bucket = _LOGIN_BUCKETS.get(ip) - if not bucket: - return 0 - while bucket and now - bucket[0] > _LOGIN_WINDOW_SECONDS: - bucket.popleft() - if len(bucket) >= _LOGIN_MAX_FAILS: - return max(1, int(_LOGIN_WINDOW_SECONDS - (now - bucket[0]))) - return 0 + return max( + _blocked_for(_LOGIN_BUCKETS.get(key), now, _LOGIN_MAX_FAILS), + _blocked_for(_LOGIN_IP_BUCKETS.get(ip), now, _LOGIN_IP_MAX_FAILS), + ) -def _clear_login_bucket(ip: str) -> None: +def _clear_login_bucket(key: tuple[str, str]) -> None: + ip, _username = key with _LOGIN_BUCKETS_LOCK: - _LOGIN_BUCKETS.pop(ip, None) + _LOGIN_BUCKETS.pop(key, None) + _LOGIN_IP_BUCKETS.pop(ip, None) @router.get("/status", response_model = AuthStatusResponse) @@ -95,14 +205,17 @@ async def auth_status() -> AuthStatusResponse: @router.post("/login", response_model = Token) async def login(payload: AuthLoginRequest, request: Request) -> Token: - """Login with username/password. Rate-limited per source IP.""" - ip = _client_key(request) - blocked_for = _login_blocked(ip) + """Login with username/password. Per-account + per-IP rate-limited.""" + key = _bucket_key(request, payload.username) + unknown_key = _unknown_user_key(request) + blocked_for = max(_login_blocked(key), _login_blocked(unknown_key)) if blocked_for > 0: raise HTTPException( status_code = status.HTTP_429_TOO_MANY_REQUESTS, + # IP is intentionally not interpolated into the body; behind a + # proxy or NAT it is either misleading or an info leak. detail = ( - f"Too many failed login attempts from {ip}. " + f"Too many failed login attempts. " f"Try again in {blocked_for} seconds." ), headers = {"Retry-After": str(blocked_for)}, @@ -110,7 +223,9 @@ async def login(payload: AuthLoginRequest, request: Request) -> Token: record = storage.get_user_and_secret(payload.username) if record is None: - _record_login_failure(ip) + # Record under a single sentinel key per IP so attacker-controlled + # username cardinality does not allocate buckets without bound. + _record_login_failure(unknown_key) raise HTTPException( status_code = status.HTTP_401_UNAUTHORIZED, detail = "Incorrect password. Run 'unsloth studio reset-password' in your terminal to reset it.", @@ -118,13 +233,14 @@ async def login(payload: AuthLoginRequest, request: Request) -> Token: salt, pwd_hash, _jwt_secret, must_change_password = record if not hashing.verify_password(payload.password, salt, pwd_hash): - _record_login_failure(ip) + _record_login_failure(key) raise HTTPException( status_code = status.HTTP_401_UNAUTHORIZED, detail = "Incorrect password. Run 'unsloth studio reset-password' in your terminal to reset it.", ) - _clear_login_bucket(ip) + _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) return Token( diff --git a/studio/backend/tests/test_login_rate_limit.py b/studio/backend/tests/test_login_rate_limit.py new file mode 100644 index 0000000000..c8498d4857 --- /dev/null +++ b/studio/backend/tests/test_login_rate_limit.py @@ -0,0 +1,285 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +"""Tests for the per-(ip, username) login rate limiter. + +Covers: + - bucket key composition is (client-ip, username.lower()) + - X-Forwarded-For is honoured only when UNSLOTH_STUDIO_TRUST_FORWARDED is set + - 429 detail body does NOT leak the client IP + - One username failing does not lock out a different user from the same IP + - One IP failing does not lock out the same user from a different IP +""" + +import os +import sys +from pathlib import Path + +import pytest + +_BACKEND_ROOT = Path(__file__).resolve().parents[1] +if str(_BACKEND_ROOT) not in sys.path: + sys.path.insert(0, str(_BACKEND_ROOT)) + + +@pytest.fixture(autouse = True) +def _reset_buckets(): + """Clear the in-memory bucket dicts between tests.""" + from routes import auth as auth_routes + + auth_routes._LOGIN_BUCKETS.clear() + auth_routes._LOGIN_IP_BUCKETS.clear() + yield + auth_routes._LOGIN_BUCKETS.clear() + auth_routes._LOGIN_IP_BUCKETS.clear() + + +@pytest.fixture +def env_no_proxy(monkeypatch): + monkeypatch.delenv("UNSLOTH_STUDIO_TRUST_FORWARDED", raising = False) + + +@pytest.fixture +def env_trust_proxy(monkeypatch): + monkeypatch.setenv("UNSLOTH_STUDIO_TRUST_FORWARDED", "1") + + +class _FakeRequest: + def __init__(self, client_host = "127.0.0.1", headers = None): + from starlette.datastructures import Headers + + self.client = type("Client", (), {"host": client_host})() + self.headers = Headers(headers or {}) + + +# ---------- _client_ip ---------- + + +class TestClientIp: + def test_uses_request_client_host_by_default(self, env_no_proxy): + from routes.auth import _client_ip + + assert _client_ip(_FakeRequest("203.0.113.5")) == "203.0.113.5" + + def test_ignores_xff_when_trust_off(self, env_no_proxy): + from routes.auth import _client_ip + + req = _FakeRequest( + "127.0.0.1", + {"x-forwarded-for": "198.51.100.7, 10.0.0.1"}, + ) + # The proxy header could be spoofed; without the opt-in we + # only trust the direct connection. + assert _client_ip(req) == "127.0.0.1" + + def test_honours_first_xff_when_trust_on(self, env_trust_proxy): + from routes.auth import _client_ip + + req = _FakeRequest( + "127.0.0.1", + {"x-forwarded-for": "198.51.100.7, 10.0.0.1"}, + ) + assert _client_ip(req) == "198.51.100.7" + + def test_falls_back_to_client_host_when_xff_missing(self, env_trust_proxy): + from routes.auth import _client_ip + + assert _client_ip(_FakeRequest("203.0.113.9")) == "203.0.113.9" + + def test_honours_forwarded_header_when_trust_on(self, env_trust_proxy): + from routes.auth import _client_ip + + req = _FakeRequest( + "127.0.0.1", + {"forwarded": 'for="198.51.100.42";proto=https'}, + ) + assert _client_ip(req) == "198.51.100.42" + + def test_unknown_when_no_client(self, env_no_proxy): + from routes.auth import _client_ip + + req = _FakeRequest() + req.client = None + assert _client_ip(req) == "_unknown" + + def test_xff_strips_ipv4_port(self, env_trust_proxy): + from routes.auth import _client_ip + + req = _FakeRequest( + "127.0.0.1", {"x-forwarded-for": "198.51.100.7:50001, 10.0.0.1"} + ) + assert _client_ip(req) == "198.51.100.7" + + def test_xff_strips_bracketed_ipv6_port(self, env_trust_proxy): + from routes.auth import _client_ip + + req = _FakeRequest( + "127.0.0.1", {"x-forwarded-for": "[2001:db8::1]:50001, 10.0.0.1"} + ) + assert _client_ip(req) == "2001:db8::1" + + def test_forwarded_strips_ipv4_port(self, env_trust_proxy): + from routes.auth import _client_ip + + req = _FakeRequest( + "127.0.0.1", {"forwarded": 'for="198.51.100.7:50001";proto=https'} + ) + assert _client_ip(req) == "198.51.100.7" + + def test_forwarded_strips_bracketed_ipv6_port(self, env_trust_proxy): + from routes.auth import _client_ip + + req = _FakeRequest( + "127.0.0.1", {"forwarded": 'for="[2001:db8::1]:50001";proto=https'} + ) + assert _client_ip(req) == "2001:db8::1" + + def test_forwarded_isolates_first_element(self, env_trust_proxy): + from routes.auth import _client_ip + + # Multi-element Forwarded must pick the first element only, + # otherwise suffix variations create attacker-controlled buckets. + req = _FakeRequest( + "127.0.0.1", + {"forwarded": "for=198.51.100.42, for=10.0.0.1;proto=https"}, + ) + assert _client_ip(req) == "198.51.100.42" + + def test_xff_invalid_ip_falls_back_to_client_host(self, env_trust_proxy): + from routes.auth import _client_ip + + # A garbage XFF must not propagate into the bucket key. + req = _FakeRequest("127.0.0.1", {"x-forwarded-for": "not-an-ip"}) + assert _client_ip(req) == "127.0.0.1" + + +# ---------- bucket compose / blocking ---------- + + +class TestBucketKeyAndBlocking: + def test_record_per_user_isolates_other_users(self, env_no_proxy): + from routes.auth import ( + _bucket_key, + _record_login_failure, + _login_blocked, + _LOGIN_MAX_FAILS, + ) + + req = _FakeRequest("203.0.113.1") + for _ in range(_LOGIN_MAX_FAILS): + _record_login_failure(_bucket_key(req, "alice")) + assert _login_blocked(_bucket_key(req, "alice")) > 0 + # bob's account from the same IP is unaffected by alice's typos. + assert _login_blocked(_bucket_key(req, "bob")) == 0 + + def test_record_per_ip_isolates_other_ips(self, env_no_proxy): + from routes.auth import ( + _bucket_key, + _record_login_failure, + _login_blocked, + _LOGIN_MAX_FAILS, + ) + + req_a = _FakeRequest("203.0.113.1") + req_b = _FakeRequest("203.0.113.2") + for _ in range(_LOGIN_MAX_FAILS): + _record_login_failure(_bucket_key(req_a, "alice")) + assert _login_blocked(_bucket_key(req_a, "alice")) > 0 + # Same username, different IP, not blocked. + assert _login_blocked(_bucket_key(req_b, "alice")) == 0 + + def test_username_lowercased_in_key(self, env_no_proxy): + from routes.auth import _bucket_key + + req = _FakeRequest("203.0.113.1") + assert _bucket_key(req, "Alice") == _bucket_key(req, "alice") + assert _bucket_key(req, "ALICE") == _bucket_key(req, "alice") + + def test_rotating_usernames_hit_ip_aggregate_cap(self, env_no_proxy, monkeypatch): + """Spraying nonexistent usernames from one IP must still be throttled.""" + from routes import auth as auth_routes + + monkeypatch.setattr(auth_routes, "_LOGIN_IP_MAX_FAILS", 5) + req = _FakeRequest("203.0.113.10") + for idx in range(5): + auth_routes._record_login_failure(auth_routes._unknown_user_key(req)) + # Different "username" each attempt would not have throttled + # under per-(ip,username) only; the IP aggregate must. + # The next missing-user attempt is blocked. + assert auth_routes._login_blocked(auth_routes._unknown_user_key(req)) > 0 + + def test_unknown_user_bucket_is_single_sentinel(self, env_no_proxy): + """Random unknown usernames from one IP collapse to one bucket.""" + from routes import auth as auth_routes + + req = _FakeRequest("203.0.113.11") + unknown_key = auth_routes._unknown_user_key(req) + for _ in range(20): + auth_routes._record_login_failure(unknown_key) + # Account bucket cardinality stays at exactly one sentinel entry + # for this IP regardless of how many distinct usernames sprayed. + ip_keys = [k for k in auth_routes._LOGIN_BUCKETS if k[0] == "203.0.113.11"] + assert len(ip_keys) == 1 + assert ip_keys[0][1].startswith("\x00") + + def test_account_bucket_cap_bounded(self, env_no_proxy, monkeypatch): + """The per-account bucket dict cannot grow without bound.""" + from routes import auth as auth_routes + + monkeypatch.setattr(auth_routes, "_LOGIN_MAX_BUCKETS", 10) + req = _FakeRequest("203.0.113.12") + for idx in range(50): + auth_routes._record_login_failure((req.client.host, f"user-{idx}")) + # Hard cap respected; further keys do not allocate. + assert len(auth_routes._LOGIN_BUCKETS) <= 10 + + +# ---------- /login 429 body ---------- + + +class TestLogin429Body: + @pytest.fixture + def login_client(self, tmp_path, monkeypatch): + from auth import storage + from fastapi import FastAPI + from fastapi.testclient import TestClient + from routes.auth import router as auth_router + import secrets as _secrets + + 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) + storage.create_initial_user( + username = storage.DEFAULT_ADMIN_USERNAME, + password = "human-password-123", + jwt_secret = _secrets.token_urlsafe(64), + must_change_password = False, + ) + + app = FastAPI() + app.include_router(auth_router, prefix = "/api/auth") + return TestClient(app) + + def test_429_detail_does_not_leak_ip(self, env_no_proxy, login_client): + from routes.auth import _LOGIN_MAX_FAILS + + # Drive 6 failures from the same client IP / username. + for _ in range(_LOGIN_MAX_FAILS): + r = login_client.post( + "/api/auth/login", + json = {"username": "unsloth", "password": "wrong"}, + ) + assert r.status_code == 401 + r = login_client.post( + "/api/auth/login", + json = {"username": "unsloth", "password": "wrong"}, + ) + assert r.status_code == 429 + detail = r.json()["detail"] + # The 429 body must not interpolate the source IP. + assert "127.0.0.1" not in detail + assert "Too many" in detail + # Retry-After header is still set for clients. + assert "Retry-After" in r.headers diff --git a/studio/backend/tests/test_middleware.py b/studio/backend/tests/test_middleware.py index 4e396db9c5..bbaf20298d 100644 --- a/studio/backend/tests/test_middleware.py +++ b/studio/backend/tests/test_middleware.py @@ -196,6 +196,28 @@ class TestSecurityHeadersMiddleware: nonced = main_module._build_csp("XYZ") assert "script-src 'self' 'nonce-XYZ';" in nonced + def test_img_src_allows_google_favicons(self, main_module): + # sources.tsx fetches https://www.google.com/s2/favicons?... ; without + # this allowlist entry citation favicons fall back to gray initials. + csp = main_module._build_csp() + img_directive = next( + chunk.strip() + for chunk in csp.split(";") + if chunk.strip().startswith("img-src ") + ) + # Tokenise and compare with `==` so CodeQL's URL-substring rule does + # not read directive-string `in` membership as URL sanitisation. + img_sources = img_directive.split() + assert any(src == "https://www.google.com" for src in img_sources) + # Pre-existing favicon CDNs stay allowed. + for host in ( + "https://t0.gstatic.com", + "https://t1.gstatic.com", + "https://t2.gstatic.com", + "https://t3.gstatic.com", + ): + assert any(src == host for src in img_sources) + # ===================================================================== # /api/health auth gate