Fix critical/high security and correctness issues from code review
Critical (all fixed): - bcrypt statt SHA-256 für Passwörter - API-Keys gehasht in DB, Plaintext nur einmalig zurückgegeben - DB-Session-Leak behoben (SessionLocal + try/finally, Depends(get_db)) - Admin-Check via is_admin-Spalte statt Hardcoded-Username - CORS: konfigurierbare Origins via ALLOWED_ORIGINS, kein Wildcard mit Credentials High (all fixed): - TOCTOU-Race: check_and_increment_quota mit SELECT FOR UPDATE atomar - Getrennte Tages-/Monatszähler in Usage + automatische Reset-Logik - Token-Zählung mit tiktoken (cl100k_base) statt .split() Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Oliver Hofmann
co-authored by
Claude Sonnet 4.6
Claude Sonnet 4.6
parent
562f6ecd9c
commit
bf694b79e2
@@ -0,0 +1 @@
|
||||
"""Tests for the Ollama Proxy."""
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,87 @@
|
||||
import pytest
|
||||
from fastapi.testclient import TestClient
|
||||
import tempfile
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
def create_test_db():
|
||||
"""Create a temporary SQLite database for tests."""
|
||||
temp_db = tempfile.NamedTemporaryFile(suffix='.db', delete=False)
|
||||
temp_db.close()
|
||||
os.environ["DATABASE_URL"] = f"sqlite:///{temp_db.name}"
|
||||
return temp_db.name
|
||||
|
||||
def cleanup_test_db(db_path):
|
||||
"""Remove the temporary database."""
|
||||
if os.path.exists(db_path):
|
||||
os.unlink(db_path)
|
||||
os.environ.pop("DATABASE_URL", None)
|
||||
|
||||
def setup_test_db():
|
||||
"""Setup test database with required data."""
|
||||
from database import Base, engine, SessionLocal
|
||||
from models import User, APIKey, Quota, Usage
|
||||
from crud import create_api_key, hash_password
|
||||
|
||||
# Create tables
|
||||
Base.metadata.create_all(bind=engine)
|
||||
|
||||
db = SessionLocal()
|
||||
|
||||
# Create test user
|
||||
test_user = User(
|
||||
username="testuser",
|
||||
email="test@example.com",
|
||||
hashed_password=hash_password("test123"),
|
||||
is_active=True
|
||||
)
|
||||
db.add(test_user)
|
||||
db.commit()
|
||||
db.refresh(test_user)
|
||||
|
||||
# Create API key for test user
|
||||
_, raw_key = create_api_key(db, test_user.id, "test-key")
|
||||
os.environ["TEST_API_KEY"] = raw_key
|
||||
|
||||
# Create admin user
|
||||
admin_user = User(
|
||||
username="admin",
|
||||
email="admin@example.com",
|
||||
hashed_password=hash_password("admin123"),
|
||||
is_active=True,
|
||||
is_admin=True,
|
||||
)
|
||||
db.add(admin_user)
|
||||
db.commit()
|
||||
db.refresh(admin_user)
|
||||
|
||||
# Create admin API key
|
||||
_, admin_raw_key = create_api_key(db, admin_user.id, "admin-key")
|
||||
os.environ["ADMIN_API_KEY"] = admin_raw_key
|
||||
|
||||
db.close()
|
||||
|
||||
return raw_key, admin_raw_key
|
||||
|
||||
def teardown_test_db():
|
||||
"""Clean up test database and environment."""
|
||||
from database import engine
|
||||
from models import Base
|
||||
Base.metadata.drop_all(bind=engine)
|
||||
|
||||
os.environ.pop("TEST_API_KEY", None)
|
||||
os.environ.pop("ADMIN_API_KEY", None)
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def test_client():
|
||||
"""Create test client with test database."""
|
||||
db_path = create_test_db()
|
||||
setup_test_db()
|
||||
|
||||
from main import app
|
||||
client = TestClient(app)
|
||||
|
||||
yield client
|
||||
|
||||
teardown_test_db()
|
||||
cleanup_test_db(db_path)
|
||||
@@ -0,0 +1,106 @@
|
||||
import pytest
|
||||
import os
|
||||
from unittest.mock import AsyncMock, patch
|
||||
from fastapi.testclient import TestClient
|
||||
from main import app
|
||||
from database import Base, engine, SessionLocal
|
||||
from models import User, APIKey, Quota
|
||||
from crud import create_api_key, hash_password
|
||||
|
||||
os.environ["OLLAMA_URL"] = "http://127.0.0.1:9999"
|
||||
|
||||
def setup_test_db():
|
||||
"""Setup test database with required data."""
|
||||
Base.metadata.drop_all(bind=engine)
|
||||
Base.metadata.create_all(bind=engine)
|
||||
|
||||
db = SessionLocal()
|
||||
|
||||
test_user = User(
|
||||
username="testuser",
|
||||
email="test@example.com",
|
||||
hashed_password=hash_password("test123"),
|
||||
is_active=True
|
||||
)
|
||||
db.add(test_user)
|
||||
db.commit()
|
||||
db.refresh(test_user)
|
||||
|
||||
quota = Quota(
|
||||
user_id=test_user.id,
|
||||
daily_tokens=1000000,
|
||||
monthly_tokens=10000000,
|
||||
daily_requests=1000,
|
||||
monthly_requests=10000
|
||||
)
|
||||
db.add(quota)
|
||||
db.commit()
|
||||
|
||||
api_key_record, raw_key = create_api_key(db, test_user.id, "test-key")
|
||||
os.environ["TEST_API_KEY"] = raw_key
|
||||
|
||||
admin_user = User(
|
||||
username="admin",
|
||||
email="admin@example.com",
|
||||
hashed_password=hash_password("admin123"),
|
||||
is_active=True,
|
||||
is_admin=True,
|
||||
)
|
||||
db.add(admin_user)
|
||||
db.commit()
|
||||
db.refresh(admin_user)
|
||||
|
||||
admin_quota = Quota(
|
||||
user_id=admin_user.id,
|
||||
daily_tokens=10000000,
|
||||
monthly_tokens=100000000,
|
||||
daily_requests=10000,
|
||||
monthly_requests=100000
|
||||
)
|
||||
db.add(admin_quota)
|
||||
db.commit()
|
||||
|
||||
_, admin_raw_key = create_api_key(db, admin_user.id, "admin-key")
|
||||
os.environ["ADMIN_API_KEY"] = admin_raw_key
|
||||
|
||||
db.close()
|
||||
|
||||
return os.environ["TEST_API_KEY"], os.environ["ADMIN_API_KEY"]
|
||||
|
||||
def teardown_test_db():
|
||||
"""Clean up test database and environment."""
|
||||
Base.metadata.drop_all(bind=engine)
|
||||
os.environ.pop("TEST_API_KEY", None)
|
||||
os.environ.pop("ADMIN_API_KEY", None)
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def test_client():
|
||||
setup_test_db()
|
||||
client = TestClient(app, raise_server_exceptions=False)
|
||||
yield client
|
||||
teardown_test_db()
|
||||
|
||||
def test_auth_middleware_missing_auth(test_client):
|
||||
response = test_client.post("/api/generate", json={"model": "llama3", "prompt": "test"})
|
||||
assert response.status_code == 401
|
||||
|
||||
def test_auth_middleware_invalid_key(test_client):
|
||||
response = test_client.post(
|
||||
"/api/generate",
|
||||
headers={"Authorization": "sk-invalid-key"},
|
||||
json={"model": "llama3", "prompt": "test"}
|
||||
)
|
||||
assert response.status_code == 401
|
||||
|
||||
@patch("main.proxy_request", new_callable=AsyncMock)
|
||||
def test_auth_middleware_valid_key(mock_proxy, test_client):
|
||||
mock_proxy.return_value.status_code = 200
|
||||
mock_proxy.return_value.json = lambda: {"response": "success"}
|
||||
mock_proxy.return_value.headers = {}
|
||||
|
||||
response = test_client.post(
|
||||
"/api/generate",
|
||||
headers={"Authorization": os.environ.get("TEST_API_KEY", "")},
|
||||
json={"model": "llama3", "prompt": "test"}
|
||||
)
|
||||
assert response.status_code == 200
|
||||
@@ -0,0 +1,182 @@
|
||||
import pytest
|
||||
import os
|
||||
from datetime import datetime, timedelta
|
||||
|
||||
os.environ.setdefault("OLLAMA_URL", "http://127.0.0.1:9999")
|
||||
|
||||
from database import Base, engine, SessionLocal
|
||||
from models import User, Quota, Usage
|
||||
from crud import check_and_increment_quota, count_tokens, hash_password
|
||||
|
||||
|
||||
def make_user_and_quota(db, daily_tokens=None, monthly_tokens=None,
|
||||
daily_requests=None, monthly_requests=None):
|
||||
user = User(
|
||||
username="quotauser",
|
||||
email="quota@example.com",
|
||||
hashed_password=hash_password("pass"),
|
||||
is_active=True,
|
||||
)
|
||||
db.add(user)
|
||||
db.commit()
|
||||
db.refresh(user)
|
||||
|
||||
quota = Quota(
|
||||
user_id=user.id,
|
||||
daily_tokens=daily_tokens,
|
||||
monthly_tokens=monthly_tokens,
|
||||
daily_requests=daily_requests,
|
||||
monthly_requests=monthly_requests,
|
||||
)
|
||||
db.add(quota)
|
||||
db.commit()
|
||||
|
||||
return user.id
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def db():
|
||||
Base.metadata.drop_all(bind=engine)
|
||||
Base.metadata.create_all(bind=engine)
|
||||
session = SessionLocal()
|
||||
yield session
|
||||
session.close()
|
||||
Base.metadata.drop_all(bind=engine)
|
||||
|
||||
|
||||
# --- count_tokens ---
|
||||
|
||||
def test_count_tokens_empty():
|
||||
assert count_tokens("") == 0
|
||||
|
||||
def test_count_tokens_returns_int():
|
||||
assert isinstance(count_tokens("hello world"), int)
|
||||
|
||||
def test_count_tokens_scales_with_length():
|
||||
assert count_tokens("a b c d e f g h") > count_tokens("a")
|
||||
|
||||
def test_count_tokens_more_accurate_than_split():
|
||||
# tiktoken counts "don't" as 2 tokens, split counts it as 1 word
|
||||
text = "don't do that"
|
||||
assert count_tokens(text) >= len(text.split())
|
||||
|
||||
|
||||
# --- check_and_increment_quota ---
|
||||
|
||||
def test_allowed_within_daily_token_limit(db):
|
||||
user_id = make_user_and_quota(db, daily_tokens=1000)
|
||||
assert check_and_increment_quota(db, user_id, tokens=100, requests=1) is True
|
||||
|
||||
def test_denied_when_daily_tokens_exceeded(db):
|
||||
user_id = make_user_and_quota(db, daily_tokens=50)
|
||||
assert check_and_increment_quota(db, user_id, tokens=100, requests=1) is False
|
||||
|
||||
def test_denied_when_monthly_tokens_exceeded(db):
|
||||
user_id = make_user_and_quota(db, monthly_tokens=50)
|
||||
assert check_and_increment_quota(db, user_id, tokens=100, requests=1) is False
|
||||
|
||||
def test_denied_when_daily_requests_exceeded(db):
|
||||
user_id = make_user_and_quota(db, daily_requests=1)
|
||||
check_and_increment_quota(db, user_id, tokens=0, requests=1)
|
||||
assert check_and_increment_quota(db, user_id, tokens=0, requests=1) is False
|
||||
|
||||
def test_denied_when_monthly_requests_exceeded(db):
|
||||
user_id = make_user_and_quota(db, monthly_requests=1)
|
||||
check_and_increment_quota(db, user_id, tokens=0, requests=1)
|
||||
assert check_and_increment_quota(db, user_id, tokens=0, requests=1) is False
|
||||
|
||||
def test_increments_both_daily_and_monthly_counters(db):
|
||||
user_id = make_user_and_quota(db, daily_tokens=1000, monthly_tokens=10000,
|
||||
daily_requests=100, monthly_requests=1000)
|
||||
check_and_increment_quota(db, user_id, tokens=50, requests=1)
|
||||
|
||||
usage = db.query(Usage).filter(Usage.user_id == user_id).first()
|
||||
assert usage.tokens_used_today == 50
|
||||
assert usage.tokens_used_month == 50
|
||||
assert usage.requests_today == 1
|
||||
assert usage.requests_month == 1
|
||||
|
||||
def test_creates_usage_record_on_first_call(db):
|
||||
user_id = make_user_and_quota(db, daily_tokens=1000)
|
||||
assert db.query(Usage).filter(Usage.user_id == user_id).first() is None
|
||||
|
||||
check_and_increment_quota(db, user_id, tokens=10, requests=1)
|
||||
|
||||
assert db.query(Usage).filter(Usage.user_id == user_id).first() is not None
|
||||
|
||||
def test_no_quota_allows_any_request(db):
|
||||
user_id = make_user_and_quota(db) # all limits None
|
||||
assert check_and_increment_quota(db, user_id, tokens=999999, requests=9999) is True
|
||||
|
||||
def test_cumulative_usage_across_calls(db):
|
||||
user_id = make_user_and_quota(db, daily_tokens=200)
|
||||
check_and_increment_quota(db, user_id, tokens=100, requests=1)
|
||||
check_and_increment_quota(db, user_id, tokens=99, requests=1)
|
||||
# 199 used, 1 remaining – exactly 1 more token should pass
|
||||
assert check_and_increment_quota(db, user_id, tokens=1, requests=1) is True
|
||||
# Now 200 used – next request must fail
|
||||
assert check_and_increment_quota(db, user_id, tokens=1, requests=1) is False
|
||||
|
||||
|
||||
# --- Reset logic ---
|
||||
|
||||
def test_daily_reset_restores_access(db):
|
||||
user_id = make_user_and_quota(db, daily_tokens=100)
|
||||
check_and_increment_quota(db, user_id, tokens=90, requests=1)
|
||||
|
||||
# Backdate daily_reset_at to yesterday
|
||||
usage = db.query(Usage).filter(Usage.user_id == user_id).first()
|
||||
usage.daily_reset_at = datetime.utcnow() - timedelta(days=1)
|
||||
db.commit()
|
||||
|
||||
# Should pass again after reset
|
||||
assert check_and_increment_quota(db, user_id, tokens=90, requests=1) is True
|
||||
|
||||
usage = db.query(Usage).filter(Usage.user_id == user_id).first()
|
||||
assert usage.tokens_used_today == 90
|
||||
|
||||
def test_daily_reset_does_not_affect_monthly_counter(db):
|
||||
user_id = make_user_and_quota(db, daily_tokens=1000, monthly_tokens=10000)
|
||||
check_and_increment_quota(db, user_id, tokens=50, requests=1)
|
||||
|
||||
usage = db.query(Usage).filter(Usage.user_id == user_id).first()
|
||||
usage.daily_reset_at = datetime.utcnow() - timedelta(days=1)
|
||||
db.commit()
|
||||
|
||||
check_and_increment_quota(db, user_id, tokens=50, requests=1)
|
||||
|
||||
usage = db.query(Usage).filter(Usage.user_id == user_id).first()
|
||||
assert usage.tokens_used_today == 50
|
||||
assert usage.tokens_used_month == 100 # cumulative across days
|
||||
|
||||
def test_monthly_reset_restores_access(db):
|
||||
user_id = make_user_and_quota(db, monthly_tokens=100)
|
||||
check_and_increment_quota(db, user_id, tokens=90, requests=1)
|
||||
|
||||
usage = db.query(Usage).filter(Usage.user_id == user_id).first()
|
||||
usage.monthly_reset_at = datetime.utcnow() - timedelta(days=32)
|
||||
db.commit()
|
||||
|
||||
assert check_and_increment_quota(db, user_id, tokens=90, requests=1) is True
|
||||
|
||||
usage = db.query(Usage).filter(Usage.user_id == user_id).first()
|
||||
assert usage.tokens_used_month == 90
|
||||
|
||||
def test_failed_quota_check_still_commits_reset(db):
|
||||
user_id = make_user_and_quota(db, daily_tokens=100, daily_requests=5)
|
||||
check_and_increment_quota(db, user_id, tokens=80, requests=1)
|
||||
|
||||
# Backdate so a reset fires, but the new request still exceeds the limit
|
||||
usage = db.query(Usage).filter(Usage.user_id == user_id).first()
|
||||
usage.daily_reset_at = datetime.utcnow() - timedelta(days=1)
|
||||
usage.tokens_used_today = 80
|
||||
db.commit()
|
||||
|
||||
# After reset tokens_used_today = 0; 200 tokens exceeds 100 limit
|
||||
result = check_and_increment_quota(db, user_id, tokens=200, requests=1)
|
||||
assert result is False
|
||||
|
||||
# Reset must still be persisted so the next request sees fresh counters
|
||||
db.expire_all()
|
||||
usage = db.query(Usage).filter(Usage.user_id == user_id).first()
|
||||
assert usage.tokens_used_today == 0
|
||||
Reference in New Issue
Block a user