Skip to content

Commit b10a5ec

Browse files
authored
Merge pull request #1 from syzayd/night-shift/2026-07-09
night-shift: gateway token auth + Origin rejection + tests
2 parents c88ef72 + 319b33c commit b10a5ec

3 files changed

Lines changed: 148 additions & 1 deletion

File tree

src/personal_llm/config.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ class Settings(BaseSettings):
2323
personal_llm_chroma_dir: str = "./data/chroma"
2424
personal_llm_workspace_dir: str = "./data/workspace"
2525
personal_llm_voice_dir: str = "./data/voice"
26+
personal_llm_gateway_token_path: str = "./data/gateway_token"
2627

2728
retrieval_top_k: int = 8
2829
retrieval_min_similarity: float = 0.25

src/personal_llm/interfaces/api.py

Lines changed: 43 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,10 +2,12 @@
22

33
from __future__ import annotations
44

5+
import secrets
56
import tempfile
67
from pathlib import Path
78

8-
from fastapi import FastAPI, File, HTTPException, UploadFile
9+
from fastapi import FastAPI, File, HTTPException, Request, UploadFile
10+
from fastapi.responses import JSONResponse
911
from pydantic import BaseModel
1012

1113
from personal_llm import __version__
@@ -24,6 +26,46 @@
2426

2527
app = FastAPI(title="Personal LLM", version=__version__)
2628

29+
GATEWAY_TOKEN_HEADER = "x-dreamos-token"
30+
_gateway_token: str | None = None
31+
32+
33+
def _load_or_create_gateway_token() -> str:
34+
"""Shared secret DreamOS (or any other local caller) must send back on every request.
35+
36+
Generated once and persisted to disk so it survives a gateway restart without the
37+
caller needing to re-pair.
38+
"""
39+
global _gateway_token
40+
if _gateway_token is not None:
41+
return _gateway_token
42+
path = Path(get_settings().personal_llm_gateway_token_path)
43+
existing = path.read_text().strip() if path.exists() else ""
44+
if existing:
45+
_gateway_token = existing
46+
else:
47+
_gateway_token = secrets.token_hex(32)
48+
path.parent.mkdir(parents=True, exist_ok=True)
49+
path.write_text(_gateway_token)
50+
return _gateway_token
51+
52+
53+
@app.middleware("http")
54+
async def gateway_auth(request: Request, call_next):
55+
"""CSRF hardening (MASTER-FIX-PLAN.md S3 / Phase 3 item 12).
56+
57+
Multipart and form-encoded POSTs are reachable cross-origin without a CORS
58+
preflight, so a browser page could otherwise submit them straight to this
59+
gateway. Any request carrying an Origin header - which every real browser
60+
request does, and no local non-browser caller does - is rejected outright,
61+
and every request must also present the shared token.
62+
"""
63+
if "origin" in request.headers:
64+
return JSONResponse(status_code=403, content={"detail": "Cross-origin requests are not allowed."})
65+
if request.headers.get(GATEWAY_TOKEN_HEADER) != _load_or_create_gateway_token():
66+
return JSONResponse(status_code=401, content={"detail": "Missing or invalid X-DreamOS-Token."})
67+
return await call_next(request)
68+
2769

2870
async def _save_upload_to_temp(upload: UploadFile) -> str:
2971
suffix = Path(upload.filename or "upload").suffix

tests/test_api.py

Lines changed: 104 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,104 @@
1+
"""Gateway auth: shared token + Origin rejection (MASTER-FIX-PLAN.md Phase 3 item 12).
2+
3+
FakeRouter/store/vectors fixtures (conftest.py) keep this offline and keyless -
4+
no network call, no API key, ever.
5+
"""
6+
7+
from __future__ import annotations
8+
9+
from fastapi.testclient import TestClient
10+
11+
from personal_llm.engine import Engine
12+
from personal_llm.interfaces import api
13+
from personal_llm.voice import SpeechToText, TextToSpeech
14+
15+
TOKEN = "test-token"
16+
17+
18+
def _fake_engine(store, vectors, router) -> Engine:
19+
return Engine(store=store, vectors=vectors, router=router, stt=SpeechToText(), tts=TextToSpeech())
20+
21+
22+
def _client(monkeypatch, store, vectors, router):
23+
monkeypatch.setattr(api, "build_engine", lambda: _fake_engine(store, vectors, router))
24+
monkeypatch.setattr(api, "_load_or_create_gateway_token", lambda: TOKEN)
25+
return TestClient(api.app)
26+
27+
28+
def test_tokenless_request_is_rejected(monkeypatch, store, vectors, router):
29+
client = _client(monkeypatch, store, vectors, router)
30+
31+
resp = client.get("/stats")
32+
33+
assert resp.status_code == 401
34+
35+
36+
def test_wrong_token_is_rejected(monkeypatch, store, vectors, router):
37+
client = _client(monkeypatch, store, vectors, router)
38+
39+
resp = client.get("/stats", headers={api.GATEWAY_TOKEN_HEADER: "not-the-token"})
40+
41+
assert resp.status_code == 401
42+
43+
44+
def test_correct_token_allows_stats(monkeypatch, store, vectors, router):
45+
client = _client(monkeypatch, store, vectors, router)
46+
47+
resp = client.get("/stats", headers={api.GATEWAY_TOKEN_HEADER: TOKEN})
48+
49+
assert resp.status_code == 200
50+
assert resp.json() == store.stats()
51+
52+
53+
def test_origin_header_is_rejected_even_with_a_valid_token(monkeypatch, store, vectors, router):
54+
client = _client(monkeypatch, store, vectors, router)
55+
56+
resp = client.get(
57+
"/stats",
58+
headers={api.GATEWAY_TOKEN_HEADER: TOKEN, "Origin": "https://evil.example"},
59+
)
60+
61+
assert resp.status_code == 403
62+
63+
64+
def test_ask_endpoint_requires_token(monkeypatch, store, vectors, router):
65+
client = _client(monkeypatch, store, vectors, router)
66+
67+
resp = client.post("/ask", json={"question": "anything"})
68+
69+
assert resp.status_code == 401
70+
71+
72+
def test_ask_endpoint_round_trips_with_a_valid_token(monkeypatch, store, vectors, router):
73+
client = _client(monkeypatch, store, vectors, router)
74+
75+
resp = client.post(
76+
"/ask",
77+
json={"question": "anything"},
78+
headers={api.GATEWAY_TOKEN_HEADER: TOKEN},
79+
)
80+
81+
assert resp.status_code == 200
82+
assert "text" in resp.json()
83+
84+
85+
def test_voice_ask_endpoint_requires_token(monkeypatch, store, vectors, router):
86+
client = _client(monkeypatch, store, vectors, router)
87+
88+
resp = client.post("/voice/ask", files={"file": ("clip.wav", b"not real audio", "audio/wav")})
89+
90+
assert resp.status_code == 401
91+
92+
93+
def test_load_or_create_gateway_token_persists_across_calls(tmp_path, monkeypatch):
94+
token_path = tmp_path / "gateway_token"
95+
monkeypatch.setattr(api, "_gateway_token", None)
96+
monkeypatch.setattr(api, "get_settings", lambda: type("S", (), {"personal_llm_gateway_token_path": str(token_path)})())
97+
98+
first = api._load_or_create_gateway_token()
99+
100+
monkeypatch.setattr(api, "_gateway_token", None)
101+
second = api._load_or_create_gateway_token()
102+
103+
assert first == second
104+
assert token_path.read_text().strip() == first

0 commit comments

Comments
 (0)