mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-04-25 11:18:22 +00:00
refactor: Remove init_token handling from admin initialization logic and related tests
This commit is contained in:
parent
839563f308
commit
0572ef44b9
@ -40,18 +40,16 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
async def _ensure_admin_user(app: FastAPI) -> None:
|
||||
"""Startup hook: generate init token on first boot; migrate orphan threads otherwise.
|
||||
"""Startup hook: handle first boot and migrate orphan threads otherwise.
|
||||
|
||||
After admin creation, migrate orphan threads from the LangGraph
|
||||
store (metadata.user_id unset) to the admin account. This is the
|
||||
"no-auth → with-auth" upgrade path: users who ran DeerFlow without
|
||||
authentication have existing LangGraph thread data that needs an
|
||||
owner assigned.
|
||||
First boot (no admin exists):
|
||||
- Generates a one-time ``init_token`` stored in ``app.state.init_token``
|
||||
- Logs the token to stdout so the operator can copy-paste it into the
|
||||
``/setup`` form to create the first admin account interactively.
|
||||
- Does NOT create any user accounts automatically.
|
||||
First boot (no admin exists):
|
||||
- Does NOT create any user accounts automatically.
|
||||
- The operator must visit ``/setup`` to create the first admin.
|
||||
|
||||
Subsequent boots (admin already exists):
|
||||
- Runs the one-time "no-auth → with-auth" orphan thread migration for
|
||||
@ -62,8 +60,6 @@ async def _ensure_admin_user(app: FastAPI) -> None:
|
||||
alongside the auth module via create_all, so freshly created tables
|
||||
never contain NULL-owner rows.
|
||||
"""
|
||||
import secrets
|
||||
|
||||
from sqlalchemy import select
|
||||
|
||||
from app.gateway.deps import get_local_provider
|
||||
@ -74,13 +70,8 @@ async def _ensure_admin_user(app: FastAPI) -> None:
|
||||
admin_count = await provider.count_admin_users()
|
||||
|
||||
if admin_count == 0:
|
||||
init_token = secrets.token_urlsafe(32)
|
||||
app.state.init_token = init_token
|
||||
logger.info("=" * 60)
|
||||
logger.info(" First boot detected — no admin account exists.")
|
||||
logger.info(" Use the one-time token below to create the admin account.")
|
||||
logger.info(" Copy it into the /setup form when prompted.")
|
||||
logger.info(" INIT TOKEN: %s", init_token)
|
||||
logger.info(" Visit /setup to complete admin account creation.")
|
||||
logger.info("=" * 60)
|
||||
return
|
||||
@ -365,11 +356,6 @@ This gateway provides custom endpoints for models, MCP configuration, skills, an
|
||||
"""
|
||||
return {"status": "healthy", "service": "deer-flow-gateway"}
|
||||
|
||||
# Ensure init_token always exists on app.state (None until lifespan sets it
|
||||
# if no admin is found). This prevents AttributeError in tests that don't
|
||||
# run the full lifespan.
|
||||
app.state.init_token = None
|
||||
|
||||
return app
|
||||
|
||||
|
||||
|
||||
@ -21,7 +21,6 @@ class AuthErrorCode(StrEnum):
|
||||
PROVIDER_NOT_FOUND = "provider_not_found"
|
||||
NOT_AUTHENTICATED = "not_authenticated"
|
||||
SYSTEM_ALREADY_INITIALIZED = "system_already_initialized"
|
||||
INVALID_INIT_TOKEN = "invalid_init_token"
|
||||
|
||||
|
||||
class TokenError(StrEnum):
|
||||
|
||||
@ -2,7 +2,6 @@
|
||||
|
||||
import logging
|
||||
import os
|
||||
import secrets
|
||||
import time
|
||||
from ipaddress import ip_address, ip_network
|
||||
|
||||
@ -389,7 +388,6 @@ class InitializeAdminRequest(BaseModel):
|
||||
|
||||
email: EmailStr
|
||||
password: str = Field(..., min_length=8)
|
||||
init_token: str | None = Field(default=None, description="One-time initialization token printed to server logs on first boot")
|
||||
|
||||
_strong_password = field_validator("password")(classmethod(lambda cls, v: _validate_strong_password(v)))
|
||||
|
||||
@ -399,31 +397,13 @@ async def initialize_admin(request: Request, response: Response, body: Initializ
|
||||
"""Create the first admin account on initial system setup.
|
||||
|
||||
Only callable when no admin exists. Returns 409 Conflict if an admin
|
||||
already exists. Requires the one-time ``init_token`` that is logged to
|
||||
stdout at startup whenever the system has no admin account.
|
||||
already exists.
|
||||
|
||||
On success the token is consumed (one-time use), the admin account is
|
||||
created with ``needs_setup=False``, and the session cookie is set.
|
||||
On success, the admin account is created with ``needs_setup=False`` and
|
||||
the session cookie is set.
|
||||
"""
|
||||
# Validate the one-time initialization token. The token is generated
|
||||
# at startup and stored in app.state.init_token; it is consumed here on
|
||||
# the first successful call so it cannot be replayed.
|
||||
# Using str | None allows a missing/null token to return 403 (not 422),
|
||||
# giving a consistent error response regardless of whether the token is
|
||||
# absent or incorrect.
|
||||
stored_token: str | None = getattr(request.app.state, "init_token", None)
|
||||
provided_token: str = body.init_token or ""
|
||||
if stored_token is None or not secrets.compare_digest(stored_token, provided_token):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_403_FORBIDDEN,
|
||||
detail=AuthErrorResponse(code=AuthErrorCode.INVALID_INIT_TOKEN, message="Invalid or expired initialization token").model_dump(),
|
||||
)
|
||||
|
||||
admin_count = await get_local_provider().count_admin_users()
|
||||
if admin_count > 0:
|
||||
# Do NOT consume the token on this error path — consuming it here
|
||||
# would allow an attacker to exhaust the token by calling with the
|
||||
# correct token when admin already exists (denial-of-service).
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_409_CONFLICT,
|
||||
detail=AuthErrorResponse(code=AuthErrorCode.SYSTEM_ALREADY_INITIALIZED, message="System already initialized").model_dump(),
|
||||
@ -433,16 +413,11 @@ async def initialize_admin(request: Request, response: Response, body: Initializ
|
||||
user = await get_local_provider().create_user(email=body.email, password=body.password, system_role="admin", needs_setup=False)
|
||||
except ValueError:
|
||||
# DB unique-constraint race: another concurrent request beat us.
|
||||
# Do NOT consume the token here for the same reason as above.
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_409_CONFLICT,
|
||||
detail=AuthErrorResponse(code=AuthErrorCode.SYSTEM_ALREADY_INITIALIZED, message="System already initialized").model_dump(),
|
||||
)
|
||||
|
||||
# Consume the token only after successful initialization — this is the
|
||||
# single place where one-time use is enforced.
|
||||
request.app.state.init_token = None
|
||||
|
||||
token = create_access_token(str(user.id), token_version=user.token_version)
|
||||
_set_session_cookie(response, token, request)
|
||||
|
||||
|
||||
@ -63,14 +63,13 @@ def _make_session_factory(admin_row=None):
|
||||
return sf
|
||||
|
||||
|
||||
# ── First boot: no admin → generate init_token, return early ─────────────
|
||||
# ── First boot: no admin → return early ──────────────────────────────────
|
||||
|
||||
|
||||
def test_first_boot_does_not_create_admin():
|
||||
"""admin_count==0 → generate init_token, do NOT create admin automatically."""
|
||||
"""admin_count==0 → do NOT create admin automatically."""
|
||||
provider = _make_provider(admin_count=0)
|
||||
app = _make_app_stub()
|
||||
app.state.init_token = None # lifespan sets this
|
||||
|
||||
with patch("app.gateway.deps.get_local_provider", return_value=provider):
|
||||
from app.gateway.app import _ensure_admin_user
|
||||
@ -78,9 +77,6 @@ def test_first_boot_does_not_create_admin():
|
||||
asyncio.run(_ensure_admin_user(app))
|
||||
|
||||
provider.create_user.assert_not_called()
|
||||
# init_token must have been set on app.state
|
||||
assert app.state.init_token is not None
|
||||
assert len(app.state.init_token) > 10
|
||||
|
||||
|
||||
def test_first_boot_skips_migration():
|
||||
@ -89,7 +85,6 @@ def test_first_boot_skips_migration():
|
||||
store = AsyncMock()
|
||||
store.asearch = AsyncMock(return_value=[])
|
||||
app = _make_app_stub(store=store)
|
||||
app.state.init_token = None # lifespan sets this
|
||||
|
||||
with patch("app.gateway.deps.get_local_provider", return_value=provider):
|
||||
from app.gateway.app import _ensure_admin_user
|
||||
|
||||
@ -1,7 +1,7 @@
|
||||
"""Tests for the POST /api/v1/auth/initialize endpoint.
|
||||
|
||||
Covers: first-boot admin creation, rejection when system already
|
||||
initialized, invalid/missing init_token, password strength validation,
|
||||
initialized, password strength validation,
|
||||
and public accessibility (no auth cookie required).
|
||||
"""
|
||||
|
||||
@ -16,7 +16,6 @@ os.environ.setdefault("AUTH_JWT_SECRET", "test-secret-key-initialize-admin-min-3
|
||||
from app.gateway.auth.config import AuthConfig, set_auth_config
|
||||
|
||||
_TEST_SECRET = "test-secret-key-initialize-admin-min-32"
|
||||
_INIT_TOKEN = "test-init-token-for-initialization-tests"
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
@ -45,9 +44,6 @@ def client(_setup_auth):
|
||||
|
||||
set_auth_config(AuthConfig(jwt_secret=_TEST_SECRET))
|
||||
app = create_app()
|
||||
# Pre-set the init token on app.state (normally done by the lifespan on
|
||||
# first boot; tests don't run the lifespan because it requires config.yaml).
|
||||
app.state.init_token = _INIT_TOKEN
|
||||
# Do NOT use TestClient as a context manager — that would trigger the
|
||||
# full lifespan which requires config.yaml. The auth endpoints work
|
||||
# without the lifespan (persistence engine is set up by _setup_auth).
|
||||
@ -55,11 +51,10 @@ def client(_setup_auth):
|
||||
|
||||
|
||||
def _init_payload(**extra):
|
||||
"""Build a valid /initialize payload with the test init_token."""
|
||||
"""Build a valid /initialize payload."""
|
||||
return {
|
||||
"email": "admin@example.com",
|
||||
"password": "Str0ng!Pass99",
|
||||
"init_token": _INIT_TOKEN,
|
||||
**extra,
|
||||
}
|
||||
|
||||
@ -85,53 +80,12 @@ def test_initialize_needs_setup_false(client):
|
||||
assert me.json()["needs_setup"] is False
|
||||
|
||||
|
||||
# ── Token validation ──────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_initialize_rejects_wrong_token(client):
|
||||
"""Wrong init_token → 403 invalid_init_token."""
|
||||
resp = client.post(
|
||||
"/api/v1/auth/initialize",
|
||||
json={**_init_payload(), "init_token": "wrong-token"},
|
||||
)
|
||||
assert resp.status_code == 403
|
||||
assert resp.json()["detail"]["code"] == "invalid_init_token"
|
||||
|
||||
|
||||
def test_initialize_rejects_empty_token(client):
|
||||
"""Empty init_token → 403 (fails constant-time comparison against stored token)."""
|
||||
resp = client.post(
|
||||
"/api/v1/auth/initialize",
|
||||
json={**_init_payload(), "init_token": ""},
|
||||
)
|
||||
assert resp.status_code == 403
|
||||
|
||||
|
||||
def test_initialize_token_consumed_after_success(client):
|
||||
"""After a successful /initialize the token is consumed and cannot be reused."""
|
||||
client.post("/api/v1/auth/initialize", json=_init_payload())
|
||||
# The token is now None; any subsequent call with the old token must be rejected (403)
|
||||
resp2 = client.post(
|
||||
"/api/v1/auth/initialize",
|
||||
json={**_init_payload(), "email": "other@example.com"},
|
||||
)
|
||||
assert resp2.status_code == 403
|
||||
|
||||
|
||||
# ── Rejection when already initialized ───────────────────────────────────
|
||||
|
||||
|
||||
def test_initialize_rejected_when_admin_exists(client):
|
||||
"""Second call to /initialize after admin exists → 409 system_already_initialized.
|
||||
|
||||
The first call consumes the token. Re-setting it on app.state simulates
|
||||
what would happen if the operator somehow restarted or manually refreshed
|
||||
the token (e.g., in testing).
|
||||
"""
|
||||
"""Second call to /initialize after admin exists → 409 system_already_initialized."""
|
||||
client.post("/api/v1/auth/initialize", json=_init_payload())
|
||||
# Re-set the token so the second attempt can pass token validation
|
||||
# and reach the admin-exists check.
|
||||
client.app.state.init_token = _INIT_TOKEN
|
||||
resp2 = client.post(
|
||||
"/api/v1/auth/initialize",
|
||||
json={**_init_payload(), "email": "other@example.com"},
|
||||
@ -141,24 +95,6 @@ def test_initialize_rejected_when_admin_exists(client):
|
||||
assert body["detail"]["code"] == "system_already_initialized"
|
||||
|
||||
|
||||
def test_initialize_token_not_consumed_on_admin_exists(client):
|
||||
"""Token is NOT consumed when the admin-exists guard rejects the request.
|
||||
|
||||
This prevents a DoS where an attacker calls with the correct token when
|
||||
admin already exists and permanently burns the init token.
|
||||
"""
|
||||
client.post("/api/v1/auth/initialize", json=_init_payload())
|
||||
# Token consumed by success above; re-simulate the scenario:
|
||||
# admin exists, token is still valid (re-set), call should 409 and NOT consume token.
|
||||
client.app.state.init_token = _INIT_TOKEN
|
||||
client.post(
|
||||
"/api/v1/auth/initialize",
|
||||
json={**_init_payload(), "email": "other@example.com"},
|
||||
)
|
||||
# Token must still be set (not consumed) after the 409 rejection.
|
||||
assert client.app.state.init_token == _INIT_TOKEN
|
||||
|
||||
|
||||
def test_initialize_register_does_not_block_initialization(client):
|
||||
"""/register creating a user before /initialize doesn't block admin creation."""
|
||||
# Register a regular user first
|
||||
|
||||
@ -26,9 +26,6 @@ export default function SetupPage() {
|
||||
const [error, setError] = useState("");
|
||||
const [loading, setLoading] = useState(false);
|
||||
|
||||
// --- Init-admin mode only ---
|
||||
const [initToken, setInitToken] = useState("");
|
||||
|
||||
// --- Change-password mode only ---
|
||||
const [currentPassword, setCurrentPassword] = useState("");
|
||||
|
||||
@ -82,7 +79,6 @@ export default function SetupPage() {
|
||||
body: JSON.stringify({
|
||||
email,
|
||||
password: newPassword,
|
||||
init_token: initToken,
|
||||
}),
|
||||
});
|
||||
|
||||
@ -190,23 +186,6 @@ export default function SetupPage() {
|
||||
required
|
||||
/>
|
||||
</div>
|
||||
<div className="flex flex-col space-y-1">
|
||||
<label htmlFor="initToken" className="text-sm font-medium">
|
||||
Initialization Token
|
||||
</label>
|
||||
<Input
|
||||
id="initToken"
|
||||
type="text"
|
||||
placeholder="Copy from server startup logs"
|
||||
value={initToken}
|
||||
onChange={(e) => setInitToken(e.target.value)}
|
||||
required
|
||||
autoComplete="off"
|
||||
/>
|
||||
<p className="text-muted-foreground text-xs">
|
||||
Find the <code>INIT TOKEN</code> printed in the server startup logs.
|
||||
</p>
|
||||
</div>
|
||||
<div className="flex flex-col space-y-1">
|
||||
<label htmlFor="password" className="text-sm font-medium">
|
||||
Password
|
||||
|
||||
@ -40,7 +40,6 @@ const AUTH_ERROR_CODES = [
|
||||
"provider_not_found",
|
||||
"not_authenticated",
|
||||
"system_already_initialized",
|
||||
"invalid_init_token",
|
||||
] as const;
|
||||
|
||||
export type AuthErrorCode = (typeof AUTH_ERROR_CODES)[number];
|
||||
|
||||
@ -322,7 +322,7 @@ export function useThreadStream({
|
||||
useEffect(() => {
|
||||
if (
|
||||
optimisticMessages.length > 0 &&
|
||||
thread.messages.length > prevMsgCountRef.current
|
||||
thread.messages.length > prevMsgCountRef.current + 1
|
||||
) {
|
||||
setOptimisticMessages([]);
|
||||
}
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user