diff --git a/backend/models/__init__.py b/backend/models/__init__.py index cd89bdd7..eec143f4 100644 --- a/backend/models/__init__.py +++ b/backend/models/__init__.py @@ -1,4 +1,4 @@ -from typing import Optional, Union, List +from typing import Optional, Union, List, Literal from pydantic import BaseModel, Field @@ -10,6 +10,7 @@ class StartSessionBody(BaseModel): mode: str = "socratic" use_shared_context: bool = True course_id: Optional[str] = None # Direct course_id lookup instead of resolving from topic + model_pref: Optional[Literal["fast", "smart"]] = None # "fast" (default, gemini-2.5-flash) or "smart" (gemini-2.5-pro) class ChatBody(BaseModel): @@ -18,6 +19,7 @@ class ChatBody(BaseModel): message: str mode: str = "socratic" use_shared_context: bool = True + model_pref: Optional[Literal["fast", "smart"]] = None # "fast" (default, gemini-2.5-flash) or "smart" (gemini-2.5-pro) class EndSessionBody(BaseModel): @@ -31,6 +33,7 @@ class ActionBody(BaseModel): action_type: str = "hint" mode: str = "socratic" use_shared_context: bool = True + model_pref: Optional[Literal["fast", "smart"]] = None # "fast" (default, gemini-2.5-flash) or "smart" (gemini-2.5-pro) # ── Quiz ────────────────────────────────────────────────────────────────────── diff --git a/backend/routes/auth.py b/backend/routes/auth.py index aa95a406..f9a83c95 100644 --- a/backend/routes/auth.py +++ b/backend/routes/auth.py @@ -9,6 +9,7 @@ import base64 import hashlib import hmac as _hmac +import re import secrets import time as _time from urllib.parse import urlencode @@ -39,6 +40,14 @@ router = APIRouter() +OAUTH_STATE_COOKIE = "sapling_oauth_state" +_OAUTH_COOKIE_MAX_AGE = 600 +_POPUP_ID_RE = re.compile(r"^[A-Za-z0-9_-]{1,128}$") + +# Fallback in-memory store for environments without SESSION_SECRET; entries +# are keyed by nonce and expire after _OAUTH_COOKIE_MAX_AGE seconds. +_OAUTH_FALLBACK_STORE: dict[str, tuple[float, dict]] = {} + def _google_client_config() -> dict: return { @@ -73,6 +82,68 @@ def _generate_pkce_pair(): return code_verifier, code_challenge +def _clean_popup_id(s: str | None) -> str | None: + if not s: + return None + return s if _POPUP_ID_RE.match(s) else None + + +def _encode_oauth_cookie(payload: dict) -> str: + raw = json.dumps(payload, separators=(",", ":")).encode() + payload_b64 = base64.urlsafe_b64encode(raw).decode().rstrip("=") + if SESSION_SECRET: + sig_bytes = _hmac.new(SESSION_SECRET.encode(), payload_b64.encode(), hashlib.sha256).digest() + sig_b64 = base64.urlsafe_b64encode(sig_bytes).decode().rstrip("=") + return f"{payload_b64}.{sig_b64}" + nonce = payload.get("n", "") + if nonce: + _OAUTH_FALLBACK_STORE[nonce] = (_time.monotonic() + _OAUTH_COOKIE_MAX_AGE, payload) + _prune_fallback_store() + return payload_b64 + + +def _decode_oauth_cookie(cookie_value: str | None) -> dict | None: + if not cookie_value: + return None + if SESSION_SECRET: + if "." not in cookie_value: + return None + try: + payload_b64, sig_b64 = cookie_value.rsplit(".", 1) + except ValueError: + return None + expected = _hmac.new(SESSION_SECRET.encode(), payload_b64.encode(), hashlib.sha256).digest() + expected_b64 = base64.urlsafe_b64encode(expected).decode().rstrip("=") + if not _hmac.compare_digest(expected_b64, sig_b64): + return None + try: + padded = payload_b64 + "=" * (-len(payload_b64) % 4) + payload = json.loads(base64.urlsafe_b64decode(padded.encode()).decode()) + except Exception: + return None + return payload if isinstance(payload, dict) else None + try: + padded = cookie_value + "=" * (-len(cookie_value) % 4) + payload = json.loads(base64.urlsafe_b64decode(padded.encode()).decode()) + except Exception: + return None + if not isinstance(payload, dict): + return None + nonce = payload.get("n") + _prune_fallback_store() + entry = _OAUTH_FALLBACK_STORE.get(nonce or "") + if not entry: + return None + return entry[1] + + +def _prune_fallback_store() -> None: + now = _time.monotonic() + expired = [k for k, (exp, _) in _OAUTH_FALLBACK_STORE.items() if exp < now] + for k in expired: + _OAUTH_FALLBACK_STORE.pop(k, None) + + @router.get("/me") def get_me(request: Request): """Return approval and onboarding status for a given user_id.""" @@ -136,36 +207,84 @@ def get_me(request: Request): @router.get("/google") -def google_login(): +def google_login(popup_id: str = Query(None)): """Redirect to Google consent screen with identity + calendar scopes.""" if not GOOGLE_AVAILABLE or not GOOGLE_CLIENT_ID: raise HTTPException(status_code=400, detail="Google OAuth not configured") code_verifier, code_challenge = _generate_pkce_pair() + nonce = secrets.token_urlsafe(32) + clean_popup = _clean_popup_id(popup_id) + flow = Flow.from_client_config(_google_client_config(), scopes=AUTH_SCOPES) flow.redirect_uri = GOOGLE_AUTH_REDIRECT_URI auth_url, _ = flow.authorization_url( prompt="consent", access_type="offline", - state=_encode_state({"action": "signin", "cv": code_verifier}), + state=_encode_state({"action": "signin", "n": nonce}), code_challenge=code_challenge, code_challenge_method="S256", ) - return RedirectResponse(auth_url) + + cookie_value = _encode_oauth_cookie({ + "n": nonce, + "cv": code_verifier, + "popup_id": clean_popup, + }) + response = RedirectResponse(auth_url) + response.set_cookie( + key=OAUTH_STATE_COOKIE, + value=cookie_value, + max_age=_OAUTH_COOKIE_MAX_AGE, + httponly=True, + secure=True, + samesite="lax", + path="/", + ) + return response @router.get("/google/callback") -def google_callback(code: str = Query(...), state: str = Query(None)): +def google_callback(request: Request, code: str = Query(...), state: str = Query(None)): """Exchange auth code for tokens, validate @bu.edu, upsert user.""" + cookie_payload = _decode_oauth_cookie(request.cookies.get(OAUTH_STATE_COOKIE)) + code_verifier = cookie_payload.get("cv") if cookie_payload else None + popup_id = _clean_popup_id(cookie_payload.get("popup_id")) if cookie_payload else None + cookie_nonce = cookie_payload.get("n") if cookie_payload else None + + def _fail_redirect(error_code: str, fallback_path: str = "/auth") -> RedirectResponse: + # In popup mode, route failures through /auth/callback so the popup + # can broadcast the error and self-close instead of stranding the opener. + if popup_id: + params = urlencode({"error": error_code, "popup_id": popup_id}) + resp = RedirectResponse(f"{FRONTEND_URL}/auth/callback?{params}") + else: + resp = RedirectResponse(f"{FRONTEND_URL}{fallback_path}?error={error_code}") + resp.set_cookie( + key=OAUTH_STATE_COOKIE, + value="", + max_age=0, + httponly=True, + secure=True, + samesite="lax", + path="/", + ) + return resp + if not GOOGLE_AVAILABLE: - return RedirectResponse(f"{FRONTEND_URL}/auth?error=google_not_configured") + return _fail_redirect("google_not_configured") state_data = _decode_state(state) if state else {} - code_verifier = state_data.get("cv") + state_nonce = state_data.get("n") + if not cookie_payload or not cookie_nonce or not state_nonce or not _hmac.compare_digest(str(state_nonce), str(cookie_nonce)): + return _fail_redirect("invalid_state") flow = Flow.from_client_config(_google_client_config(), scopes=AUTH_SCOPES) flow.redirect_uri = GOOGLE_AUTH_REDIRECT_URI - flow.fetch_token(code=code, code_verifier=code_verifier) + try: + flow.fetch_token(code=code, code_verifier=code_verifier) + except Exception: + return _fail_redirect("oauth_exchange_failed") creds = flow.credentials # Fetch user info from Google @@ -184,9 +303,7 @@ def google_callback(code: str = Query(...), state: str = Query(None)): # Restrict to @bu.edu accounts if not email.endswith("@bu.edu"): - return RedirectResponse( - f"{FRONTEND_URL}/auth?error=invalid_domain" - ) + return _fail_redirect("invalid_domain") # Determine user_id: check if this Google ID already exists existing = table("users").select("id,is_approved", filters={"google_id": f"eq.{google_id}"}) @@ -233,7 +350,19 @@ def google_callback(code: str = Query(...), state: str = Query(None)): ) if not is_approved: - return RedirectResponse(f"{FRONTEND_URL}/pending") + if popup_id: + return _fail_redirect("not_approved") + resp = RedirectResponse(f"{FRONTEND_URL}/pending") + resp.set_cookie( + key=OAUTH_STATE_COOKIE, + value="", + max_age=0, + httponly=True, + secure=True, + samesite="lax", + path="/", + ) + return resp # Build a short-lived HMAC token so the frontend can verify this redirect # without a second round-trip to the backend. @@ -250,5 +379,16 @@ def google_callback(code: str = Query(...), state: str = Query(None)): "avatar": avatar_url, "is_approved": "true", **({"auth_token": auth_token} if auth_token else {}), + **({"popup_id": popup_id} if popup_id else {}), }) - return RedirectResponse(f"{FRONTEND_URL}/auth/callback?{params}") + resp = RedirectResponse(f"{FRONTEND_URL}/auth/callback?{params}") + resp.set_cookie( + key=OAUTH_STATE_COOKIE, + value="", + max_age=0, + httponly=True, + secure=True, + samesite="lax", + path="/", + ) + return resp diff --git a/backend/routes/learn.py b/backend/routes/learn.py index 37df0b3a..c289969c 100644 --- a/backend/routes/learn.py +++ b/backend/routes/learn.py @@ -11,7 +11,12 @@ from models import StartSessionBody, ChatBody, EndSessionBody, ActionBody, ModeSwitchBody from services.auth_guard import require_self, get_session_user_id from services.encryption import encrypt_if_present, encrypt_json, decrypt_if_present, decrypt_json -from services.gemini_service import MODEL_SMART, call_gemini_multiturn, extract_graph_update +from services.gemini_service import ( + MODEL_DEFAULT, + MODEL_SMART, + call_gemini_multiturn, + extract_graph_update, +) from services.graph_service import get_graph, apply_graph_update router = APIRouter() @@ -28,6 +33,18 @@ "teachback": "Teach-back (you explain to me)", } +# User-facing speed/quality knob for the tutor chat. +# "fast" = flash (default, faster), "smart" = pro (opt-in, slower but stronger reasoning). +# Anything unrecognized falls back to fast so the default is the snappy one. +_MODEL_PREF_TO_MODEL = { + "fast": MODEL_DEFAULT, + "smart": MODEL_SMART, +} + + +def _resolve_tutor_model(model_pref: str | None) -> str: + return _MODEL_PREF_TO_MODEL.get(model_pref or "", MODEL_DEFAULT) + def _load_prompt(name: str) -> str: with open(os.path.join(PROMPTS_DIR, name)) as f: @@ -293,13 +310,15 @@ def start_session(body: StartSessionBody, request: Request): ) try: - raw = call_gemini_multiturn(system_prompt, [], user_message, model=MODEL_SMART) + raw = call_gemini_multiturn( + system_prompt, [], user_message, model=_resolve_tutor_model(body.model_pref) + ) except Exception as e: raise HTTPException(status_code=502, detail=f"Gemini error: {e}") reply, graph_update = extract_graph_update(raw) apply_graph_update(body.user_id, graph_update, course_id=course_id) - + PENDING_SESSIONS[session_id] = { "user_id": body.user_id, "mode": body.mode, @@ -339,7 +358,9 @@ def chat(body: ChatBody, request: Request): ) try: - raw = call_gemini_multiturn(system_prompt, history, body.message, model=MODEL_SMART) + raw = call_gemini_multiturn( + system_prompt, history, body.message, model=_resolve_tutor_model(body.model_pref) + ) except Exception as e: raise HTTPException(status_code=502, detail=f"Gemini error: {e}") @@ -558,7 +579,9 @@ def action(body: ActionBody, request: Request): action_message = f"[ACTION: {action_prompts.get(body.action_type, '')}]" try: - raw = call_gemini_multiturn(system_prompt, history, action_message, model=MODEL_SMART) + raw = call_gemini_multiturn( + system_prompt, history, action_message, model=_resolve_tutor_model(body.model_pref) + ) except Exception as e: raise HTTPException(status_code=502, detail=f"Gemini error: {e}") diff --git a/backend/tests/test_auth_state.py b/backend/tests/test_auth_state.py new file mode 100644 index 00000000..bdc876c2 --- /dev/null +++ b/backend/tests/test_auth_state.py @@ -0,0 +1,207 @@ +""" +Unit tests for routes/auth.py OAuth state hardening. + +Covers: +- HMAC cookie round-trip (encode -> decode) +- Tampered cookie rejection (MAC and payload) +- _clean_popup_id charset/length validation +- _decode_state happy path +- Callback rejects when state nonce doesn't match cookie nonce +- Callback handles missing/malformed cookie via popup-aware error redirect +""" +import json +import base64 + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient + +import routes.auth as auth_module + +# Build a minimal app that mounts ONLY the auth router so these tests don't +# pull in main.py (which imports logfire and the full router stack). +_app = FastAPI() +_app.include_router(auth_module.router, prefix="/api/auth") +client = TestClient(_app) + + +# ── HMAC cookie round-trip ──────────────────────────────────────────────────── + + +class TestOAuthCookieRoundTrip: + def test_round_trip_returns_same_payload(self, monkeypatch): + monkeypatch.setattr(auth_module, "SESSION_SECRET", "test-secret-key") + payload = {"n": "abc123", "cv": "verifier", "popup_id": "popup-1"} + cookie = auth_module._encode_oauth_cookie(payload) + decoded = auth_module._decode_oauth_cookie(cookie) + assert decoded == payload + + def test_tampered_mac_rejected(self, monkeypatch): + monkeypatch.setattr(auth_module, "SESSION_SECRET", "test-secret-key") + cookie = auth_module._encode_oauth_cookie( + {"n": "abc123", "cv": "verifier", "popup_id": None} + ) + payload_b64, sig_b64 = cookie.rsplit(".", 1) + # flip a character in the signature + bad_sig = ("A" if sig_b64[0] != "A" else "B") + sig_b64[1:] + tampered = f"{payload_b64}.{bad_sig}" + assert auth_module._decode_oauth_cookie(tampered) is None + + def test_tampered_payload_rejected(self, monkeypatch): + monkeypatch.setattr(auth_module, "SESSION_SECRET", "test-secret-key") + cookie = auth_module._encode_oauth_cookie( + {"n": "abc123", "cv": "verifier", "popup_id": None} + ) + _, sig_b64 = cookie.rsplit(".", 1) + bogus_payload = base64.urlsafe_b64encode( + json.dumps({"n": "evil", "cv": "x", "popup_id": None}).encode() + ).decode().rstrip("=") + tampered = f"{bogus_payload}.{sig_b64}" + assert auth_module._decode_oauth_cookie(tampered) is None + + def test_missing_cookie_returns_none(self, monkeypatch): + monkeypatch.setattr(auth_module, "SESSION_SECRET", "test-secret-key") + assert auth_module._decode_oauth_cookie(None) is None + assert auth_module._decode_oauth_cookie("") is None + + def test_malformed_cookie_returns_none(self, monkeypatch): + monkeypatch.setattr(auth_module, "SESSION_SECRET", "test-secret-key") + assert auth_module._decode_oauth_cookie("not-a-valid-cookie") is None + assert auth_module._decode_oauth_cookie("only_payload_no_dot") is None + + def test_fallback_in_memory_store_round_trip(self, monkeypatch): + # Without SESSION_SECRET, encode stashes the payload in an in-memory + # dict keyed by nonce. Decode retrieves it. + monkeypatch.setattr(auth_module, "SESSION_SECRET", "") + auth_module._OAUTH_FALLBACK_STORE.clear() + payload = {"n": "fallback-nonce", "cv": "verifier", "popup_id": "p"} + cookie = auth_module._encode_oauth_cookie(payload) + decoded = auth_module._decode_oauth_cookie(cookie) + assert decoded == payload + + def test_fallback_unknown_nonce_returns_none(self, monkeypatch): + monkeypatch.setattr(auth_module, "SESSION_SECRET", "") + auth_module._OAUTH_FALLBACK_STORE.clear() + # Build a cookie payload whose nonce was never registered + bogus = base64.urlsafe_b64encode( + json.dumps({"n": "ghost", "cv": "x", "popup_id": None}).encode() + ).decode().rstrip("=") + assert auth_module._decode_oauth_cookie(bogus) is None + + +# ── _clean_popup_id ─────────────────────────────────────────────────────────── + + +class TestCleanPopupId: + def test_valid_uuid_returned_as_is(self): + uid = "550e8400-e29b-41d4-a716-446655440000" + assert auth_module._clean_popup_id(uid) == uid + + def test_alphanumeric_with_underscore_ok(self): + assert auth_module._clean_popup_id("popup_abc_123") == "popup_abc_123" + + def test_slash_rejected(self): + assert auth_module._clean_popup_id("abc/def") is None + + def test_empty_string_rejected(self): + assert auth_module._clean_popup_id("") is None + + def test_none_returns_none(self): + assert auth_module._clean_popup_id(None) is None + + def test_overlong_rejected(self): + assert auth_module._clean_popup_id("a" * 129) is None + + def test_max_length_accepted(self): + s = "a" * 128 + assert auth_module._clean_popup_id(s) == s + + def test_special_chars_rejected(self): + assert auth_module._clean_popup_id("abc def") is None + assert auth_module._clean_popup_id("abc