xsl/hair
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(gateway): 接入42和72并为42/52/72启用可刷新令牌

Browse Source
This commit is contained in:
xsl
2026-11-07 20:50:22 +08:00
parent dc53a0ecfe
commit 51baf96d78
24 changed files with 1312 additions and 126 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_gateway_client_auth.py
+95
View File
@@ -0,0 +1,95 @@
import importlib
import json
from types import SimpleNamespace
from unittest.mock import AsyncMock, Mock
import pytest
from fastapi.testclient import TestClient
from gateway import client_auth
from gateway.logging_middleware import _should_log
gateway = importlib.import_module("gateway.app")
LOGIN = "/api/v1/auth/login"
REFRESH = "/api/v1/auth/refresh"
FEATURES_42 = "/api/v1/face/features-42"
PROTECTED = [FEATURES_42, "/api/v1/hairline/generate-52", "/api/v1/face/b-end-points"]
@pytest.fixture
def client(monkeypatch):
monkeypatch.setenv("HAIR_API_TOKEN_SECRET", "gateway-test-signing-secret-at-least-32")
session = TestClient(gateway.app)
yield session
session.close()
def login(client):
response = client.post(LOGIN, json={"username": "yonghe", "password": "yonghe1999"})
assert response.status_code == 200
assert response.headers["cache-control"] == "no-store"
return response.json()["data"]
@pytest.mark.parametrize("path", PROTECTED)
def test_all_three_routes_reject_missing_static_and_refresh_tokens(client, monkeypatch, path):
proxy = AsyncMock()
features = AsyncMock()
monkeypatch.setattr(gateway, "_proxy", proxy)
monkeypatch.setattr(gateway, "_face_features_response", features)
refresh_token = login(client)["refresh_token"]
for headers in ({}, {"X-API-Token": "yonghe1999"}, {"Authorization": "Bearer " + refresh_token}):
result = client.post(path, headers=headers)
assert result.status_code == 401
assert result.json()["code"] == 1009
proxy.assert_not_awaited()
features.assert_not_awaited()
def test_login_expiry_tampering_and_refresh(client, monkeypatch):
clock = SimpleNamespace(now=1_000_000)
monkeypatch.setattr(client_auth, "time", SimpleNamespace(time=lambda: clock.now))
proxy = AsyncMock(return_value={"code": 0})
monkeypatch.setattr(gateway, "_proxy", proxy)
path = PROTECTED[-1]
assert client.post(LOGIN, json={"username": "yonghe", "password": "wrong"}).status_code == 401
data = login(client)
assert data["expires_in"] == 900 and data["refresh_expires_in"] == 604800
access, refresh = data["access_token"], data["refresh_token"]
headers = {"Authorization": f"Bearer {access}"}
assert client.post(path, headers=headers).status_code == 200
assert client.post(path, headers={"Authorization": f"Bearer {access}x"}).status_code == 401
assert client.post(REFRESH, json={"refresh_token": access}).status_code == 401
clock.now += data["expires_in"]
assert client.post(path, headers=headers).status_code == 401
response = client.post(REFRESH, json={"refresh_token": refresh})
assert response.status_code == 200 and response.headers["cache-control"] == "no-store"
refreshed = response.json()["data"]
assert "refresh_token" not in refreshed and refreshed["access_token"] != access
assert client.post(path, headers={"Authorization": "Bearer " + refreshed["access_token"]}).status_code == 200
clock.now += data["refresh_expires_in"] - data["expires_in"]
assert client.post(REFRESH, json={"refresh_token": refresh}).status_code == 401
assert proxy.await_count == 2
def test_unconfigured_secret_fails_closed_and_token_responses_are_not_logged(client, monkeypatch):
monkeypatch.setenv("HAIR_API_TOKEN_SECRET", "")
result = client.post(LOGIN, json={"username": "yonghe", "password": "yonghe1999"})
assert result.status_code == 503 and result.json()["code"] == 1007
assert not _should_log(LOGIN) and not _should_log(REFRESH)
def test_features42_combines_processed_photo_with_main_features(client, monkeypatch):
analyzer = Mock(return_value={"face_shape": "鹅蛋脸", "contour_type": "轮廓适中"})
monkeypatch.setattr("face_features.analyze_features", analyzer)
prepared = AsyncMock(return_value=(b"processed", {"zygoma_expansion": True, "temple_hollowing": True}))
monkeypatch.setattr(gateway, "_prepare_features_42", prepared)
headers = {"Authorization": "Bearer " + login(client)["access_token"]}
response = client.post(FEATURES_42, headers=headers, files={"image_file": ("face.jpg", b"raw")})
assert response.status_code == 200 and response.json()["code"] == 0
assert json.loads(response.json()["data"]["features"]) == {
"face_shape": "鹅蛋脸", "contour_type": "轮廓适中",
"zygoma_expansion": True, "temple_hollowing": True,
}
prepared.assert_awaited_once_with(b"raw", None)
analyzer.assert_called_once_with(b"processed", None)
tests/test_gateway_interface52.py
+15 -74
View File
@@ -15,75 +15,15 @@ from gateway import config, forward
gateway = importlib.import_module("gateway.app")
PATH = "/api/v1/hairline/generate-52"
CLIENT_HEADERS = {"X-API-Token": "yonghe1999"}
@pytest.mark.parametrize("headers", [
{}, {"X-API-Token": ""}, {"X-API-Token": "wrong"},
{"X-API-Token": "yonghe1999 "}, {"X-Internal-Token": "yonghe1999"},
])
def test_invalid_client_token_is_rejected_before_dispatch(monkeypatch, headers):
monkeypatch.setattr(config, "get_config", lambda: {})
proxy = AsyncMock()
monkeypatch.setattr(gateway, "_proxy", proxy)
@pytest.fixture
def client_headers(monkeypatch):
monkeypatch.setenv("HAIR_API_TOKEN_SECRET", "gateway-test-signing-secret-at-least-32")
client = TestClient(gateway.app)
try:
response = client.post(PATH, headers=headers)
assert response.status_code == 401
body = response.json()
assert body["code"] == 1009 and body["data"] is None
assert body["request_id"].startswith("gw-")
response = client.post("/api/v1/auth/login", json={"username": "yonghe", "password": "yonghe1999"})
finally:
client.close()
proxy.assert_not_awaited()
def test_custom_token_replaces_default_only_on_interface52(monkeypatch):
monkeypatch.setattr(config, "get_config", lambda: {"interface52_token": "custom-client-token"})
proxy = AsyncMock(return_value={"code": 0, "data": None})
monkeypatch.setattr(gateway, "_proxy", proxy)
client = TestClient(gateway.app)
try:
assert client.post(PATH, headers=CLIENT_HEADERS).status_code == 401
assert client.post(PATH, headers={"X-API-Token": "custom-client-token"}).status_code == 200
assert client.post("/api/v1/hairline/generate").status_code == 200
finally:
client.close()
assert proxy.await_count == 2
@pytest.mark.parametrize("token", ["", None, 123, []])
def test_invalid_configured_token_cannot_disable_auth(monkeypatch, token):
monkeypatch.setattr(config, "get_config", lambda: {"interface52_token": token})
proxy = AsyncMock()
monkeypatch.setattr(gateway, "_proxy", proxy)
client = TestClient(gateway.app)
try:
assert client.post(PATH, headers=CLIENT_HEADERS).status_code == 401
finally:
client.close()
proxy.assert_not_awaited()
@pytest.mark.parametrize("token", ["", None, 123, []])
def test_startup_rejects_invalid_client_token_configuration(tmp_path, monkeypatch, token):
import json
path = tmp_path / "config.json"
path.write_text(json.dumps({"interface52_token": token}))
monkeypatch.setattr(config, "CONFIG_PATH", path)
monkeypatch.setattr(config, "_config_cache", None)
with pytest.raises(ValueError, match="interface52_token"):
config.load_config()
def test_existing_config_gets_default_client_token(tmp_path, monkeypatch):
import json
path = tmp_path / "config.json"
path.write_text(json.dumps({"workers": ["http://worker"], "shared_password": "test-internal-token",
"static_dir": str(tmp_path)}))
monkeypatch.setattr(config, "CONFIG_PATH", path)
monkeypatch.setattr(config, "_config_cache", None)
assert config.load_config()["interface52_token"] == "yonghe1999"
return {"Authorization": "Bearer " + response.json()["data"]["access_token"]}
def image_bytes(fmt):
@@ -113,14 +53,14 @@ def proxy_setup(tmp_path, monkeypatch):
return cfg, acquire, release, release_worker, unhealthy
@pytest.mark.parametrize("path", ["/api/v1/hairline/generate", PATH])
def test_hairline_routes_forward_to_their_own_worker_path(monkeypatch, path):
@pytest.mark.parametrize("path", ["/api/v1/hairline/generate", PATH, "/api/v1/face/b-end-points"])
def test_routes_forward_to_their_own_worker_path(monkeypatch, path, client_headers):
monkeypatch.setattr(config, "get_config", lambda: {})
proxy = AsyncMock(return_value={"code": 0, "data": None})
monkeypatch.setattr(gateway, "_proxy", proxy)
client = TestClient(gateway.app, client=("203.0.113.9", 50000))
try:
assert client.post(path, headers=CLIENT_HEADERS if path == PATH else {}).json() == {"code": 0, "data": None}
assert client.post(path, headers=client_headers if path != "/api/v1/hairline/generate" else {}).json() == {"code": 0, "data": None}
finally:
client.close()
proxy.assert_awaited_once()
@@ -128,7 +68,7 @@ def test_hairline_routes_forward_to_their_own_worker_path(monkeypatch, path):
@pytest.mark.parametrize("generate_grow_image", [False, True])
def test_interface52_raw_forwarding_and_served_image_urls(proxy_setup, monkeypatch, generate_grow_image):
def test_interface52_raw_forwarding_and_served_image_urls(proxy_setup, monkeypatch, generate_grow_image, client_headers):
cfg, acquire, release, release_worker, unhealthy = proxy_setup
png, jpg = image_bytes("PNG"), image_bytes("JPEG")
metadata = {"coordinate_space": "processed_image", "width": 6, "height": 8,
@@ -163,7 +103,7 @@ def test_interface52_raw_forwarding_and_served_image_urls(proxy_setup, monkeypat
client = TestClient(gateway.app, client=("203.0.113.9", 50000))
try:
result = client.post(PATH, content=raw, headers={
**CLIENT_HEADERS,
**client_headers,
"Content-Type": source.headers["content-type"], "X-Internal-Token": "untrusted-client-token"})
assert result.status_code == 200
data = result.json()["data"]
@@ -190,21 +130,21 @@ def test_interface52_raw_forwarding_and_served_image_urls(proxy_setup, monkeypat
assert calls[0].headers["content-type"] == source.headers["content-type"]
assert calls[0].headers["x-internal-token"] == cfg["shared_password"]
assert calls[0].headers.get("x-forwarded-for") is None
assert calls[0].headers.get("x-api-token") is None
assert calls[0].headers.get("authorization") is None
acquire.assert_awaited_once_with(cfg, 0)
release.assert_called_once()
release_worker.assert_awaited_once()
unhealthy.assert_not_awaited()
def test_worker_ip_rejection_is_preserved(proxy_setup, monkeypatch):
def test_worker_ip_rejection_is_preserved(proxy_setup, monkeypatch, client_headers):
_, _, release, release_worker, unhealthy = proxy_setup
body = {"code": 1010, "message": "来源IP不在接口白名单", "data": None}
upstream = httpx.AsyncClient(transport=httpx.MockTransport(lambda request: httpx.Response(403, json=body)))
monkeypatch.setattr(forward, "get_client", lambda: upstream)
client = TestClient(gateway.app)
try:
response = client.post(PATH, headers=CLIENT_HEADERS, data={"gender": "male", "hair_style": "1"})
response = client.post(PATH, headers=client_headers, data={"gender": "male", "hair_style": "1"})
assert response.status_code == 403 and response.json() == body
finally:
client.close()
@@ -220,7 +160,8 @@ def test_public_interface52_page_has_no_internal_password():
page = client.get("/static/test_interface52.html")
assert page.status_code == 200
assert PATH in page.text and "X-Internal-Token" not in page.text
assert 'id="apiToken"' in page.text and "X-API-Token" in page.text
assert 'id="client-auth"' in page.text and "X-API-Token" not in page.text
assert "yonghe1999" not in page.text
assert "localhost:8187" not in page.text
assert "test_interface52_compare.html" not in page.text
assert client.get("/static/hairline_reference_box.js").status_code == 200