diff --git a/.env.example b/.env.example index d4a09bf..10b1c64 100644 --- a/.env.example +++ b/.env.example @@ -33,6 +33,12 @@ SNOWFLAKE_WAREHOUSE= OLLAMA_HOST=http://ollama:11434 LLM_MODEL=llama2 +# --------------------------------------------------------------------------- +# Inference API authentication +# --------------------------------------------------------------------------- +# Required to enable POST /infer. Use a unique local secret; do not commit it. +API_BEARER_TOKEN= + # --------------------------------------------------------------------------- # Internal service URLs # --------------------------------------------------------------------------- diff --git a/README.md b/README.md index ffaa454..1615ae8 100644 --- a/README.md +++ b/README.md @@ -151,6 +151,7 @@ All configuration is via environment variables. Copy `.env.example` to `.env` a | `POSTGRES_DB` | `sentinel` | Postgres database | | `OLLAMA_HOST` | `http://ollama:11434` | Ollama endpoint (optional) | | `LLM_MODEL` | `llama2` | LLM model name | +| `API_BEARER_TOKEN` | *(unset)* | Required shared bearer token for `POST /infer`; the endpoint returns 503 until configured | **Snowflake (optional):** set `WAREHOUSE_MODE=snowflake` and fill in `SNOWFLAKE_ACCOUNT`, `SNOWFLAKE_USER`, `SNOWFLAKE_PASSWORD`, `SNOWFLAKE_DATABASE`, `SNOWFLAKE_SCHEMA`, `SNOWFLAKE_WAREHOUSE`. diff --git a/api/main.py b/api/main.py index c13dcb4..6a45b71 100644 --- a/api/main.py +++ b/api/main.py @@ -1,3 +1,6 @@ +import hmac +import os + from fastapi import Depends, FastAPI, HTTPException, Header, Request from slowapi import Limiter, _rate_limit_exceeded_handler from slowapi.util import get_remote_address @@ -22,10 +25,21 @@ class PromptRequest(BaseModel): def _require_bearer(authorization: Optional[str] = Header(default=None)) -> str: - """Dependency that validates a Bearer token is present.""" + """Require a configured bearer token for the public inference stub.""" + + expected_token = os.getenv("API_BEARER_TOKEN") + if not expected_token: + raise HTTPException( + status_code=503, + detail="Inference authentication is not configured.", + ) if not authorization or not authorization.startswith("Bearer "): raise HTTPException(status_code=401, detail="Missing or invalid Authorization header") - return authorization.removeprefix("Bearer ").strip() + + token = authorization.removeprefix("Bearer ").strip() + if not hmac.compare_digest(token, expected_token): + raise HTTPException(status_code=401, detail="Missing or invalid Authorization header") + return token @app.get("/") diff --git a/tests/test_auth.py b/tests/test_auth.py index e8b40f4..2850bbe 100644 --- a/tests/test_auth.py +++ b/tests/test_auth.py @@ -1,3 +1,17 @@ -def test_missing_api_key(client): +def test_infer_fails_closed_when_bearer_token_is_unconfigured(client, monkeypatch): + monkeypatch.delenv("API_BEARER_TOKEN", raising=False) + response = client.post("/infer", json={"prompt": "Hello"}) - assert response.status_code == 401 + + assert response.status_code == 503 + + +def test_infer_rejects_missing_and_invalid_bearer_tokens(client, monkeypatch): + monkeypatch.setenv("API_BEARER_TOKEN", "test-token") + + assert client.post("/infer", json={"prompt": "Hello"}).status_code == 401 + assert client.post( + "/infer", + json={"prompt": "Hello"}, + headers={"Authorization": "Bearer wrong-token"}, + ).status_code == 401 diff --git a/tests/test_inference.py b/tests/test_inference.py index 65cd0b2..85ca306 100644 --- a/tests/test_inference.py +++ b/tests/test_inference.py @@ -1,5 +1,6 @@ -def test_inference_success(client): - headers = {"Authorization": "Bearer test-key"} +def test_inference_success(client, monkeypatch): + monkeypatch.setenv("API_BEARER_TOKEN", "test-token") + headers = {"Authorization": "Bearer test-token"} payload = {"prompt": "Explain AI in one sentence."} response = client.post("/infer", json=payload, headers=headers) diff --git a/tests/test_rate_limit.py b/tests/test_rate_limit.py index c5723da..d329cfa 100644 --- a/tests/test_rate_limit.py +++ b/tests/test_rate_limit.py @@ -1,5 +1,6 @@ -def test_rate_limit(client): - headers = {"Authorization": "Bearer test-key"} +def test_rate_limit(client, monkeypatch): + monkeypatch.setenv("API_BEARER_TOKEN", "test-token") + headers = {"Authorization": "Bearer test-token"} payload = {"prompt": "rate limit test"} for _ in range(5):