From 3a1c4566d437edc37bab42941b8981d3007473f9 Mon Sep 17 00:00:00 2001 From: Rerowros Date: Mon, 10 Aug 2026 11:30:53 +0400 Subject: [PATCH 1/2] Harden distributed node and user lifecycle --- app/app_factory.py | 11 + app/db/crud/user.py | 58 +- ...5_bind_xray_client_inbounds_to_loopback.py | 89 + ...d12f6a8b9c30_add_bridge_sync_namespaces.py | 48 + app/db/models.py | 11 + app/jobs/node_checker.py | 57 +- app/jobs/record_usages.py | 47 +- app/jobs/remove_expired_users.py | 14 +- app/models/node.py | 17 + app/models/user.py | 7 +- app/nats/rpc_service.py | 14 +- app/node/__init__.py | 850 ++++++- app/node/errors.py | 4 + app/node/manager_sync.py | 54 +- app/node/nats_memory.py | 1091 ++++++++- app/node/sync.py | 608 ++++- app/node/user.py | 16 +- app/node/worker.py | 97 +- app/notification/webhook/__init__.py | 6 +- app/operation/admin.py | 18 +- app/operation/node.py | 443 +++- app/operation/subscription.py | 22 +- app/operation/user.py | 55 +- app/routers/admin.py | 10 +- app/routers/node.py | 18 +- app/routers/user.py | 12 +- app/utils/jwt.py | 13 +- app/utils/responses.py | 6 + config.py | 4 +- dashboard/src/service/api/index.ts | 429 ++-- install_service.sh | 25 +- pyproject.toml | 7 +- tests/api/test_node.py | 15 +- tests/api/test_usage_functions_timezone.py | 3 + tests/api/test_user.py | 30 + tests/test_app_factory.py | 16 +- tests/test_connect_concurrency.py | 9 +- tests/test_nats_node_memory.py | 769 +++++- tests/test_node_bridge_namespace.py | 79 + tests/test_node_manager.py | 5 +- tests/test_node_manager_sync.py | 194 +- tests/test_node_sync.py | 2105 ++++++++++++++++- tests/test_record_usages.py | 108 +- tests/test_rpc_service.py | 122 + tests/test_security_hardening.py | 557 +++++ tests/test_sync_namespace_migration.py | 56 + tests/test_xray_loopback_migration.py | 60 + uv.lock | 10 +- 48 files changed, 7840 insertions(+), 459 deletions(-) create mode 100644 app/db/migrations/versions/a8c2d491e705_bind_xray_client_inbounds_to_loopback.py create mode 100644 app/db/migrations/versions/d12f6a8b9c30_add_bridge_sync_namespaces.py create mode 100644 app/node/errors.py create mode 100644 tests/test_node_bridge_namespace.py create mode 100644 tests/test_rpc_service.py create mode 100644 tests/test_security_hardening.py create mode 100644 tests/test_sync_namespace_migration.py create mode 100644 tests/test_xray_loopback_migration.py diff --git a/app/app_factory.py b/app/app_factory.py index 4ec508667..6894ebf8e 100644 --- a/app/app_factory.py +++ b/app/app_factory.py @@ -12,6 +12,7 @@ from app.nats import is_multi_worker, require_nats_if_multiworker from app.nats.message import MessageTopic from app.nats.router import router +from app.node.errors import NodeRevocationError from app.settings import handle_settings_message from app.subscription.client_templates import handle_client_template_message from app.utils.logger import get_logger @@ -35,6 +36,15 @@ async def database_operational_error_handler(request: Request, exc: DBAPIError): ) +async def node_revocation_error_handler(request: Request, exc: NodeRevocationError): + logger.warning("Node revocation unavailable while handling %s %s: %s", request.method, request.url.path, exc) + return JSONResponse( + status_code=status.HTTP_503_SERVICE_UNAVAILABLE, + content={"detail": "User removal was not confirmed by all runtime nodes. Retry when nodes are available."}, + headers={"Retry-After": "1"}, + ) + + def _use_route_names_as_operation_ids(app: FastAPI) -> None: def _simplify_operation_ids(routes): for route in routes: @@ -275,6 +285,7 @@ def validation_exception_handler(request: Request, exc: RequestValidationError): ) app.add_exception_handler(DBAPIError, database_operational_error_handler) + app.add_exception_handler(NodeRevocationError, node_revocation_error_handler) from app.operation.permissions import LimitExceeded, PermissionDenied diff --git a/app/db/crud/user.py b/app/db/crud/user.py index 1f325c611..8d8722d98 100644 --- a/app/db/crud/user.py +++ b/app/db/crud/user.py @@ -953,7 +953,7 @@ async def _delete_user_dependencies(db: AsyncSession, user_ids: list[int]): await db.execute(users_groups_association.delete().where(users_groups_association.c.user_id.in_(user_ids))) -async def remove_user(db: AsyncSession, db_user: User) -> User: +async def remove_user(db: AsyncSession, db_user: User, *, commit: bool = True) -> User: """ Removes a user from the database. @@ -967,11 +967,14 @@ async def remove_user(db: AsyncSession, db_user: User) -> User: await release_users_allocations(db, [db_user]) await _delete_user_dependencies(db, [db_user.id]) await db.execute(delete(User).where(User.id == db_user.id)) - await db.commit() + if commit: + await db.commit() + else: + await db.flush() return db_user -async def remove_users(db: AsyncSession, db_users: list[User]): +async def remove_users(db: AsyncSession, db_users: list[User], *, commit: bool = True): """ Removes multiple users from the database. @@ -987,7 +990,10 @@ async def remove_users(db: AsyncSession, db_users: list[User]): await release_users_allocations(db, db_users) await _delete_user_dependencies(db, user_ids) await db.execute(delete(User).where(User.id.in_(user_ids))) - await db.commit() + if commit: + await db.commit() + else: + await db.flush() async def modify_user( @@ -1424,20 +1430,10 @@ async def get_users_subscription_agent_stats( return rows -async def autodelete_expired_users( +async def get_autodelete_expired_users( db: AsyncSession, include_limited_users: bool = False -) -> list[UserNotificationResponse]: - """ - Deletes expired (optionally also limited) users whose auto-delete time has passed. - - Args: - db (AsyncSession): Database session - include_limited_users (bool, optional): Whether to delete limited users as well. - Defaults to False. - - Returns: - list[UserNotificationResponse]: List of deleted users. - """ +) -> tuple[list[User], list[UserNotificationResponse]]: + """Return auto-delete targets and their node-removal snapshots without deleting them.""" target_status = [UserStatus.expired] if not include_limited_users else [UserStatus.expired, UserStatus.limited] auto_delete = func.coalesce(User.auto_delete_in_days, literal(user_cleanup_settings.autodelete_days)) @@ -1449,6 +1445,7 @@ async def autodelete_expired_users( ) .where( auto_delete >= 0, # Negative values prevent auto-deletion + auto_delete <= 36500, # Keep persisted legacy values within datetime arithmetic bounds User.status.in_(target_status), ) .options(joinedload(User.admin)) @@ -1461,12 +1458,29 @@ async def autodelete_expired_users( ] result: list[UserNotificationResponse] = [] - if expired_users: - for user in expired_users: - await load_user_attrs(user) - result.append(UserNotificationResponse.model_validate(user)) + for user in expired_users: + await load_user_attrs(user) + result.append(UserNotificationResponse.model_validate(user)) + return expired_users, result + + +async def autodelete_expired_users( + db: AsyncSession, include_limited_users: bool = False, *, commit: bool = True +) -> list[UserNotificationResponse]: + """ + Delete expired (optionally also limited) users whose auto-delete time has passed. - await remove_users(db, expired_users) + Args: + db (AsyncSession): Database session + include_limited_users (bool, optional): Whether to delete limited users as well. + Defaults to False. + + Returns: + list[UserNotificationResponse]: List of deleted users. + """ + expired_users, result = await get_autodelete_expired_users(db, include_limited_users) + if expired_users: + await remove_users(db, expired_users, commit=commit) return result diff --git a/app/db/migrations/versions/a8c2d491e705_bind_xray_client_inbounds_to_loopback.py b/app/db/migrations/versions/a8c2d491e705_bind_xray_client_inbounds_to_loopback.py new file mode 100644 index 000000000..23d074fd4 --- /dev/null +++ b/app/db/migrations/versions/a8c2d491e705_bind_xray_client_inbounds_to_loopback.py @@ -0,0 +1,89 @@ +"""Bind generated Xray client proxy inbounds to loopback. + +Revision ID: a8c2d491e705 +Revises: fb32155473c1 +Create Date: 2026-08-09 +""" + +import re + +import sqlalchemy as sa +from alembic import op + +revision = "a8c2d491e705" +down_revision = "fb32155473c1" +branch_labels = None +depends_on = None + + +client_templates = sa.table( + "client_templates", + sa.column("id", sa.Integer()), + sa.column("template_type", sa.String()), + sa.column("content", sa.Text()), + sa.column("is_system", sa.Boolean()), +) + +EXPOSED_LISTENER = re.compile(r'("listen"\s*:\s*)"0\.0\.0\.0"') +INBOUNDS_ARRAY = re.compile(r'"inbounds"\s*:\s*\[') + + +def _json_array_end(content: str, start: int) -> int | None: + depth = 0 + in_string = False + escaped = False + for index in range(start, len(content)): + character = content[index] + if in_string: + if escaped: + escaped = False + elif character == "\\": + escaped = True + elif character == '"': + in_string = False + continue + if character == '"': + in_string = True + elif character == "[": + depth += 1 + elif character == "]": + depth -= 1 + if depth == 0: + return index + 1 + return None + + +def _bind_client_listeners_to_loopback(content: str) -> str: + """Rewrite only exposed Xray client listener values, preserving the template.""" + match = INBOUNDS_ARRAY.search(content) + if not match: + return content + start = match.end() - 1 + end = _json_array_end(content, start) + if end is None: + return content + inbounds = EXPOSED_LISTENER.sub(r'\1"127.0.0.1"', content[start:end]) + return content[:start] + inbounds + content[end:] + + +def upgrade() -> None: + connection = op.get_bind() + rows = connection.execute( + sa.select(client_templates.c.id, client_templates.c.content).where( + client_templates.c.template_type == "xray_subscription", + client_templates.c.is_system.is_(True), + ) + ).mappings() + for row in rows: + content = row["content"] + updated_content = _bind_client_listeners_to_loopback(content) + if updated_content == content: + continue + connection.execute( + client_templates.update().where(client_templates.c.id == row["id"]).values(content=updated_content) + ) + + +def downgrade() -> None: + # Do not reintroduce an unauthenticated network listener on downgrade. + pass diff --git a/app/db/migrations/versions/d12f6a8b9c30_add_bridge_sync_namespaces.py b/app/db/migrations/versions/d12f6a8b9c30_add_bridge_sync_namespaces.py new file mode 100644 index 000000000..a5608f50e --- /dev/null +++ b/app/db/migrations/versions/d12f6a8b9c30_add_bridge_sync_namespaces.py @@ -0,0 +1,48 @@ +"""Add stable Bridge node and user sync namespaces. + +Revision ID: d12f6a8b9c30 +Revises: a8c2d491e705 +Create Date: 2026-08-10 +""" + +import sqlalchemy as sa +from alembic import op + +revision = "d12f6a8b9c30" +down_revision = "a8c2d491e705" +branch_labels = None +depends_on = None + + +def _backfill_legacy_namespace(connection, table: str, column: str) -> None: + # Preserve the namespace already used by running Bridge/core processes and + # NATS KV. Newly created ORM rows use UUID defaults, so a later reused + # numeric database id cannot inherit this incarnation's state or stats. + table_ref = sa.table(table, sa.column("id"), sa.column(column)) + connection.execute( + sa.update(table_ref).values({column: sa.cast(table_ref.c.id, sa.String(36))}) + ) + + +def upgrade() -> None: + connection = op.get_bind() + op.add_column("nodes", sa.Column("bridge_id", sa.String(length=36), nullable=True)) + op.add_column("users", sa.Column("sync_id", sa.String(length=36), nullable=True)) + _backfill_legacy_namespace(connection, "nodes", "bridge_id") + _backfill_legacy_namespace(connection, "users", "sync_id") + + with op.batch_alter_table("nodes") as batch_op: + batch_op.alter_column("bridge_id", existing_type=sa.String(length=36), nullable=False) + batch_op.create_unique_constraint("uq_nodes_bridge_id", ["bridge_id"]) + with op.batch_alter_table("users") as batch_op: + batch_op.alter_column("sync_id", existing_type=sa.String(length=36), nullable=False) + batch_op.create_unique_constraint("uq_users_sync_id", ["sync_id"]) + + +def downgrade() -> None: + with op.batch_alter_table("users") as batch_op: + batch_op.drop_constraint("uq_users_sync_id", type_="unique") + batch_op.drop_column("sync_id") + with op.batch_alter_table("nodes") as batch_op: + batch_op.drop_constraint("uq_nodes_bridge_id", type_="unique") + batch_op.drop_column("bridge_id") diff --git a/app/db/models.py b/app/db/models.py index 7588a31a5..2d16aae10 100644 --- a/app/db/models.py +++ b/app/db/models.py @@ -2,6 +2,7 @@ from datetime import UTC, datetime as dt from enum import Enum from typing import Any +from uuid import uuid4 from sqlalchemy import ( JSON, @@ -193,10 +194,14 @@ class DataLimitResetStrategy(str, Enum): class User(Base, CreatedAtUTCMixin): __tablename__ = "users" __table_args__ = ( + UniqueConstraint("sync_id", name="uq_users_sync_id"), Index("idx_users_admin_online", "admin_id", "online_at"), Index("idx_users_admin_status", "admin_id", "status"), Index("idx_users_admin_created", "admin_id", "created_at"), ) + sync_id: Mapped[str] = mapped_column( + String(36), default_factory=lambda: str(uuid4()), nullable=False, init=False + ) username: Mapped[str] = mapped_column(CaseSensitiveString(128), unique=True, index=True) node_usages: Mapped[list[NodeUserUsage]] = relationship( back_populates="user", @@ -596,6 +601,12 @@ class NodeStatus(str, Enum): class Node(Base, CreatedAtUTCMixin): __tablename__ = "nodes" + __table_args__ = (UniqueConstraint("bridge_id", name="uq_nodes_bridge_id"),) + # Stable internal namespace for distributed Bridge/KV state. Public APIs + # continue to route by numeric id, which SQLite may reuse after deletion. + bridge_id: Mapped[str] = mapped_column( + String(36), default_factory=lambda: str(uuid4()), nullable=False, init=False + ) name: Mapped[str] = mapped_column(CaseSensitiveString(256), unique=True) address: Mapped[str] = mapped_column(String(256), unique=False, nullable=False) port: Mapped[int] = mapped_column(unique=False, nullable=False) diff --git a/app/jobs/node_checker.py b/app/jobs/node_checker.py index 32e8c1b45..36f08a520 100644 --- a/app/jobs/node_checker.py +++ b/app/jobs/node_checker.py @@ -5,7 +5,7 @@ from app import notification, on_shutdown, on_startup, scheduler from app.db import GetDB -from app.db.crud.node import get_limited_nodes, get_nodes +from app.db.crud.node import get_limited_nodes, get_node_by_id, get_nodes from app.db.models import Node, NodeStatus from app.models.node import NodeListQuery, NodeNotification from app.nats import is_multi_worker @@ -92,7 +92,28 @@ async def process_node_health_check(db_node: Node, node: PasarGuardNode): - For other errors (code > -1): Reconnect (connection works but has another issue) - For NOT_CONNECTED/INVALID: Reconnect immediately """ + _, coordinator, _ = get_bridge_memory() + bridge_namespace = str(getattr(db_node, "bridge_id", None) or db_node.id) + if coordinator is not None and await coordinator.is_deleted(bridge_namespace): + if node is not None: + await node_manager.remove_node( + db_node.id, + remote_stop=False, + expected_bridge_namespace=bridge_namespace, + ) + return + if node is None: + async with GetDB() as db: + await node_operator.connect_single_node(db, db_node.id) + return + + # Broadcast delivery is best-effort. A worker that missed an upsert heals + # from the authoritative DB row during its ordinary health pass. + if getattr(node, "_extra", {}).get("config_signature") is not None and not node_manager.runtime_config_matches( + node, db_node + ): + await node_manager.update_node(db_node) return # Limit concurrent health checks to prevent DB/API overload @@ -147,7 +168,7 @@ async def process_node_health_check(db_node: Node, node: PasarGuardNode): return _, coordinator, _ = get_bridge_memory() - if coordinator is not None and await coordinator.has_active_lease(str(db_node.id)): + if coordinator is not None and await coordinator.has_active_lease(bridge_namespace): logger.debug( "[%s] Shared lifecycle HEALTHY with active lease; waiting for owner", db_node.name, @@ -249,6 +270,19 @@ async def node_health_check(): await asyncio.gather(*check_tasks, return_exceptions=True) +async def reconcile_orphaned_user_sync(): + """Periodically resolve durable NATS barriers after worker/process crashes.""" + if not runtime_settings.role.runs_node: + return + node_ids = list((await node_manager.get_nodes()).keys()) + for node_id in node_ids: + async with GetDB() as db: + db_node = await get_node_by_id(db, node_id, load_usage_logs=False) + if db_node is None: + continue + await node_operator.reconcile_orphaned_user_sync(db, db_node) + + _node_loop_tasks: list[asyncio.Task] = [] @@ -290,6 +324,16 @@ async def initialize_nodes(): name="node_health_loop", ) ) + _node_loop_tasks.append( + asyncio.create_task( + _interval_loop( + reconcile_orphaned_user_sync, + job_settings.core_health_check_interval, + "user-sync-recovery", + ), + name="node_user_sync_recovery_loop", + ) + ) else: scheduler.add_job( node_health_check, @@ -300,6 +344,15 @@ async def initialize_nodes(): id="node_health_check", replace_existing=True, ) + scheduler.add_job( + reconcile_orphaned_user_sync, + "interval", + seconds=job_settings.core_health_check_interval, + coalesce=True, + max_instances=1, + id="reconcile_orphaned_user_sync", + replace_existing=True, + ) # Limit checks mutate node status / disconnect; run only on the leader scheduler. scheduler.add_job( diff --git a/app/jobs/record_usages.py b/app/jobs/record_usages.py index 5ad109082..a0ea1f710 100644 --- a/app/jobs/record_usages.py +++ b/app/jobs/record_usages.py @@ -76,7 +76,7 @@ def _process_node_chunk(chunk_data: tuple) -> dict: _node_id, params, coeff = chunk_data users_usage = defaultdict(int) for param in params: - uid = int(param["uid"]) + uid = str(param["uid"]) value = int(param["value"] * coeff) users_usage[uid] += value return dict(users_usage) @@ -504,9 +504,9 @@ def _process_users_stats_response(stats_response): validated_params = [] invalid_uids = [] for uid, value in params.items(): - try: - validated_params.append({"uid": int(uid), "value": value}) - except ValueError, TypeError: + if isinstance(uid, str) and uid: + validated_params.append({"uid": uid, "value": value}) + else: invalid_uids.append(uid) return validated_params, invalid_uids @@ -578,30 +578,29 @@ async def get_outbounds_stats(node: PasarGuardNode): return [] -async def calculate_admin_usage(users_usage: list) -> tuple[dict, set[int]]: +async def calculate_admin_usage(users_usage: list) -> tuple[dict, dict[str, int]]: if not users_usage: - return {}, set() + return {}, {} - # Get unique user IDs from users_usage - uids = {int(user_usage["uid"]) for user_usage in users_usage} + sync_ids = {str(user_usage["uid"]) for user_usage in users_usage} async with GetDB() as db: - # Query only relevant users' admin IDs - user_admin_pairs = [] - for uid_batch in _chunked(list(uids), USER_ADMIN_LOOKUP_BATCH_SIZE): - stmt = select(User.id, User.admin_id).where(User.id.in_(uid_batch)) + user_rows = [] + for sync_id_batch in _chunked(list(sync_ids), USER_ADMIN_LOOKUP_BATCH_SIZE): + stmt = select(User.sync_id, User.id, User.admin_id).where(User.sync_id.in_(sync_id_batch)) result = await db.execute(stmt) - user_admin_pairs.extend(result.fetchall()) + user_rows.extend(result.fetchall()) - user_admin_map = {uid: admin_id for uid, admin_id in user_admin_pairs} + sync_to_user_id = {sync_id: user_id for sync_id, user_id, _ in user_rows} + sync_to_admin_id = {sync_id: admin_id for sync_id, _, admin_id in user_rows} admin_usage = defaultdict(int) for user_usage in users_usage: - admin_id = user_admin_map.get(int(user_usage["uid"])) + admin_id = sync_to_admin_id.get(str(user_usage["uid"])) if admin_id: admin_usage[admin_id] += user_usage["value"] - return admin_usage, set(user_admin_map.keys()) + return admin_usage, sync_to_user_id async def calculate_users_usage(api_params: dict, usage_coefficient: dict) -> list: @@ -618,7 +617,7 @@ def _process_usage_sync(chunks_data: list[tuple[int, list[dict], float]]): users_usage = defaultdict(int) for _, params, coeff in chunks_data: for param in params: - uid = int(param["uid"]) + uid = str(param["uid"]) value = int(param["value"] * coeff) users_usage[uid] += value return [{"uid": uid, "value": value} for uid, value in users_usage.items()] @@ -711,14 +710,16 @@ async def _record_user_usages_impl(): logger.debug("No user usage to record") return - admin_usage, valid_user_ids = await calculate_admin_usage(users_usage) - if not valid_user_ids: + admin_usage, sync_to_user_id = await calculate_admin_usage(users_usage) + if not sync_to_user_id: logger.warning("Skipping user usage recording; no matching users found for received stats") return # Filter valid users - only include users with actual non-zero traffic valid_users_usage = [ - usage for usage in users_usage if int(usage["uid"]) in valid_user_ids and usage["value"] > 0 + {"uid": sync_to_user_id[str(usage["uid"])], "value": usage["value"]} + for usage in users_usage + if str(usage["uid"]) in sync_to_user_id and usage["value"] > 0 ] # Update User table with concurrency control @@ -756,7 +757,11 @@ async def _record_user_usages_impl(): # Filter params to only valid users filtered_node_params = {} for node_id, params in api_params.items(): - filtered_params = [param for param in params if int(param["uid"]) in valid_user_ids] + filtered_params = [ + {"uid": sync_to_user_id[str(param["uid"])], "value": param["value"]} + for param in params + if str(param["uid"]) in sync_to_user_id + ] if filtered_params: filtered_node_params[node_id] = filtered_params diff --git a/app/jobs/remove_expired_users.py b/app/jobs/remove_expired_users.py index caaa131df..1c9ca39fe 100644 --- a/app/jobs/remove_expired_users.py +++ b/app/jobs/remove_expired_users.py @@ -2,8 +2,9 @@ from app import notification, scheduler from app.db import GetDB -from app.db.crud.user import autodelete_expired_users +from app.db.crud.user import get_autodelete_expired_users, remove_users from app.jobs.dependencies import SYSTEM_ADMIN +from app.node.sync import finalize_users_removal, remove_users_and_wait, resolve_user_removal_after_db_error from app.utils.logger import get_logger from config import job_settings, runtime_settings, user_cleanup_settings @@ -12,8 +13,15 @@ async def remove_expired_users(): async with GetDB() as db: - deleted_users = await autodelete_expired_users(db, user_cleanup_settings.include_limited_accounts) - + db_users, deleted_users = await get_autodelete_expired_users(db, user_cleanup_settings.include_limited_accounts) + revocation = await remove_users_and_wait(db_users) + try: + await remove_users(db, db_users) + except BaseException: + await resolve_user_removal_after_db_error(revocation, db) + raise + if revocation is not None: + await finalize_users_removal(revocation) for user in deleted_users: asyncio.create_task(notification.remove_user(user=user, by=SYSTEM_ADMIN)) logger.info(f"User `{user.username}` has been deleted due to expiration.") diff --git a/app/models/node.py b/app/models/node.py index 231eef5fb..8f4636e27 100644 --- a/app/models/node.py +++ b/app/models/node.py @@ -1,9 +1,11 @@ import re from enum import Enum from ipaddress import ip_address +from typing import Literal from uuid import UUID from cryptography.x509 import load_pem_x509_certificate +from PasarGuardNodeBridge.storage import LifecycleStatus from pydantic import BaseModel, ConfigDict, Field, computed_field, field_validator, model_validator from app.db.models import DataLimitResetStrategy, NodeConnectionType, NodeStatus @@ -36,6 +38,20 @@ class NodeSettings(BaseModel): min_node_version: str = "v1.0.0" +class NodeLifecycleRecovery(BaseModel): + """Explicit acknowledgement for an expired, outcome-unknown operation.""" + + observed: LifecycleStatus + acknowledge_expired_operation: Literal[True] + + @field_validator("observed") + @classmethod + def validate_terminal_observed_state(cls, value: LifecycleStatus) -> LifecycleStatus: + if value not in (LifecycleStatus.HEALTHY, LifecycleStatus.STOPPED, LifecycleStatus.BROKEN): + raise ValueError("observed must be a terminal inspected state") + return value + + class Node(BaseModel): name: str address: str @@ -364,6 +380,7 @@ class RemoveNodesResponse(BaseModel): nodes: list[str] count: int + failed: dict[int, str] = Field(default_factory=dict) class BulkNodesActionResponse(BaseModel): diff --git a/app/models/user.py b/app/models/user.py index 7af5d0756..135ef6ccc 100644 --- a/app/models/user.py +++ b/app/models/user.py @@ -48,7 +48,12 @@ class User(BaseModel): ) on_hold_timeout: dt | int | None = Field(default=None) group_ids: list[int] | None = Field(default_factory=list) - auto_delete_in_days: int | None = Field(default=None) + auto_delete_in_days: int | None = Field( + ge=-1, + le=36500, + default=None, + description="Per-user cleanup delay in days; -1 disables automatic deletion", + ) hwid_limit: int | None = Field(default=None) next_plan: NextPlanModel | None = Field(default=None) diff --git a/app/nats/rpc_service.py b/app/nats/rpc_service.py index 921173108..a8f9ef8da 100644 --- a/app/nats/rpc_service.py +++ b/app/nats/rpc_service.py @@ -19,12 +19,14 @@ def __init__( start_msg: str | None = None, stop_msg: str | None = None, rpc_concurrency: int = 20, + queue_group: str | None = None, ): self._rpc_subject = subject self._logger = logger self._role_check = role_check self._start_msg = start_msg self._stop_msg = stop_msg + self._queue_group = queue_group self._nc: nats.NATS | None = None self._rpc_sub: Subscription | None = None self._rpc_semaphore = asyncio.Semaphore(rpc_concurrency) @@ -44,7 +46,14 @@ async def start(self): if not self._nc: return - self._rpc_sub = await self._nc.subscribe(self._rpc_subject, cb=self._handle_rpc) + if self._queue_group is None: + self._rpc_sub = await self._nc.subscribe(self._rpc_subject, cb=self._handle_rpc) + else: + self._rpc_sub = await self._nc.subscribe( + self._rpc_subject, + queue=self._queue_group, + cb=self._handle_rpc, + ) if self._start_msg: self._logger.info(self._start_msg) @@ -79,7 +88,8 @@ async def _run_rpc(self, msg, action: str | None, data: dict): await msg.respond(json.dumps({"ok": True, "data": result}).encode()) except Exception as exc: error_msg = str(exc) - await msg.respond(json.dumps({"ok": False, "error": error_msg, "code": 500}).encode()) + error_code = getattr(exc, "code", 500) + await msg.respond(json.dumps({"ok": False, "error": error_msg, "code": error_code}).encode()) async def _dispatch_rpc(self, action: str | None, data: dict): if not action: diff --git a/app/node/__init__.py b/app/node/__init__.py index 479fbfa43..338329e7c 100644 --- a/app/node/__init__.py +++ b/app/node/__init__.py @@ -1,10 +1,13 @@ import asyncio +import hashlib from aiorwlock import RWLock -from PasarGuardNodeBridge import Health, NodeType, PasarGuardNode, create_node +from PasarGuardNodeBridge import Health, NodeAPIError, NodeType, PasarGuardNode, create_node from PasarGuardNodeBridge.common.service_pb2 import User as ProtoUser from app.db.models import Node, NodeConnectionType +from app.nats import needs_shared_bridge_memory +from app.node.errors import NodeRevocationError from app.node.nats_memory import ensure_bridge_memory, get_bridge_memory from app.node.user import core_users from app.utils.logger import get_logger @@ -19,7 +22,24 @@ class NodeManager: def __init__(self): self._nodes: dict[int, PasarGuardNode] = {} + self._retiring_nodes: dict[int, list[PasarGuardNode]] = {} + self._removing_node_ids: set[int] = set() + self._replacing_node_ids: set[int] = set() + self._deleted_node_namespaces: set[str] = set() + self._runtime_transition_tasks: set[asyncio.Task] = set() + self._runtime_transition_locks: dict[int, asyncio.Lock] = {} self._user_sync_locks: dict[int, asyncio.Lock] = {} + # A permanent deletion fence. User IDs are generated monotonically and + # are not reused, so retaining the tombstone prevents an already queued + # update from re-admitting a deleted user after revocation completes. + self._deleted_user_keys: set[str] = set() + self._deletion_fence_owners: dict[str, set[str]] = {} + self._revocation_nodes: dict[str, list[tuple[int, PasarGuardNode, frozenset[str]]]] = {} + self._revocations_idle = asyncio.Event() + self._revocations_idle.set() + # Deployment mode is immutable for this process. Never build a hybrid + # manager containing private and NATS-backed bridge stores. + self._uses_shared_revocation_store = needs_shared_bridge_memory() self._lock = RWLock(fast=True) self.logger = get_logger("node-manager") @@ -36,55 +56,308 @@ def _create_node_kwargs(self, node: Node) -> dict: "default_timeout": node.default_timeout, "internal_timeout": node.internal_timeout, "proxy": node.proxy_url, - "extra": {"id": node.id, "usage_coefficient": node.usage_coefficient}, - "node_id": str(node.id), + "extra": { + "id": node.id, + "usage_coefficient": node.usage_coefficient, + "config_signature": self.node_config_signature(node), + }, + "node_id": self.bridge_namespace(node), } store, coordinator, worker_id = get_bridge_memory() + if self._uses_shared_revocation_store and (store is None or coordinator is None): + raise NodeAPIError(503, "shared node bridge memory is unavailable") if store is not None and coordinator is not None: kwargs["user_sync_store"] = store kwargs["lifecycle_coordinator"] = coordinator kwargs["worker_id"] = worker_id return kwargs - async def _shutdown_node(self, node: PasarGuardNode | None, *, remote_stop: bool = True): + @staticmethod + def bridge_namespace(node: Node) -> str: + """Return the persisted Bridge namespace, with a test/legacy fallback.""" + return str(getattr(node, "bridge_id", None) or node.id) + + def is_bridge_namespace_deleted(self, namespace: str) -> bool: + return str(namespace) in self._deleted_node_namespaces + + @staticmethod + def node_config_signature(node: Node) -> str: + values = ( + node.connection_type, + node.address, + node.port, + node.api_port, + node.server_ca, + node.api_key, + node.name, + node.default_timeout, + node.internal_timeout, + node.proxy_url, + node.usage_coefficient, + ) + return hashlib.sha256(repr(values).encode()).hexdigest() + + async def runtime_matches(self, node: Node) -> bool: + async with self._lock.reader_lock: + runtime = self._nodes.get(node.id) + return runtime is not None and self.runtime_config_matches(runtime, node) + + @classmethod + def runtime_config_matches(cls, runtime: PasarGuardNode, node: Node) -> bool: + return bool( + str(runtime.node_id) == cls.bridge_namespace(node) + and getattr(runtime, "_extra", {}).get("config_signature") + == cls.node_config_signature(node) + ) + + @property + def uses_shared_revocation_store(self) -> bool: + return self._uses_shared_revocation_store + + async def _shutdown_node(self, node: PasarGuardNode | None, *, remote_stop: bool = True) -> bool: if node is None: - return + return True try: await node.set_health(Health.INVALID) - if remote_stop: - await node.stop() except Exception: pass + try: + if remote_stop: + await node.stop() + else: + # Another worker already owns the remote Stop. This local + # controller must still cancel and await its background sync + # and stats tasks before shared state can be purged. + await node.disconnect() + return True + except Exception as exc: + self.logger.error("Failed to quiesce retiring node runtime: %s", exc) + return False + + async def _finish_retiring_node( + self, + node_id: int, + node: PasarGuardNode | None, + *, + remote_stop: bool = True, + ) -> None: + if node is None or not await self._shutdown_node(node, remote_stop=remote_stop): + return + async with self._lock.writer_lock: + retiring = self._retiring_nodes.get(node_id) + if retiring is None: + return + self._retiring_nodes[node_id] = [candidate for candidate in retiring if candidate is not node] + if not self._retiring_nodes[node_id]: + self._retiring_nodes.pop(node_id, None) async def update_node(self, node: Node) -> PasarGuardNode: await ensure_bridge_memory() - + namespace = self.bridge_namespace(node) + if namespace in self._deleted_node_namespaces: + await self.remove_node( + node.id, + remote_stop=False, + expected_bridge_namespace=namespace, + ) + raise NodeAPIError(410, f"node {node.id} incarnation is permanently deleted") + _, coordinator, _ = get_bridge_memory() + if self._uses_shared_revocation_store: + if coordinator is None: + raise NodeAPIError(503, "shared node bridge memory is unavailable") + if await coordinator.is_deleted(namespace): + # A missed broadcast must not leave a stale local controller + # capable of attaching to or restarting a deleted incarnation. + await self.remove_node( + node.id, + remote_stop=False, + expected_bridge_namespace=namespace, + ) + raise NodeAPIError(410, f"node {node.id} incarnation is permanently deleted") + # Validate shared-memory availability and construct the replacement + # before mutating the active/retiring topology. A fail-closed 503 must + # leave the currently working runtime untouched. + new_node = create_node(**self._create_node_kwargs(node)) + transition = asyncio.create_task(self._replace_runtime(node.id, new_node, coordinator, namespace)) + self._runtime_transition_tasks.add(transition) + + def _transition_done(task: asyncio.Task) -> None: + self._runtime_transition_tasks.discard(task) + if task.cancelled(): + return + # Retrieve failures when a cancelled caller no longer awaits us. + task.exception() + + transition.add_done_callback(_transition_done) + return await asyncio.shield(transition) + + async def _replace_runtime( + self, + node_id: int, + new_node: PasarGuardNode, + coordinator, + namespace: str, + ) -> PasarGuardNode: + transition_lock = self._runtime_transition_locks.setdefault(node_id, asyncio.Lock()) + async with transition_lock: + return await self._replace_runtime_locked(node_id, new_node, coordinator, namespace) + + async def _replace_runtime_locked( + self, + node_id: int, + new_node: PasarGuardNode, + coordinator, + namespace: str, + ) -> PasarGuardNode: + """Quiesce every old runtime before making a replacement active.""" + deleted_after_shutdown = False async with self._lock.writer_lock: - old_node: PasarGuardNode | None = self._nodes.pop(node.id, None) - - new_node = create_node(**self._create_node_kwargs(node)) - - self._nodes[node.id] = new_node - self._user_sync_locks.setdefault(node.id, asyncio.Lock()) + if self.is_bridge_namespace_deleted(namespace) or ( + coordinator is not None and await coordinator.is_deleted(namespace) + ): + deleted = True + retiring: list[PasarGuardNode] = [] + else: + deleted = False + if node_id in self._removing_node_ids and node_id not in self._replacing_node_ids: + raise NodeAPIError(409, f"node {node_id} removal is in progress") + self._replacing_node_ids.add(node_id) + self._removing_node_ids.add(node_id) + old_node = self._nodes.pop(node_id, None) + if old_node is not None: + current = self._retiring_nodes.setdefault(node_id, []) + if old_node not in current: + current.append(old_node) + retiring = list(self._retiring_nodes.get(node_id, ())) + + if deleted: + await self.remove_node( + node_id, + remote_stop=False, + expected_bridge_namespace=namespace, + ) + raise NodeAPIError(410, f"node {node_id} incarnation is permanently deleted") - # Stop the old node after releasing the lock. - await self._shutdown_node(old_node) + results = await asyncio.gather( + *(self._shutdown_node(runtime) for runtime in retiring), + return_exceptions=True, + ) + succeeded = [runtime for runtime, result in zip(retiring, results) if result is True] + failed = [runtime for runtime, result in zip(retiring, results) if result is not True] + async with self._lock.writer_lock: + current = self._retiring_nodes.get(node_id, []) + current = [runtime for runtime in current if runtime not in succeeded] + if current: + self._retiring_nodes[node_id] = current + else: + self._retiring_nodes.pop(node_id, None) + if failed or current: + # Keep the replacement barrier retryable. A later update will + # drain all accumulated retirees before installing anything. + raise NodeAPIError(503, f"cannot confirm old node {node_id} runtime shutdown") + if self.is_bridge_namespace_deleted(namespace) or ( + coordinator is not None and await coordinator.is_deleted(namespace) + ): + self._replacing_node_ids.discard(node_id) + self._removing_node_ids.discard(node_id) + deleted_after_shutdown = True + else: + self._nodes[node_id] = new_node + self._user_sync_locks.setdefault(node_id, asyncio.Lock()) + self._replacing_node_ids.discard(node_id) + self._removing_node_ids.discard(node_id) + + if deleted_after_shutdown: + try: + await new_node.disconnect() + except Exception: + pass + raise NodeAPIError(410, f"node {node_id} incarnation is permanently deleted") return new_node - async def remove_node(self, id: int, *, remote_stop: bool = True) -> None: + async def remove_node( + self, + id: int, + *, + remote_stop: bool = True, + expected_bridge_namespace: str | None = None, + permanent_delete: bool = False, + ) -> None: + await ensure_bridge_memory() + _, coordinator, _ = get_bridge_memory() async with self._lock.writer_lock: + active_node = self._nodes.get(id) + if ( + expected_bridge_namespace is not None + and active_node is not None + and str(active_node.node_id) != str(expected_bridge_namespace) + ): + # A delayed cross-worker message for a deleted row must not + # remove a replacement row that reused the public numeric id. + return + namespace = str( + expected_bridge_namespace + or (active_node.node_id if active_node is not None else "") + ) + if permanent_delete and namespace: + self._deleted_node_namespaces.add(namespace) + if self._uses_shared_revocation_store and permanent_delete: + if coordinator is None: + raise NodeAPIError(503, "shared node lifecycle memory is unavailable") + if not namespace: + raise NodeAPIError(503, f"cannot durably fence node {id} without its Bridge namespace") + # This durable marker is the correctness mechanism. The later + # broadcast is only a prompt for sibling-local cleanup. + await coordinator.mark_deleted(namespace) + self._removing_node_ids.add(id) old_node: PasarGuardNode | None = self._nodes.pop(id, None) - self._user_sync_locks.pop(id, None) - - # Do cleanup without holding the lock to avoid slow delete operations. - asyncio.create_task(self._shutdown_node(old_node, remote_stop=remote_stop)) + if old_node is not None: + self._retiring_nodes.setdefault(id, []).append(old_node) + retiring = list(self._retiring_nodes.get(id, ())) + + # Stop outside the topology lock, but do not return success (and do + # not let callers purge shared fencing state) until every runtime is + # confirmed quiescent. A failed/ambiguous Stop remains visible in + # _retiring_nodes so revocation preflight continues to fail closed. + results = await asyncio.gather( + *(self._shutdown_node(node, remote_stop=remote_stop) for node in retiring), + return_exceptions=True, + ) + failed = [node for node, result in zip(retiring, results) if result is not True] + succeeded = [node for node, result in zip(retiring, results) if result is True] + async with self._lock.writer_lock: + current = self._retiring_nodes.get(id, []) + if succeeded: + current = [node for node in current if node not in succeeded] + if current: + self._retiring_nodes[id] = current + else: + self._retiring_nodes.pop(id, None) + self._user_sync_locks.pop(id, None) + self._removing_node_ids.discard(id) + if failed: + raise NodeAPIError(503, f"cannot confirm node {id} runtime shutdown") async def get_node(self, id: int) -> PasarGuardNode | None: async with self._lock.reader_lock: return self._nodes.get(id, None) + async def get_lifecycle_recovery_node(self, id: int) -> PasarGuardNode | None: + """Return the active controller or a retained ambiguous retiree.""" + async with self._lock.reader_lock: + active = self._nodes.get(id) + if active is not None: + return active + retiring = self._retiring_nodes.get(id, ()) + return retiring[-1] if retiring else None + + async def get_bridge_namespace(self, id: int) -> str | None: + async with self._lock.reader_lock: + node = self._nodes.get(id) + return None if node is None else str(node.node_id) + async def get_nodes(self) -> dict[int, PasarGuardNode]: async with self._lock.reader_lock: return self._nodes @@ -116,13 +389,22 @@ async def _snapshot_nodes(self) -> list[PasarGuardNode]: async def _snapshot_node_items(self) -> list[tuple[int, PasarGuardNode]]: async with self._lock.reader_lock: - return list(self._nodes.items()) + return [ + *self._nodes.items(), + *((node_id, node) for node_id, retiring in self._retiring_nodes.items() for node in retiring), + ] @staticmethod def _chunk_users(users: list[ProtoUser], size: int) -> list[list[ProtoUser]]: return [users[start : start + size] for start in range(0, len(users), size)] - async def _sync_user_batch_to_node(self, node: PasarGuardNode, batch: list[ProtoUser]) -> int: + async def _sync_user_batch_to_node( + self, + node: PasarGuardNode, + batch: list[ProtoUser], + *, + revocation_id: str | None = None, + ) -> int: users_to_sync = batch supports_chunked = True supports_chunked_check = getattr(node, "_supports_chunked_sync", None) @@ -130,53 +412,545 @@ async def _sync_user_batch_to_node(self, node: PasarGuardNode, batch: list[Proto supports_chunked, _ = await supports_chunked_check() if supports_chunked: - users_to_sync = await node.sync_users_chunked( - batch, - chunk_size=len(batch), - flush_pending=False, - ) + kwargs = { + "chunk_size": len(batch), + "flush_pending": False, + } + if revocation_id is not None: + # Passing this only for the new revocation path preserves + # compatibility with older custom bridge subclasses used by + # ordinary background sync. + kwargs["revocation_id"] = revocation_id + users_to_sync = await node.sync_users_chunked(batch, **kwargs) if not users_to_sync: return 0 + # A revocation permit covers the public direct request. Falling + # through to the bridge's private batch helper would bypass that + # permit and reopen the stale-update race. + if revocation_id is not None: + return len(users_to_sync) + + elif revocation_id is not None: + # SyncUsers is a full replacement on the Go node. Sending a + # one-user removal/restoration through it would erase every other + # account. Permanent revocation therefore requires the guaranteed + # partial chunked endpoint and fails closed on legacy nodes. + raise NodeRevocationError("node does not support partial chunked sync required for user revocation") + sync_batch_users = getattr(node, "_sync_batch_users", None) if callable(sync_batch_users): users_to_sync = await sync_batch_users(users_to_sync) return len(users_to_sync) - async def _sync_users_to_node(self, node_id: int, node: PasarGuardNode, users: list[ProtoUser]): + @staticmethod + def _user_key(user: ProtoUser) -> str: + # PasarGuardNodeBridge's protobuf has no panel database `id` field. + # The serializer intentionally stores that stable ID in `email`. + return user.email + + def _without_deleted_users(self, users: list[ProtoUser]) -> list[ProtoUser]: + return [user for user in users if self._user_key(user) not in self._deleted_user_keys] + + async def _sync_users_to_node( + self, + node_id: int, + node: PasarGuardNode, + users: list[ProtoUser], + *, + allow_deleted: bool = False, + revocation_id: str | None = None, + allowed_user_keys: frozenset[str] | None = None, + ): batch_size = max(1, nats_settings.node_update_users_batch_size) lock = self._user_sync_locks.setdefault(node_id, asyncio.Lock()) failed_count = 0 async with lock: for batch in self._chunk_users(users, batch_size): - failed_count += await self._sync_user_batch_to_node(node, batch) + current_batch = ( + batch if allow_deleted or self._uses_shared_revocation_store else self._without_deleted_users(batch) + ) + if allowed_user_keys is not None: + current_batch = [user for user in current_batch if self._user_key(user) in allowed_user_keys] + if current_batch: + failed_count += await self._sync_user_batch_to_node( + node, + current_batch, + revocation_id=revocation_id, + ) if failed_count: raise RuntimeError(f"failed to sync {failed_count}/{len(users)} users to node {node_id}") - async def _update_users(self, users: list[ProtoUser]): - nodes = await self._snapshot_node_items() + async def _update_users( + self, + users: list[ProtoUser], + *, + raise_on_failure: bool = False, + allow_deleted: bool = False, + revocation_id: str | None = None, + node_items: list[tuple[int, PasarGuardNode]] | None = None, + node_user_keys: dict[int, frozenset[str]] | None = None, + ): + if not allow_deleted and not self._uses_shared_revocation_store: + users = self._without_deleted_users(users) + if not users: + return + + nodes = node_items if node_items is not None else await self._snapshot_node_items() if not nodes: + # There are no active runtime nodes that could admit this user. return results = await asyncio.gather( - *(self._sync_users_to_node(node_id, node, users) for node_id, node in nodes), return_exceptions=True + *( + self._sync_users_to_node( + node_id, + node, + users, + allow_deleted=allow_deleted, + revocation_id=revocation_id, + allowed_user_keys=None if node_user_keys is None else node_user_keys[node_id], + ) + for node_id, node in nodes + ), + return_exceptions=True, ) - for result in results: - if isinstance(result, Exception): - self.logger.error("Failed to sync users to one of the nodes: %s", result) + # ``asyncio.CancelledError`` is a ``BaseException``. With + # ``return_exceptions=True`` a node task which cancels itself is + # returned as a result, rather than cancelling this caller. Treat that + # as an unconfirmed revocation: silently accepting it could let the + # database delete commit while one node still has the user. + failures = [ + (node_id, result) for (node_id, _), result in zip(nodes, results) if isinstance(result, BaseException) + ] + for node_id, failure in failures: + self.logger.error("Failed to sync users to node %s: %s", node_id, failure) + if failures and raise_on_failure: + failed_node_ids = ", ".join(str(node_id) for node_id, _ in failures) + raise NodeRevocationError( + f"failed to sync users to {len(failures)}/{len(nodes)} nodes (node ids: {failed_node_ids})" + ) from failures[0][1] async def update_users(self, users: list[ProtoUser]) -> None: asyncio.create_task(self._update_users(users)) + async def update_users_and_wait(self, users: list[ProtoUser]) -> None: + """Synchronize a removal-sensitive batch and surface node failures.""" + await self._update_users(users, raise_on_failure=True) + + async def wait_for_user_revocations(self) -> None: + """Delay node startup until every provisional delete is resolved.""" + await self._revocations_idle.wait() + + def filter_permanently_deleted_users(self, users: list[ProtoUser]) -> list[ProtoUser]: + """Keep an old transaction snapshot from reintroducing tombstoned users.""" + return [user for user in users if self._user_key(user) not in self._deleted_user_keys] + + def _acquire_deletion_fences(self, user_keys: set[str], revocation_id: str) -> None: + self._revocations_idle.clear() + for user_key in user_keys: + # An absent owner entry for an already deleted key denotes a + # committed tombstone. Do not make it provisional again if a stale + # or duplicate revoke arrives after the database commit. + if user_key in self._deleted_user_keys and user_key not in self._deletion_fence_owners: + continue + self._deletion_fence_owners.setdefault(user_key, set()).add(revocation_id) + self._deleted_user_keys.update(user_keys) + + def _release_deletion_fences(self, user_keys: set[str], revocation_id: str) -> None: + for user_key in user_keys: + owners = self._deletion_fence_owners.get(user_key) + if owners is None: + continue + owners.discard(revocation_id) + if not owners: + self._deletion_fence_owners.pop(user_key, None) + self._deleted_user_keys.discard(user_key) + if not self._deletion_fence_owners: + self._revocations_idle.set() + + def _finalize_deletion_fences(self, user_keys: set[str], revocation_id: str) -> None: + for user_key in user_keys: + owners = self._deletion_fence_owners.get(user_key) + if owners is None or revocation_id not in owners: + continue + # Once any overlapping delete commits, the tombstone is permanent. + # Drop every per-operation owner so successful deletes do not leak + # revocation IDs and a later abort cannot undo the committed fence. + self._deletion_fence_owners.pop(user_key, None) + if not self._deletion_fence_owners: + self._revocations_idle.set() + + def _record_revocation_nodes( + self, + revocation_id: str, + nodes: list[tuple[int, PasarGuardNode, frozenset[str]]], + ) -> None: + """Merge separately transported chunks without losing earlier keys.""" + current = self._revocation_nodes.setdefault(revocation_id, []) + by_identity = {(node_id, id(node)): index for index, (node_id, node, _) in enumerate(current)} + for node_id, node, active_user_keys in nodes: + if not active_user_keys: + continue + identity = (node_id, id(node)) + index = by_identity.get(identity) + if index is None: + by_identity[identity] = len(current) + current.append((node_id, node, active_user_keys)) + continue + old_node_id, old_node, old_keys = current[index] + current[index] = (old_node_id, old_node, old_keys | active_user_keys) + + def _peek_revocation_nodes( + self, + revocation_id: str, + user_keys: set[str], + ) -> list[tuple[int, PasarGuardNode, frozenset[str]]] | None: + """Read this RPC chunk without losing retry state on close failure.""" + current = self._revocation_nodes.get(revocation_id) + if current is None: + return None + selected: list[tuple[int, PasarGuardNode, frozenset[str]]] = [] + for node_id, node, active_user_keys in current: + selected_keys = active_user_keys & user_keys + if selected_keys: + selected.append((node_id, node, frozenset(selected_keys))) + return selected + + def _discard_revocation_nodes(self, revocation_id: str, user_keys: set[str]) -> None: + """Forget a chunk only after every bridge close has been acknowledged.""" + current = self._revocation_nodes.get(revocation_id) + if current is None: + return + remaining: list[tuple[int, PasarGuardNode, frozenset[str]]] = [] + for node_id, node, active_user_keys in current: + remaining_keys = active_user_keys - user_keys + if remaining_keys: + remaining.append((node_id, node, frozenset(remaining_keys))) + if remaining: + self._revocation_nodes[revocation_id] = remaining + else: + self._revocation_nodes.pop(revocation_id, None) + + def _restorable_users(self, users: list[ProtoUser], revocation_id: str) -> list[ProtoUser]: + """Return users which are still owned exclusively by this failed delete.""" + restorable = [] + for user in users: + user_key = self._user_key(user) + owners = self._deletion_fence_owners.get(user_key) + if owners == {revocation_id}: + restorable.append(user) + return restorable + + async def _restore_users_to_node( + self, + node_id: int, + node: PasarGuardNode, + users: list[ProtoUser], + revocation_id: str, + active_user_keys: frozenset[str], + ) -> bool: + """Restore authoritative DB state and leave an ordered retry behind.""" + lock = self._user_sync_locks.setdefault(node_id, asyncio.Lock()) + async with lock: + candidate_users = ( + users if self._uses_shared_revocation_store else self._restorable_users(users, revocation_id) + ) + current_users = [user for user in candidate_users if self._user_key(user) in active_user_keys] + direct_succeeded = True + for batch in self._chunk_users(current_users, max(1, nats_settings.node_update_users_batch_size)): + try: + if await self._sync_user_batch_to_node( + node, + batch, + revocation_id=revocation_id, + ): + direct_succeeded = False + except BaseException as exc: + direct_succeeded = False + self.logger.error("Failed to restore users immediately on node %s: %s", node_id, exc) + + abort = getattr(node, "abort_user_revocation", None) + if not callable(abort): + self.logger.error("Node %s bridge does not support revocation abort", node_id) + return False + try: + await abort(sorted(active_user_keys), revocation_id) + except BaseException as exc: + self.logger.error("Failed to abort user revocation fence on node %s: %s", node_id, exc) + return False + + queue_succeeded = False + if current_users: + try: + # Once the distributed fence is released, leave the + # authoritative original state in the retry queue as well. + await node.update_users(current_users) + queue_succeeded = True + except BaseException as exc: + self.logger.error("Failed to queue user restoration for node %s: %s", node_id, exc) + else: + queue_succeeded = True + + return direct_succeeded or queue_succeeded + + async def _restore_failed_revocation( + self, + nodes: list[tuple[int, PasarGuardNode, frozenset[str]]], + users: list[ProtoUser], + revocation_id: str, + ) -> list[int]: + if not users or not nodes: + return [] + results = await asyncio.gather( + *( + self._restore_users_to_node(node_id, node, users, revocation_id, active_user_keys) + for node_id, node, active_user_keys in nodes + ), + return_exceptions=True, + ) + return [ + node_id + for (node_id, _, _), result in zip(nodes, results) + if isinstance(result, BaseException) or result is not True + ] + + @staticmethod + async def _unavailable_revocation_nodes(nodes: list[tuple[int, PasarGuardNode]]) -> list[int]: + health = await asyncio.gather(*(node.get_health() for _, node in nodes), return_exceptions=True) + return [ + node_id + for (node_id, node), result in zip(nodes, health) + if ( + isinstance(result, BaseException) + or result != Health.HEALTHY + or not callable(getattr(node, "begin_user_revocation", None)) + or not callable(getattr(node, "abort_user_revocation", None)) + or not callable(getattr(node, "finalize_user_revocation", None)) + ) + ] + + @staticmethod + def _require_complete_revocation_topology( + nodes: list[tuple[int, PasarGuardNode]], + expected_node_ids: set[int] | None, + ) -> None: + if expected_node_ids is None: + return + missing_node_ids = expected_node_ids - {node_id for node_id, _ in nodes} + if missing_node_ids: + node_ids = ", ".join(str(node_id) for node_id in sorted(missing_node_ids)) + raise NodeRevocationError( + f"runtime topology is incomplete for user revocation (missing node ids: {node_ids})" + ) + + async def _begin_node_revocations( + self, + nodes: list[tuple[int, PasarGuardNode]], + user_keys: set[str], + revocation_id: str, + ) -> list[tuple[int, PasarGuardNode, frozenset[str]]]: + prepared: list[tuple[int, PasarGuardNode, frozenset[str]]] = [] + for node_id, node in nodes: + begin = node.begin_user_revocation + try: + result = await begin(sorted(user_keys), revocation_id) + active_user_keys = frozenset(result.active_user_keys) + finalized_user_keys = frozenset(result.finalized_user_keys) + if active_user_keys & finalized_user_keys or active_user_keys | finalized_user_keys != user_keys: + raise NodeRevocationError("node bridge returned an invalid user revocation result") + except BaseException: + nodes_to_unwind = [ + *((prepared_node_id, prepared_node) for prepared_node_id, prepared_node, _ in prepared), + (node_id, node), + ] + + async def unwind(nodes_to_unwind=nodes_to_unwind) -> None: + # begin may persist its fences before waiting for active + # leases, so the current node is ambiguous as well. + for _, prepared_node in reversed(nodes_to_unwind): + try: + await prepared_node.abort_user_revocation(sorted(user_keys), revocation_id) + except BaseException as abort_exc: + self.logger.error("Failed to unwind prepared user revocation: %s", abort_exc) + + unwind_task = asyncio.create_task(unwind()) + try: + await asyncio.shield(unwind_task) + except asyncio.CancelledError: + await unwind_task + raise + raise + prepared.append((node_id, node, active_user_keys)) + return prepared + + @classmethod + def _resolve_revocation_id(cls, users: list[ProtoUser], revocation_id: str | None) -> str: + if revocation_id: + return revocation_id + # Rolling upgrades can pair legacy revoke/abort calls only through data + # visible to both requests. A stable digest avoids an unobservable UUID. + user_keys = "\0".join(sorted({cls._user_key(user) for user in users})) + return f"legacy:{hashlib.sha256(user_keys.encode()).hexdigest()}" + + async def revoke_users_and_wait( + self, + users: list[ProtoUser], + revocation_id: str | None = None, + restore_users: list[ProtoUser] | None = None, + *, + expected_node_ids: set[int] | None = None, + ) -> str: + """Fence permanent deletions before removing users from every node.""" + revocation_id = self._resolve_revocation_id(users, revocation_id) + user_keys = {self._user_key(user) for user in users} + restore_users = restore_users or list(users) + restore_user_keys = {self._user_key(user) for user in restore_users} + if restore_user_keys != user_keys: + raise NodeRevocationError("removal and restoration users do not match") + + if not self._uses_shared_revocation_store: + self._acquire_deletion_fences(user_keys, revocation_id) + nodes: list[tuple[int, PasarGuardNode]] = [] + prepared_nodes: list[tuple[int, PasarGuardNode, frozenset[str]]] = [] + revocation_started = False + try: + nodes = await self._snapshot_node_items() + self._require_complete_revocation_topology(nodes, expected_node_ids) + unavailable_node_ids = await self._unavailable_revocation_nodes(nodes) + if unavailable_node_ids: + node_ids = ", ".join(str(node_id) for node_id in unavailable_node_ids) + raise NodeRevocationError(f"runtime nodes are not ready for user revocation (node ids: {node_ids})") + + prepared_nodes = await self._begin_node_revocations(nodes, user_keys, revocation_id) + revocation_started = True + await self._update_users( + users, + raise_on_failure=True, + allow_deleted=True, + revocation_id=revocation_id, + node_items=[(node_id, node) for node_id, node, _ in prepared_nodes], + node_user_keys={node_id: active_user_keys for node_id, _, active_user_keys in prepared_nodes}, + ) + except BaseException as exc: + # The database row is kept when revocation is not confirmed, so + # every possibly-mutated node must converge back to the original + # authoritative state before normal updates are admitted again. + restoration_failures: list[int] = [] + try: + if revocation_started: + task = asyncio.create_task( + self._restore_failed_revocation(prepared_nodes, restore_users, revocation_id) + ) + try: + restoration_failures = await asyncio.shield(task) + except asyncio.CancelledError: + # A second cancellation must not let the caller tear + # down the local fence while compensation is still + # mutating the nodes. Finish it before propagating the + # original cancellation. + restoration_failures = await task + raise + elif prepared_nodes: + await asyncio.gather( + *( + node.abort_user_revocation(sorted(active_user_keys), revocation_id) + for _, node, active_user_keys in prepared_nodes + if active_user_keys + ), + return_exceptions=True, + ) + finally: + if not self._uses_shared_revocation_store: + self._release_deletion_fences(user_keys, revocation_id) + + if restoration_failures and not isinstance(exc, asyncio.CancelledError): + failed_node_ids = ", ".join(str(node_id) for node_id in restoration_failures) + raise NodeRevocationError( + f"user revocation failed and restoration was not accepted by nodes: {failed_node_ids}" + ) from exc + if not isinstance(exc, (NodeRevocationError, asyncio.CancelledError)): + raise NodeRevocationError(f"cannot confirm user revocation: {exc}") from exc + raise + if not self._uses_shared_revocation_store: + self._record_revocation_nodes(revocation_id, prepared_nodes) + return revocation_id + + async def abort_user_revocations( + self, + users: list[ProtoUser], + revocation_id: str | None = None, + restore_users: list[ProtoUser] | None = None, + *, + expected_node_ids: set[int] | None = None, + ) -> None: + """Restore users and release provisional fences after a DB rollback.""" + revocation_id = self._resolve_revocation_id(users, revocation_id) + user_keys = {self._user_key(user) for user in users} + nodes = self._peek_revocation_nodes(revocation_id, user_keys) + if not nodes: + topology = await self._snapshot_node_items() + self._require_complete_revocation_topology(topology, expected_node_ids) + nodes = [(node_id, node, frozenset(user_keys)) for node_id, node in topology] + failures = await self._restore_failed_revocation(nodes, restore_users or [], revocation_id) + if failures: + failed_node_ids = ", ".join(str(node_id) for node_id in failures) + raise NodeRevocationError(f"failed to restore users after database rollback on nodes: {failed_node_ids}") + if not self._uses_shared_revocation_store: + self._discard_revocation_nodes(revocation_id, user_keys) + self._release_deletion_fences(user_keys, revocation_id) + + async def finalize_user_revocations( + self, + users: list[ProtoUser], + revocation_id: str | None = None, + *, + expected_node_ids: set[int] | None = None, + ) -> None: + """Commit deletion tombstones and discard per-operation ownership.""" + revocation_id = self._resolve_revocation_id(users, revocation_id) + user_keys = {self._user_key(user) for user in users} + nodes = self._peek_revocation_nodes(revocation_id, user_keys) + if not nodes: + topology = await self._snapshot_node_items() + self._require_complete_revocation_topology(topology, expected_node_ids) + nodes = [(node_id, node, frozenset(user_keys)) for node_id, node in topology] + results = await asyncio.gather( + *( + node.finalize_user_revocation(sorted(active_user_keys), revocation_id) + for _, node, active_user_keys in nodes + if active_user_keys + if callable(getattr(node, "finalize_user_revocation", None)) + ), + return_exceptions=True, + ) + failures = [result for result in results if isinstance(result, BaseException)] + if failures: + raise NodeRevocationError(f"failed to finalize revocation fences on {len(failures)} nodes") from failures[0] + if not self._uses_shared_revocation_store: + self._discard_revocation_nodes(revocation_id, user_keys) + self._finalize_deletion_fences(user_keys, revocation_id) + async def update_user(self, user: ProtoUser) -> None: - nodes = await self._snapshot_nodes() + user_key = self._user_key(user) + if not self._uses_shared_revocation_store and user_key in self._deleted_user_keys: + return + + nodes = await self._snapshot_node_items() if not nodes: return - results = await asyncio.gather(*(node.update_user(user) for node in nodes), return_exceptions=True) + async def sync_one(node_id: int, node: PasarGuardNode) -> None: + lock = self._user_sync_locks.setdefault(node_id, asyncio.Lock()) + async with lock: + if not self._uses_shared_revocation_store and user_key in self._deleted_user_keys: + return + await node.update_user(user) + + results = await asyncio.gather(*(sync_one(node_id, node) for node_id, node in nodes), return_exceptions=True) for result in results: if isinstance(result, Exception): raise result diff --git a/app/node/errors.py b/app/node/errors.py new file mode 100644 index 000000000..988f77ea1 --- /dev/null +++ b/app/node/errors.py @@ -0,0 +1,4 @@ +class NodeRevocationError(RuntimeError): + """A user removal was not acknowledged by every configured runtime node.""" + + code = 503 diff --git a/app/node/manager_sync.py b/app/node/manager_sync.py index f7a949b94..0597de8ee 100644 --- a/app/node/manager_sync.py +++ b/app/node/manager_sync.py @@ -8,17 +8,20 @@ from app.nats.message import MessageTopic from app.nats.router import router from app.node import node_manager -from app.node.nats_memory import WORKER_ID, clear_bridge_memory_for_node +from app.node.nats_memory import WORKER_ID from app.utils.logger import get_logger logger = get_logger("node-manager-sync") -async def publish_node_sync(action: str, node_id: int) -> None: +async def publish_node_sync(action: str, node_id: int, bridge_id: str | None = None) -> None: try: + data = {"action": action, "node_id": node_id, "origin": WORKER_ID} + if bridge_id is not None: + data["bridge_id"] = bridge_id await router.publish( MessageTopic.NODE, - {"action": action, "node_id": node_id, "origin": WORKER_ID}, + data, ) except Exception as exc: logger.warning("Failed to publish node sync action=%s node_id=%s: %s", action, node_id, exc) @@ -33,14 +36,22 @@ async def handle_node_message(data: dict) -> None: if not action or node_id is None: return node_id = int(node_id) + announced_bridge_id = data.get("bridge_id") if action == "remove": - await node_manager.remove_node(node_id, remote_stop=False) - await clear_bridge_memory_for_node(node_id) + if announced_bridge_id is None: + logger.warning("Ignoring unsafe legacy node remove without bridge_id node_id=%s", node_id) + return + remove_kwargs = {"remote_stop": False, "permanent_delete": True} + remove_kwargs["expected_bridge_namespace"] = str(announced_bridge_id) + await node_manager.remove_node(node_id, **remove_kwargs) return if action == "disconnect": - await node_manager.remove_node(node_id, remote_stop=False) + remove_kwargs = {"remote_stop": False} + if announced_bridge_id is not None: + remove_kwargs["expected_bridge_namespace"] = str(announced_bridge_id) + await node_manager.remove_node(node_id, **remove_kwargs) return if action == "upsert": @@ -48,7 +59,11 @@ async def handle_node_message(data: dict) -> None: db_node = await get_node_by_id(db, node_id, load_usage_logs=False) if db_node is None: return - await node_manager.update_node(db_node) + if announced_bridge_id is not None and str(db_node.bridge_id) != str(announced_bridge_id): + logger.warning("Ignoring stale node sync action=%s node_id=%s", action, node_id) + return + if not await node_manager.runtime_matches(db_node): + await node_manager.update_node(db_node) return if action == "connect": @@ -59,17 +74,26 @@ async def handle_node_message(data: dict) -> None: db_node = await get_node_by_id(db, node_id, load_usage_logs=False) if db_node is None or db_node.status in (NodeStatus.disabled, NodeStatus.limited): return + if announced_bridge_id is not None and str(db_node.bridge_id) != str(announced_bridge_id): + logger.warning("Ignoring stale node connect node_id=%s", node_id) + return + # Match the local startup ordering. Register this worker's runtime + # before locking the authoritative user snapshot, then keep those + # row locks until the epoch-fenced full apply completes. A delete + # either sees this runtime in topology or commits before snapshot. + try: + if not await node_manager.runtime_matches(db_node): + await node_manager.update_node(db_node) + except Exception: + logger.exception("Node sync connect runtime registration failed for node_id=%s", node_id) + return core_id = db_node.core_config_id or 1 - cores_by_id, users_by_core = await NodeOperation._get_core_users_map(db, {core_id}) + cores_by_id, users_by_core, authoritative_user_keys = await NodeOperation._get_core_users_map( + db, {core_id} + ) core = cores_by_id.get(core_id) users = users_by_core.get(core_id, []) - - try: - await node_manager.update_node(db_node) - except Exception: - logger.exception("Node sync connect update_node failed for node_id=%s", node_id) - return - await NodeOperation.connect_node(db_node, core, users) + await NodeOperation.connect_node(db_node, core, users, authoritative_user_keys) return logger.warning("Unknown node sync action: %s", action) diff --git a/app/node/nats_memory.py b/app/node/nats_memory.py index 61e5aafe0..b2a1b92cf 100644 --- a/app/node/nats_memory.py +++ b/app/node/nats_memory.py @@ -9,6 +9,7 @@ import json import os import time +from contextvars import ContextVar, Token from typing import Any from uuid import uuid4 @@ -22,16 +23,23 @@ LifecycleOperation, LifecycleStatus, NodeLifecycleState, + StartupUserSyncLease, + UserRevocationConflictError, + UserRevocationResult, + UserSyncLease, + UserSyncLeaseLostError, ) from app.nats import needs_shared_bridge_memory from app.nats.client import create_nats_client, get_jetstream_context, get_or_create_kv_bucket -from app.nats.kv_cas import CasKv, kv_cas_json, kv_get_json, kv_list_keys, kv_put_json +from app.nats.kv_cas import CasKv, kv_cas_json, kv_get_json, kv_list_keys from app.utils.logger import get_logger from config import nats_settings logger = get_logger("node-nats-memory") +_REVOCATION_METADATA_LEASE_SECONDS = 300.0 + WORKER_ID = f"{os.getpid()}:{uuid4().hex[:8]}" # Stay under default NATS max_payload (1MiB) with headroom for JSON framing. _MAX_USER_SYNC_VALUE_BYTES = min(900_000, nats_settings.node_command_max_payload_bytes) @@ -92,10 +100,27 @@ def _state_to_dict(state: NodeLifecycleState) -> dict[str, Any]: class NatsUserSyncStore: - """Per-user pending / per-token claimed keys so each KV value stays payload-safe.""" + """Generation-fenced user queue shared by every Panel/node worker.""" def __init__(self, kv: CasKv): self._kv = kv + self._authoritative_reconciliation_membership: ContextVar[ + tuple[str, str, frozenset[str]] | None + ] = ContextVar(f"node-reconciliation-membership-{id(self)}", default=None) + + async def set_authoritative_reconciliation_membership( + self, node_id: str, worker_id: str, user_keys: list[str] + ) -> Token[tuple[str, str, frozenset[str]] | None]: + """Supply row-locked global DB membership for one reconciliation scope.""" + return self._authoritative_reconciliation_membership.set( + (node_id, worker_id, frozenset(user_keys)) + ) + + def reset_authoritative_reconciliation_membership( + self, token: Token[tuple[str, str, frozenset[str]] | None] + ) -> None: + """Synchronously end a DB-locked scope, including cancellation paths.""" + self._authoritative_reconciliation_membership.reset(token) def _pending_prefix(self, node_id: str) -> str: return f"p.{node_id}." @@ -109,6 +134,168 @@ def _pending_key(self, node_id: str, email: str) -> str: def _claimed_key(self, node_id: str, token: str) -> str: return f"{self._claimed_prefix(node_id)}{_digest(token)}" + def _barrier_prefix(self, node_id: str) -> str: + return f"b.{node_id}." + + def _barrier_key(self, node_id: str, user_key: str) -> str: + return f"{self._barrier_prefix(node_id)}{_digest(user_key)}" + + def _execution_prefix(self, node_id: str) -> str: + return f"x.{node_id}." + + def _execution_key(self, node_id: str, token: str) -> str: + return f"{self._execution_prefix(node_id)}{_digest(token)}" + + def _revocation_lock_key(self, node_id: str) -> str: + return f"m.{node_id}" + + def _epoch_key(self, node_id: str) -> str: + return f"e.{node_id}" + + async def _next_user_sync_epoch(self, node_id: str) -> int: + key = self._epoch_key(node_id) + for _ in range(64): + doc, rev = await kv_get_json(self._kv, key) + epoch = int((doc or {}).get("epoch", 0)) + 1 + if await kv_cas_json(self._kv, key, {"epoch": epoch}, rev): + return epoch + raise RuntimeError(f"failed to reserve user-sync epoch key={key} after CAS retries") + + async def advance_user_sync_epoch(self, node_id: str, minimum_epoch: int) -> None: + if minimum_epoch < 0: + raise ValueError("minimum_epoch must be non-negative") + key = self._epoch_key(node_id) + for _ in range(64): + doc, rev = await kv_get_json(self._kv, key) + current = int((doc or {}).get("epoch", 0)) + if current >= minimum_epoch: + return + if await kv_cas_json(self._kv, key, {"epoch": minimum_epoch}, rev): + return + raise RuntimeError(f"failed to advance user-sync epoch key={key} after CAS retries") + + @contextlib.asynccontextmanager + async def _revocation_lock(self, node_id: str, user_keys: list[str]): + """Serialize multi-key revocation transitions for one node. + + The lease is long relative to the bounded metadata transition and is + renewed while the owner is alive. Expiry remains fail-closed until the + periodic authoritative DB-locked recovery resolves it. + """ + key = self._revocation_lock_key(node_id) + token = str(uuid4()) + acquired = False + for _ in range(32): + now = time.time() + doc, rev = await kv_get_json(self._kv, key) + # Never steal an expired metadata lock. Its owner may merely be + # paused and can still resume. Only authoritative recovery, while + # holding the target DB row locks, may clear this evidence. + if doc is not None: + break + value = {"token": token, "expires_at": now + _REVOCATION_METADATA_LEASE_SECONDS} + if await kv_cas_json(self._kv, key, value, rev): + acquired = True + break + if not acquired: + raise UserRevocationConflictError(tuple(sorted(set(user_keys)))) + + lease_lost = asyncio.Event() + + async def _assert_owned() -> None: + if lease_lost.is_set(): + raise UserSyncLeaseLostError(f"revocation metadata lease was lost for node_id={node_id}") + doc, _ = await kv_get_json(self._kv, key) + if ( + doc is None + or doc.get("token") != token + or float(doc.get("expires_at", 0)) <= time.time() + ): + lease_lost.set() + raise UserSyncLeaseLostError(f"revocation metadata lease was lost for node_id={node_id}") + + async def _renew() -> None: + try: + while True: + await asyncio.sleep(_REVOCATION_METADATA_LEASE_SECONDS / 3) + for _ in range(32): + doc, rev = await kv_get_json(self._kv, key) + if doc is None or doc.get("token") != token: + raise UserSyncLeaseLostError( + f"revocation metadata lease was lost for node_id={node_id}" + ) + if float(doc.get("expires_at", 0)) <= time.time(): + raise UserSyncLeaseLostError( + f"revocation metadata lease expired for node_id={node_id}" + ) + doc["expires_at"] = time.time() + _REVOCATION_METADATA_LEASE_SECONDS + if await kv_cas_json(self._kv, key, doc, rev): + break + else: + raise UserSyncLeaseLostError( + f"revocation metadata lease heartbeat failed for node_id={node_id}" + ) + except asyncio.CancelledError: + raise + except BaseException: + lease_lost.set() + raise + + heartbeat = asyncio.create_task(_renew()) + body_failed = False + try: + yield _assert_owned + await _assert_owned() + except BaseException: + body_failed = True + raise + finally: + heartbeat.cancel() + heartbeat_error: BaseException | None = None + try: + await heartbeat + except asyncio.CancelledError: + pass + except BaseException as exc: + heartbeat_error = exc + + async def _release() -> None: + doc, rev = await kv_get_json(self._kv, key) + if doc is not None and doc.get("token") == token: + await self._kv.delete(key, last=rev) + + release = asyncio.create_task(_release()) + try: + await asyncio.shield(release) + except asyncio.CancelledError: + await release + raise + if heartbeat_error is not None and not body_failed: + raise UserSyncLeaseLostError( + f"revocation metadata lease heartbeat failed for node_id={node_id}" + ) from heartbeat_error + + async def _clear_expired_revocation_lock_for_recovery( + self, + node_id: str, + authorization: tuple[str, str, frozenset[str]], + ) -> None: + """Clear stale metadata only for a row-locked authoritative recovery.""" + if authorization[0] != node_id: + raise UserSyncLeaseLostError("authoritative database locks are required for metadata recovery") + key = self._revocation_lock_key(node_id) + for _ in range(32): + doc, rev = await kv_get_json(self._kv, key) + if doc is None or float(doc.get("expires_at", 0)) > time.time(): + return + try: + await self._kv.delete(key, last=rev) + return + except Exception as exc: + logger.debug("Stale revocation metadata lock delete raced node_id=%s: %s", node_id, exc) + continue + raise RuntimeError(f"failed to clear stale revocation metadata lock node_id={node_id}") + def _ensure_value_size(self, key: str, value: dict[str, Any]) -> None: size = len(json.dumps(value, separators=(",", ":")).encode()) if size > _MAX_USER_SYNC_VALUE_BYTES: @@ -116,16 +303,99 @@ def _ensure_value_size(self, key: str, value: dict[str, Any]) -> None: f"user sync KV value for key={key} is {size} bytes; exceeds limit {_MAX_USER_SYNC_VALUE_BYTES}" ) + @staticmethod + def _empty_barrier(user_key: str) -> dict[str, Any]: + return { + "user_key": user_key, + "generation": 0, + "active_owner": None, + "closing": False, + "permanent": False, + } + + async def _get_barrier(self, node_id: str, user_key: str) -> tuple[dict[str, Any], int]: + doc, rev = await kv_get_json(self._kv, self._barrier_key(node_id, user_key)) + if doc is None or doc.get("user_key") != user_key: + return self._empty_barrier(user_key), rev + return doc, rev + + @staticmethod + def _barrier_allows(barrier: dict[str, Any], generation: int | None = None) -> bool: + if barrier.get("permanent") or barrier.get("active_owner") is not None or barrier.get("closing"): + return False + return generation is None or int(barrier.get("generation", 0)) == generation + + @staticmethod + def _barrier_allows_lease(barrier: dict[str, Any], generation: int, revocation_id: str | None) -> bool: + if barrier.get("permanent") or int(barrier.get("generation", 0)) != generation: + return False + owner = barrier.get("active_owner") + if revocation_id is not None: + return owner == revocation_id and not barrier.get("closing") + return owner is None and not barrier.get("closing") + + async def _delete_revision(self, key: str, revision: int) -> None: + try: + await self._kv.delete(key, last=revision) + except Exception as exc: + logger.debug("Failed to delete stale user sync key=%s: %s", key, exc) + + async def _put_pending_if_current( + self, + node_id: str, + email: str, + user_b64: str, + generation: int, + *, + overwrite: bool, + ) -> bool: + """Put pending work only while the matching generation stays unfenced.""" + key = self._pending_key(node_id, email) + value = {"email": email, "user": user_b64, "generation": generation} + self._ensure_value_size(key, value) + + for _ in range(32): + barrier, _ = await self._get_barrier(node_id, email) + if not self._barrier_allows(barrier, generation): + return False + + current, rev = await kv_get_json(self._kv, key) + if current is not None and not overwrite: + current_generation = int(current.get("generation", 0)) + if current_generation == generation: + return True + if current_generation > generation: + return False + if not await kv_cas_json(self._kv, key, value, rev): + continue + + # A revoke may have fenced the key between the first read and the + # pending CAS. Delete only our exact revision; a newer enqueue wins. + written, written_rev = await kv_get_json(self._kv, key) + barrier, _ = await self._get_barrier(node_id, email) + if self._barrier_allows(barrier, generation): + return True + if written == value: + await self._delete_revision(key, written_rev) + return False + raise RuntimeError(f"failed to enqueue user sync key={key} after CAS retries") + async def enqueue_users(self, node_id: str, users: list[User]) -> None: if not users: return # Latest payload per email wins (dedupe across the batch first). by_email = {user.email: user for user in users} for email, user in by_email.items(): - key = self._pending_key(node_id, email) - value = {"email": email, "user": _b64_user(user)} - self._ensure_value_size(key, value) - await kv_put_json(self._kv, key, value) + barrier, _ = await self._get_barrier(node_id, email) + if not self._barrier_allows(barrier): + continue + await self._put_pending_if_current( + node_id, + email, + _b64_user(user), + int(barrier.get("generation", 0)), + overwrite=True, + ) async def _requeue_expired_claims(self, node_id: str) -> None: now = time.time() @@ -137,10 +407,15 @@ async def _requeue_expired_claims(self, node_id: str) -> None: continue email = doc.get("email") user_b64 = doc.get("user") + generation = int(doc.get("generation", 0)) if isinstance(email, str) and isinstance(user_b64, str): - pending_key = self._pending_key(node_id, email) - pending_value = {"email": email, "user": user_b64} - await kv_put_json(self._kv, pending_key, pending_value) + await self._put_pending_if_current( + node_id, + email, + user_b64, + generation, + overwrite=False, + ) try: await self._kv.delete(key, last=rev) except Exception as exc: @@ -163,6 +438,11 @@ async def claim_users(self, node_id: str, worker_id: str, limit: int, lease_seco user_b64 = doc.get("user") if not isinstance(email, str) or not isinstance(user_b64, str): continue + generation = int(doc.get("generation", 0)) + barrier, _ = await self._get_barrier(node_id, email) + if not self._barrier_allows(barrier, generation): + await self._delete_revision(pending_key, rev) + continue token = f"{worker_id}:{uuid4()}" claimed_key = self._claimed_key(node_id, token) @@ -170,6 +450,7 @@ async def claim_users(self, node_id: str, worker_id: str, limit: int, lease_seco "token": token, "email": email, "user": user_b64, + "generation": generation, "expires_at": now + lease_seconds, } self._ensure_value_size(claimed_key, claimed_value) @@ -189,9 +470,57 @@ async def claim_users(self, node_id: str, worker_id: str, limit: int, lease_seco cleanup_exc, ) continue - result.append(ClaimedUser(token=token, user=_user_from_b64(user_b64))) + # A fence installed after the pending delete invalidates this + # claim. The worker would reject it at lease acquisition too, but + # avoid returning known-stale work in the first place. + barrier, _ = await self._get_barrier(node_id, email) + if not self._barrier_allows(barrier, generation): + claimed_doc, claimed_rev = await kv_get_json(self._kv, claimed_key) + if claimed_doc is not None: + await self._delete_revision(claimed_key, claimed_rev) + continue + result.append(ClaimedUser(token=token, user=_user_from_b64(user_b64), generation=generation)) return result + async def next_claim_delay(self, node_id: str) -> float | None: + """Return when valid queued work can next be claimed.""" + now = time.time() + has_pending = False + for key in await kv_list_keys(self._kv, self._pending_prefix(node_id)): + doc, rev = await kv_get_json(self._kv, key) + if doc is None: + continue + email = doc.get("email") + generation = int(doc.get("generation", 0)) + if not isinstance(email, str): + await self._delete_revision(key, rev) + continue + barrier, _ = await self._get_barrier(node_id, email) + if not self._barrier_allows(barrier, generation): + await self._delete_revision(key, rev) + continue + has_pending = True + if has_pending: + return 0.0 + + delay: float | None = None + for key in await kv_list_keys(self._kv, self._claimed_prefix(node_id)): + doc, rev = await kv_get_json(self._kv, key) + if doc is None: + continue + email = doc.get("email") + generation = int(doc.get("generation", 0)) + if not isinstance(email, str): + await self._delete_revision(key, rev) + continue + barrier, _ = await self._get_barrier(node_id, email) + if not self._barrier_allows(barrier, generation): + await self._delete_revision(key, rev) + continue + claim_delay = max(0.0, float(doc.get("expires_at", 0)) - now) + delay = claim_delay if delay is None else min(delay, claim_delay) + return delay + async def ack_users(self, node_id: str, tokens: list[str]) -> None: if not tokens: return @@ -209,36 +538,690 @@ async def requeue_users(self, node_id: str, claimed_users: list[ClaimedUser]) -> if not claimed_users: return for item in claimed_users: - pending_key = self._pending_key(node_id, item.user.email) - pending_value = {"email": item.user.email, "user": _b64_user(item.user)} - self._ensure_value_size(pending_key, pending_value) - await kv_put_json(self._kv, pending_key, pending_value) claimed_key = self._claimed_key(node_id, item.token) doc, rev = await kv_get_json(self._kv, claimed_key) if doc is None: continue + generation = int(getattr(item, "generation", 0)) + if ( + doc.get("token") != item.token + or doc.get("email") != item.user.email + or int(doc.get("generation", 0)) != generation + ): + continue + user_b64 = doc.get("user") + if not isinstance(user_b64, str): + continue + await self._put_pending_if_current( + node_id, + item.user.email, + user_b64, + generation, + overwrite=False, + ) try: await self._kv.delete(claimed_key, last=rev) except Exception as exc: logger.debug("Failed to delete requeued claim key=%s: %s", claimed_key, exc) - async def clear(self, node_id: str) -> None: - for key in await kv_list_keys(self._kv, self._pending_prefix(node_id)): - doc, rev = await kv_get_json(self._kv, key) - if doc is None: + async def _set_revocation_owner(self, node_id: str, user_key: str, revocation_id: str) -> bool: + """Acquire one key, returning False when it is already finalized.""" + key = self._barrier_key(node_id, user_key) + for _ in range(32): + doc, rev = await self._get_barrier(node_id, user_key) + if doc.get("permanent"): + return False + owner = doc.get("active_owner") + if owner == revocation_id and not doc.get("closing"): + return True + if owner is not None or doc.get("closing"): + raise UserRevocationConflictError((user_key,)) + doc["active_owner"] = revocation_id + doc["closing"] = False + doc["generation"] = int(doc.get("generation", 0)) + 1 + if await kv_cas_json(self._kv, key, doc, rev): + return True + raise RuntimeError(f"failed to fence user sync key={user_key} after CAS retries") + + async def _clear_revocation_owner(self, node_id: str, user_key: str, revocation_id: str) -> None: + key = self._barrier_key(node_id, user_key) + for _ in range(32): + doc, rev = await self._get_barrier(node_id, user_key) + if doc.get("permanent") or doc.get("active_owner") != revocation_id: + return + doc["active_owner"] = None + doc["closing"] = False + if await kv_cas_json(self._kv, key, doc, rev): + return + raise RuntimeError(f"failed to release user revocation key={user_key} after CAS retries") + + async def _restore_barrier_snapshot( + self, + node_id: str, + user_key: str, + snapshot: dict[str, Any], + existed: bool, + ) -> None: + """Restore a pre-begin barrier while the node metadata lock is held.""" + key = self._barrier_key(node_id, user_key) + for _ in range(32): + current, rev = await kv_get_json(self._kv, key) + if current is None: + if not existed: + return + if await kv_cas_json(self._kv, key, snapshot, 0): + return continue - try: + if not existed: await self._kv.delete(key, last=rev) - except Exception as exc: - logger.debug("Failed to clear pending key=%s: %s", key, exc) - for key in await kv_list_keys(self._kv, self._claimed_prefix(node_id)): - doc, rev = await kv_get_json(self._kv, key) - if doc is None: + return + if await kv_cas_json(self._kv, key, snapshot, rev): + return + raise RuntimeError(f"failed to restore user revocation key={user_key} after CAS retries") + + async def _reopen_revocation_owner(self, node_id: str, user_key: str, revocation_id: str) -> None: + """Allow the owner to perform an authoritative write after a failed close.""" + key = self._barrier_key(node_id, user_key) + for _ in range(32): + doc, rev = await self._get_barrier(node_id, user_key) + if doc.get("permanent") or doc.get("active_owner") != revocation_id: + return + if not doc.get("closing"): + return + doc["closing"] = False + if await kv_cas_json(self._kv, key, doc, rev): + return + raise RuntimeError(f"failed to reopen user revocation key={user_key} after CAS retries") + + async def _reopen_revocation_owners(self, node_id: str, user_keys: list[str], revocation_id: str) -> None: + await asyncio.gather( + *(self._reopen_revocation_owner(node_id, user_key, revocation_id) for user_key in user_keys) + ) + + async def _purge_user_work(self, node_id: str, user_keys: set[str]) -> None: + for prefix in (self._pending_prefix(node_id), self._claimed_prefix(node_id)): + for key in await kv_list_keys(self._kv, prefix): + doc, rev = await kv_get_json(self._kv, key) + if doc is not None and doc.get("email") in user_keys: + await self._delete_revision(key, rev) + + async def _drain_execution_leases(self, node_id: str, user_keys: set[str]) -> None: + prefix = self._execution_prefix(node_id) + while True: + now = time.time() + wait_seconds: float | None = None + found_active = False + for key in await kv_list_keys(self._kv, prefix): + doc, _ = await kv_get_json(self._kv, key) + if doc is None or not ( + doc.get("covers_all_users") or user_keys.intersection(doc.get("user_keys") or []) + ): + continue + expires_at = float(doc.get("expires_at", 0)) + if expires_at <= now: + # The remote request may have completed after its permit + # expired. Keep both the execution record and the newly + # installed fence so an operator/reconnect can reconcile + # the ambiguous outcome instead of deleting the DB row. + raise UserSyncLeaseLostError("an expired user-sync execution lease has an unknown remote outcome") + found_active = True + remaining = expires_at - now + wait_seconds = remaining if wait_seconds is None else min(wait_seconds, remaining) + if not found_active: + return + await asyncio.sleep(min(0.05, max(0.001, wait_seconds or 0.001))) + + async def begin_user_revocation( + self, node_id: str, user_keys: list[str], revocation_id: str + ) -> UserRevocationResult: + if not revocation_id: + raise ValueError("revocation_id must not be empty") + keys = sorted(set(user_keys)) + async with self._revocation_lock(node_id, keys) as assert_metadata_owned: + # Validate the complete set before changing any generation. A + # conflicting bulk begin must leave unrelated queued work intact. + barrier_snapshots = {user_key: await self._get_barrier(node_id, user_key) for user_key in keys} + barriers = {user_key: snapshot[0] for user_key, snapshot in barrier_snapshots.items()} + conflicts = tuple( + user_key + for user_key, barrier in barriers.items() + if not barrier.get("permanent") + and (barrier.get("closing") or barrier.get("active_owner") not in (None, revocation_id)) + ) + if conflicts: + raise UserRevocationConflictError(conflicts) + + acquired: list[str] = [] + finalized: list[str] = [] + attempted: list[str] = [] + try: + for user_key in keys: + await assert_metadata_owned() + if barriers[user_key].get("permanent"): + finalized.append(user_key) + else: + attempted.append(user_key) + if await self._set_revocation_owner(node_id, user_key, revocation_id): + acquired.append(user_key) + except BaseException: + await assert_metadata_owned() + + async def restore_snapshots() -> None: + await asyncio.gather( + *( + self._restore_barrier_snapshot( + node_id, + user_key, + barrier_snapshots[user_key][0], + barrier_snapshots[user_key][1] != 0, + ) + for user_key in attempted + ) + ) + + restore = asyncio.create_task(restore_snapshots()) + try: + await asyncio.shield(restore) + except asyncio.CancelledError: + await restore + raise + raise + + active_keys = set(acquired) + # Fencing happens before cleanup. Any enqueue/claim racing with + # cleanup observes the new generation and self-deletes or rejects. + await self._purge_user_work(node_id, active_keys) + await assert_metadata_owned() + await self._drain_execution_leases(node_id, active_keys) + await assert_metadata_owned() + return UserRevocationResult(tuple(acquired), tuple(finalized)) + + async def abort_user_revocation(self, node_id: str, user_keys: list[str], revocation_id: str) -> None: + if not revocation_id: + raise ValueError("revocation_id must not be empty") + keys = sorted(set(user_keys)) + async with self._revocation_lock(node_id, keys) as assert_metadata_owned: + affected: list[str] = [] + try: + for user_key in keys: + key = self._barrier_key(node_id, user_key) + for _ in range(32): + await assert_metadata_owned() + doc, rev = await self._get_barrier(node_id, user_key) + if doc.get("permanent") or doc.get("active_owner") != revocation_id: + break + doc["closing"] = True + if await kv_cas_json(self._kv, key, doc, rev): + affected.append(user_key) + break + else: + raise RuntimeError(f"failed to abort user revocation key={user_key} after CAS retries") + if not affected: + return + await self._drain_execution_leases(node_id, set(affected)) + await assert_metadata_owned() + except BaseException: + await assert_metadata_owned() + # Include every requested key: the current CAS may have + # succeeded even if its transport reply was lost. + reopen = asyncio.create_task(self._reopen_revocation_owners(node_id, keys, revocation_id)) + try: + await asyncio.shield(reopen) + except asyncio.CancelledError: + await reopen + raise + raise + for user_key in affected: + await assert_metadata_owned() + await self._clear_revocation_owner(node_id, user_key, revocation_id) + + async def finalize_user_revocation(self, node_id: str, user_keys: list[str], revocation_id: str) -> None: + if not revocation_id: + raise ValueError("revocation_id must not be empty") + keys = sorted(set(user_keys)) + async with self._revocation_lock(node_id, keys) as assert_metadata_owned: + affected: list[str] = [] + try: + for user_key in keys: + key = self._barrier_key(node_id, user_key) + for _ in range(32): + await assert_metadata_owned() + doc, rev = await self._get_barrier(node_id, user_key) + if doc.get("permanent") or doc.get("active_owner") != revocation_id: + break + doc["closing"] = True + if await kv_cas_json(self._kv, key, doc, rev): + affected.append(user_key) + break + else: + raise RuntimeError(f"failed to finalize user revocation key={user_key} after CAS retries") + if not affected: + return + await self._drain_execution_leases(node_id, set(affected)) + await assert_metadata_owned() + except BaseException: + await assert_metadata_owned() + reopen = asyncio.create_task(self._reopen_revocation_owners(node_id, keys, revocation_id)) + try: + await asyncio.shield(reopen) + except asyncio.CancelledError: + await reopen + raise + raise + for user_key in affected: + key = self._barrier_key(node_id, user_key) + for _ in range(32): + await assert_metadata_owned() + doc, rev = await self._get_barrier(node_id, user_key) + if doc.get("permanent"): + break + if doc.get("active_owner") != revocation_id or not doc.get("closing"): + raise RuntimeError(f"lost ownership while finalizing user revocation key={user_key}") + doc["permanent"] = True + doc["active_owner"] = None + doc["closing"] = False + if await kv_cas_json(self._kv, key, doc, rev): + break + else: + raise RuntimeError(f"failed to finalize user revocation key={user_key} after CAS retries") + await self._purge_user_work(node_id, set(keys)) + await assert_metadata_owned() + + async def acquire_user_sync_lease( + self, + node_id: str, + worker_id: str, + user_keys: list[str], + lease_seconds: float, + expected_generations: dict[str, int] | None = None, + revocation_id: str | None = None, + ) -> UserSyncLease: + # A startup snapshot replaces the complete node user set. Do not admit + # a per-user write while a live wildcard startup lease is in flight. + for key in await kv_list_keys(self._kv, self._execution_prefix(node_id)): + doc, _ = await kv_get_json(self._kv, key) + if doc is not None and doc.get("covers_all_users") and float(doc.get("expires_at", 0)) > time.time(): + return UserSyncLease(node_id, worker_id, "", (), {}, lease_seconds) + + generations: dict[str, int] = {} + for user_key in dict.fromkeys(user_keys): + barrier, _ = await self._get_barrier(node_id, user_key) + generation = int(barrier.get("generation", 0)) + expected = None if expected_generations is None else expected_generations.get(user_key) + if self._barrier_allows_lease(barrier, generation, revocation_id) and ( + expected_generations is None or expected == generation + ): + generations[user_key] = generation + + if not generations: + return UserSyncLease(node_id, worker_id, "", (), {}, lease_seconds) + + token = f"{worker_id}:{uuid4()}" + lease = UserSyncLease( + node_id=node_id, + worker_id=worker_id, + token=token, + user_keys=tuple(generations), + generations=generations, + lease_seconds=lease_seconds, + revocation_id=revocation_id, + epoch=await self._next_user_sync_epoch(node_id), + ) + key = self._execution_key(node_id, token) + value = { + "token": token, + "worker_id": worker_id, + "user_keys": list(lease.user_keys), + "generations": generations, + "revocation_id": revocation_id, + "covers_all_users": False, + "epoch": lease.epoch, + "expires_at": time.time() + lease_seconds, + } + self._ensure_value_size(key, value) + if not await kv_cas_json(self._kv, key, value, 0): + raise RuntimeError(f"failed to create user sync execution lease key={key}") + + # Close the acquire/begin race: begin either sees this lease and waits, + # or the post-check sees its fence and discards the lease before use. + for user_key, generation in generations.items(): + barrier, _ = await self._get_barrier(node_id, user_key) + if not self._barrier_allows_lease(barrier, generation, revocation_id): + doc, rev = await kv_get_json(self._kv, key) + if doc is not None: + await self._delete_revision(key, rev) + return UserSyncLease(node_id, worker_id, "", (), {}, lease_seconds) + for other_key in await kv_list_keys(self._kv, self._execution_prefix(node_id)): + if other_key == key: continue + other, _ = await kv_get_json(self._kv, other_key) + if other is not None and other.get("covers_all_users") and float(other.get("expires_at", 0)) > time.time(): + doc, rev = await kv_get_json(self._kv, key) + if doc is not None: + await self._delete_revision(key, rev) + return UserSyncLease(node_id, worker_id, "", (), {}, lease_seconds) + return lease + + async def acquire_startup_user_sync_lease( + self, + node_id: str, + worker_id: str, + user_keys: list[str], + lease_seconds: float, + ) -> StartupUserSyncLease: + """Take a node-wide replacement permit and drain prior writes.""" + unique_keys = tuple(dict.fromkeys(user_keys)) + while True: try: - await self._kv.delete(key, last=rev) - except Exception as exc: - logger.debug("Failed to clear claimed key=%s: %s", key, exc) + async with self._revocation_lock(node_id, list(unique_keys)) as assert_metadata_owned: + barriers: dict[str, dict[str, Any]] = {} + has_provisional = False + for key in await kv_list_keys(self._kv, self._barrier_prefix(node_id)): + doc, _ = await kv_get_json(self._kv, key) + if doc is None: + continue + user_key = doc.get("user_key") + if isinstance(user_key, str): + barriers[user_key] = doc + if doc.get("active_owner") is not None or doc.get("closing"): + has_provisional = True + if has_provisional: + pass + else: + now = time.time() + for key in await kv_list_keys(self._kv, self._execution_prefix(node_id)): + doc, _ = await kv_get_json(self._kv, key) + if doc is not None and float(doc.get("expires_at", 0)) <= now: + raise UserSyncLeaseLostError( + "an expired user-sync execution lease has an unknown remote outcome" + ) + + generations = { + user_key: int(barriers.get(user_key, {}).get("generation", 0)) for user_key in unique_keys + } + included_keys = tuple( + user_key for user_key in unique_keys if not barriers.get(user_key, {}).get("permanent") + ) + token = f"{worker_id}:{uuid4()}" + lease = UserSyncLease( + node_id=node_id, + worker_id=worker_id, + token=token, + user_keys=unique_keys, + generations=generations, + lease_seconds=lease_seconds, + covers_all_users=True, + epoch=await self._next_user_sync_epoch(node_id), + ) + lease_key = self._execution_key(node_id, token) + value = { + "token": token, + "worker_id": worker_id, + "user_keys": list(unique_keys), + "generations": generations, + "revocation_id": None, + "covers_all_users": True, + "epoch": lease.epoch, + "expires_at": now + lease_seconds, + } + self._ensure_value_size(lease_key, value) + await assert_metadata_owned() + if not await kv_cas_json(self._kv, lease_key, value, 0): + raise RuntimeError(f"failed to create startup execution lease key={lease_key}") + break + except UserRevocationConflictError: + await asyncio.sleep(0.01) + continue + await asyncio.sleep(0.01) + + try: + while True: + now = time.time() + active_prior = False + for key in await kv_list_keys(self._kv, self._execution_prefix(node_id)): + if key == lease_key: + continue + doc, _ = await kv_get_json(self._kv, key) + if doc is None: + continue + if float(doc.get("expires_at", 0)) <= now: + raise UserSyncLeaseLostError( + "an expired user-sync execution lease has an unknown remote outcome" + ) + active_prior = True + if not active_prior: + return StartupUserSyncLease(lease=lease, included_user_keys=included_keys) + if not await self.heartbeat_user_sync_lease(lease): + raise UserSyncLeaseLostError("startup execution lease expired while draining prior writes") + await asyncio.sleep(0.01) + except BaseException: + await self.release_user_sync_lease(lease) + raise + + async def acquire_user_sync_reconciliation_lease( + self, + node_id: str, + worker_id: str, + user_keys: list[str], + lease_seconds: float, + ) -> StartupUserSyncLease: + """Resolve orphan barriers and take an authoritative snapshot lease. + + ``user_keys`` is the current, row-locked database membership. Once no + live transport lease remains, provisional barriers left by a crashed + delete can be decided safely: present rows are restored/unfenced and + absent rows are finalized permanently. The full replacement is then + protected by a newly allocated monotonic epoch. + """ + unique_keys = tuple(dict.fromkeys(user_keys)) + membership = self._authoritative_reconciliation_membership.get() + if membership is None or membership[:2] != (node_id, worker_id): + raise UserSyncLeaseLostError("authoritative database membership was not supplied for reconciliation") + authoritative_keys = set(membership[2]) + while True: + await self._clear_expired_revocation_lock_for_recovery(node_id, membership) + try: + async with self._revocation_lock(node_id, list(unique_keys)) as assert_metadata_owned: + now = time.time() + active_prior = False + for key in await kv_list_keys(self._kv, self._execution_prefix(node_id)): + doc, rev = await kv_get_json(self._kv, key) + if doc is None: + continue + if float(doc.get("expires_at", 0)) <= now: + await self._delete_revision(key, rev) + else: + active_prior = True + if active_prior: + await asyncio.sleep(0.01) + continue + + absent_orphans: set[str] = set() + for barrier_key in await kv_list_keys(self._kv, self._barrier_prefix(node_id)): + for _ in range(32): + await assert_metadata_owned() + doc, rev = await kv_get_json(self._kv, barrier_key) + if doc is None: + break + user_key = doc.get("user_key") + provisional = doc.get("active_owner") is not None or bool(doc.get("closing")) + if not provisional or not isinstance(user_key, str): + break + if user_key in authoritative_keys: + doc["active_owner"] = None + doc["closing"] = False + else: + doc["permanent"] = True + doc["active_owner"] = None + doc["closing"] = False + absent_orphans.add(user_key) + if await kv_cas_json(self._kv, barrier_key, doc, rev): + break + else: + raise RuntimeError( + f"failed to resolve orphan user revocation barrier key={barrier_key}" + ) + if absent_orphans: + await self._purge_user_work(node_id, absent_orphans) + + barriers: dict[str, dict[str, Any]] = {} + for key in await kv_list_keys(self._kv, self._barrier_prefix(node_id)): + doc, _ = await kv_get_json(self._kv, key) + if doc is not None and isinstance(doc.get("user_key"), str): + barriers[doc["user_key"]] = doc + + generations = { + user_key: int(barriers.get(user_key, {}).get("generation", 0)) for user_key in unique_keys + } + included_keys = tuple( + user_key for user_key in unique_keys if not barriers.get(user_key, {}).get("permanent") + ) + token = f"{worker_id}:{uuid4()}" + lease = UserSyncLease( + node_id=node_id, + worker_id=worker_id, + token=token, + user_keys=unique_keys, + generations=generations, + lease_seconds=lease_seconds, + covers_all_users=True, + epoch=await self._next_user_sync_epoch(node_id), + ) + lease_key = self._execution_key(node_id, token) + value = { + "token": token, + "worker_id": worker_id, + "user_keys": list(unique_keys), + "generations": generations, + "revocation_id": None, + "covers_all_users": True, + "epoch": lease.epoch, + "expires_at": now + lease_seconds, + } + self._ensure_value_size(lease_key, value) + await assert_metadata_owned() + if not await kv_cas_json(self._kv, lease_key, value, 0): + raise RuntimeError(f"failed to create reconciliation execution lease key={lease_key}") + return StartupUserSyncLease(lease=lease, included_user_keys=included_keys) + except UserRevocationConflictError: + await asyncio.sleep(0.01) + continue + await asyncio.sleep(0.01) + + async def retain_user_sync_lease_keys(self, lease: UserSyncLease, retained_user_keys: list[str]) -> UserSyncLease: + retained = tuple(dict.fromkeys(retained_user_keys)) + if lease.covers_all_users: + raise ValueError("a node-wide startup lease cannot be narrowed") + if not retained or not set(retained).issubset(lease.user_keys): + raise ValueError("retained_user_keys must be a non-empty subset of the lease") + key = self._execution_key(lease.node_id, lease.token) + for _ in range(32): + doc, rev = await kv_get_json(self._kv, key) + if doc is None or not self._lease_matches(doc, lease): + raise UserSyncLeaseLostError("user-sync execution lease is no longer owned") + narrowed = UserSyncLease( + node_id=lease.node_id, + worker_id=lease.worker_id, + token=lease.token, + user_keys=retained, + generations={user_key: lease.generations[user_key] for user_key in retained}, + lease_seconds=lease.lease_seconds, + revocation_id=lease.revocation_id, + epoch=lease.epoch, + ) + doc["user_keys"] = list(retained) + doc["generations"] = narrowed.generations + if await kv_cas_json(self._kv, key, doc, rev): + return narrowed + raise RuntimeError("failed to narrow user-sync execution lease after CAS retries") + + @staticmethod + def _lease_matches(doc: dict[str, Any], lease: UserSyncLease) -> bool: + return ( + doc.get("token") == lease.token + and doc.get("worker_id") == lease.worker_id + and tuple(doc.get("user_keys") or ()) == lease.user_keys + and doc.get("generations") == lease.generations + and doc.get("revocation_id") == lease.revocation_id + and bool(doc.get("covers_all_users")) == lease.covers_all_users + and int(doc.get("epoch", 0)) == lease.epoch + ) + + async def heartbeat_user_sync_lease(self, lease: UserSyncLease) -> bool: + if not lease.token: + return False + key = self._execution_key(lease.node_id, lease.token) + for _ in range(32): + doc, rev = await kv_get_json(self._kv, key) + if doc is None or not self._lease_matches(doc, lease): + return False + if float(doc.get("expires_at", 0)) <= time.time(): + # Keep the record as fail-closed evidence of an unknown remote + # outcome. Only an explicit release of this exact lease may + # reconcile it. + return False + doc["expires_at"] = time.time() + lease.lease_seconds + if await kv_cas_json(self._kv, key, doc, rev): + return True + logger.warning("User sync lease heartbeat CAS exhausted for node_id=%s", lease.node_id) + return False + + async def release_user_sync_lease(self, lease: UserSyncLease) -> None: + if not lease.token: + return + key = self._execution_key(lease.node_id, lease.token) + doc, rev = await kv_get_json(self._kv, key) + if doc is None or not self._lease_matches(doc, lease): + return + await self._delete_revision(key, rev) + + async def needs_authoritative_recovery(self, node_id: str) -> bool: + """Detect durable orphan evidence that requires a DB-locked reconcile.""" + now = time.time() + lock_doc, _ = await kv_get_json(self._kv, self._revocation_lock_key(node_id)) + if lock_doc is not None and float(lock_doc.get("expires_at", 0)) <= now: + return True + for key in await kv_list_keys(self._kv, self._barrier_prefix(node_id)): + doc, _ = await kv_get_json(self._kv, key) + if doc is not None and (doc.get("active_owner") is not None or bool(doc.get("closing"))): + return True + for key in await kv_list_keys(self._kv, self._execution_prefix(node_id)): + doc, _ = await kv_get_json(self._kv, key) + if doc is not None and float(doc.get("expires_at", 0)) <= now: + return True + return False + + async def clear(self, node_id: str) -> None: + """Flush pending work without erasing revocation fences or poison.""" + for prefix in ( + self._pending_prefix(node_id), + self._claimed_prefix(node_id), + ): + for key in await kv_list_keys(self._kv, prefix): + doc, rev = await kv_get_json(self._kv, key) + if doc is not None: + await self._delete_revision(key, rev) + + async def purge_node(self, node_id: str) -> None: + """Delete all memory only after the node itself is removed.""" + for prefix in ( + self._pending_prefix(node_id), + self._claimed_prefix(node_id), + self._barrier_prefix(node_id), + self._execution_prefix(node_id), + ): + for key in await kv_list_keys(self._kv, prefix): + doc, rev = await kv_get_json(self._kv, key) + if doc is not None: + await self._delete_revision(key, rev) + lock_key = self._revocation_lock_key(node_id) + lock_doc, lock_rev = await kv_get_json(self._kv, lock_key) + if lock_doc is not None: + await self._delete_revision(lock_key, lock_rev) + epoch_key = self._epoch_key(node_id) + epoch_doc, epoch_rev = await kv_get_json(self._kv, epoch_key) + if epoch_doc is not None: + await self._delete_revision(epoch_key, epoch_rev) class NatsNodeLifecycleCoordinator: @@ -248,6 +1231,25 @@ def __init__(self, kv: CasKv): def _key(self, node_id: str) -> str: return f"lifecycle.{node_id}" + async def mark_deleted(self, node_id: str) -> None: + """Permanently fence a Bridge incarnation before any remote Stop.""" + key = self._key(node_id) + for _ in range(32): + doc, rev = await kv_get_json(self._kv, key) + if doc is None: + doc = _empty_lifecycle_doc() + if doc.get("deleted") is True: + return + doc["deleted"] = True + doc["deleted_at"] = time.time() + if await kv_cas_json(self._kv, key, doc, rev): + return + raise RuntimeError(f"failed to persist node deletion tombstone node_id={node_id}") + + async def is_deleted(self, node_id: str) -> bool: + doc, _ = await kv_get_json(self._kv, self._key(node_id)) + return bool(doc and doc.get("deleted") is True) + async def try_acquire( self, node_id: str, worker_id: str, operation: LifecycleOperation, lease_seconds: float ) -> LifecycleLease | None: @@ -258,8 +1260,13 @@ async def try_acquire( if doc is None: doc = _empty_lifecycle_doc() + # A deleted incarnation may only be stopped. START remains fenced + # even when a worker missed the best-effort cleanup broadcast. + if doc.get("deleted") is True and operation is not LifecycleOperation.STOP: + return None + lease_data = doc.get("lease") - if lease_data is not None and float(lease_data.get("expires_at", 0)) > now: + if lease_data is not None: return None state = _state_from_dict(doc.get("state")) @@ -329,6 +1336,8 @@ async def heartbeat(self, lease: LifecycleLease) -> bool: lease_data = doc.get("lease") if lease_data is None or lease_data.get("token") != lease.token: return False + if float(lease_data.get("expires_at", 0)) <= now: + return False lease_data["expires_at"] = now + lease.lease_seconds doc["lease"] = lease_data if await kv_cas_json(self._kv, key, doc, rev): @@ -336,6 +1345,30 @@ async def heartbeat(self, lease: LifecycleLease) -> bool: logger.warning("Lifecycle heartbeat CAS exhausted for node_id=%s key=%s", lease.node_id, key) return False + async def reconcile(self, node_id: str, observed: LifecycleStatus) -> bool: + """Clear only an expired unknown lifecycle lease after a state probe.""" + key = self._key(node_id) + for _ in range(32): + now = time.time() + doc, rev = await kv_get_json(self._kv, key) + if doc is None: + doc = _empty_lifecycle_doc() + lease_data = doc.get("lease") + if lease_data is not None and float(lease_data.get("expires_at", 0)) > now: + return False + state = _state_from_dict(doc.get("state")) + state.epoch += 1 + state.desired = observed + state.observed = observed + state.operation = None + state.owner = None + state.updated_at = now + doc["state"] = _state_to_dict(state) + doc["lease"] = None + if await kv_cas_json(self._kv, key, doc, rev): + return True + return False + async def get_state(self, node_id: str) -> NodeLifecycleState | None: doc, _ = await kv_get_json(self._kv, self._key(node_id)) if doc is None or doc.get("state") is None: @@ -384,7 +1417,7 @@ async def clear_bridge_memory_for_node(node_id: int | str) -> None: store, coordinator, _ = get_bridge_memory() nid = str(node_id) if store is not None: - await store.clear(nid) + await store.purge_node(nid) if coordinator is not None: await coordinator.clear(nid) diff --git a/app/node/sync.py b/app/node/sync.py index 75d383f57..4c5fc05e0 100644 --- a/app/node/sync.py +++ b/app/node/sync.py @@ -1,18 +1,179 @@ import asyncio +import uuid +from dataclasses import dataclass +from PasarGuardNodeBridge.common.service_pb2 import User as ProtoUser from sqlalchemy import select from sqlalchemy.ext.asyncio import async_object_session +from sqlalchemy.orm.exc import UnmappedInstanceError -from app.db.models import Admin, AdminRole, AdminStatus, User -from app.models.user import UserNotificationResponse +from app.db import GetDB +from app.db.models import Admin, AdminRole, AdminStatus, Node, NodeStatus, User from app.nats.node_rpc import encode_node_command, node_nats_client from app.nats.proto_utils import serialize_proto_message, serialize_proto_messages from app.node import node_manager +from app.node.errors import NodeRevocationError from app.node.user import _serialize_user_for_node, serialize_user, serialize_users_for_node from app.utils.logger import get_logger from config import nats_settings, runtime_settings logger = get_logger("node-sync") +_abort_retry_tasks: dict[str, asyncio.Task] = {} +_finalize_retry_tasks: dict[str, asyncio.Task] = {} +_resolution_retry_tasks: dict[str, asyncio.Task] = {} + + +@dataclass(frozen=True, slots=True) +class UserRevocation: + revocation_id: str + removal_users: tuple[ProtoUser, ...] + original_users: tuple[ProtoUser, ...] + expected_node_ids: frozenset[int] | None = None + + +def _subset_revocation(revocation: UserRevocation, user_keys: set[str]) -> UserRevocation | None: + removal = tuple(user for user in revocation.removal_users if user.email in user_keys) + if not removal: + return None + originals = {user.email: user for user in revocation.original_users} + return UserRevocation( + revocation.revocation_id, + removal, + tuple(originals[user.email] for user in removal), + revocation.expected_node_ids, + ) + + +async def _resolve_user_removal_from_fresh_db(revocation: UserRevocation) -> None: + """Resolve an ambiguous delete commit using a new DB transaction.""" + user_keys = {user.email for user in revocation.removal_users} + async with GetDB() as db: + present_keys = set( + (await db.execute(select(User.sync_id).where(User.sync_id.in_(user_keys)))).scalars().all() + ) + present = _subset_revocation(revocation, present_keys) + absent = _subset_revocation(revocation, user_keys - present_keys) + if present is not None: + await _dispatch_abort_with_topology_retry(present) + if absent is not None: + await _dispatch_finalize_with_topology_retry(absent) + + +def _schedule_resolution_retry(revocation: UserRevocation) -> None: + if revocation.revocation_id in _resolution_retry_tasks: + return + + async def retry() -> None: + delay = 0.25 + try: + while True: + try: + await _resolve_user_removal_from_fresh_db(revocation) + return + except asyncio.CancelledError: + raise + except BaseException as exc: + logger.error( + "Retrying ambiguous database user-removal resolution %s: %s", + revocation.revocation_id, + exc, + ) + await asyncio.sleep(delay) + delay = min(delay * 2, 30.0) + finally: + _resolution_retry_tasks.pop(revocation.revocation_id, None) + + _resolution_retry_tasks[revocation.revocation_id] = asyncio.create_task(retry()) + + +async def resolve_user_removal_after_db_error(revocation: UserRevocation | None, failed_db=None) -> None: + """Resolve commit ambiguity; unknown membership remains fenced for retry.""" + if revocation is None: + return + if failed_db is not None: + try: + # Release locks from the failed transaction before opening the + # authoritative read. If rollback itself is ambiguous, defer the + # read until the request context has closed this session. + await failed_db.rollback() + except asyncio.CancelledError: + _schedule_resolution_retry(revocation) + raise + except BaseException as exc: + logger.error("Cannot close failed database delete transaction %s: %s", revocation.revocation_id, exc) + _schedule_resolution_retry(revocation) + return + try: + await _resolve_user_removal_from_fresh_db(revocation) + except asyncio.CancelledError: + _schedule_resolution_retry(revocation) + raise + except BaseException as exc: + logger.error("Cannot resolve ambiguous database delete %s: %s", revocation.revocation_id, exc) + _schedule_resolution_retry(revocation) + + +async def _lock_users_for_revocation(users: list[User]) -> set[int] | None: + """Hold target DB rows until revoke and delete commit/rollback finish. + + Node startup locks its user snapshot after registering the runtime node. + This complementary lock makes the ordering safe across Panel processes: + startup either finishes first and is included in the topology snapshot, or + observes the committed deletion and cannot start with a stale user. + """ + user_ids = sorted({user.id for user in users}) + if not user_ids: + return set() + sessions = set() + for user in users: + try: + sessions.add(async_object_session(user)) + except UnmappedInstanceError: + # Lightweight DTO-like test/custom integrations have no ORM + # session. Production CRUD passes mapped User instances. + continue + sessions.discard(None) + if not sessions: + return None + if len(sessions) != 1: + raise NodeRevocationError("users scheduled for removal belong to different database sessions") + session = sessions.pop() + await session.execute(select(User.id).where(User.id.in_(user_ids)).with_for_update()) + return set( + ( + await session.execute( + select(Node.id).where(Node.status.not_in([NodeStatus.disabled, NodeStatus.limited])).with_for_update() + ) + ) + .scalars() + .all() + ) + + +async def _refresh_expected_node_ids() -> frozenset[int]: + """Refresh topology only after a close rejected the preflight snapshot. + + The normal rollback path never depends on a second database connection. + This fallback handles a node which was legitimately removed after the + original transaction released its row locks. + """ + async with GetDB() as db: + return frozenset( + (await db.execute(select(Node.id).where(Node.status.not_in([NodeStatus.disabled, NodeStatus.limited])))) + .scalars() + .all() + ) + + +def _is_incomplete_topology_error(exc: BaseException) -> bool: + seen: set[int] = set() + current: BaseException | None = exc + while current is not None and id(current) not in seen: + seen.add(id(current)) + if "runtime topology is incomplete for user revocation" in str(current): + return True + current = current.__cause__ or current.__context__ + return False def _chunk_serialized_users_for_nats(users: list[dict]) -> list[list[dict]]: @@ -47,6 +208,57 @@ def _chunk_serialized_users_for_nats(users: list[dict]) -> list[list[dict]]: return chunks +def _chunk_serialized_revocations_for_nats( + users: list[dict], original_users: list[dict], revocation_id: str +) -> list[tuple[list[dict], list[dict]]]: + """Chunk paired removal/original payloads without exceeding NATS limits.""" + originals_by_key = {user["email"]: user for user in original_users} + if set(originals_by_key) != {user["email"] for user in users}: + raise NodeRevocationError("removal and restoration users do not match") + + max_payload_bytes = max(1024, nats_settings.node_command_max_payload_bytes) + max_batch_size = max(1, nats_settings.node_update_users_batch_size) + chunks: list[tuple[list[dict], list[dict]]] = [] + current_users: list[dict] = [] + current_originals: list[dict] = [] + + for user in users: + original_user = originals_by_key[user["email"]] + candidate_users = [*current_users, user] + candidate_originals = [*current_originals, original_user] + payload = { + "users": candidate_users, + "original_users": candidate_originals, + "revocation_id": revocation_id, + } + if current_users and ( + len(candidate_users) > max_batch_size + or len(encode_node_command("revoke_users", payload)) > max_payload_bytes + ): + chunks.append((current_users, current_originals)) + current_users = [user] + current_originals = [original_user] + else: + current_users = candidate_users + current_originals = candidate_originals + + if len(current_users) == 1: + single_payload = { + "users": current_users, + "original_users": current_originals, + "revocation_id": revocation_id, + } + if len(encode_node_command("revoke_users", single_payload)) > max_payload_bytes: + logger.warning( + "Single serialized user revocation exceeds configured NATS node command payload limit: user=%s", + user.get("email") or "unknown", + ) + + if current_users: + chunks.append((current_users, current_originals)) + return chunks + + def _loaded_admin_sync_blocked(admin: Admin) -> bool | None: state = getattr(admin, "__dict__", {}) status = state.get("status") @@ -145,6 +357,323 @@ async def _dispatch_users_update(proto_users): await node_nats_client.publish("update_users", {"users": users_chunk}) +async def _dispatch_user_removal(proto_user, original_user, expected_node_ids: set[int] | None = None) -> str: + revocation_id = uuid.uuid4().hex + if runtime_settings.role.runs_node: + return await node_manager.revoke_users_and_wait( + [proto_user], + revocation_id, + [original_user], + expected_node_ids=expected_node_ids, + ) + + user_dict = serialize_proto_message(proto_user) + original_user_dict = serialize_proto_message(original_user) + payload = {"user": user_dict, "original_user": original_user_dict, "revocation_id": revocation_id} + if expected_node_ids is not None: + payload["expected_node_ids"] = sorted(expected_node_ids) + try: + await _request_node_revocation("revoke_user", payload) + except BaseException as revoke_exc: + # The worker may have applied the revoke even if its reply was lost. + try: + await _compensate_remote_revocation( + [user_dict], + [original_user_dict], + revocation_id, + expected_node_ids, + ) + except BaseException as compensation_exc: + raise NodeRevocationError( + f"{revoke_exc}; revocation compensation also failed: {compensation_exc}" + ) from compensation_exc + raise + return revocation_id + + +async def _dispatch_users_removal( + proto_users, + original_users, + expected_node_ids: set[int] | None = None, +) -> str: + revocation_id = uuid.uuid4().hex + if runtime_settings.role.runs_node: + return await node_manager.revoke_users_and_wait( + proto_users, + revocation_id, + original_users, + expected_node_ids=expected_node_ids, + ) + + serialized_users = serialize_proto_messages(proto_users) + serialized_original_users = serialize_proto_messages(original_users) + possibly_applied_chunks: list[tuple[list[dict], list[dict]]] = [] + try: + for users_chunk, original_users_chunk in _chunk_serialized_revocations_for_nats( + serialized_users, serialized_original_users, revocation_id + ): + # Include the in-flight chunk before awaiting: timeout/cancellation + # is ambiguous because the worker may have applied it already. + possibly_applied_chunks.append((users_chunk, original_users_chunk)) + payload = { + "users": users_chunk, + "original_users": original_users_chunk, + "revocation_id": revocation_id, + } + if expected_node_ids is not None: + payload["expected_node_ids"] = sorted(expected_node_ids) + await _request_node_revocation( + "revoke_users", + payload, + ) + except BaseException as revoke_exc: + try: + await _compensate_remote_revocation_chunks( + possibly_applied_chunks, + revocation_id, + expected_node_ids, + ) + except BaseException as compensation_exc: + raise NodeRevocationError( + f"{revoke_exc}; revocation compensation also failed: {compensation_exc}" + ) from compensation_exc + raise + return revocation_id + + +async def _run_bounded_shielded(operation, description: str) -> None: + task = asyncio.create_task(operation) + + async def wait_once() -> None: + try: + await asyncio.wait_for(asyncio.shield(task), timeout=nats_settings.node_rpc_timeout) + except TimeoutError: + task.cancel() + await asyncio.gather(task, return_exceptions=True) + raise NodeRevocationError(f"timed out while {description}") from None + + try: + await wait_once() + except asyncio.CancelledError: + # A second cancellation must not cancel the compensating RPC. Give it + # one bounded chance, then restore cancellation to the caller. + try: + await wait_once() + except BaseException: + if not task.done(): + task.cancel() + raise + + +async def _abort_remote_chunks( + users_chunks: list[tuple[list[dict], list[dict]]], + revocation_id: str, + expected_node_ids: set[int] | frozenset[int] | None = None, +) -> None: + failures: list[BaseException] = [] + cancellation: asyncio.CancelledError | None = None + for users_chunk, original_users_chunk in users_chunks: + try: + payload = { + "users": users_chunk, + "original_users": original_users_chunk, + "revocation_id": revocation_id, + } + if expected_node_ids is not None: + payload["expected_node_ids"] = sorted(expected_node_ids) + await _run_bounded_shielded( + _request_node_revocation("abort_revoke_users", payload), + "aborting a user revocation chunk", + ) + except asyncio.CancelledError as exc: + # Finish every chunk before restoring cancellation; otherwise a + # later chunk can remain removed while its DB row survives. + cancellation = exc + except BaseException as exc: + failures.append(exc) + + if failures: + raise NodeRevocationError( + f"failed to compensate {len(failures)}/{len(users_chunks)} user revocation chunks" + ) from failures[0] + if cancellation is not None: + raise cancellation + + +async def _compensate_remote_revocation_chunks( + users_chunks: list[tuple[list[dict], list[dict]]], + revocation_id: str, + expected_node_ids: set[int] | None = None, +) -> None: + await _abort_remote_chunks(users_chunks, revocation_id, expected_node_ids) + + +async def _compensate_remote_revocation( + users: list[dict], + original_users: list[dict], + revocation_id: str, + expected_node_ids: set[int] | None = None, +) -> None: + await _compensate_remote_revocation_chunks( + [(users, original_users)], + revocation_id, + expected_node_ids, + ) + + +async def _dispatch_users_removal_abort( + proto_users, + original_users, + revocation_id: str, + expected_node_ids: frozenset[int] | None = None, +) -> None: + if runtime_settings.role.runs_node: + await _run_bounded_shielded( + node_manager.abort_user_revocations( + proto_users, + revocation_id, + original_users, + expected_node_ids=set(expected_node_ids) if expected_node_ids is not None else None, + ), + "aborting local user removals", + ) + return + + await _abort_remote_chunks( + _chunk_serialized_revocations_for_nats( + serialize_proto_messages(proto_users), + serialize_proto_messages(original_users), + revocation_id, + ), + revocation_id, + expected_node_ids, + ) + + +async def _dispatch_users_removal_finalize( + proto_users, + revocation_id: str, + expected_node_ids: frozenset[int] | None = None, +) -> None: + if runtime_settings.role.runs_node: + await node_manager.finalize_user_revocations( + proto_users, + revocation_id, + expected_node_ids=set(expected_node_ids) if expected_node_ids is not None else None, + ) + return + + failures = [] + for users_chunk in _chunk_serialized_users_for_nats(serialize_proto_messages(proto_users)): + try: + payload = {"users": users_chunk, "revocation_id": revocation_id} + if expected_node_ids is not None: + payload["expected_node_ids"] = sorted(expected_node_ids) + await _request_node_revocation("finalize_revoke_users", payload) + except Exception as exc: + failures.append(exc) + if failures: + raise NodeRevocationError(f"failed to finalize {len(failures)} user revocation chunks") from failures[0] + + +async def _dispatch_abort_with_topology_retry(revocation: UserRevocation) -> None: + async def dispatch(expected_node_ids: frozenset[int] | None) -> None: + await _dispatch_users_removal_abort( + list(revocation.removal_users), + list(revocation.original_users), + revocation.revocation_id, + expected_node_ids, + ) + + try: + await dispatch(revocation.expected_node_ids) + except NodeRevocationError as exc: + if not _is_incomplete_topology_error(exc): + raise + await dispatch(await _refresh_expected_node_ids()) + + +async def _dispatch_finalize_with_topology_retry(revocation: UserRevocation) -> None: + async def dispatch(expected_node_ids: frozenset[int] | None) -> None: + await _run_bounded_shielded( + _dispatch_users_removal_finalize( + list(revocation.removal_users), + revocation.revocation_id, + expected_node_ids, + ), + "finalizing user removals", + ) + + try: + await dispatch(revocation.expected_node_ids) + except NodeRevocationError as exc: + if not _is_incomplete_topology_error(exc): + raise + await dispatch(await _refresh_expected_node_ids()) + + +def _schedule_finalize_retry(revocation: UserRevocation) -> None: + if revocation.revocation_id in _finalize_retry_tasks: + return + + async def retry() -> None: + delay = 0.25 + try: + while True: + try: + await _dispatch_finalize_with_topology_retry(revocation) + return + except asyncio.CancelledError: + raise + except BaseException as exc: + logger.error( + "Retrying incomplete user revocation finalize %s: %s", + revocation.revocation_id, + exc, + ) + await asyncio.sleep(delay) + delay = min(delay * 2, 30.0) + finally: + _finalize_retry_tasks.pop(revocation.revocation_id, None) + + _finalize_retry_tasks[revocation.revocation_id] = asyncio.create_task(retry()) + + +def _schedule_abort_retry(revocation: UserRevocation) -> None: + if revocation.revocation_id in _abort_retry_tasks: + return + + async def retry() -> None: + delay = 0.25 + try: + while True: + try: + await _dispatch_abort_with_topology_retry(revocation) + return + except asyncio.CancelledError: + raise + except BaseException as exc: + logger.error( + "Retrying incomplete user revocation abort %s: %s", + revocation.revocation_id, + exc, + ) + await asyncio.sleep(delay) + delay = min(delay * 2, 30.0) + finally: + _abort_retry_tasks.pop(revocation.revocation_id, None) + + _abort_retry_tasks[revocation.revocation_id] = asyncio.create_task(retry()) + + +async def _request_node_revocation(action: str, payload: dict) -> None: + """Translate a remote node-worker failure into a retryable local API error.""" + try: + await node_nats_client.request(action, payload) + except Exception as exc: + raise NodeRevocationError(f"cannot confirm user revocation: {exc}") from exc + + async def sync_user(db_user: User) -> None: if await _user_sync_blocked(db_user): return @@ -153,16 +682,83 @@ async def sync_user(db_user: User) -> None: asyncio.create_task(_dispatch_user_update(proto_user)) -async def remove_user(user: UserNotificationResponse) -> None: - proto_user = _serialize_user_for_node(user.id, user.proxy_settings.dict()) - asyncio.create_task(_dispatch_user_update(proto_user)) +async def remove_user(user: User) -> UserRevocation: + expected_node_ids = await _lock_users_for_revocation([user]) + removal_user = _serialize_user_for_node(user.sync_id, user.proxy_settings) + original_user = await serialize_user(user) + revocation_id = await _dispatch_user_removal(removal_user, original_user, expected_node_ids) + return UserRevocation( + revocation_id, + (removal_user,), + (original_user,), + frozenset(expected_node_ids) if expected_node_ids is not None else None, + ) + + +async def remove_users_and_wait(users: list[User]) -> UserRevocation | None: + """Publish a batch user removal before reporting cleanup as completed.""" + if not users: + return + expected_node_ids = await _lock_users_for_revocation(users) + removal_users = [_serialize_user_for_node(user.sync_id, user.proxy_settings) for user in users] + original_users = await serialize_users_for_node(users) + revocation_id = await _dispatch_users_removal(removal_users, original_users, expected_node_ids) + return UserRevocation( + revocation_id, + tuple(removal_users), + tuple(original_users), + frozenset(expected_node_ids) if expected_node_ids is not None else None, + ) + + +async def abort_user_removal(revocation: UserRevocation) -> None: + """Abort a provisional fence when the following database delete fails.""" + try: + await _dispatch_abort_with_topology_retry(revocation) + except BaseException: + _schedule_abort_retry(revocation) + raise + + +async def abort_users_removal(revocation: UserRevocation) -> None: + """Abort provisional fences when a following bulk database delete fails.""" + try: + await _dispatch_abort_with_topology_retry(revocation) + except BaseException: + _schedule_abort_retry(revocation) + raise + + +async def finalize_user_removal(revocation: UserRevocation) -> None: + """Make a successful single-user revocation tombstone permanent.""" + try: + await _dispatch_finalize_with_topology_retry(revocation) + except BaseException as finalize_exc: + # The database delete has committed and cannot be rolled back here. + # Retain NodeManager/store ownership and retry until the idempotent + # finalize is acknowledged instead of silently poisoning startup. + logger.error("Failed to finalize user removal; scheduling retry: %s", finalize_exc) + _schedule_finalize_retry(revocation) + if isinstance(finalize_exc, asyncio.CancelledError): + raise + + +async def finalize_users_removal(revocation: UserRevocation) -> None: + """Make successful bulk revocation tombstones permanent.""" + try: + await _dispatch_finalize_with_topology_retry(revocation) + except BaseException as finalize_exc: + logger.error("Failed to finalize user removals; scheduling retry: %s", finalize_exc) + _schedule_finalize_retry(revocation) + if isinstance(finalize_exc, asyncio.CancelledError): + raise async def remove_users(users: list[User]) -> None: """Batch-remove users from nodes (serialized without inbounds so nodes drop them).""" if not users: return - proto_users = [_serialize_user_for_node(u.id, u.proxy_settings) for u in users] + proto_users = [_serialize_user_for_node(u.sync_id, u.proxy_settings) for u in users] asyncio.create_task(_dispatch_users_update(proto_users)) diff --git a/app/node/user.py b/app/node/user.py index 1a0ca1000..e067a4abe 100644 --- a/app/node/user.py +++ b/app/node/user.py @@ -51,11 +51,11 @@ async def serialize_user(user: User, allowed_protocols: frozenset[ProxyProtocol] if inbounds is None: inbounds = await user.inbounds() - return _serialize_user_for_node(user.id, user_settings, inbounds, allowed_protocols) + return _serialize_user_for_node(user.sync_id, user_settings, inbounds, allowed_protocols) def _serialize_user_for_node( - id: int, + sync_id: str | int, user_settings: dict, inbounds: list[str] | None = None, allowed_protocols: frozenset[ProxyProtocol] | None = None, @@ -81,7 +81,7 @@ def _serialize_user_for_node( proxy_kwargs["hysteria_auth"] = user_settings.get("hysteria", {}).get("auth") return create_user( - str(id), + str(sync_id), create_proxy(**proxy_kwargs), inbounds, ) @@ -104,7 +104,7 @@ async def core_users( stmt = ( select( - User.id, + User.sync_id, User.proxy_settings, inbound_agg, ) @@ -136,7 +136,7 @@ async def core_users( and_(Admin.status == AdminStatus.disabled, AdminRole.disconnect_users_when_disabled.is_not(True)), ) ) - .group_by(User.id) + .group_by(User.id, User.sync_id) ) results = (await db.execute(stmt)).all() @@ -147,7 +147,7 @@ async def core_users( if inbound_tags: bridge_users.append( _serialize_user_for_node( - row.id, + row.sync_id, row.proxy_settings, inbound_tags, allowed_protocols, @@ -172,6 +172,8 @@ async def serialize_users_for_node( else: inbounds_list = loaded_inbounds - bridge_users.append(_serialize_user_for_node(user.id, user.proxy_settings, inbounds_list, allowed_protocols)) + bridge_users.append( + _serialize_user_for_node(user.sync_id, user.proxy_settings, inbounds_list, allowed_protocols) + ) return bridge_users diff --git a/app/node/worker.py b/app/node/worker.py index 80689136f..60423439e 100644 --- a/app/node/worker.py +++ b/app/node/worker.py @@ -12,7 +12,7 @@ from app.db import GetDB from app.db.crud.node import get_node_by_id, get_nodes from app.db.models import NodeStatus -from app.models.node import NodeCoreUpdate, NodeGeoFilesUpdate, NodeListQuery +from app.models.node import NodeCoreUpdate, NodeGeoFilesUpdate, NodeLifecycleRecovery, NodeListQuery from app.nats.proto_utils import deserialize_proto_message, deserialize_proto_messages from app.nats.rpc_service import BaseRpcService from app.node import node_manager @@ -22,6 +22,7 @@ from config import nats_settings, runtime_settings logger = get_logger("node-worker") +NODE_RPC_QUEUE_GROUP = f"{nats_settings.node_rpc_subject}.workers" class NodeWorkerService(BaseRpcService): @@ -30,6 +31,7 @@ def __init__(self): subject=nats_settings.node_rpc_subject, logger=logger, role_check=lambda: runtime_settings.role.runs_node, + queue_group=NODE_RPC_QUEUE_GROUP, ) self._command_sub: Subscription | None = None self._log_tasks: dict[str, asyncio.Task] = {} @@ -46,7 +48,7 @@ def _register_handlers(self): self.register_command_handler("update_user", self._update_user) self.register_command_handler("update_users", self._update_users) self.register_command_handler("update_node", self._update_node) - self.register_command_handler("remove_node", self._remove_node) + self.register_rpc_handler("remove_node", self._remove_node) self.register_command_handler("connect_node", self._connect_node) self.register_command_handler("connect_nodes_bulk", self._connect_nodes_bulk) self.register_command_handler("disconnect_node", self._disconnect_node) @@ -62,6 +64,11 @@ def _register_handlers(self): self.register_rpc_handler("update_core", self._update_core) self.register_rpc_handler("update_geofiles", self._update_geofiles) self.register_rpc_handler("start_logs", self._start_logs) + self.register_rpc_handler("revoke_user", self._rpc_revoke_user) + self.register_rpc_handler("revoke_users", self._rpc_revoke_users) + self.register_rpc_handler("abort_revoke_users", self._rpc_abort_revoke_users) + self.register_rpc_handler("finalize_revoke_users", self._rpc_finalize_revoke_users) + self.register_rpc_handler("recover_node_lifecycle", self._rpc_recover_node_lifecycle) async def start(self): await super().start() @@ -118,9 +125,11 @@ async def _run_rpc(self, msg, action: str | None, data: dict): result = await self._dispatch_rpc(action, data) await msg.respond(json.dumps({"ok": True, "data": result}).encode()) except Exception as exc: - error_msg = str(exc) + error_msg = exc.detail if isinstance(exc, NodeAPIError) else str(exc) # Determine error code based on error message content - if "NotFound" in error_msg or "not found" in error_msg.lower(): + if isinstance(exc, NodeAPIError): + error_code = exc.code + elif "NotFound" in error_msg or "not found" in error_msg.lower(): error_code = 404 elif "not allowed" in error_msg.lower() or "permission" in error_msg.lower(): error_code = 403 @@ -149,6 +158,74 @@ async def _update_users(self, data: dict): proto_users = deserialize_proto_messages(users_dicts, ProtoUser) await node_manager.update_users(proto_users) + async def _rpc_revoke_user(self, data: dict) -> dict: + user_dict = data.get("user") + original_user_dict = data.get("original_user") + if user_dict and not original_user_dict: + raise RuntimeError("original_user is required for compensatable revocation") + if user_dict: + await node_manager.revoke_users_and_wait( + [deserialize_proto_message(user_dict, ProtoUser)], + data.get("revocation_id"), + [deserialize_proto_message(original_user_dict, ProtoUser)], + expected_node_ids=( + {int(node_id) for node_id in data["expected_node_ids"]} + if data.get("expected_node_ids") is not None + else None + ), + ) + return {} + + async def _rpc_revoke_users(self, data: dict) -> dict: + users_dicts = data.get("users") or [] + original_users_dicts = data.get("original_users") or [] + if users_dicts and not original_users_dicts: + raise RuntimeError("original_users are required for compensatable revocation") + if users_dicts: + await node_manager.revoke_users_and_wait( + deserialize_proto_messages(users_dicts, ProtoUser), + data.get("revocation_id"), + deserialize_proto_messages(original_users_dicts, ProtoUser), + expected_node_ids=( + {int(node_id) for node_id in data["expected_node_ids"]} + if data.get("expected_node_ids") is not None + else None + ), + ) + return {} + + async def _rpc_abort_revoke_users(self, data: dict) -> dict: + users_dicts = data.get("users") or [] + original_users_dicts = data.get("original_users") or [] + if users_dicts and not original_users_dicts: + raise RuntimeError("original_users are required for compensatable revocation abort") + if users_dicts: + await node_manager.abort_user_revocations( + deserialize_proto_messages(users_dicts, ProtoUser), + data.get("revocation_id"), + deserialize_proto_messages(original_users_dicts, ProtoUser), + expected_node_ids=( + {int(node_id) for node_id in data["expected_node_ids"]} + if data.get("expected_node_ids") is not None + else None + ), + ) + return {} + + async def _rpc_finalize_revoke_users(self, data: dict) -> dict: + users_dicts = data.get("users") or [] + if users_dicts: + await node_manager.finalize_user_revocations( + deserialize_proto_messages(users_dicts, ProtoUser), + data.get("revocation_id"), + expected_node_ids=( + {int(node_id) for node_id in data["expected_node_ids"]} + if data.get("expected_node_ids") is not None + else None + ), + ) + return {} + async def _update_node(self, data: dict): node_id = data.get("node_id") if not node_id: @@ -162,7 +239,17 @@ async def _remove_node(self, data: dict): node_id = data.get("node_id") if not node_id: return - await node_manager.remove_node(node_id) + await self._node_operator._remove_node_impl( + node_id, + data.get("bridge_id"), + force=bool(data.get("force", False)), + ) + return {} + + async def _rpc_recover_node_lifecycle(self, data: dict) -> dict: + node_id = int(data["node_id"]) + recovery = NodeLifecycleRecovery.model_validate(data.get("recovery") or {}) + return await self._node_operator._recover_node_lifecycle_local(node_id, recovery) async def _connect_node(self, data: dict): node_id = data.get("node_id") diff --git a/app/notification/webhook/__init__.py b/app/notification/webhook/__init__.py index 672169884..b1cefd6b2 100644 --- a/app/notification/webhook/__init__.py +++ b/app/notification/webhook/__init__.py @@ -9,6 +9,8 @@ from app.notification.queue_manager import enqueue_webhook from app.settings import webhook_settings +_SENSITIVE_USER_FIELDS = {"proxy_settings", "subscription_url"} + def get_current_timestamp() -> float: """Factory function to get current timestamp""" @@ -123,10 +125,10 @@ async def status_change(user: UserNotificationResponse): async def notify(message: type[Notification]) -> None: if (await webhook_settings()).enable: - await enqueue_webhook(jsonable_encoder(message)) + await enqueue_webhook(jsonable_encoder(message, exclude={"user": _SENSITIVE_USER_FIELDS})) async def bulk_notify(messages: list[type[Notification]]) -> None: if (await webhook_settings()).enable: for message in messages: - await enqueue_webhook(jsonable_encoder(message)) + await enqueue_webhook(jsonable_encoder(message, exclude={"user": _SENSITIVE_USER_FIELDS})) diff --git a/app/operation/admin.py b/app/operation/admin.py index 0d7766a89..70a51022c 100644 --- a/app/operation/admin.py +++ b/app/operation/admin.py @@ -38,7 +38,12 @@ ) from app.models.stats import Period, UserUsageStatsList from app.models.user import UserListQuery -from app.node.sync import remove_user as sync_remove_user, sync_users +from app.node.sync import ( + finalize_users_removal, + remove_users_and_wait, + resolve_user_removal_after_db_error, + sync_users, +) from app.operation import BaseOperation from app.operation.admin_sync import admin_users_sync_blocked, sync_admin_users_for_block_transition from app.operation.permissions import PermissionDenied, enforce_permission @@ -268,9 +273,14 @@ async def _remove_all_users_for_admin(self, db: AsyncSession, db_admin: Admin, a user_operation = UserOperation(self.operator_type) serialized_users = [await user_operation.validate_user(user) for user in users] - await remove_users(db, users) - for user in serialized_users: - await sync_remove_user(user) + revocation = await remove_users_and_wait(users) + try: + await remove_users(db, users) + except BaseException: + await resolve_user_removal_after_db_error(revocation, db) + raise + if revocation is not None: + await finalize_users_removal(revocation) for user in serialized_users: asyncio.create_task(notification.remove_user(user, admin)) diff --git a/app/operation/node.py b/app/operation/node.py index bd72008c2..94808f399 100644 --- a/app/operation/node.py +++ b/app/operation/node.py @@ -1,10 +1,12 @@ import asyncio from collections.abc import AsyncIterator, Callable +from contextlib import asynccontextmanager from fastapi import HTTPException -from PasarGuardNodeBridge import NodeAPIError, PasarGuardNode +from PasarGuardNodeBridge import Health, NodeAPIError, PasarGuardNode from PasarGuardNodeBridge.common import service_pb2 as service from PasarGuardNodeBridge.storage import LifecycleStatus +from sqlalchemy import select from sqlalchemy.exc import IntegrityError from app import notification @@ -22,12 +24,11 @@ get_nodes_usage, modify_node, remove_node, - remove_nodes, reset_node_usage, update_node_status, ) from app.db.crud.user import get_user_by_id, get_user_count_metric_stats -from app.db.models import Node, NodeStatus +from app.db.models import Node, NodeStatus, User from app.models.admin import AdminDetails from app.models.core import CoreType from app.models.node import ( @@ -37,6 +38,7 @@ NodeCoreUpdate, NodeCreate, NodeGeoFilesUpdate, + NodeLifecycleRecovery, NodeListQuery, NodeModify, NodeNotification, @@ -101,6 +103,7 @@ def __init__(self, operator_type: OperatorType): self._update_geofiles_impl = self._update_geofiles_local self._get_logs_impl = self._get_logs_local self._restart_all_impl = self._restart_all_nodes_local + self._recover_lifecycle_impl = self._recover_node_lifecycle_local else: self._update_node_impl = self._update_node_remote self._remove_node_impl = self._remove_node_remote @@ -119,6 +122,7 @@ def __init__(self, operator_type: OperatorType): self._update_geofiles_impl = self._update_geofiles_remote self._get_logs_impl = self._get_logs_remote self._restart_all_impl = self._restart_all_nodes_remote + self._recover_lifecycle_impl = self._recover_node_lifecycle_remote async def get_db_nodes( self, @@ -206,15 +210,26 @@ async def _update_single_node_status( @staticmethod async def _get_core_users_map( db: AsyncSession, core_ids: set[int] - ) -> tuple[dict[int, object | None], dict[int, list]]: + ) -> tuple[dict[int, object | None], dict[int, list], set[str]]: if not core_ids: - return {}, {} + return {}, {}, set() resolved_cores = await core_manager.get_cores(core_ids | {1}) default_core = resolved_cores.get(1) cores_by_id: dict[int, object | None] = {} users_by_core: dict[int, list] = {} + # Register runtime nodes before taking this lock. A concurrent delete + # then has exactly two safe outcomes: it sees/revokes these nodes, or it + # commits first and this current-read snapshot no longer contains the + # deleted user. In shared mode Bridge 0.10's node-wide startup lease is + # the cross-worker admission barrier; local mode retains its event. + if not node_manager.uses_shared_revocation_store: + await node_manager.wait_for_user_revocations() + locked_user_keys = set( + (await db.execute(select(User.sync_id).with_for_update())).scalars().all() + ) + for core_id in core_ids: core = resolved_cores.get(core_id) or default_core cores_by_id[core_id] = core @@ -222,17 +237,81 @@ async def _get_core_users_map( users_by_core[core_id] = [] continue - users_by_core[core_id] = await core_users( + users = await core_users( db=db, inbound_tags=core.inbounds, allowed_protocols=core.protocols, ) + # The locking query is a current read even under MySQL's default + # REPEATABLE READ. In shared mode Bridge startup filters permanent + # tombstones atomically with full-snapshot apply; local mode keeps + # its in-process tombstone filter. + current_users = [user for user in users if user.email in locked_user_keys] + users_by_core[core_id] = ( + current_users + if node_manager.uses_shared_revocation_store + else node_manager.filter_permanently_deleted_users(current_users) + ) + + return cores_by_id, users_by_core, locked_user_keys + + @staticmethod + async def _prepare_authoritative_user_reconciliation( + pg_node: PasarGuardNode, authoritative_user_keys: set[str] | None + ): + store = getattr(pg_node, "_user_sync_store", None) + prepare = getattr(store, "set_authoritative_reconciliation_membership", None) + if callable(prepare): + if authoritative_user_keys is None: + raise NodeAPIError(503, "authoritative database membership is required for reconciliation") + token = await prepare(pg_node.node_id, pg_node.worker_id, sorted(authoritative_user_keys)) + return store, token + return None - return cores_by_id, users_by_core + @staticmethod + @asynccontextmanager + async def _authoritative_user_reconciliation_scope( + pg_node: PasarGuardNode, authoritative_user_keys: set[str] | None + ): + authorization = await NodeOperation._prepare_authoritative_user_reconciliation( + pg_node, authoritative_user_keys + ) + try: + yield + finally: + if authorization is not None: + store, token = authorization + # ContextVar.reset is deliberately synchronous: even repeated + # task cancellation cannot interrupt cleanup or leak the + # row-lock authorization into a later reconciliation. + store.reset_authoritative_reconciliation_membership(token) + + @staticmethod + async def _assert_node_incarnation_active( + pg_node: PasarGuardNode, + *, + stop_remote: bool = False, + ) -> None: + locally_deleted = node_manager.is_bridge_namespace_deleted(str(pg_node.node_id)) + coordinator = getattr(pg_node, "_lifecycle_coordinator", None) + is_deleted = getattr(type(coordinator), "is_deleted", None) + remotely_deleted = callable(is_deleted) and await coordinator.is_deleted(pg_node.node_id) + if not locally_deleted and not remotely_deleted: + return + try: + if stop_remote: + await pg_node.stop() + else: + await pg_node.set_health(Health.INVALID) + await pg_node.disconnect() + except Exception as exc: + logger.error("Failed to quiesce deleted node incarnation %s: %s", pg_node.node_id, exc) + raise NodeAPIError(410, "Node incarnation is permanently deleted") @staticmethod async def _attach_if_running(pg_node: PasarGuardNode, node_name: str): """Attach to an already-started remote core without calling Start RPC.""" + await NodeOperation._assert_node_incarnation_active(pg_node) try: state = await pg_node.get_lifecycle_state() if ( @@ -243,41 +322,82 @@ async def _attach_if_running(pg_node: PasarGuardNode, node_name: str): return None info = await pg_node.info() - if info is None or not info.node_version or not info.core_version: + if info is None or not info.started or not info.node_version or not info.core_version: return None await pg_node.connect(info.node_version, info.core_version) + await NodeOperation._assert_node_incarnation_active(pg_node, stop_remote=True) if state is not None: await pg_node.update_observed_lifecycle(LifecycleStatus.HEALTHY, expected_epoch=state.epoch) logger.info( f'Attached to already-running "{node_name}" node v{info.node_version}, core v{info.core_version}' ) return info + except NodeAPIError as exc: + if exc.code == 410: + raise + logger.debug(f'Attach skipped for "{node_name}": {exc}') + return None except Exception as exc: logger.debug(f'Attach skipped for "{node_name}": {exc}') return None @staticmethod - async def _start_or_attach_node(pg_node: PasarGuardNode, db_node: Node, core, users: list, backend_type): + async def _start_or_attach_node( + pg_node: PasarGuardNode, + db_node: Node, + core, + users: list, + backend_type, + authoritative_user_keys: set[str] | None = None, + ): + await NodeOperation._assert_node_incarnation_active(pg_node) state = await pg_node.get_lifecycle_state() + if state is not None and state.operation is not None: + # A probe cannot prove that an old timed-out request will not finish + # later. Keep the distributed lease fail-closed; reconciliation is + # an explicit operator action after the old worker/request is known + # to be gone. + await pg_node.info() + raise NodeAPIError( + 503, + "Node has an unresolved lifecycle operation; explicit reconciliation is required", + ) if state is not None and state.observed is LifecycleStatus.HEALTHY: attached = await NodeOperation._attach_if_running(pg_node, db_node.name) if attached is not None: + # Attach only proves that the core process is running. Apply an + # authoritative user snapshot before reporting the node ready; + # this also reconciles poison left by a crashed sync worker. + async with NodeOperation._authoritative_user_reconciliation_scope( + pg_node, authoritative_user_keys + ): + await pg_node.reconcile_users(users) return attached + capability = await pg_node.info() + if capability is None or not getattr(capability, "user_sync_epoch_supported", False): + raise NodeAPIError(426, "Node must support monotonic user-sync epoch fencing before startup") + start_kwargs = { "config": core.to_str(), "backend_type": backend_type, "users": users, "keep_alive": db_node.keep_alive, + "reconcile_user_sync": True, } if core.type == CoreType.xray: start_kwargs["exclude_inbounds"] = core.exclude_inbound_tags - return await pg_node.start(**start_kwargs) + async with NodeOperation._authoritative_user_reconciliation_scope(pg_node, authoritative_user_keys): + result = await pg_node.start(**start_kwargs) + await NodeOperation._assert_node_incarnation_active(pg_node, stop_remote=True) + return result @staticmethod - async def connect_node(db_node: Node, core, users: list) -> dict | None: + async def connect_node( + db_node: Node, core, users: list, authoritative_user_keys: set[str] | None = None + ) -> dict | None: """ Connect to a node and return status result (does NOT update database). @@ -295,13 +415,16 @@ async def connect_node(db_node: Node, core, users: list) -> dict | None: return None if core is None: return None + await NodeOperation._assert_node_incarnation_active(pg_node) old_status = db_node.status logger.info(f'Connecting to "{db_node.name}" node') type = service.BackendType.WIREGUARD if core.type == CoreType.wg else service.BackendType.XRAY try: - info = await NodeOperation._start_or_attach_node(pg_node, db_node, core, users, type) + info = await NodeOperation._start_or_attach_node( + pg_node, db_node, core, users, type, authoritative_user_keys + ) if info is None: return None @@ -322,6 +445,15 @@ async def connect_node(db_node: Node, core, users: list) -> dict | None: # Another worker holds the lifecycle lease; try attach once more. attached = await NodeOperation._attach_if_running(pg_node, db_node.name) if attached is not None: + # The competing Start may have used an older/empty + # snapshot. Attachment only proves process liveness, so + # do not publish CONNECTED until this worker has applied + # the authoritative database snapshot under an epoch + # fenced reconciliation lease. + async with NodeOperation._authoritative_user_reconciliation_scope( + pg_node, authoritative_user_keys + ): + await pg_node.reconcile_users(users) return { "node_id": db_node.id, "status": NodeStatus.connected, @@ -397,12 +529,30 @@ async def modify_node(self, db: AsyncSession, node_id: int, modified_node: NodeM return node - async def remove_node(self, db: AsyncSession, node_id: int, admin: AdminDetails) -> None: + async def remove_node( + self, + db: AsyncSession, + node_id: int, + admin: AdminDetails, + *, + force: bool = False, + ) -> None: db_node: Node = await self.get_validated_node(db=db, node_id=node_id) node_response = NodeResponse.model_validate(db_node) - await self._remove_node_impl(db_node.id) - await remove_node(db=db, db_node=db_node) + await self._remove_node_impl(db_node.id, db_node.bridge_id, force=force) + try: + await remove_node(db=db, db_node=db_node) + except Exception as exc: + if not await self._resolve_node_delete_after_db_error(db, db_node): + raise NodeAPIError(503, "Node deletion was not committed; durable tombstone remains active") from exc + + if force: + logger.warning( + 'Node "%s" force-deleted by admin "%s"; remote core may remain and requires manual decommission', + node_response.name, + admin.username, + ) logger.info(f'Node "{node_response.name}" with id "{node_response.id}" deleted by admin "{admin.username}"') @@ -590,6 +740,69 @@ async def _get_node_user_ip_list_safe(self, node_id: int, email: str) -> dict[st async def sync_node_users(self, db: AsyncSession, node_id: int, flush_users: bool = False) -> NodeResponse: return await self._sync_node_users_impl(db, node_id, flush_users) + async def reconcile_orphaned_user_sync(self, db: AsyncSession, db_node: Node) -> bool: + """Resolve persisted Bridge poison while DB user locks remain held.""" + pg_node = await node_manager.get_node(db_node.id) + if pg_node is None: + return False + store = getattr(pg_node, "_user_sync_store", None) + needs_recovery = getattr(store, "needs_authoritative_recovery", None) + if not callable(needs_recovery) or not await needs_recovery(pg_node.node_id): + return False + + core_id = db_node.core_config_id or 1 + _, users_by_core, authoritative_user_keys = await self._get_core_users_map(db, {core_id}) + async with self._authoritative_user_reconciliation_scope(pg_node, authoritative_user_keys): + await pg_node.reconcile_users(users_by_core.get(core_id, [])) + return True + + async def recover_node_lifecycle( + self, db: AsyncSession, node_id: int, recovery: NodeLifecycleRecovery + ) -> dict: + await self.get_validated_node(db, node_id, load_usage_logs=False) + return await self._recover_lifecycle_impl(node_id, recovery) + + @staticmethod + async def _probe_recovery_state(pg_node: PasarGuardNode) -> LifecycleStatus: + try: + info = await pg_node.info() + except Exception: + return LifecycleStatus.BROKEN + if info is not None and info.started and info.node_version and info.core_version: + return LifecycleStatus.HEALTHY + return LifecycleStatus.STOPPED + + async def _recover_node_lifecycle_local( + self, node_id: int, recovery: NodeLifecycleRecovery + ) -> dict: + pg_node = await node_manager.get_node(node_id) + if pg_node is None: + pg_node = await node_manager.get_lifecycle_recovery_node(node_id) + if pg_node is None: + raise NodeAPIError(409, "Node runtime is not registered on this worker") + state = await pg_node.get_lifecycle_state() + if state is None or state.operation is None: + raise NodeAPIError(409, "Node has no unresolved lifecycle operation") + detected = await self._probe_recovery_state(pg_node) + if detected is not recovery.observed: + raise NodeAPIError( + 409, + f"Observed node state is {detected.value}, not {recovery.observed.value}", + ) + await pg_node.reconcile_lifecycle(recovery.observed) + return {"node_id": node_id, "observed": recovery.observed.value, "reconciled": True} + + async def _recover_node_lifecycle_remote( + self, node_id: int, recovery: NodeLifecycleRecovery + ) -> dict: + try: + return await node_nats_client.request( + "recover_node_lifecycle", + {"node_id": node_id, "recovery": recovery.model_dump(mode="json")}, + ) + except RuntimeError as exc: + await self.handle_rpc_error(exc) + async def clear_usage_data(self, db: AsyncSession, table: UsageTable, query: NodeClearUsageQuery): if query.start and query.end and query.start >= query.end: await self.raise_error(code=400, message="Start time must be before end time.") @@ -617,39 +830,67 @@ async def _update_node_local(self, db_node: Node) -> None: async def _update_node_sync(self, db_node: Node) -> None: await self._update_node_local(db_node) - await publish_node_sync("upsert", db_node.id) + await publish_node_sync("upsert", db_node.id, db_node.bridge_id) async def _update_node_remote(self, db_node: Node) -> None: await node_nats_client.publish("update_node", {"node_id": db_node.id}) - async def _remove_node_local(self, node_id: int) -> None: - await node_manager.remove_node(node_id) - await clear_bridge_memory_for_node(node_id) + async def _remove_node_local( + self, node_id: int, bridge_id: str | None = None, *, force: bool = False + ) -> None: + await node_manager.remove_node( + node_id, + remote_stop=not force, + expected_bridge_namespace=bridge_id, + permanent_delete=True, + ) + if not force: + await clear_bridge_memory_for_node(bridge_id or node_id) - async def _remove_node_sync(self, node_id: int) -> None: - await self._remove_node_local(node_id) - await publish_node_sync("remove", node_id) + async def _remove_node_sync( + self, node_id: int, bridge_id: str | None = None, *, force: bool = False + ) -> None: + # Broadcast NATS has no all-worker acknowledgement barrier. Quiesce + # this runtime and retain shared KV state fail-closed; a future + # acknowledged cleanup protocol may purge it after every worker has + # confirmed disconnect. + await node_manager.remove_node( + node_id, + remote_stop=not force, + expected_bridge_namespace=bridge_id, + permanent_delete=True, + ) + await publish_node_sync("remove", node_id, bridge_id) - async def _remove_node_remote(self, node_id: int) -> None: - await node_nats_client.publish("remove_node", {"node_id": node_id}) + async def _remove_node_remote( + self, node_id: int, bridge_id: str | None = None, *, force: bool = False + ) -> None: + # Deletion must not commit while the worker outcome is ambiguous: the + # shared KV fences are safe to purge only after Stop is acknowledged + # and the local runtime has quiesced. + await node_nats_client.request( + "remove_node", + {"node_id": node_id, "bridge_id": bridge_id, "force": force}, + ) async def _connect_nodes_bulk_local(self, db: AsyncSession, nodes: list[Node]) -> None: if not nodes: return - core_ids = {node.core_config_id or 1 for node in nodes} - cores_by_id, users_by_core = await self._get_core_users_map(db, core_ids) - sem = asyncio.Semaphore(CONNECT_CONCURRENCY) - - async def connect_single(node: Node) -> dict | None: + ready_nodes: list[Node] = [] + results: list[dict | None] = [] + # Phase one registers every usable runtime node. Only then snapshot and + # lock users once for the whole bulk start, avoiding stale snapshots and + # concurrent use of the same AsyncSession. + for node in nodes: if node is None or node.status in (NodeStatus.disabled, NodeStatus.limited): - return - - async with sem: - try: - await node_manager.update_node(node) - except NodeAPIError as e: - return { + continue + try: + await node_manager.update_node(node) + ready_nodes.append(node) + except NodeAPIError as e: + results.append( + { "node_id": node.id, "status": NodeStatus.error, "message": e.detail, @@ -657,11 +898,23 @@ async def connect_single(node: Node) -> dict | None: "node_version": "", "old_status": node.status, } + ) + + core_ids = {node.core_config_id or 1 for node in ready_nodes} + cores_by_id, users_by_core, authoritative_user_keys = await self._get_core_users_map(db, core_ids) + sem = asyncio.Semaphore(CONNECT_CONCURRENCY) + async def connect_single(node: Node) -> dict | None: + async with sem: core_id = node.core_config_id or 1 - return await self.connect_node(node, cores_by_id.get(core_id), users_by_core.get(core_id, [])) + return await self.connect_node( + node, + cores_by_id.get(core_id), + users_by_core.get(core_id, []), + authoritative_user_keys, + ) - results = await asyncio.gather(*[connect_single(node) for node in nodes]) + results.extend(await asyncio.gather(*[connect_single(node) for node in ready_nodes])) # Filter out None results valid_results = [r for r in results if r is not None] @@ -705,7 +958,7 @@ async def _connect_nodes_bulk_sync(self, db: AsyncSession, nodes: list[Node]) -> await self._connect_nodes_bulk_local(db, nodes) for node in nodes: if node is not None and node.status not in (NodeStatus.disabled, NodeStatus.limited): - await publish_node_sync("connect", node.id) + await publish_node_sync("connect", node.id, node.bridge_id) async def _connect_nodes_bulk_remote(self, db: AsyncSession, nodes: list[Node]) -> None: if not nodes: @@ -717,14 +970,11 @@ async def _connect_single_node_local(self, db: AsyncSession, node_id: int) -> No if db_node is None or db_node.status in (NodeStatus.disabled, NodeStatus.limited): return - core_id = db_node.core_config_id or 1 - cores_by_id, users_by_core = await self._get_core_users_map(db, {core_id}) - core = cores_by_id.get(core_id) - users = users_by_core.get(core_id, []) - - # Update node manager + # Add the runtime node before locking/snapshotting users. This ordering + # pairs with the removal row lock and closes add-node vs delete races. try: - await node_manager.update_node(db_node) + if not await node_manager.runtime_matches(db_node): + await node_manager.update_node(db_node) except NodeAPIError as e: # Update status to error using simple CRUD await update_node_status( @@ -743,8 +993,13 @@ async def _connect_single_node_local(self, db: AsyncSession, node_id: int) -> No asyncio.create_task(notification.error_node(node_notif)) return + core_id = db_node.core_config_id or 1 + cores_by_id, users_by_core, authoritative_user_keys = await self._get_core_users_map(db, {core_id}) + core = cores_by_id.get(core_id) + users = users_by_core.get(core_id, []) + # Connect the node - result = await NodeOperation.connect_node(db_node, core, users) + result = await NodeOperation.connect_node(db_node, core, users, authoritative_user_keys) if not result: return @@ -778,7 +1033,9 @@ async def _connect_single_node_local(self, db: AsyncSession, node_id: int) -> No async def _connect_single_node_sync(self, db: AsyncSession, node_id: int) -> None: await self._connect_single_node_local(db, node_id) - await publish_node_sync("connect", node_id) + db_node = await get_node_by_id(db, node_id, load_usage_logs=False) + if db_node is not None: + await publish_node_sync("connect", node_id, db_node.bridge_id) async def _connect_single_node_remote(self, db: AsyncSession, node_id: int) -> None: await node_nats_client.publish("connect_node", {"node_id": node_id}) @@ -787,8 +1044,9 @@ async def _disconnect_single_node_local(self, node_id: int) -> None: await node_manager.remove_node(node_id) async def _disconnect_single_node_sync(self, node_id: int) -> None: + bridge_id = await node_manager.get_bridge_namespace(node_id) await self._disconnect_single_node_local(node_id) - await publish_node_sync("disconnect", node_id) + await publish_node_sync("disconnect", node_id, bridge_id) async def _disconnect_single_node_remote(self, node_id: int) -> None: await node_nats_client.publish("disconnect_node", {"node_id": node_id}) @@ -927,7 +1185,7 @@ async def _get_user_online_stats_local(self, db: AsyncSession, node_id: int, use await self.raise_error(message="Node not found", code=404) try: - stats = await node.get_user_online_stats(email=f"{db_user.id}") + stats = await node.get_user_online_stats(email=db_user.sync_id) except NodeAPIError as e: await self.raise_error(message=e.detail, code=e.code) @@ -947,7 +1205,7 @@ async def _get_user_ip_list_local(self, db: AsyncSession, node_id: int, user_id: if db_user is None: await self.raise_error(message="User not found", code=404) - email = f"{db_user.id}" + email = db_user.sync_id ips = await self._get_node_user_ip_list_safe(node_id, email) if ips is None: @@ -968,7 +1226,7 @@ async def _get_user_ip_list_all_local(self, db: AsyncSession, user_id: int) -> U await self.raise_error(message="User not found", code=404) nodes = await node_manager.get_healthy_nodes() - email = f"{db_user.id}" + email = db_user.sync_id ip_list_tasks = {id: asyncio.create_task(self._get_node_user_ip_list_safe(id, email)) for id, _ in nodes} @@ -1002,9 +1260,10 @@ async def _sync_node_users_local(self, db: AsyncSession, node_id: int, flush_use try: core_id = db_node.core_config_id or 1 - _, users_by_core = await self._get_core_users_map(db, {core_id}) + _, users_by_core, authoritative_user_keys = await self._get_core_users_map(db, {core_id}) users = users_by_core.get(core_id, []) - await pg_node.sync_users(users, flush_pending=flush_users) + async with self._authoritative_user_reconciliation_scope(pg_node, authoritative_user_keys): + await pg_node.reconcile_users(users, flush_pending=flush_users) except NodeAPIError as e: await update_node_status(db=db, db_node=db_node, status=NodeStatus.error, message=e.detail) await self.raise_error(message=e.detail, code=e.code) @@ -1067,31 +1326,57 @@ async def _update_geofiles_remote(self, node_id: int, node_geofiles_update: Node ) async def bulk_remove_nodes( - self, db: AsyncSession, bulk_nodes: BulkNodeSelection, admin: AdminDetails + self, + db: AsyncSession, + bulk_nodes: BulkNodeSelection, + admin: AdminDetails, + *, + force: bool = False, ) -> RemoveNodesResponse: """Remove multiple nodes by ID""" db_nodes = [] - for node_id in bulk_nodes.ids: + for node_id in sorted(bulk_nodes.ids): db_node = await self.get_validated_node(db, node_id) db_nodes.append(db_node) - node_ids = [n.id for n in db_nodes] - node_names = [n.name for n in db_nodes] - node_responses = [NodeResponse.model_validate(n) for n in db_nodes] - - # Remove nodes from RPC first - for node_id in node_ids: - await self._remove_node_impl(node_id) - - # Batch delete using CRUD function - await remove_nodes(db, node_ids) - - # Notify - for node_response in node_responses: + removed_names: list[str] = [] + failed: dict[int, str] = {} + # Each confirmed tombstone/Stop is committed to the DB immediately. + # If a later node fails, an earlier tombstoned node is never left as a + # normal visible row that health workers can no longer start. + for db_node in db_nodes: + node_response = NodeResponse.model_validate(db_node) + try: + await self._remove_node_impl(db_node.id, db_node.bridge_id, force=force) + try: + await remove_node(db, db_node) + except Exception as exc: + if not await self._resolve_node_delete_after_db_error(db, db_node): + raise NodeAPIError( + 503, + "Node deletion was not committed; durable tombstone remains active", + ) from exc + except Exception as exc: + detail = exc.detail if isinstance(exc, NodeAPIError) else str(exc) + failed[db_node.id] = detail + logger.error( + 'Bulk node deletion failed for "%s" (id=%s): %s', + node_response.name, + db_node.id, + detail, + ) + continue + removed_names.append(node_response.name) + if force: + logger.warning( + 'Node "%s" force-deleted by admin "%s"; remote core may remain and requires manual decommission', + node_response.name, + admin.username, + ) logger.info(f'Node "{node_response.name}" with id "{node_response.id}" deleted by admin "{admin.username}"') asyncio.create_task(notification.remove_node(node_response, admin.username)) - return RemoveNodesResponse(nodes=node_names, count=len(db_nodes)) + return RemoveNodesResponse(nodes=removed_names, count=len(removed_names), failed=failed) async def _get_validated_nodes(self, db: AsyncSession, node_ids: list[int] | set[int]) -> list[Node]: if not node_ids: @@ -1209,3 +1494,21 @@ async def bulk_update_nodes( raise errors[0] return self._build_bulk_action_response(updated_nodes) + @staticmethod + async def _resolve_node_delete_after_db_error(db: AsyncSession, db_node: Node) -> bool: + """Resolve a commit-ACK loss using a fresh authoritative session.""" + try: + await db.rollback() + except Exception: + pass + try: + async with GetDB() as fresh_db: + current_bridge_id = await fresh_db.scalar( + select(Node.bridge_id).where(Node.id == db_node.id) + ) + except Exception as exc: + raise NodeAPIError( + 503, + "Node deletion database outcome is unknown; durable tombstone remains active", + ) from exc + return current_bridge_id is None or str(current_bridge_id) != str(db_node.bridge_id) diff --git a/app/operation/subscription.py b/app/operation/subscription.py index 4d345dce2..8c20d0235 100644 --- a/app/operation/subscription.py +++ b/app/operation/subscription.py @@ -1,4 +1,5 @@ import re +from datetime import timedelta as td from json import dumps as json_dumps from typing import Any, ClassVar @@ -12,7 +13,7 @@ register_user_hwid, ) from app.db.crud.user import get_user_usages, user_sub_update -from app.db.models import User +from app.db.models import User, UserStatus from app.models.admin import AdminDetails from app.models.settings import Application, ConfigFormat, HWIDSettings, SubRule, Subscription as SubSettings from app.models.stats import UserUsageStatsList @@ -87,6 +88,12 @@ class SubscriptionOperation(BaseOperation): _ENCODED_RULE_RESPONSE_HEADERS: ClassVar[set[str]] = {"announce", "profile-title"} + _CONFIG_ELIGIBLE_STATUSES: ClassVar[set[UserStatus]] = {UserStatus.active, UserStatus.on_hold} + _MAX_USAGE_RANGE_DAYS: ClassVar[int] = 31 + + async def require_config_eligible(self, user: User | UsersResponseWithInbounds) -> None: + if user.status not in self._CONFIG_ELIGIBLE_STATUSES: + await self.raise_error(message="Subscription is not active", code=403) @staticmethod async def validated_user(db_user: User) -> UsersResponseWithInbounds: @@ -434,7 +441,7 @@ async def user_subscription( not is_hwid_enabled or not global_hwid_conf.require_hwid_for_manual_sub ) links = [] - if is_allow_browser_config: + if is_allow_browser_config and user.status in self._CONFIG_ELIGIBLE_STATUSES: conf, media_type = await self.fetch_config( user, ConfigFormat.links, @@ -453,6 +460,7 @@ async def user_subscription( ) ) else: + await self.require_config_eligible(user) await self.validate_and_register_hwid( db, db_user.id, @@ -542,6 +550,7 @@ async def user_subscription_with_client_type( await self.raise_error(message="Client not supported", code=406) db_user = await self.get_validated_sub(db, token=token, load_admin_role=True) user = await self.validated_user(db_user) + await self.require_config_eligible(user) await self.validate_and_register_hwid( db, @@ -614,6 +623,7 @@ async def user_subscription_raw(self, db: AsyncSession, token: str, request_url: sub_settings: SubSettings = await subscription_settings() db_user = await self.get_validated_sub(db, token, load_admin_role=True) user = await self.validated_user(db_user) + await self.require_config_eligible(user) is_hwid_enabled = await self.is_user_hwid_enabled(db_user) links = [] @@ -684,6 +694,7 @@ async def user_subscription_info( sub_settings: SubSettings = await subscription_settings() db_user = await self.get_validated_sub(db, token=token) user = await self.validated_user(db_user) + await self.require_config_eligible(user) response_headers = self.create_info_response_headers(user, sub_settings) try: @@ -701,6 +712,7 @@ async def user_subscription_apps(self, db: AsyncSession, token: str) -> list[App """ db_user = await self.get_validated_sub(db, token=token, load_admin_role=True) user = await self.validated_user(db_user) + await self.require_config_eligible(user) is_hwid_enabled = await self.is_user_hwid_enabled(db_user) sub_settings: SubSettings = await subscription_settings() format_variables = await self.get_format_variables(user) @@ -747,6 +759,7 @@ async def user_subscription_headers( "content-type": "text/html; charset=utf-8", } else: + await self.require_config_eligible(user) matched_rule = self.detect_client_rule(user_agent, sub_settings.rules) client_type = matched_rule.target if matched_rule else None if client_type == ConfigFormat.block or not client_type: @@ -790,6 +803,11 @@ async def get_user_usage( ) -> UserUsageStatsList: """Fetches the usage statistics for the user within a specified date range.""" start, end = await self.validate_dates(query.start, query.end, True) + if end - start > td(days=self._MAX_USAGE_RANGE_DAYS): + await self.raise_error( + message=f"Subscription usage range cannot exceed {self._MAX_USAGE_RANGE_DAYS} days", + code=400, + ) db_user = await self.get_validated_sub(db, token=token) diff --git a/app/operation/user.py b/app/operation/user.py index cf0f4b8a9..5bc6ed257 100644 --- a/app/operation/user.py +++ b/app/operation/user.py @@ -45,7 +45,6 @@ load_user_attrs, lock_admin_quota_row, modify_user as crud_modify_user, - remove_expired_users, remove_user, remove_users, reset_user_by_next, @@ -96,7 +95,15 @@ UsersUsageQuery, UserUsageQuery, ) -from app.node.sync import remove_user as sync_remove_user, sync_user, sync_users +from app.node.sync import ( + finalize_user_removal, + finalize_users_removal, + remove_user as sync_remove_user, + remove_users_and_wait, + resolve_user_removal_after_db_error, + sync_user, + sync_users, +) from app.operation import BaseOperation, OperatorType from app.operation.permissions import ( PermissionDenied, @@ -940,8 +947,13 @@ async def set_user_disabled_by_id( async def _remove_user(self, db: AsyncSession, db_user: User, admin: AdminDetails) -> dict: user = await self.validate_user(db_user, include_subscription_url=False) - await remove_user(db, db_user) - await sync_remove_user(user) + revocation = await sync_remove_user(db_user) + try: + await remove_user(db, db_user) + except BaseException: + await resolve_user_removal_after_db_error(revocation, db) + raise + await finalize_user_removal(revocation) asyncio.create_task(notification.remove_user(user, admin)) logger.info(f'User "{db_user.username}" with id "{db_user.id}" deleted by admin "{admin.username}"') @@ -1013,10 +1025,16 @@ async def bulk_remove_users( db_users = await self._get_validated_users_by_ids(db, bulk_users.ids, admin, scope_action="delete") users = [await self.validate_user(db_user, include_subscription_url=False) for db_user in db_users] - await remove_users(db, db_users) + revocation = await remove_users_and_wait(db_users) + try: + await remove_users(db, db_users) + except BaseException: + await resolve_user_removal_after_db_error(revocation, db) + raise + if revocation is not None: + await finalize_users_removal(revocation) for user in users: - await sync_remove_user(user) asyncio.create_task(notification.remove_user(user, admin)) logger.info(f'User "{user.username}" with id "{user.id}" deleted by admin "{admin.username}"') @@ -1586,16 +1604,25 @@ async def delete_expired_users( admin_id = (await self.get_validated_admin(db, query.admin_username)).id else: admin_id = None - username_list = await remove_expired_users( + cleanup_query = query.model_copy(update={"expired_after": expired_after, "expired_before": expired_before}) + db_users = await get_expired_users( db, - expired_after, - expired_before, - admin_id, - target=query.target, - dry_run=query.dry_run, + query=cleanup_query, + admin_id=admin_id, ) - if not query.dry_run: - await self.remove_users_logger(users=username_list, by=admin.username) + username_list = [user.username for user in db_users] + if query.dry_run: + return RemoveUsersResponse(users=username_list, count=len(username_list)) + + revocation = await remove_users_and_wait(db_users) + try: + await remove_users(db, db_users) + except BaseException: + await resolve_user_removal_after_db_error(revocation, db) + raise + if revocation is not None: + await finalize_users_removal(revocation) + await self.remove_users_logger(users=username_list, by=admin.username) return RemoveUsersResponse(users=username_list, count=len(username_list)) diff --git a/app/routers/admin.py b/app/routers/admin.py index 7ae023427..13457e2e1 100644 --- a/app/routers/admin.py +++ b/app/routers/admin.py @@ -310,7 +310,7 @@ async def activate_all_disabled_users_by_id( return {} -@router.delete("/{username}/users", responses={403: responses._403, 404: responses._404}) +@router.delete("/{username}/users", responses={403: responses._403, 404: responses._404, 503: responses._503}) async def remove_all_users( username: str, db: AsyncSession = Depends(get_db), @@ -321,7 +321,9 @@ async def remove_all_users( return {"detail": f"operation has been successfuly done {deleted} users deleted"} -@router.delete("/by-username/{username}/users", responses={403: responses._403, 404: responses._404}) +@router.delete( + "/by-username/{username}/users", responses={403: responses._403, 404: responses._404, 503: responses._503} +) async def remove_all_users_by_username( username: str, db: AsyncSession = Depends(get_db), @@ -331,7 +333,7 @@ async def remove_all_users_by_username( return {"detail": f"operation has been successfuly done {deleted} users deleted"} -@router.delete("/by-id/{admin_id}/users", responses={403: responses._403, 404: responses._404}) +@router.delete("/by-id/{admin_id}/users", responses={403: responses._403, 404: responses._404, 503: responses._503}) async def remove_all_users_by_id( admin_id: int, db: AsyncSession = Depends(get_db), @@ -456,7 +458,7 @@ async def bulk_activate_all_disabled_users( @router.delete( "s/bulk/users", response_model=BulkAdminsActionResponse, - responses={400: responses._400, 403: responses._403, 404: responses._404}, + responses={400: responses._400, 403: responses._403, 404: responses._404, 503: responses._503}, ) async def bulk_remove_all_users( bulk_admins: BulkAdminSelection, diff --git a/app/routers/node.py b/app/routers/node.py index 505a84fa4..152fcbd12 100644 --- a/app/routers/node.py +++ b/app/routers/node.py @@ -16,6 +16,7 @@ NodeCoreUpdate, NodeCreate, NodeGeoFilesUpdate, + NodeLifecycleRecovery, NodeListQuery, NodeModify, NodeResponse, @@ -292,6 +293,17 @@ async def reconnect_node( return {} +@router.post("/{node_id}/lifecycle/recover") +async def recover_node_lifecycle( + node_id: int, + recovery: NodeLifecycleRecovery, + db: AsyncSession = Depends(get_db), + _: AdminDetails = Depends(require_permission("nodes", "reconnect")), +): + """Explicitly resolve an expired lifecycle operation after inspection.""" + return await node_operator.recover_node_lifecycle(db, node_id, recovery) + + @router.put("/{node_id}/sync") async def sync_node( node_id: int, @@ -305,11 +317,12 @@ async def sync_node( @router.delete("/{node_id}", status_code=status.HTTP_204_NO_CONTENT) async def remove_node( node_id: int, + force: bool = False, db: AsyncSession = Depends(get_db), admin: AdminDetails = Depends(require_permission("nodes", "delete")), ): """Remove a node and remove it from xray in the background.""" - await node_operator.remove_node(db=db, node_id=node_id, admin=admin) + await node_operator.remove_node(db=db, node_id=node_id, admin=admin, force=force) return {} @@ -420,11 +433,12 @@ async def clear_usage_data( ) async def bulk_delete_nodes( bulk_nodes: BulkNodeSelection, + force: bool = False, db: AsyncSession = Depends(get_db), admin: AdminDetails = Depends(require_permission("nodes", "delete")), ): """Delete selected nodes by ID.""" - return await node_operator.bulk_remove_nodes(db, bulk_nodes, admin) + return await node_operator.bulk_remove_nodes(db, bulk_nodes, admin, force=force) @router.post( diff --git a/app/routers/user.py b/app/routers/user.py index 305e13d0c..e3b989824 100644 --- a/app/routers/user.py +++ b/app/routers/user.py @@ -191,7 +191,9 @@ async def set_user_disabled_by_id( @router.delete( - "/{username}", responses={403: responses._403, 404: responses._404}, status_code=status.HTTP_204_NO_CONTENT + "/{username}", + responses={403: responses._403, 404: responses._404, 503: responses._503}, + status_code=status.HTTP_204_NO_CONTENT, ) async def remove_user( username: str, @@ -204,7 +206,7 @@ async def remove_user( @router.delete( "/by-username/{username}", - responses={403: responses._403, 404: responses._404}, + responses={403: responses._403, 404: responses._404, 503: responses._503}, status_code=status.HTTP_204_NO_CONTENT, ) async def remove_user_by_username( @@ -217,7 +219,7 @@ async def remove_user_by_username( @router.delete( "/by-id/{user_id}", - responses={403: responses._403, 404: responses._404}, + responses={403: responses._403, 404: responses._404, 503: responses._503}, status_code=status.HTTP_204_NO_CONTENT, ) async def remove_user_by_id( @@ -614,7 +616,7 @@ async def get_expired_users( return await user_operator.get_expired_users(db, query=query) -@router.delete("s/expired", response_model=RemoveUsersResponse) +@router.delete("s/expired", response_model=RemoveUsersResponse, responses={503: responses._503}) async def delete_expired_users( query: Annotated[ExpiredUsersQuery, Depends(get_expired_users_query)], db: AsyncSession = Depends(get_db), @@ -635,7 +637,7 @@ async def delete_expired_users( @router.post( "s/bulk/delete", response_model=RemoveUsersResponse, - responses={400: responses._400, 403: responses._403, 404: responses._404}, + responses={400: responses._400, 403: responses._403, 404: responses._404, 503: responses._503}, ) async def bulk_delete_users( bulk_users: BulkUsersSelection, diff --git a/app/utils/jwt.py b/app/utils/jwt.py index 277314963..7af5f57e7 100644 --- a/app/utils/jwt.py +++ b/app/utils/jwt.py @@ -24,16 +24,21 @@ async def create_admin_token(admin_id: int | None, username: str) -> str: data = {"sub": username, "access": "admin", "iat": datetime.now(UTC)} if admin_id is not None: data["aid"] = int(admin_id) - if jwt_settings.access_token_expire_minutes > 0: - expire = datetime.now(UTC) + timedelta(minutes=jwt_settings.access_token_expire_minutes) - data["exp"] = expire + expire = datetime.now(UTC) + timedelta(minutes=jwt_settings.access_token_expire_minutes) + data["exp"] = expire encoded_jwt = jwt.encode(data, await get_secret_key(), algorithm="HS256") return encoded_jwt async def get_admin_payload(token: str) -> dict | None: try: - payload = jwt.decode(token, await get_secret_key(), algorithms=["HS256"], leeway=5) + payload = jwt.decode( + token, + await get_secret_key(), + algorithms=["HS256"], + leeway=5, + options={"require": ["exp", "iat", "sub"]}, + ) username: str = payload.get("sub") access: str = payload.get("access") admin_id = payload.get("aid") diff --git a/app/utils/responses.py b/app/utils/responses.py index bcd8a5fb3..0bcdd8a9b 100644 --- a/app/utils/responses.py +++ b/app/utils/responses.py @@ -23,6 +23,10 @@ class Conflict(HTTPException): detail: str = "Entity already exists" +class ServiceUnavailable(HTTPException): + detail: str = "Service temporarily unavailable" + + _400 = {"description": "BadRequest Error", "model": HTTPException} _401 = { @@ -41,3 +45,5 @@ class Conflict(HTTPException): _404 = {"description": "NotFound Error", "model": NotFound} _409 = {"description": "Conflict Error", "model": Conflict} + +_503 = {"description": "ServiceUnavailable Error", "model": ServiceUnavailable} diff --git a/config.py b/config.py index 8fb04ab51..eb64398f2 100644 --- a/config.py +++ b/config.py @@ -130,7 +130,7 @@ def path(self) -> str: class JwtSettings(EnvSettings): - access_token_expire_minutes: int = Field(default=1440, validation_alias="JWT_ACCESS_TOKEN_EXPIRE_MINUTES") + access_token_expire_minutes: int = Field(gt=0, default=1440, validation_alias="JWT_ACCESS_TOKEN_EXPIRE_MINUTES") class TemplateSettings(EnvSettings): @@ -142,7 +142,7 @@ class TemplateSettings(EnvSettings): class UserCleanupSettings(EnvSettings): - autodelete_days: int = Field(default=-1, validation_alias="USERS_AUTODELETE_DAYS") + autodelete_days: int = Field(ge=-1, le=36500, default=-1, validation_alias="USERS_AUTODELETE_DAYS") include_limited_accounts: bool = Field(default=False, validation_alias="USER_AUTODELETE_INCLUDE_LIMITED_ACCOUNTS") diff --git a/dashboard/src/service/api/index.ts b/dashboard/src/service/api/index.ts index 3371339fc..533820a77 100644 --- a/dashboard/src/service/api/index.ts +++ b/dashboard/src/service/api/index.ts @@ -139,6 +139,7 @@ export type GetUsersParams = { admin?: string[] | null admin_ids?: number[] | null group?: number[] | null + no_group?: boolean search?: string | null status?: UserStatus | UserStatus[] | null sort?: string | null @@ -192,6 +193,10 @@ export type GetUsersSubUpdateChartParams = { end?: string | null } +export type BulkDeleteNodesParams = { + force?: boolean +} + export type ClearUsageDataParams = { start?: string | null end?: string | null @@ -216,6 +221,10 @@ export type SyncNodeParams = { flush_users?: boolean } +export type RemoveNodeParams = { + force?: boolean +} + export type ReconnectAllNodeParams = { core_id?: number | null } @@ -426,13 +435,6 @@ export type XrayMuxSettingsInputXudpConcurrency = number | null export type XrayMuxSettingsInputConcurrency = number | null -export interface XrayMuxSettingsInput { - enabled?: boolean - concurrency?: XrayMuxSettingsInputConcurrency - xudp_concurrency?: XrayMuxSettingsInputXudpConcurrency - xudp_proxy_udp_443?: Xudp -} - export type XrayFragmentSettingsMaxSplit = string | null export interface XrayFragmentSettings { @@ -455,6 +457,13 @@ export const Xudp = { skip: 'skip', } as const +export interface XrayMuxSettingsInput { + enabled?: boolean + concurrency?: XrayMuxSettingsInputConcurrency + xudp_concurrency?: XrayMuxSettingsInputXudpConcurrency + xudp_proxy_udp_443?: Xudp +} + export type XMuxSettingsHKeepAlivePeriod = number | null export type XMuxSettingsHMaxRequestTimes = string | null @@ -518,6 +527,18 @@ export type XHttpSettingsXPaddingBytes = string | null export type XHttpSettingsNoGrpcHeader = boolean | null +export type XHttpModes = (typeof XHttpModes)[keyof typeof XHttpModes] + +// eslint-disable-next-line @typescript-eslint/no-redeclare +export const XHttpModes = { + auto: 'auto', + 'packet-up': 'packet-up', + 'stream-up': 'stream-up', + 'stream-one': 'stream-one', +} as const + +export type XHttpSettingsMode = XHttpModes | null + export interface XHttpSettings { mode?: XHttpSettingsMode no_grpc_header?: XHttpSettingsNoGrpcHeader @@ -543,33 +564,21 @@ export interface XHttpSettings { download_settings?: XHttpSettingsDownloadSettings } -export type XHttpModes = (typeof XHttpModes)[keyof typeof XHttpModes] - -// eslint-disable-next-line @typescript-eslint/no-redeclare -export const XHttpModes = { - auto: 'auto', - 'packet-up': 'packet-up', - 'stream-up': 'stream-up', - 'stream-one': 'stream-one', -} as const +export type WorkerHealthError = string | null -export type XHttpSettingsMode = XHttpModes | null +export interface WorkerHealth { + status: string + response_time_ms?: WorkerHealthResponseTimeMs + error?: WorkerHealthError +} export interface WorkersHealth { scheduler: WorkerHealth node: WorkerHealth } -export type WorkerHealthError = string | null - export type WorkerHealthResponseTimeMs = number | null -export interface WorkerHealth { - status: string - response_time_ms?: WorkerHealthResponseTimeMs - error?: WorkerHealthError -} - export interface WireGuardSubnetUsage { subnet: string interface_tags: string[] @@ -672,18 +681,6 @@ export type UsersPermissionsActivateNextPlanAnyOf = { [key: string]: PermissionS export type UsersPermissionsActivateNextPlan = boolean | UsersPermissionsActivateNextPlanAnyOf | null -export type UsersPermissionsSetOwnerAnyOf = { [key: string]: PermissionScope | number } - -export type UsersPermissionsSetOwner = boolean | UsersPermissionsSetOwnerAnyOf | null - -export type UsersPermissionsRevokeSubAnyOf = { [key: string]: PermissionScope | number } - -export type UsersPermissionsRevokeSub = boolean | UsersPermissionsRevokeSubAnyOf | null - -export type UsersPermissionsResetUsageAnyOf = { [key: string]: PermissionScope | number } - -export type UsersPermissionsResetUsage = boolean | UsersPermissionsResetUsageAnyOf | null - export interface UsersPermissions { create?: UsersPermissionsCreate read?: UsersPermissionsRead @@ -696,6 +693,18 @@ export interface UsersPermissions { activate_next_plan?: UsersPermissionsActivateNextPlan } +export type UsersPermissionsSetOwnerAnyOf = { [key: string]: PermissionScope | number } + +export type UsersPermissionsSetOwner = boolean | UsersPermissionsSetOwnerAnyOf | null + +export type UsersPermissionsRevokeSubAnyOf = { [key: string]: PermissionScope | number } + +export type UsersPermissionsRevokeSub = boolean | UsersPermissionsRevokeSubAnyOf | null + +export type UsersPermissionsResetUsageAnyOf = { [key: string]: PermissionScope | number } + +export type UsersPermissionsResetUsage = boolean | UsersPermissionsResetUsageAnyOf | null + export type UsersPermissionsDeleteAnyOf = { [key: string]: PermissionScope | number } export type UsersPermissionsDelete = boolean | UsersPermissionsDeleteAnyOf | null @@ -726,13 +735,6 @@ export const UsernameGenerationStrategy = { export type UserUsageStatsListPeriod = Period | null -export interface UserUsageStat { - period_start: string - total_traffic: number -} - -export type UserUsageStatsListStats = { [key: string]: UserUsageStat[] } - export interface UserUsageStatsList { period?: UserUsageStatsListPeriod start: string @@ -740,6 +742,13 @@ export interface UserUsageStatsList { stats: UserUsageStatsListStats } +export interface UserUsageStat { + period_start: string + total_traffic: number +} + +export type UserUsageStatsListStats = { [key: string]: UserUsageStat[] } + export type UserTemplateSimpleName = string | null /** @@ -979,6 +988,9 @@ export type UserResponseNextPlan = NextPlanModel | null export type UserResponseHwidLimit = number | null +/** + * Per-user cleanup delay in days; -1 disables automatic deletion + */ export type UserResponseAutoDeleteInDays = number | null export type UserResponseGroupIds = number[] | null @@ -1012,6 +1024,7 @@ export interface UserResponse { on_hold_expire_duration?: UserResponseOnHoldExpireDuration on_hold_timeout?: UserResponseOnHoldTimeout group_ids?: UserResponseGroupIds + /** Per-user cleanup delay in days; -1 disables automatic deletion */ auto_delete_in_days?: UserResponseAutoDeleteInDays hwid_limit?: UserResponseHwidLimit next_plan?: UserResponseNextPlan @@ -1043,6 +1056,9 @@ export type UserModifyNextPlan = NextPlanModel | null export type UserModifyHwidLimit = number | null +/** + * Per-user cleanup delay in days; -1 disables automatic deletion + */ export type UserModifyAutoDeleteInDays = number | null export type UserModifyGroupIds = number[] | null @@ -1078,19 +1094,13 @@ export interface UserModify { on_hold_expire_duration?: UserModifyOnHoldExpireDuration on_hold_timeout?: UserModifyOnHoldTimeout group_ids?: UserModifyGroupIds + /** Per-user cleanup delay in days; -1 disables automatic deletion */ auto_delete_in_days?: UserModifyAutoDeleteInDays hwid_limit?: UserModifyHwidLimit next_plan?: UserModifyNextPlan status?: UserModifyStatus } -/** - * User IP lists for all nodes - */ -export interface UserIPListAll { - nodes: UserIPListAllNodes -} - export type UserIPListIps = { [key: string]: number } /** @@ -1102,6 +1112,13 @@ export interface UserIPList { export type UserIPListAllNodes = { [key: string]: UserIPList | null } +/** + * User IP lists for all nodes + */ +export interface UserIPListAll { + nodes: UserIPListAllNodes +} + export type UserHWIDResponseDeviceModel = string | null export type UserHWIDResponseOsVersion = string | null @@ -1129,6 +1146,9 @@ export type UserCreateNextPlan = NextPlanModel | null export type UserCreateHwidLimit = number | null +/** + * Per-user cleanup delay in days; -1 disables automatic deletion + */ export type UserCreateAutoDeleteInDays = number | null export type UserCreateGroupIds = number[] | null @@ -1162,6 +1182,7 @@ export interface UserCreate { on_hold_expire_duration?: UserCreateOnHoldExpireDuration on_hold_timeout?: UserCreateOnHoldTimeout group_ids?: UserCreateGroupIds + /** Per-user cleanup delay in days; -1 disables automatic deletion */ auto_delete_in_days?: UserCreateAutoDeleteInDays hwid_limit?: UserCreateHwidLimit next_plan?: UserCreateNextPlan @@ -1522,6 +1543,10 @@ export type SettingsPermissionsReadAnyOf = { [key: string]: PermissionScope | nu export type SettingsPermissionsRead = boolean | SettingsPermissionsReadAnyOf | null +export interface ServiceUnavailable { + detail?: string +} + export type RunMethod = (typeof RunMethod)[keyof typeof RunMethod] // eslint-disable-next-line @typescript-eslint/no-redeclare @@ -1650,12 +1675,15 @@ export interface RemoveUserTemplatesResponse { count: number } +export type RemoveNodesResponseFailed = { [key: string]: string } + /** * Response model for bulk node deletion */ export interface RemoveNodesResponse { nodes: string[] count: number + failed?: RemoveNodesResponseFailed } /** @@ -1986,6 +2014,14 @@ export type NodesPermissionsReadSimpleAnyOf = { [key: string]: PermissionScope | export type NodesPermissionsReadSimple = boolean | NodesPermissionsReadSimpleAnyOf | null +export type NodesPermissionsReadAnyOf = { [key: string]: PermissionScope | number } + +export type NodesPermissionsRead = boolean | NodesPermissionsReadAnyOf | null + +export type NodesPermissionsCreateAnyOf = { [key: string]: PermissionScope | number } + +export type NodesPermissionsCreate = boolean | NodesPermissionsCreateAnyOf | null + export interface NodesPermissions { create?: NodesPermissionsCreate read?: NodesPermissionsRead @@ -1998,24 +2034,10 @@ export interface NodesPermissions { stats?: NodesPermissionsStats } -export type NodesPermissionsReadAnyOf = { [key: string]: PermissionScope | number } - -export type NodesPermissionsRead = boolean | NodesPermissionsReadAnyOf | null - -export type NodesPermissionsCreateAnyOf = { [key: string]: PermissionScope | number } - -export type NodesPermissionsCreate = boolean | NodesPermissionsCreateAnyOf | null +export type NodeUsageStatsListStats = { [key: string]: NodeUsageStat[] } export type NodeUsageStatsListPeriod = Period | null -export interface NodeUsageStat { - period_start: string - uplink: number - downlink: number -} - -export type NodeUsageStatsListStats = { [key: string]: NodeUsageStat[] } - export interface NodeUsageStatsList { period?: NodeUsageStatsListPeriod start: string @@ -2023,6 +2045,12 @@ export interface NodeUsageStatsList { stats: NodeUsageStatsListStats } +export interface NodeUsageStat { + period_start: string + uplink: number + downlink: number +} + export type NodeStatus = (typeof NodeStatus)[keyof typeof NodeStatus] // eslint-disable-next-line @typescript-eslint/no-redeclare @@ -2207,12 +2235,33 @@ export interface NodeModify { status?: NodeModifyStatus } +/** + * Explicit acknowledgement for an expired, outcome-unknown operation. + */ +export interface NodeLifecycleRecovery { + observed: LifecycleStatus + acknowledge_expired_operation: true +} + export interface NodeGeoFilesUpdate { region?: GeoFilseRegion } export type NodeCreateProxyUrl = string | null +export interface NodeCoreUpdate { + /** @pattern ^(latest|v?\d+\.\d+\.\d+)$ */ + core_version?: string +} + +export type NodeConnectionType = (typeof NodeConnectionType)[keyof typeof NodeConnectionType] + +// eslint-disable-next-line @typescript-eslint/no-redeclare +export const NodeConnectionType = { + grpc: 'grpc', + rest: 'rest', +} as const + export interface NodeCreate { name: string address: string @@ -2241,19 +2290,6 @@ export interface NodeCreate { proxy_url?: NodeCreateProxyUrl } -export interface NodeCoreUpdate { - /** @pattern ^(latest|v?\d+\.\d+\.\d+)$ */ - core_version?: string -} - -export type NodeConnectionType = (typeof NodeConnectionType)[keyof typeof NodeConnectionType] - -// eslint-disable-next-line @typescript-eslint/no-redeclare -export const NodeConnectionType = { - grpc: 'grpc', - rest: 'rest', -} as const - export type NextPlanModelExpire = number | null export type NextPlanModelDataLimit = number | null @@ -2307,6 +2343,18 @@ export interface ModifyUserByTemplate { note?: ModifyUserByTemplateNote } +export type LifecycleStatus = (typeof LifecycleStatus)[keyof typeof LifecycleStatus] + +// eslint-disable-next-line @typescript-eslint/no-redeclare +export const LifecycleStatus = { + unknown: 'unknown', + starting: 'starting', + healthy: 'healthy', + stopping: 'stopping', + stopped: 'stopped', + broken: 'broken', +} as const + export type Language = (typeof Language)[keyof typeof Language] // eslint-disable-next-line @typescript-eslint/no-redeclare @@ -2358,19 +2406,25 @@ export type HwidsPermissionsDeleteAnyOf = { [key: string]: PermissionScope | num export type HwidsPermissionsDelete = boolean | HwidsPermissionsDeleteAnyOf | null -export type HwidsPermissionsReadAnyOf = { [key: string]: PermissionScope | number } - -export type HwidsPermissionsRead = boolean | HwidsPermissionsReadAnyOf | null - export interface HwidsPermissions { read?: HwidsPermissionsRead delete?: HwidsPermissionsDelete } +export type HwidsPermissionsReadAnyOf = { [key: string]: PermissionScope | number } + +export type HwidsPermissionsRead = boolean | HwidsPermissionsReadAnyOf | null + export type HostsPermissionsUpdateAnyOf = { [key: string]: PermissionScope | number } export type HostsPermissionsUpdate = boolean | HostsPermissionsUpdateAnyOf | null +export interface HostsPermissions { + create?: HostsPermissionsCreate + read?: HostsPermissionsRead + update?: HostsPermissionsUpdate +} + export type HostsPermissionsReadAnyOf = { [key: string]: PermissionScope | number } export type HostsPermissionsRead = boolean | HostsPermissionsReadAnyOf | null @@ -2379,12 +2433,6 @@ export type HostsPermissionsCreateAnyOf = { [key: string]: PermissionScope | num export type HostsPermissionsCreate = boolean | HostsPermissionsCreateAnyOf | null -export interface HostsPermissions { - create?: HostsPermissionsCreate - read?: HostsPermissionsRead - update?: HostsPermissionsUpdate -} - export interface HostNotificationEnable { create?: boolean modify?: boolean @@ -2625,14 +2673,27 @@ export const FinalMaskUdpType = { 'mkcp-aes128gcm': 'mkcp-aes128gcm', } as const +export type FinalMaskUdpLayerSettingsAnyOf = { [key: string]: unknown } + +export type FinalMaskUdpLayerSettings = + | FinalMaskUdpHeaderCustomSettings + | FinalMaskPasswordSettings + | FinalMaskSudokuSettings + | FinalMaskDomainSettings + | FinalMaskXdnsSettings + | FinalMaskXicmpSettings + | FinalMaskNoiseSettings + | FinalMaskSalamanderSettings + | FinalMaskRealmSettings + | FinalMaskMkcpLegacySettings + | FinalMaskUdpLayerSettingsAnyOf + export interface FinalMaskUdpLayer { type: FinalMaskUdpType settings?: FinalMaskUdpLayerSettings [key: string]: unknown } -export type FinalMaskUdpLayerSettingsAnyOf = { [key: string]: unknown } - export type FinalMaskUdpHopInterval = string | number | null export type FinalMaskUdpHopPorts = string | null @@ -2663,14 +2724,16 @@ export const FinalMaskTcpType = { xmc: 'xmc', } as const +export type FinalMaskTcpLayerSettingsAnyOf = { [key: string]: unknown } + +export type FinalMaskTcpLayerSettings = FinalMaskTcpHeaderCustomSettings | FinalMaskFragmentSettings | FinalMaskSudokuSettings | FinalMaskXmcSettings | FinalMaskTcpLayerSettingsAnyOf + export interface FinalMaskTcpLayer { type: FinalMaskTcpType settings?: FinalMaskTcpLayerSettings [key: string]: unknown } -export type FinalMaskTcpLayerSettingsAnyOf = { [key: string]: unknown } - export type FinalMaskTcpHeaderCustomSettingsErrors = XrayNoiseSettings[][] | null export type FinalMaskTcpHeaderCustomSettingsServers = XrayNoiseSettings[][] | null @@ -2684,8 +2747,6 @@ export interface FinalMaskTcpHeaderCustomSettings { [key: string]: unknown } -export type FinalMaskTcpLayerSettings = FinalMaskTcpHeaderCustomSettings | FinalMaskFragmentSettings | FinalMaskSudokuSettings | FinalMaskXmcSettings | FinalMaskTcpLayerSettingsAnyOf - export type FinalMaskSudokuSettingsPaddingMax = number | null export type FinalMaskSudokuSettingsPaddingMin = number | null @@ -2718,19 +2779,6 @@ export interface FinalMaskSalamanderSettings { [key: string]: unknown } -export type FinalMaskUdpLayerSettings = - | FinalMaskUdpHeaderCustomSettings - | FinalMaskPasswordSettings - | FinalMaskSudokuSettings - | FinalMaskDomainSettings - | FinalMaskXdnsSettings - | FinalMaskXicmpSettings - | FinalMaskNoiseSettings - | FinalMaskSalamanderSettings - | FinalMaskRealmSettings - | FinalMaskMkcpLegacySettings - | FinalMaskUdpLayerSettingsAnyOf - export type FinalMaskRealmSettingsTlsConfigAnyOf = { [key: string]: unknown } export type FinalMaskRealmSettingsTlsConfig = FinalMaskRealmSettingsTlsConfigAnyOf | null @@ -3170,6 +3218,10 @@ export type CRUDPermissionsReadSimpleAnyOf = { [key: string]: PermissionScope | export type CRUDPermissionsReadSimple = boolean | CRUDPermissionsReadSimpleAnyOf | null +export type CRUDPermissionsReadAnyOf = { [key: string]: PermissionScope | number } + +export type CRUDPermissionsRead = boolean | CRUDPermissionsReadAnyOf | null + /** * Standard create/read/read_simple/update/delete permissions. Used directly by: groups, templates, client_templates, cores, admin_roles. @@ -3183,10 +3235,6 @@ export interface CRUDPermissions { delete?: CRUDPermissionsDelete } -export type CRUDPermissionsReadAnyOf = { [key: string]: PermissionScope | number } - -export type CRUDPermissionsRead = boolean | CRUDPermissionsReadAnyOf | null - export type CRUDPermissionsCreateAnyOf = { [key: string]: PermissionScope | number } export type CRUDPermissionsCreate = boolean | CRUDPermissionsCreateAnyOf | null @@ -3338,6 +3386,7 @@ export interface BulkGroupSelection { export interface BulkGroup { group_ids: number[] has_group_ids?: number[] + has_no_group?: boolean admins?: number[] users?: number[] dry_run?: boolean @@ -5209,7 +5258,7 @@ export const removeAllUsers = (username: string) => { export const getRemoveAllUsersMutationOptions = < TData = Awaited>, - TError = ErrorType, + TError = ErrorType, TContext = unknown, >(options?: { mutation?: UseMutationOptions @@ -5232,12 +5281,16 @@ export const getRemoveAllUsersMutationOptions = < export type RemoveAllUsersMutationResult = NonNullable>> -export type RemoveAllUsersMutationError = ErrorType +export type RemoveAllUsersMutationError = ErrorType /** * @summary Remove All Users */ -export const useRemoveAllUsers = >, TError = ErrorType, TContext = unknown>(options?: { +export const useRemoveAllUsers = < + TData = Awaited>, + TError = ErrorType, + TContext = unknown, +>(options?: { mutation?: UseMutationOptions }): UseMutationResult => { const mutationOptions = getRemoveAllUsersMutationOptions(options) @@ -5254,7 +5307,7 @@ export const removeAllUsersByUsername = (username: string) => { export const getRemoveAllUsersByUsernameMutationOptions = < TData = Awaited>, - TError = ErrorType, + TError = ErrorType, TContext = unknown, >(options?: { mutation?: UseMutationOptions @@ -5277,14 +5330,14 @@ export const getRemoveAllUsersByUsernameMutationOptions = < export type RemoveAllUsersByUsernameMutationResult = NonNullable>> -export type RemoveAllUsersByUsernameMutationError = ErrorType +export type RemoveAllUsersByUsernameMutationError = ErrorType /** * @summary Remove All Users By Username */ export const useRemoveAllUsersByUsername = < TData = Awaited>, - TError = ErrorType, + TError = ErrorType, TContext = unknown, >(options?: { mutation?: UseMutationOptions @@ -5303,7 +5356,7 @@ export const removeAllUsersById = (adminId: number) => { export const getRemoveAllUsersByIdMutationOptions = < TData = Awaited>, - TError = ErrorType, + TError = ErrorType, TContext = unknown, >(options?: { mutation?: UseMutationOptions @@ -5326,14 +5379,14 @@ export const getRemoveAllUsersByIdMutationOptions = < export type RemoveAllUsersByIdMutationResult = NonNullable>> -export type RemoveAllUsersByIdMutationError = ErrorType +export type RemoveAllUsersByIdMutationError = ErrorType /** * @summary Remove All Users By Id */ export const useRemoveAllUsersById = < TData = Awaited>, - TError = ErrorType, + TError = ErrorType, TContext = unknown, >(options?: { mutation?: UseMutationOptions @@ -5797,7 +5850,7 @@ export const bulkRemoveAllUsers = (bulkAdminSelection: BodyType>, - TError = ErrorType, + TError = ErrorType, TContext = unknown, >(options?: { mutation?: UseMutationOptions }, TContext> @@ -5820,14 +5873,14 @@ export const getBulkRemoveAllUsersMutationOptions = < export type BulkRemoveAllUsersMutationResult = NonNullable>> export type BulkRemoveAllUsersMutationBody = BodyType -export type BulkRemoveAllUsersMutationError = ErrorType +export type BulkRemoveAllUsersMutationError = ErrorType /** * @summary Bulk Remove All Users */ export const useBulkRemoveAllUsers = < TData = Awaited>, - TError = ErrorType, + TError = ErrorType, TContext = unknown, >(options?: { mutation?: UseMutationOptions }, TContext> @@ -9599,12 +9652,12 @@ export const useModifyNode = >, TE * Remove a node and remove it from xray in the background. * @summary Remove Node */ -export const removeNode = (nodeId: number) => { - return orvalFetcher({ url: `/api/node/${nodeId}`, method: 'DELETE' }) +export const removeNode = (nodeId: number, params?: RemoveNodeParams) => { + return orvalFetcher({ url: `/api/node/${nodeId}`, method: 'DELETE', params }) } export const getRemoveNodeMutationOptions = >, TError = ErrorType, TContext = unknown>(options?: { - mutation?: UseMutationOptions + mutation?: UseMutationOptions }) => { const mutationKey = ['removeNode'] const { mutation: mutationOptions } = options @@ -9613,13 +9666,13 @@ export const getRemoveNodeMutationOptions = >, { nodeId: number }> = props => { - const { nodeId } = props ?? {} + const mutationFn: MutationFunction>, { nodeId: number; params?: RemoveNodeParams }> = props => { + const { nodeId, params } = props ?? {} - return removeNode(nodeId) + return removeNode(nodeId, params) } - return { mutationFn, ...mutationOptions } as UseMutationOptions + return { mutationFn, ...mutationOptions } as UseMutationOptions } export type RemoveNodeMutationResult = NonNullable>> @@ -9630,8 +9683,8 @@ export type RemoveNodeMutationError = ErrorType>, TError = ErrorType, TContext = unknown>(options?: { - mutation?: UseMutationOptions -}): UseMutationResult => { + mutation?: UseMutationOptions +}): UseMutationResult => { const mutationOptions = getRemoveNodeMutationOptions(options) return useMutation(mutationOptions) @@ -9846,6 +9899,52 @@ export const useReconnectNode = , signal?: AbortSignal) => { + return orvalFetcher({ url: `/api/node/${nodeId}/lifecycle/recover`, method: 'POST', headers: { 'Content-Type': 'application/json' }, data: nodeLifecycleRecovery, signal }) +} + +export const getRecoverNodeLifecycleMutationOptions = < + TData = Awaited>, + TError = ErrorType, + TContext = unknown, +>(options?: { + mutation?: UseMutationOptions }, TContext> +}) => { + const mutationKey = ['recoverNodeLifecycle'] + const { mutation: mutationOptions } = options + ? options.mutation && 'mutationKey' in options.mutation && options.mutation.mutationKey + ? options + : { ...options, mutation: { ...options.mutation, mutationKey } } + : { mutation: { mutationKey } } + + const mutationFn: MutationFunction>, { nodeId: number; data: BodyType }> = props => { + const { nodeId, data } = props ?? {} + + return recoverNodeLifecycle(nodeId, data) + } + + return { mutationFn, ...mutationOptions } as UseMutationOptions }, TContext> +} + +export type RecoverNodeLifecycleMutationResult = NonNullable>> +export type RecoverNodeLifecycleMutationBody = BodyType +export type RecoverNodeLifecycleMutationError = ErrorType + +/** + * @summary Recover Node Lifecycle + */ +export const useRecoverNodeLifecycle = >, TError = ErrorType, TContext = unknown>(options?: { + mutation?: UseMutationOptions }, TContext> +}): UseMutationResult }, TContext> => { + const mutationOptions = getRecoverNodeLifecycleMutationOptions(options) + + return useMutation(mutationOptions) +} + /** * @summary Sync Node */ @@ -10461,8 +10560,8 @@ export const useClearUsageData = , signal?: AbortSignal) => { - return orvalFetcher({ url: `/api/nodes/bulk/delete`, method: 'POST', headers: { 'Content-Type': 'application/json' }, data: bulkNodeSelection, signal }) +export const bulkDeleteNodes = (bulkNodeSelection: BodyType, params?: BulkDeleteNodesParams, signal?: AbortSignal) => { + return orvalFetcher({ url: `/api/nodes/bulk/delete`, method: 'POST', headers: { 'Content-Type': 'application/json' }, data: bulkNodeSelection, params, signal }) } export const getBulkDeleteNodesMutationOptions = < @@ -10470,7 +10569,7 @@ export const getBulkDeleteNodesMutationOptions = < TError = ErrorType, TContext = unknown, >(options?: { - mutation?: UseMutationOptions }, TContext> + mutation?: UseMutationOptions; params?: BulkDeleteNodesParams }, TContext> }) => { const mutationKey = ['bulkDeleteNodes'] const { mutation: mutationOptions } = options @@ -10479,13 +10578,13 @@ export const getBulkDeleteNodesMutationOptions = < : { ...options, mutation: { ...options.mutation, mutationKey } } : { mutation: { mutationKey } } - const mutationFn: MutationFunction>, { data: BodyType }> = props => { - const { data } = props ?? {} + const mutationFn: MutationFunction>, { data: BodyType; params?: BulkDeleteNodesParams }> = props => { + const { data, params } = props ?? {} - return bulkDeleteNodes(data) + return bulkDeleteNodes(data, params) } - return { mutationFn, ...mutationOptions } as UseMutationOptions }, TContext> + return { mutationFn, ...mutationOptions } as UseMutationOptions; params?: BulkDeleteNodesParams }, TContext> } export type BulkDeleteNodesMutationResult = NonNullable>> @@ -10500,8 +10599,8 @@ export const useBulkDeleteNodes = < TError = ErrorType, TContext = unknown, >(options?: { - mutation?: UseMutationOptions }, TContext> -}): UseMutationResult }, TContext> => { + mutation?: UseMutationOptions; params?: BulkDeleteNodesParams }, TContext> +}): UseMutationResult; params?: BulkDeleteNodesParams }, TContext> => { const mutationOptions = getBulkDeleteNodesMutationOptions(options) return useMutation(mutationOptions) @@ -10889,7 +10988,7 @@ export const removeUser = (username: string) => { export const getRemoveUserMutationOptions = < TData = Awaited>, - TError = ErrorType, + TError = ErrorType, TContext = unknown, >(options?: { mutation?: UseMutationOptions @@ -10912,12 +11011,16 @@ export const getRemoveUserMutationOptions = < export type RemoveUserMutationResult = NonNullable>> -export type RemoveUserMutationError = ErrorType +export type RemoveUserMutationError = ErrorType /** * @summary Remove User */ -export const useRemoveUser = >, TError = ErrorType, TContext = unknown>(options?: { +export const useRemoveUser = < + TData = Awaited>, + TError = ErrorType, + TContext = unknown, +>(options?: { mutation?: UseMutationOptions }): UseMutationResult => { const mutationOptions = getRemoveUserMutationOptions(options) @@ -11042,7 +11145,7 @@ export const removeUserByUsername = (username: string) => { export const getRemoveUserByUsernameMutationOptions = < TData = Awaited>, - TError = ErrorType, + TError = ErrorType, TContext = unknown, >(options?: { mutation?: UseMutationOptions @@ -11065,14 +11168,14 @@ export const getRemoveUserByUsernameMutationOptions = < export type RemoveUserByUsernameMutationResult = NonNullable>> -export type RemoveUserByUsernameMutationError = ErrorType +export type RemoveUserByUsernameMutationError = ErrorType /** * @summary Remove User By Username */ export const useRemoveUserByUsername = < TData = Awaited>, - TError = ErrorType, + TError = ErrorType, TContext = unknown, >(options?: { mutation?: UseMutationOptions @@ -11202,7 +11305,7 @@ export const removeUserById = (userId: number) => { export const getRemoveUserByIdMutationOptions = < TData = Awaited>, - TError = ErrorType, + TError = ErrorType, TContext = unknown, >(options?: { mutation?: UseMutationOptions @@ -11225,12 +11328,16 @@ export const getRemoveUserByIdMutationOptions = < export type RemoveUserByIdMutationResult = NonNullable>> -export type RemoveUserByIdMutationError = ErrorType +export type RemoveUserByIdMutationError = ErrorType /** * @summary Remove User By Id */ -export const useRemoveUserById = >, TError = ErrorType, TContext = unknown>(options?: { +export const useRemoveUserById = < + TData = Awaited>, + TError = ErrorType, + TContext = unknown, +>(options?: { mutation?: UseMutationOptions }): UseMutationResult => { const mutationOptions = getRemoveUserByIdMutationOptions(options) @@ -12936,7 +13043,11 @@ export const deleteExpiredUsers = (params?: DeleteExpiredUsersParams) => { return orvalFetcher({ url: `/api/users/expired`, method: 'DELETE', params }) } -export const getDeleteExpiredUsersMutationOptions = >, TError = ErrorType, TContext = unknown>(options?: { +export const getDeleteExpiredUsersMutationOptions = < + TData = Awaited>, + TError = ErrorType, + TContext = unknown, +>(options?: { mutation?: UseMutationOptions }) => { const mutationKey = ['deleteExpiredUsers'] @@ -12957,12 +13068,16 @@ export const getDeleteExpiredUsersMutationOptions = >> -export type DeleteExpiredUsersMutationError = ErrorType +export type DeleteExpiredUsersMutationError = ErrorType /** * @summary Delete Expired Users */ -export const useDeleteExpiredUsers = >, TError = ErrorType, TContext = unknown>(options?: { +export const useDeleteExpiredUsers = < + TData = Awaited>, + TError = ErrorType, + TContext = unknown, +>(options?: { mutation?: UseMutationOptions }): UseMutationResult => { const mutationOptions = getDeleteExpiredUsersMutationOptions(options) @@ -12980,7 +13095,7 @@ export const bulkDeleteUsers = (bulkUsersSelection: BodyType export const getBulkDeleteUsersMutationOptions = < TData = Awaited>, - TError = ErrorType, + TError = ErrorType, TContext = unknown, >(options?: { mutation?: UseMutationOptions }, TContext> @@ -13003,14 +13118,14 @@ export const getBulkDeleteUsersMutationOptions = < export type BulkDeleteUsersMutationResult = NonNullable>> export type BulkDeleteUsersMutationBody = BodyType -export type BulkDeleteUsersMutationError = ErrorType +export type BulkDeleteUsersMutationError = ErrorType /** * @summary Bulk Delete Users */ export const useBulkDeleteUsers = < TData = Awaited>, - TError = ErrorType, + TError = ErrorType, TContext = unknown, >(options?: { mutation?: UseMutationOptions }, TContext> diff --git a/install_service.sh b/install_service.sh index 97695fd43..0c8c04afb 100755 --- a/install_service.sh +++ b/install_service.sh @@ -3,20 +3,37 @@ SERVICE_NAME="pasarguard" SERVICE_DESCRIPTION="PasarGuard Service" SERVICE_DOCUMENTATION="https://github.com/pasarguard/panel" -MAIN_PY_PATH="$PWD/main.py" SERVICE_FILE="/etc/systemd/system/$SERVICE_NAME.service" +install_dir_with_sentinel="$(pwd -P && printf x)" || { + echo "Failed to resolve the installation directory." >&2 + exit 1 +} +# The sentinel prevents command substitution from stripping a trailing newline +# from a valid (but unsupported) filesystem path before control-character +# validation below can reject it. +INSTALL_DIR=${install_dir_with_sentinel%x} +if printf '%s' "$INSTALL_DIR" | LC_ALL=C grep -q '[[:cntrl:]]'; then + echo "Refusing to install from a path containing control characters." >&2 + exit 1 +fi + +# Escape characters that are significant inside systemd's double-quoted values. +SYSTEMD_INSTALL_DIR=${INSTALL_DIR//\\/\\\\} +SYSTEMD_INSTALL_DIR=${SYSTEMD_INSTALL_DIR//\"/\\\"} +SYSTEMD_INSTALL_DIR=${SYSTEMD_INSTALL_DIR//%/%%} + # Create the service file -cat > $SERVICE_FILE < "$SERVICE_FILE" <=0.22.1; sys_platform != 'win32'", "aiorwlock>=1.5.0", "typer>=0.27.1", - "pasarguard-node-bridge>=0.9.0", + "pasarguard-node-bridge>=0.10.0", "pip-system-certs>=5.3", "nats-py>=2.15.0", ] +[tool.uv.sources] +# Temporary release bridge for the coordinated Node/Panel rollout. Replace +# this immutable commit with the PyPI source after pasarguard-node-bridge 0.10.0 ships. +pasarguard-node-bridge = { git = "https://github.com/PasarGuard/node_bridge_py.git", rev = "bb503222e373135e8166ddd025fc51348f0806b6" } + [tool.ruff] line-length = 120 indent-width = 4 diff --git a/tests/api/test_node.py b/tests/api/test_node.py index 546e4823e..ec3b9a177 100644 --- a/tests/api/test_node.py +++ b/tests/api/test_node.py @@ -443,15 +443,15 @@ async def test_core_users_only_excludes_admins_with_blocking_sync_roles(monkeypa ) await session.commit() - expected_user_ids = {active_user.id, nonblocked_user.id} - blocked_user_id = blocked_user.id + expected_user_sync_ids = {active_user.sync_id, nonblocked_user.sync_id} + blocked_user_sync_id = blocked_user.sync_id async with TestSession() as session: users = await node_user_module.core_users(session, inbound_tags=[inbound_tag]) synced_user_ids = {user["id"] for user in users} - assert synced_user_ids == expected_user_ids - assert blocked_user_id not in synced_user_ids + assert synced_user_ids == expected_user_sync_ids + assert blocked_user_sync_id not in synced_user_ids assert all(user["inbounds"] == [inbound_tag] for user in users) @@ -831,6 +831,13 @@ def test_remove_node(access_token, node_operator_mock): assert response.status_code == status.HTTP_204_NO_CONTENT awaited_kwargs = node_operator_mock.remove_node.await_args.kwargs assert awaited_kwargs["node_id"] == 6 + assert awaited_kwargs["force"] is False + + +def test_force_remove_node_requires_explicit_query(access_token, node_operator_mock): + response = client.delete("/api/node/6?force=true", headers=auth_headers(access_token)) + assert response.status_code == status.HTTP_204_NO_CONTENT + assert node_operator_mock.remove_node.await_args.kwargs["force"] is True def test_bulk_disable_nodes(access_token, node_operator_mock): diff --git a/tests/api/test_usage_functions_timezone.py b/tests/api/test_usage_functions_timezone.py index db4be3456..723b9c437 100644 --- a/tests/api/test_usage_functions_timezone.py +++ b/tests/api/test_usage_functions_timezone.py @@ -17,6 +17,7 @@ from app.db.models import ( Admin, Node, + NodeStatus, NodeUsage, NodeUserUsage, User, @@ -61,6 +62,7 @@ async def setup_test_data(session, test_suffix=""): server_ca="ca", api_key="key", core_config_id=None, + status=NodeStatus.disabled, ) session.add(node) await session.flush() @@ -1009,6 +1011,7 @@ async def test_node_grouping_node_filter_and_admin_filter(self): server_ca="ca", api_key="key", core_config_id=None, + status=NodeStatus.disabled, ) session.add_all([admin_two, node_two]) await session.flush() diff --git a/tests/api/test_user.py b/tests/api/test_user.py index f5ce78f1c..cb7532d76 100644 --- a/tests/api/test_user.py +++ b/tests/api/test_user.py @@ -854,6 +854,36 @@ def test_user_subscriptions(access_token): cleanup_groups(access_token, core, groups) +def test_disabled_user_cannot_download_subscription_configs(access_token): + core, groups = setup_groups(access_token, 1) + hosts = create_hosts_for_inbounds(access_token) + user = create_user( + access_token, + group_ids=[group["id"] for group in groups], + payload={"username": unique_name("disabled_subscription")}, + ) + try: + disable_response = client.put( + f"/api/user/by-id/{user['id']}/disabled", + headers=auth_headers(access_token), + json={"disabled": True}, + ) + assert disable_response.status_code == status.HTTP_200_OK + + for config_format in ("links", "xray", "clash", "sing_box", "raw", "info", "apps"): + response = client.get(f"{user['subscription_url']}/{config_format}") + assert response.status_code == status.HTTP_403_FORBIDDEN + + page_response = client.get(user["subscription_url"], headers={"Accept": "text/html"}) + assert page_response.status_code == status.HTTP_200_OK + assert 'class="link-input"' not in page_response.text + finally: + delete_user(access_token, user["username"]) + for host in hosts: + client.delete(f"/api/host/{host['id']}", headers=auth_headers(access_token)) + cleanup_groups(access_token, core, groups) + + def test_user_subscription_head_route(access_token): """Test that HEAD /{token} returns headers without a body.""" core, groups = setup_groups(access_token, 1) diff --git a/tests/test_app_factory.py b/tests/test_app_factory.py index b6abb0f6b..57530bff7 100644 --- a/tests/test_app_factory.py +++ b/tests/test_app_factory.py @@ -4,7 +4,8 @@ from sqlalchemy.exc import DBAPIError, OperationalError from starlette.requests import Request -from app.app_factory import database_operational_error_handler +from app.app_factory import database_operational_error_handler, node_revocation_error_handler +from app.node.errors import NodeRevocationError @pytest.mark.asyncio @@ -27,3 +28,16 @@ async def test_database_operational_error_handler_handles_dbapi_errors(): assert response.status_code == 503 assert json.loads(response.body) == {"detail": "Database temporarily unavailable"} + + +@pytest.mark.asyncio +async def test_node_revocation_error_handler_returns_retryable_503(): + request = Request({"type": "http", "method": "DELETE", "path": "/api/user/1", "headers": []}) + + response = await node_revocation_error_handler(request, NodeRevocationError("node unavailable")) + + assert response.status_code == 503 + assert json.loads(response.body) == { + "detail": "User removal was not confirmed by all runtime nodes. Retry when nodes are available." + } + assert response.headers["retry-after"] == "1" diff --git a/tests/test_connect_concurrency.py b/tests/test_connect_concurrency.py index eac83a3c5..59f933ab8 100644 --- a/tests/test_connect_concurrency.py +++ b/tests/test_connect_concurrency.py @@ -15,7 +15,8 @@ async def test_connect_nodes_bulk_local_caps_concurrency(monkeypatch: pytest.Mon current = 0 peak = 0 - async def _connect_node(db_node, core, users): + async def _connect_node(db_node, core, users, authoritative_user_keys): + assert authoritative_user_keys == set() nonlocal current, peak current += 1 peak = max(peak, current) @@ -31,7 +32,11 @@ async def _connect_node(db_node, core, users): } monkeypatch.setattr(node_op_module, "node_manager", MagicMock(update_node=AsyncMock())) - monkeypatch.setattr(NodeOperation, "_get_core_users_map", AsyncMock(return_value=({1: object()}, {1: []}))) + monkeypatch.setattr( + NodeOperation, + "_get_core_users_map", + AsyncMock(return_value=({1: object()}, {1: []}, set())), + ) monkeypatch.setattr(NodeOperation, "connect_node", staticmethod(_connect_node)) monkeypatch.setattr(node_op_module, "bulk_update_node_status", AsyncMock()) monkeypatch.setattr(node_op_module.notification, "connect_node", AsyncMock()) diff --git a/tests/test_nats_node_memory.py b/tests/test_nats_node_memory.py index bbf0d9fa7..c6bb2b839 100644 --- a/tests/test_nats_node_memory.py +++ b/tests/test_nats_node_memory.py @@ -1,13 +1,44 @@ """CAS semantics for NATS-backed bridge user-sync + lifecycle memory.""" import asyncio +import json import pytest from PasarGuardNodeBridge.common.service_pb2 import User from PasarGuardNodeBridge.storage import LifecycleOperation, LifecycleStatus, NodeLifecycleState from app.nats.kv_cas import MemoryCasKv -from app.node.nats_memory import NatsNodeLifecycleCoordinator, NatsUserSyncStore +from app.node import nats_memory +from app.node.nats_memory import ( + ClaimedUser, + NatsNodeLifecycleCoordinator, + NatsUserSyncStore, + UserRevocationConflictError, + UserSyncLease, + UserSyncLeaseLostError, +) + + +class BlockingCreateKv(MemoryCasKv): + """Deterministic CAS race: pause one create for the selected key prefix.""" + + def __init__(self, prefix: str): + super().__init__() + self.prefix = prefix + self.started = asyncio.Event() + self.proceed = asyncio.Event() + self.armed = True + + async def create(self, key: str, value: bytes) -> int: + if self.armed and key.startswith(self.prefix): + self.armed = False + self.started.set() + await self.proceed.wait() + return await super().create(key, value) + + +def _json_doc(kv: MemoryCasKv, key: str) -> dict: + return json.loads(kv._data[key][0]) def _user(email: str, inbound: str = "in") -> User: @@ -56,6 +87,22 @@ async def test_user_sync_expired_claim_becomes_available(): assert [item.user.email for item in claimed] == ["a@example.com"] +@pytest.mark.asyncio +async def test_next_claim_delay_tracks_pending_claimed_and_fenced_work(monkeypatch): + now = [100.0] + monkeypatch.setattr(nats_memory.time, "time", lambda: now[0]) + store = NatsUserSyncStore(MemoryCasKv()) + assert await store.next_claim_delay("1") is None + + await store.enqueue_users("1", [_user("a@example.com")]) + assert await store.next_claim_delay("1") == 0.0 + await store.claim_users("1", "worker-a", limit=1, lease_seconds=5) + assert await store.next_claim_delay("1") == 5.0 + + await store.begin_user_revocation("1", ["a@example.com"], "revoke-a") + assert await store.next_claim_delay("1") is None + + @pytest.mark.asyncio async def test_user_sync_enqueue_shards_per_email_key(): kv = MemoryCasKv() @@ -72,6 +119,31 @@ async def test_user_sync_enqueue_shards_per_email_key(): assert kv._data == {} +@pytest.mark.asyncio +async def test_clear_only_flushes_queue_and_preserves_revocation_safety_state(): + store = NatsUserSyncStore(MemoryCasKv()) + await store.begin_user_revocation("1", ["a@example.com"], "revoke-a") + authorized = await store.acquire_user_sync_lease( + "1", + "worker", + ["a@example.com"], + 30, + revocation_id="revoke-a", + ) + + await store.clear("1") + + assert await store.heartbeat_user_sync_lease(authorized) is True + ordinary = await store.acquire_user_sync_lease("1", "worker", ["a@example.com"], 30) + assert ordinary.token == "" + await store.release_user_sync_lease(authorized) + + await store.purge_node("1") + after_removal = await store.acquire_user_sync_lease("1", "worker", ["a@example.com"], 30) + assert after_removal.token + await store.release_user_sync_lease(after_removal) + + @pytest.mark.asyncio async def test_claim_cleans_up_claimed_key_when_pending_delete_fails(): kv = MemoryCasKv() @@ -93,6 +165,605 @@ async def _delete(key: str, last: int | None = None) -> bool: assert any(key.startswith("p.1.") for key in kv._data) +@pytest.mark.asyncio +async def test_revocation_purges_queue_and_abort_advances_generation(): + kv = MemoryCasKv() + store = NatsUserSyncStore(kv) + await store.enqueue_users("1", [_user("a@example.com")]) + claimed = await store.claim_users("1", "worker-a", limit=10, lease_seconds=30) + assert claimed[0].generation == 0 + + await store.begin_user_revocation("1", ["a@example.com"], "revoke-a") + assert not any(key.startswith(("p.1.", "c.1.")) for key in kv._data) + + await store.enqueue_users("1", [_user("a@example.com", "blocked")]) + assert await store.claim_users("1", "worker-b", limit=10, lease_seconds=30) == [] + + await store.abort_user_revocation("1", ["a@example.com"], "revoke-a") + await store.enqueue_users("1", [_user("a@example.com", "restored")]) + fresh = await store.claim_users("1", "worker-b", limit=10, lease_seconds=30) + assert len(fresh) == 1 + assert fresh[0].generation == 1 + assert list(fresh[0].user.inbounds) == ["restored"] + + +@pytest.mark.asyncio +async def test_stale_claim_cannot_requeue_or_acquire_after_fast_abort(): + kv = MemoryCasKv() + store = NatsUserSyncStore(kv) + await store.enqueue_users("1", [_user("a@example.com")]) + stale = (await store.claim_users("1", "worker-a", limit=1, lease_seconds=30))[0] + + await store.begin_user_revocation("1", ["a@example.com"], "revoke-a") + await store.abort_user_revocation("1", ["a@example.com"], "revoke-a") + + lease = await store.acquire_user_sync_lease( + "1", + "worker-a", + ["a@example.com"], + 30, + expected_generations={"a@example.com": stale.generation}, + ) + assert lease.token == "" + + await store.requeue_users("1", [stale]) + assert await store.claim_users("1", "worker-b", limit=1, lease_seconds=30) == [] + + +@pytest.mark.asyncio +async def test_overlapping_revocation_fails_fast_and_finalize_is_permanent(): + kv = MemoryCasKv() + store = NatsUserSyncStore(kv) + key = "a@example.com" + + await store.begin_user_revocation("1", [key], "revoke-a") + with pytest.raises(UserRevocationConflictError): + await store.begin_user_revocation("1", [key], "revoke-b") + await store.abort_user_revocation("1", [key], "revoke-a") + + await store.begin_user_revocation("1", [key], "revoke-b") + await store.finalize_user_revocation("1", [key], "revoke-b") + await store.abort_user_revocation("1", [key], "revoke-b") + await store.enqueue_users("1", [_user(key)]) + lease = await store.acquire_user_sync_lease("1", "worker", [key], 30) + assert lease.token == "" + assert await store.claim_users("1", "worker", 1, 30) == [] + + barrier_key = next(key for key in kv._data if key.startswith("b.1.")) + assert _json_doc(kv, barrier_key)["permanent"] is True + assert _json_doc(kv, barrier_key)["active_owner"] is None + + +@pytest.mark.asyncio +async def test_bulk_conflict_unwinds_keys_acquired_before_the_conflict(): + kv = MemoryCasKv() + store = NatsUserSyncStore(kv) + await store.enqueue_users("1", [_user("a@example.com", "queued")]) + claimed = await store.claim_users("1", "worker-a", 1, 30) + assert claimed[0].generation == 0 + await store.begin_user_revocation("1", ["b@example.com"], "revoke-b") + + with pytest.raises(UserRevocationConflictError): + await store.begin_user_revocation( + "1", + ["a@example.com", "b@example.com"], + "revoke-a", + ) + + # The all-key conflict is a strict no-op for a: its generation and claimed + # authoritative payload remain valid instead of being silently discarded. + lease = await store.acquire_user_sync_lease( + "1", + "worker", + ["a@example.com"], + 30, + expected_generations={"a@example.com": claimed[0].generation}, + ) + assert lease.token + await store.release_user_sync_lease(lease) + await store.requeue_users("1", claimed) + requeued = await store.claim_users("1", "worker-b", 1, 30) + assert list(requeued[0].user.inbounds) == ["queued"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("close_method", ["abort_user_revocation", "finalize_user_revocation"]) +async def test_partial_close_cas_failure_reopens_every_marked_key(monkeypatch, close_method): + kv = MemoryCasKv() + store = NatsUserSyncStore(kv) + keys = ["a@example.com", "b@example.com"] + await store.begin_user_revocation("1", keys, "revoke-a") + original_cas = nats_memory.kv_cas_json + blocked_key = store._barrier_key("1", "b@example.com") + + async def fail_second_closing(kv_arg, key, value, revision): + if key == blocked_key and value.get("closing") is True: + return False + return await original_cas(kv_arg, key, value, revision) + + monkeypatch.setattr(nats_memory, "kv_cas_json", fail_second_closing) + with pytest.raises(RuntimeError, match="after CAS retries"): + await getattr(store, close_method)("1", keys, "revoke-a") + + for user_key in keys: + barrier, _ = await store._get_barrier("1", user_key) + assert barrier["active_owner"] == "revoke-a" + assert barrier["closing"] is False + + owner = await store.acquire_user_sync_lease( + "1", + "worker", + keys, + 30, + revocation_id="revoke-a", + ) + assert owner.user_keys == tuple(keys) + await store.release_user_sync_lease(owner) + + +@pytest.mark.asyncio +async def test_partial_begin_cas_failure_restores_generations_and_claims(monkeypatch): + kv = MemoryCasKv() + store = NatsUserSyncStore(kv) + keys = ["a@example.com", "b@example.com"] + await store.enqueue_users("1", [_user(keys[0], "queued")]) + claimed = await store.claim_users("1", "worker-a", 1, 30) + original_cas = nats_memory.kv_cas_json + blocked_key = store._barrier_key("1", keys[1]) + + async def fail_second_fence(kv_arg, key, value, revision): + if key == blocked_key and value.get("active_owner") == "delete": + return False + return await original_cas(kv_arg, key, value, revision) + + monkeypatch.setattr(nats_memory, "kv_cas_json", fail_second_fence) + with pytest.raises(RuntimeError, match="after CAS retries"): + await store.begin_user_revocation("1", keys, "delete") + + for user_key in keys: + barrier, _ = await store._get_barrier("1", user_key) + assert barrier["generation"] == 0 + assert barrier["active_owner"] is None + lease = await store.acquire_user_sync_lease( + "1", + "worker-b", + [keys[0]], + 30, + expected_generations={keys[0]: claimed[0].generation}, + ) + assert lease.token + await store.release_user_sync_lease(lease) + await store.requeue_users("1", claimed) + assert (await store.claim_users("1", "worker-b", 1, 30))[0].user.email == keys[0] + + +@pytest.mark.asyncio +async def test_begin_reports_already_finalized_keys_without_reopening_them(): + store = NatsUserSyncStore(MemoryCasKv()) + await store.begin_user_revocation("1", ["a@example.com"], "revoke-a") + await store.finalize_user_revocation("1", ["a@example.com"], "revoke-a") + + result = await store.begin_user_revocation( + "1", + ["a@example.com", "b@example.com"], + "revoke-b", + ) + + assert result.active_user_keys == ("b@example.com",) + assert result.finalized_user_keys == ("a@example.com",) + + +@pytest.mark.asyncio +async def test_begin_waits_for_intersecting_execution_lease(): + store = NatsUserSyncStore(MemoryCasKv()) + lease = await store.acquire_user_sync_lease("1", "worker-a", ["a@example.com", "b@example.com"], 30) + assert lease.token + + begin = asyncio.create_task(store.begin_user_revocation("1", ["b@example.com"], "revoke-a")) + for _ in range(5): + await asyncio.sleep(0) + assert not begin.done() + + await store.release_user_sync_lease(lease) + await asyncio.wait_for(begin, timeout=1) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("close_method", ["abort_user_revocation", "finalize_user_revocation"]) +async def test_revocation_close_blocks_new_admission_and_waits_authorized_lease(close_method): + store = NatsUserSyncStore(MemoryCasKv()) + key = "a@example.com" + await store.begin_user_revocation("1", [key], "revoke-a") + authorized = await store.acquire_user_sync_lease( + "1", + "worker-a", + [key], + 30, + revocation_id="revoke-a", + ) + + close = asyncio.create_task(getattr(store, close_method)("1", [key], "revoke-a")) + for _ in range(5): + await asyncio.sleep(0) + assert not close.done() + late_authorized = await store.acquire_user_sync_lease( + "1", + "worker-b", + [key], + 30, + revocation_id="revoke-a", + ) + assert late_authorized.token == "" + + await store.release_user_sync_lease(authorized) + await asyncio.wait_for(close, timeout=1) + ordinary = await store.acquire_user_sync_lease("1", "worker-b", [key], 30) + if close_method == "abort_user_revocation": + assert ordinary.token + await store.release_user_sync_lease(ordinary) + else: + assert ordinary.token == "" + + +@pytest.mark.asyncio +async def test_execution_lease_on_other_node_does_not_block_begin(): + store = NatsUserSyncStore(MemoryCasKv()) + lease = await store.acquire_user_sync_lease("1", "worker-a", ["a@example.com"], 30) + assert lease.token + + result = await store.begin_user_revocation("2", ["a@example.com"], "revoke-a") + + assert result.active_user_keys == ("a@example.com",) + await store.release_user_sync_lease(lease) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("close_method", ["abort_user_revocation", "finalize_user_revocation"]) +async def test_failed_close_reopens_only_owner_admission_for_authoritative_reconcile( + monkeypatch, + close_method, +): + now = [100.0] + monkeypatch.setattr(nats_memory.time, "time", lambda: now[0]) + store = NatsUserSyncStore(MemoryCasKv()) + key = "a@example.com" + await store.begin_user_revocation("1", [key], "revoke-a") + unknown = await store.acquire_user_sync_lease( + "1", + "worker-a", + [key], + 5, + revocation_id="revoke-a", + ) + now[0] = 106.0 + + with pytest.raises(UserSyncLeaseLostError): + await getattr(store, close_method)("1", [key], "revoke-a") + + ordinary = await store.acquire_user_sync_lease("1", "worker-b", [key], 30) + owner = await store.acquire_user_sync_lease( + "1", + "worker-b", + [key], + 30, + revocation_id="revoke-a", + ) + assert ordinary.token == "" + assert owner.token + await store.release_user_sync_lease(owner) + await store.release_user_sync_lease(unknown) + await getattr(store, close_method)("1", [key], "revoke-a") + + +@pytest.mark.asyncio +async def test_expired_execution_lease_fails_begin_closed_until_reconciled(monkeypatch): + now = [100.0] + monkeypatch.setattr(nats_memory.time, "time", lambda: now[0]) + store = NatsUserSyncStore(MemoryCasKv()) + lease = await store.acquire_user_sync_lease("1", "worker-a", ["a@example.com"], 5) + assert lease.token + now[0] = 102.0 + assert await store.heartbeat_user_sync_lease(lease) is True + now[0] = 108.0 + + with pytest.raises(UserSyncLeaseLostError): + await store.begin_user_revocation("1", ["a@example.com"], "revoke-a") + + # begin installed the fence before detecting the unknown remote outcome. + await store.enqueue_users("1", [_user("a@example.com", "blocked")]) + assert await store.claim_users("1", "worker-b", 1, 30) == [] + + await store.release_user_sync_lease(lease) + await store.abort_user_revocation("1", ["a@example.com"], "revoke-a") + await store.enqueue_users("1", [_user("a@example.com", "restored")]) + assert len(await store.claim_users("1", "worker-b", 1, 30)) == 1 + + +@pytest.mark.asyncio +async def test_authoritative_reconciliation_replaces_crashed_execution_lease(monkeypatch): + now = [100.0] + monkeypatch.setattr(nats_memory.time, "time", lambda: now[0]) + kv = MemoryCasKv() + store = NatsUserSyncStore(kv) + lost = await store.acquire_user_sync_lease("1", "dead-worker", ["a@example.com"], 5) + now[0] = 106.0 + await store.set_authoritative_reconciliation_membership( + "1", "recovery-worker", ["a@example.com", "b@example.com"] + ) + + recovery = await store.acquire_user_sync_reconciliation_lease( + "1", "recovery-worker", ["a@example.com", "b@example.com"], 30 + ) + + assert recovery.lease.covers_all_users is True + assert recovery.included_user_keys == ("a@example.com", "b@example.com") + assert not any(lost.token in json.loads(raw).get("token", "") for raw, _ in kv._data.values()) + await store.release_user_sync_lease(recovery.lease) + + +@pytest.mark.asyncio +async def test_expired_revocation_metadata_lock_requires_authoritative_recovery(monkeypatch): + now = [100.0] + monkeypatch.setattr(nats_memory.time, "time", lambda: now[0]) + kv = MemoryCasKv() + store = NatsUserSyncStore(kv) + assert await nats_memory.kv_cas_json( + kv, + store._revocation_lock_key("1"), + {"token": "dead-worker", "expires_at": 99.0}, + 0, + ) + + with pytest.raises(UserRevocationConflictError): + await store.begin_user_revocation("1", ["a@example.com"], "revoke-a") + + await store.set_authoritative_reconciliation_membership("1", "recovery", ["a@example.com"]) + recovery = await store.acquire_user_sync_reconciliation_lease( + "1", "recovery", ["a@example.com"], 30 + ) + await store.release_user_sync_lease(recovery.lease) + result = await store.begin_user_revocation("1", ["a@example.com"], "revoke-a") + + assert result.active_user_keys == ("a@example.com",) + + +@pytest.mark.asyncio +async def test_paused_expired_metadata_owner_is_not_stolen_and_fails_closed(monkeypatch): + now = [100.0] + monkeypatch.setattr(nats_memory.time, "time", lambda: now[0]) + store = NatsUserSyncStore(MemoryCasKv()) + + with pytest.raises(UserSyncLeaseLostError): + async with store._revocation_lock("1", ["a@example.com"]) as assert_owned: + now[0] += nats_memory._REVOCATION_METADATA_LEASE_SECONDS + 1 + with pytest.raises(UserRevocationConflictError): + await store.begin_user_revocation("1", ["a@example.com"], "other-owner") + await assert_owned() + + +@pytest.mark.asyncio +async def test_only_matching_revocation_owner_can_acquire_fenced_direct_sync_lease(): + store = NatsUserSyncStore(MemoryCasKv()) + key = "a@example.com" + await store.begin_user_revocation("1", [key], "revoke-a") + + ordinary = await store.acquire_user_sync_lease("1", "worker", [key], 30) + wrong_owner = await store.acquire_user_sync_lease("1", "worker", [key], 30, revocation_id="revoke-b") + authorized = await store.acquire_user_sync_lease("1", "worker", [key], 30, revocation_id="revoke-a") + + assert ordinary.token == "" + assert wrong_owner.token == "" + assert authorized.token + await store.release_user_sync_lease(authorized) + + +@pytest.mark.asyncio +async def test_startup_lease_is_node_wide_filters_finalized_and_blocks_new_writes(): + store = NatsUserSyncStore(MemoryCasKv()) + await store.begin_user_revocation("1", ["deleted@example.com"], "delete") + await store.finalize_user_revocation("1", ["deleted@example.com"], "delete") + + startup = await store.acquire_startup_user_sync_lease( + "1", + "starter", + ["active@example.com", "deleted@example.com"], + 30, + ) + + assert startup.lease.token + assert startup.lease.covers_all_users is True + assert startup.included_user_keys == ("active@example.com",) + ordinary = await store.acquire_user_sync_lease("1", "worker", ["active@example.com"], 30) + assert ordinary.token == "" + await store.release_user_sync_lease(startup.lease) + + +@pytest.mark.asyncio +async def test_startup_waits_for_prior_write_and_then_begin_waits_for_startup(): + store = NatsUserSyncStore(MemoryCasKv()) + ordinary = await store.acquire_user_sync_lease("1", "worker", ["a@example.com"], 30) + startup_task = asyncio.create_task(store.acquire_startup_user_sync_lease("1", "starter", ["a@example.com"], 30)) + for _ in range(5): + await asyncio.sleep(0) + assert not startup_task.done() + + await store.release_user_sync_lease(ordinary) + startup = await asyncio.wait_for(startup_task, timeout=1) + begin = asyncio.create_task(store.begin_user_revocation("1", ["a@example.com"], "delete")) + for _ in range(5): + await asyncio.sleep(0) + assert not begin.done() + + await store.release_user_sync_lease(startup.lease) + await asyncio.wait_for(begin, timeout=1) + + +@pytest.mark.asyncio +async def test_ordinary_lease_losing_startup_create_race_self_rejects(): + kv = BlockingCreateKv("x.1.") + store = NatsUserSyncStore(kv) + ordinary_task = asyncio.create_task(store.acquire_user_sync_lease("1", "worker", ["a@example.com"], 30)) + await kv.started.wait() + + startup = await store.acquire_startup_user_sync_lease( + "1", + "starter", + ["a@example.com"], + 30, + ) + kv.proceed.set() + ordinary = await asyncio.wait_for(ordinary_task, timeout=1) + + assert ordinary.token == "" + assert startup.lease.token + await store.release_user_sync_lease(startup.lease) + + +@pytest.mark.asyncio +async def test_startup_waiting_on_provisional_owner_does_not_block_authorized_restore(): + store = NatsUserSyncStore(MemoryCasKv()) + await store.begin_user_revocation("1", ["a@example.com"], "delete") + startup_task = asyncio.create_task(store.acquire_startup_user_sync_lease("1", "starter", ["a@example.com"], 30)) + for _ in range(5): + await asyncio.sleep(0) + assert not startup_task.done() + + restore = await store.acquire_user_sync_lease( + "1", + "worker", + ["a@example.com"], + 30, + revocation_id="delete", + ) + assert restore.token + await store.release_user_sync_lease(restore) + await store.abort_user_revocation("1", ["a@example.com"], "delete") + + startup = await asyncio.wait_for(startup_task, timeout=1) + assert startup.included_user_keys == ("a@example.com",) + await store.release_user_sync_lease(startup.lease) + + +@pytest.mark.asyncio +async def test_retain_user_sync_lease_keys_is_atomic_and_preserves_token(): + store = NatsUserSyncStore(MemoryCasKv()) + lease = await store.acquire_user_sync_lease( + "1", + "worker", + ["known@example.com", "unknown@example.com"], + 30, + ) + + narrowed = await store.retain_user_sync_lease_keys(lease, ["unknown@example.com"]) + + assert narrowed.token == lease.token + assert narrowed.user_keys == ("unknown@example.com",) + assert narrowed.epoch == lease.epoch + assert await store.heartbeat_user_sync_lease(lease) is False + assert await store.heartbeat_user_sync_lease(narrowed) is True + await store.release_user_sync_lease(narrowed) + + +@pytest.mark.asyncio +async def test_user_sync_epochs_are_unique_and_monotonic_across_workers(): + store = NatsUserSyncStore(MemoryCasKv()) + leases = await asyncio.gather( + *( + store.acquire_user_sync_lease( + "1", + f"worker-{index}", + [f"user-{index}@example.com"], + 30, + ) + for index in range(16) + ) + ) + + epochs = sorted(lease.epoch for lease in leases) + assert epochs == list(range(1, 17)) + await asyncio.gather(*(store.release_user_sync_lease(lease) for lease in leases)) + + +@pytest.mark.asyncio +async def test_user_sync_epoch_handshake_advances_floor_under_concurrency(): + store = NatsUserSyncStore(MemoryCasKv()) + await asyncio.gather( + store.advance_user_sync_epoch("1", 50), + store.advance_user_sync_epoch("1", 75), + store.advance_user_sync_epoch("1", 60), + ) + + lease = await store.acquire_user_sync_lease("1", "worker", ["a@example.com"], 30) + assert lease.epoch == 76 + await store.release_user_sync_lease(lease) + + +@pytest.mark.asyncio +async def test_forged_execution_lease_cannot_heartbeat_or_release_owner(): + store = NatsUserSyncStore(MemoryCasKv()) + lease = await store.acquire_user_sync_lease("1", "worker-a", ["a@example.com"], 30) + forged = UserSyncLease( + node_id=lease.node_id, + worker_id="worker-b", + token=lease.token, + user_keys=lease.user_keys, + generations=lease.generations, + lease_seconds=lease.lease_seconds, + ) + + assert await store.heartbeat_user_sync_lease(forged) is False + await store.release_user_sync_lease(forged) + assert await store.heartbeat_user_sync_lease(lease) is True + await store.release_user_sync_lease(lease) + + +@pytest.mark.asyncio +async def test_enqueue_racing_begin_cannot_leave_stale_pending_work(): + kv = BlockingCreateKv("p.1.") + store = NatsUserSyncStore(kv) + + enqueue = asyncio.create_task(store.enqueue_users("1", [_user("a@example.com")])) + await kv.started.wait() + await store.begin_user_revocation("1", ["a@example.com"], "revoke-a") + kv.proceed.set() + await enqueue + + assert not any(key.startswith("p.1.") for key in kv._data) + + +@pytest.mark.asyncio +async def test_lease_create_racing_begin_is_rejected_after_fence(): + kv = BlockingCreateKv("x.1.") + store = NatsUserSyncStore(kv) + + acquire = asyncio.create_task(store.acquire_user_sync_lease("1", "worker-a", ["a@example.com"], 30)) + await kv.started.wait() + await store.begin_user_revocation("1", ["a@example.com"], "revoke-a") + kv.proceed.set() + lease = await acquire + + assert lease.token == "" + assert not any(key.startswith("x.1.") for key in kv._data) + + +@pytest.mark.asyncio +async def test_requeue_requires_owned_claim_and_matching_generation(): + kv = MemoryCasKv() + store = NatsUserSyncStore(kv) + await store.enqueue_users("1", [_user("a@example.com")]) + claim = (await store.claim_users("1", "worker-a", 1, 30))[0] + + forged = ClaimedUser(token=claim.token, user=_user("other@example.com"), generation=claim.generation) + await store.requeue_users("1", [forged]) + assert not any(key.startswith("p.1.") for key in kv._data) + assert any(key.startswith("c.1.") for key in kv._data) + + wrong_generation = ClaimedUser(token=claim.token, user=claim.user, generation=claim.generation + 1) + await store.requeue_users("1", [wrong_generation]) + assert not any(key.startswith("p.1.") for key in kv._data) + assert any(key.startswith("c.1.") for key in kv._data) + + @pytest.mark.asyncio async def test_lifecycle_has_active_lease_tracks_expiry(): coordinator = NatsNodeLifecycleCoordinator(MemoryCasKv()) @@ -111,6 +782,79 @@ async def test_lifecycle_has_active_lease_tracks_expiry(): assert expired is not None await asyncio.sleep(0.01) assert await coordinator.has_active_lease("1") is False + assert await coordinator.try_acquire("1", "worker-b", LifecycleOperation.STOP, 30) is None + assert await coordinator.reconcile("1", LifecycleStatus.HEALTHY) is True + replacement = await coordinator.try_acquire("1", "worker-b", LifecycleOperation.STOP, 30) + assert replacement is not None + + +@pytest.mark.asyncio +async def test_lifecycle_reconcile_rejects_live_lease(): + coordinator = NatsNodeLifecycleCoordinator(MemoryCasKv()) + lease = await coordinator.try_acquire("1", "worker-a", LifecycleOperation.START, 30) + assert lease is not None + assert await coordinator.reconcile("1", LifecycleStatus.HEALTHY) is False + assert await coordinator.try_acquire("1", "worker-b", LifecycleOperation.STOP, 30) is None + + +@pytest.mark.asyncio +async def test_authoritative_reconcile_resolves_crash_after_db_commit_barriers(): + kv = MemoryCasKv() + store = NatsUserSyncStore(kv) + await store.begin_user_revocation("1", ["present", "deleted"], "delete-both") + assert await store.needs_authoritative_recovery("1") is True + + # Simulate a full worker/process restart: only persisted KV survives. + restarted_store = NatsUserSyncStore(kv) + await restarted_store.set_authoritative_reconciliation_membership("1", "recovery-worker", ["present"]) + + recovery = await restarted_store.acquire_user_sync_reconciliation_lease( + # The row exists globally but is not assigned to this node's core. + "1", "recovery-worker", [], 30 + ) + + assert recovery.included_user_keys == () + present, _ = await restarted_store._get_barrier("1", "present") + deleted, _ = await restarted_store._get_barrier("1", "deleted") + assert present["active_owner"] is None + assert present["closing"] is False + assert present["permanent"] is False + assert deleted["active_owner"] is None + assert deleted["closing"] is False + assert deleted["permanent"] is True + assert recovery.lease.covers_all_users is True + assert recovery.lease.epoch > 0 + assert await restarted_store.needs_authoritative_recovery("1") is False + await restarted_store.release_user_sync_lease(recovery.lease) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("close_method", ["abort_user_revocation", "finalize_user_revocation"]) +async def test_idempotent_close_without_affected_keys_ignores_unrelated_wildcard(close_method): + kv = MemoryCasKv() + store = NatsUserSyncStore(kv) + await store.set_authoritative_reconciliation_membership("1", "startup", ["other"]) + wildcard = await store.acquire_user_sync_reconciliation_lease("1", "startup", ["other"], 30) + + await asyncio.wait_for( + getattr(store, close_method)("1", ["missing"], "already-closed"), + timeout=0.1, + ) + + await store.release_user_sync_lease(wildcard.lease) + + +@pytest.mark.asyncio +async def test_purge_node_removes_epoch_fence_too(): + kv = MemoryCasKv() + store = NatsUserSyncStore(kv) + lease = await store.acquire_user_sync_lease("7", "worker", ["user"], 30) + await store.release_user_sync_lease(lease) + assert "e.7" in kv._data + + await store.purge_node("7") + + assert not any(key == "e.7" or key.startswith(("p.7.", "c.7.", "b.7.", "x.7.")) for key in kv._data) @pytest.mark.asyncio @@ -140,9 +884,32 @@ async def test_lifecycle_lease_exclusive_and_epoch_fenced(): await asyncio.sleep(0.01) newer = await coordinator.try_acquire("1", "worker-b", LifecycleOperation.STOP, 30) assert stale is not None + assert newer is None + + assert await coordinator.reconcile("1", LifecycleStatus.HEALTHY) is True + newer = await coordinator.try_acquire("1", "worker-b", LifecycleOperation.STOP, 30) assert newer is not None await coordinator.update_observed("1", LifecycleStatus.BROKEN, expected_epoch=stale.epoch) state = await coordinator.get_state("1") assert state.epoch == newer.epoch assert state.observed is not LifecycleStatus.BROKEN + + +@pytest.mark.asyncio +async def test_deleted_lifecycle_namespace_permanently_rejects_start_but_allows_stop(): + coordinator = NatsNodeLifecycleCoordinator(MemoryCasKv()) + + await coordinator.mark_deleted("stable-bridge-id") + + assert await coordinator.is_deleted("stable-bridge-id") is True + assert ( + await coordinator.try_acquire( + "stable-bridge-id", "late-worker", LifecycleOperation.START, 30 + ) + is None + ) + stop = await coordinator.try_acquire( + "stable-bridge-id", "cleanup-worker", LifecycleOperation.STOP, 30 + ) + assert stop is not None diff --git a/tests/test_node_bridge_namespace.py b/tests/test_node_bridge_namespace.py new file mode 100644 index 000000000..c94931f97 --- /dev/null +++ b/tests/test_node_bridge_namespace.py @@ -0,0 +1,79 @@ +from pathlib import Path + +import pytest +from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine + +from app.db.models import Base, Node, User +from app.models.node import NodeResponse +from app.models.user import UserResponse + + +def _node(name: str) -> Node: + return Node( + name=name, + address="127.0.0.1", + port=62050, + api_port=62051, + server_ca="ca", + api_key=None, + core_config_id=None, + ) + + +@pytest.mark.asyncio +async def test_sqlite_reused_numeric_id_gets_new_bridge_namespace(tmp_path: Path): + engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'nodes.db'}") + try: + async with engine.begin() as connection: + await connection.run_sync(Base.metadata.create_all) + sessions = async_sessionmaker(engine, expire_on_commit=False) + async with sessions() as db: + first = _node("first") + db.add(first) + await db.commit() + first_id = first.id + first_bridge_id = first.bridge_id + + await db.delete(first) + await db.commit() + + replacement = _node("replacement") + db.add(replacement) + await db.commit() + + assert replacement.id == first_id + assert replacement.bridge_id != first_bridge_id + finally: + await engine.dispose() + + +@pytest.mark.asyncio +async def test_sqlite_reused_user_id_gets_new_sync_incarnation(tmp_path: Path): + engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'users.db'}") + try: + async with engine.begin() as connection: + await connection.run_sync(Base.metadata.create_all) + sessions = async_sessionmaker(engine, expire_on_commit=False) + async with sessions() as db: + first = User(username="first") + db.add(first) + await db.commit() + first_id = first.id + first_sync_id = first.sync_id + + await db.delete(first) + await db.commit() + + replacement = User(username="replacement") + db.add(replacement) + await db.commit() + + assert replacement.id == first_id + assert replacement.sync_id != first_sync_id + finally: + await engine.dispose() + + +def test_internal_sync_namespaces_are_not_part_of_public_api_schemas(): + assert "bridge_id" not in NodeResponse.model_json_schema()["properties"] + assert "sync_id" not in UserResponse.model_json_schema()["properties"] diff --git a/tests/test_node_manager.py b/tests/test_node_manager.py index 407a97abe..1fd4b776a 100644 --- a/tests/test_node_manager.py +++ b/tests/test_node_manager.py @@ -1,4 +1,5 @@ import pytest +from PasarGuardNodeBridge.common.service_pb2 import User as ProtoUser from app.node import NodeManager @@ -6,7 +7,7 @@ @pytest.mark.asyncio async def test_node_manager_bulk_user_sync_uses_bounded_chunked_batches(monkeypatch: pytest.MonkeyPatch): manager = NodeManager() - users = [object() for _ in range(5)] + users = [ProtoUser(email=str(index)) for index in range(5)] class FakeNode: def __init__(self): @@ -32,7 +33,7 @@ async def test_node_manager_bulk_user_sync_falls_back_when_chunked_is_not_suppor monkeypatch: pytest.MonkeyPatch, ): manager = NodeManager() - users = [object() for _ in range(3)] + users = [ProtoUser(email=str(index)) for index in range(3)] class FakeNode: def __init__(self): diff --git a/tests/test_node_manager_sync.py b/tests/test_node_manager_sync.py index 9359efb0b..049e98251 100644 --- a/tests/test_node_manager_sync.py +++ b/tests/test_node_manager_sync.py @@ -1,3 +1,5 @@ +from unittest.mock import AsyncMock + import pytest from app.nats.router import _router_enabled @@ -22,20 +24,69 @@ def test_router_disabled_without_nats(monkeypatch: pytest.MonkeyPatch): @pytest.mark.asyncio async def test_handle_node_remove(monkeypatch: pytest.MonkeyPatch): removed: list[tuple[int, bool]] = [] - cleared: list[int] = [] - async def _remove(node_id: int, *, remote_stop: bool = True): + async def _remove( + node_id: int, + *, + remote_stop: bool = True, + permanent_delete: bool = False, + expected_bridge_namespace: str | None = None, + ): + assert permanent_delete is True + assert expected_bridge_namespace == "bridge-7" removed.append((node_id, remote_stop)) - async def _clear(node_id): - cleared.append(int(node_id)) - monkeypatch.setattr("app.node.manager_sync.node_manager.remove_node", _remove) - monkeypatch.setattr("app.node.manager_sync.clear_bridge_memory_for_node", _clear) - await handle_node_message({"action": "remove", "node_id": 7, "origin": "other-worker"}) + await handle_node_message( + {"action": "remove", "node_id": 7, "bridge_id": "bridge-7", "origin": "other-worker"} + ) assert removed == [(7, False)] - assert cleared == [7] + + +@pytest.mark.asyncio +async def test_handle_node_remove_failure_retains_shared_memory(monkeypatch: pytest.MonkeyPatch): + async def _remove( + _node_id: int, + *, + remote_stop: bool = True, + permanent_delete: bool = False, + expected_bridge_namespace: str | None = None, + ): + assert permanent_delete is True + assert expected_bridge_namespace == "bridge-7" + raise RuntimeError("local runtime is not quiescent") + + monkeypatch.setattr("app.node.manager_sync.node_manager.remove_node", _remove) + + with pytest.raises(RuntimeError, match="not quiescent"): + await handle_node_message( + {"action": "remove", "node_id": 7, "bridge_id": "bridge-7", "origin": "other-worker"} + ) + + +@pytest.mark.asyncio +async def test_remove_message_propagates_stable_bridge_namespace(monkeypatch: pytest.MonkeyPatch): + remove = [] + + async def _remove( + node_id: int, + *, + remote_stop: bool = True, + expected_bridge_namespace: str | None = None, + permanent_delete: bool = False, + ): + assert permanent_delete is True + remove.append((node_id, remote_stop, expected_bridge_namespace)) + + monkeypatch.setattr("app.node.manager_sync.node_manager.remove_node", _remove) + + await handle_node_message( + {"action": "remove", "node_id": 7, "bridge_id": "old-bridge-id", "origin": "other-worker"} + ) + + assert remove == [(7, False, "old-bridge-id")] + @pytest.mark.asyncio @@ -52,6 +103,16 @@ async def _remove(node_id: int, *, remote_stop: bool = True): assert removed == [] +@pytest.mark.asyncio +async def test_legacy_remove_without_namespace_cannot_delete_reused_numeric_id(monkeypatch): + remove = AsyncMock() + monkeypatch.setattr("app.node.manager_sync.node_manager.remove_node", remove) + + await handle_node_message({"action": "remove", "node_id": 7, "origin": "old-worker"}) + + remove.assert_not_awaited() + + @pytest.mark.asyncio async def test_publish_node_sync_includes_origin(monkeypatch: pytest.MonkeyPatch): from app.node.manager_sync import publish_node_sync @@ -71,20 +132,14 @@ async def _publish(topic, data): @pytest.mark.asyncio async def test_handle_node_disconnect_no_memory_clear(monkeypatch: pytest.MonkeyPatch): removed: list[tuple[int, bool]] = [] - cleared: list[int] = [] async def _remove(node_id: int, *, remote_stop: bool = True): removed.append((node_id, remote_stop)) - async def _clear(node_id): - cleared.append(int(node_id)) - monkeypatch.setattr("app.node.manager_sync.node_manager.remove_node", _remove) - monkeypatch.setattr("app.node.manager_sync.clear_bridge_memory_for_node", _clear) await handle_node_message({"action": "disconnect", "node_id": 3, "origin": "other"}) assert removed == [(3, False)] - assert cleared == [] @pytest.mark.asyncio @@ -111,7 +166,118 @@ async def _update_node(db_node): monkeypatch.setattr("app.node.manager_sync.GetDB", lambda: _DB()) monkeypatch.setattr("app.node.manager_sync.get_node_by_id", _get_node_by_id) monkeypatch.setattr("app.node.manager_sync.node_manager.update_node", _update_node) + monkeypatch.setattr("app.node.manager_sync.node_manager.runtime_matches", AsyncMock(return_value=False)) await handle_node_message({"action": "upsert", "node_id": 9, "origin": "other"}) assert len(updated) == 1 assert updated[0].id == 9 + + +@pytest.mark.asyncio +async def test_cross_worker_connect_holds_snapshot_transaction_through_apply(monkeypatch: pytest.MonkeyPatch): + from app.operation.node import NodeOperation + + events: list[str] = [] + transaction_active = False + + class _Node: + id = 11 + bridge_id = "bridge-11" + status = "connecting" + core_config_id = 1 + + class _DB: + async def __aenter__(self): + nonlocal transaction_active + transaction_active = True + events.append("db-enter") + return self + + async def __aexit__(self, *_args): + nonlocal transaction_active + events.append("db-exit") + transaction_active = False + return False + + async def _get_node(*_args, **_kwargs): + events.append("load-node") + return _Node() + + async def _register(_node): + assert transaction_active + events.append("register-runtime") + + async def _snapshot(_db, _core_ids): + assert transaction_active + assert events[-1] == "register-runtime" + events.append("lock-snapshot") + return {1: object()}, {1: []}, {"user-sync-id"} + + async def _apply(*_args): + assert transaction_active + assert events[-1] == "lock-snapshot" + events.append("authoritative-apply") + + monkeypatch.setattr("app.node.manager_sync.GetDB", lambda: _DB()) + monkeypatch.setattr("app.node.manager_sync.get_node_by_id", _get_node) + monkeypatch.setattr("app.node.manager_sync.node_manager.update_node", _register) + monkeypatch.setattr("app.node.manager_sync.node_manager.runtime_matches", AsyncMock(return_value=False)) + monkeypatch.setattr(NodeOperation, "_get_core_users_map", _snapshot) + monkeypatch.setattr(NodeOperation, "connect_node", _apply) + + await handle_node_message( + {"action": "connect", "node_id": 11, "bridge_id": "bridge-11", "origin": "other"} + ) + + assert events == [ + "db-enter", + "load-node", + "register-runtime", + "lock-snapshot", + "authoritative-apply", + "db-exit", + ] + + +@pytest.mark.asyncio +async def test_duplicate_connect_keeps_matching_runtime_but_reconciles_users(monkeypatch): + from app.operation.node import NodeOperation + + class _Node: + id = 12 + bridge_id = "bridge-12" + status = "connected" + core_config_id = 1 + + class _DB: + async def __aenter__(self): + return self + + async def __aexit__(self, *_args): + return False + + update = AsyncMock() + apply = AsyncMock() + monkeypatch.setattr("app.node.manager_sync.GetDB", lambda: _DB()) + monkeypatch.setattr( + "app.node.manager_sync.get_node_by_id", + AsyncMock(return_value=_Node()), + ) + monkeypatch.setattr( + "app.node.manager_sync.node_manager.runtime_matches", + AsyncMock(return_value=True), + ) + monkeypatch.setattr("app.node.manager_sync.node_manager.update_node", update) + monkeypatch.setattr( + NodeOperation, + "_get_core_users_map", + AsyncMock(return_value=({1: object()}, {1: []}, {"sync-user"})), + ) + monkeypatch.setattr(NodeOperation, "connect_node", apply) + + await handle_node_message( + {"action": "connect", "node_id": 12, "bridge_id": "bridge-12", "origin": "other"} + ) + + update.assert_not_awaited() + apply.assert_awaited_once() diff --git a/tests/test_node_sync.py b/tests/test_node_sync.py index 4fe3238bc..60689b29c 100644 --- a/tests/test_node_sync.py +++ b/tests/test_node_sync.py @@ -1,7 +1,428 @@ +import asyncio +from types import SimpleNamespace +from unittest.mock import ANY, AsyncMock + import pytest +from PasarGuardNodeBridge import Health, NodeAPIError +from PasarGuardNodeBridge.common import service_pb2 as service +from PasarGuardNodeBridge.common.service_pb2 import User as ProtoUser +from PasarGuardNodeBridge.rest import Node as RestBridgeNode +from PasarGuardNodeBridge.storage import LifecycleOperation, LifecycleStatus, NodeLifecycleState +from sqlalchemy.dialects import mysql +import app.node as node_module +from app.db.models import NodeConnectionType, NodeStatus +from app.models.core import CoreType +from app.models.node import BulkNodeSelection, NodeLifecycleRecovery from app.nats.node_rpc import encode_node_command -from app.node import sync as node_sync_module +from app.node import NodeManager, sync as node_sync_module, worker as node_worker_module +from app.node.user import _serialize_user_for_node +from app.operation import OperatorType +from app.operation.node import NodeOperation +from role import Role + + +def _healthy_runtime_node() -> AsyncMock: + node = AsyncMock() + node.get_health.return_value = Health.HEALTHY + node._supports_chunked_sync.return_value = (True, "0.2.0") + node.sync_users_chunked.return_value = [] + node.begin_user_revocation = AsyncMock( + side_effect=lambda user_keys, _revocation_id: SimpleNamespace( + active_user_keys=tuple(user_keys), + finalized_user_keys=(), + ) + ) + node.abort_user_revocation = AsyncMock() + node.finalize_user_revocation = AsyncMock() + return node + + +def _db_node(node_id: int = 1, bridge_id: str = "bridge-1") -> SimpleNamespace: + return SimpleNamespace( + id=node_id, + bridge_id=bridge_id, + connection_type=NodeConnectionType.rest, + address="127.0.0.1", + port=62050, + api_port=62051, + server_ca="ca", + api_key="key", + name=f"node-{node_id}", + default_timeout=10, + internal_timeout=10, + proxy_url=None, + usage_coefficient=1.0, + ) + + +@pytest.mark.asyncio +async def test_new_node_start_requires_epoch_capability_and_reconciles_snapshot(): + pg_node = AsyncMock() + pg_node._user_sync_store = None + pg_node.get_lifecycle_state.return_value = None + pg_node.info.return_value = service.BaseInfoResponse( + started=False, + user_sync_epoch_supported=True, + ) + started = service.BaseInfoResponse( + started=True, + node_version="0.4.0", + core_version="1.0.0", + user_sync_epoch_supported=True, + user_sync_epoch=7, + ) + pg_node.start.return_value = started + db_node = SimpleNamespace(name="node-1", keep_alive=30) + core = SimpleNamespace(type=CoreType.xray, to_str=lambda: "{}", exclude_inbound_tags=[]) + users = [ProtoUser(email="42")] + + result = await NodeOperation._start_or_attach_node( + pg_node, + db_node, + core, + users, + service.BackendType.XRAY, + ) + + assert result is started + pg_node.start.assert_awaited_once_with( + config="{}", + backend_type=service.BackendType.XRAY, + users=users, + keep_alive=30, + reconcile_user_sync=True, + exclude_inbounds=[], + ) + + +@pytest.mark.asyncio +async def test_new_node_start_rejects_legacy_node_before_transport(): + pg_node = AsyncMock() + pg_node._user_sync_store = None + pg_node.get_lifecycle_state.return_value = None + pg_node.info.return_value = service.BaseInfoResponse(started=False) + db_node = SimpleNamespace(name="node-1", keep_alive=30) + core = SimpleNamespace(type=CoreType.wg, to_str=lambda: "{}") + + with pytest.raises(NodeAPIError, match="monotonic user-sync epoch fencing"): + await NodeOperation._start_or_attach_node( + pg_node, + db_node, + core, + [], + service.BackendType.WIREGUARD, + ) + pg_node.start.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_unknown_lifecycle_is_not_auto_reconciled_from_racy_probe(): + pg_node = AsyncMock() + pg_node._user_sync_store = None + pg_node.get_lifecycle_state.return_value = NodeLifecycleState( + operation=LifecycleOperation.START, + observed=LifecycleStatus.STARTING, + ) + pg_node.info.return_value = service.BaseInfoResponse( + started=True, + user_sync_epoch_supported=True, + ) + + with pytest.raises(NodeAPIError, match="explicit reconciliation is required"): + await NodeOperation._start_or_attach_node( + pg_node, + SimpleNamespace(name="node-1", keep_alive=30), + SimpleNamespace(type=CoreType.wg, to_str=lambda: "{}"), + [], + service.BackendType.WIREGUARD, + ) + pg_node.reconcile_lifecycle.assert_not_awaited() + pg_node.start.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_connect_409_reconciles_authoritative_users_before_connected(monkeypatch): + pg_node = AsyncMock() + users = [ProtoUser(email="42")] + info = service.BaseInfoResponse(node_version="0.4.0", core_version="1.0.0") + monkeypatch.setattr( + NodeOperation, + "_start_or_attach_node", + AsyncMock(side_effect=NodeAPIError(409, "lease held")), + ) + monkeypatch.setattr(NodeOperation, "_attach_if_running", AsyncMock(return_value=info)) + monkeypatch.setattr(node_module.node_manager, "get_node", AsyncMock(return_value=pg_node)) + + result = await NodeOperation.connect_node( + SimpleNamespace(id=1, name="node-1", status="connecting"), + SimpleNamespace(type=CoreType.wg), + users, + {"42"}, + ) + + assert result["status"].value == "connected" + pg_node.reconcile_users.assert_awaited_once_with(users) + + +@pytest.mark.asyncio +async def test_bb503_reconcile_stale_epoch_retry_reuses_db_locked_authorization(): + from app.nats.kv_cas import MemoryCasKv + from app.node.nats_memory import NatsUserSyncStore + + store = NatsUserSyncStore(MemoryCasKv()) + pg_node = object.__new__(RestBridgeNode) + pg_node.node_id = "node-namespace" + pg_node.worker_id = "worker-a" + pg_node._user_sync_store = store + pg_node._sync_lease_seconds = 30 + pg_node._default_timeout = 10 + pg_node._node_lock = asyncio.Lock() + pg_node._user_sync_epoch_supported = True + pg_node._user_sync_epoch_capability_probed = True + epochs: list[int] = [] + + async def _request(**kwargs): + epochs.append(kwargs["proto_message"].user_sync_epoch) + if len(epochs) == 1: + raise NodeAPIError(412, "known stale epoch") + return service.Empty() + + pg_node._make_request = _request + users = [ProtoUser(email="user-sync-id")] + async with NodeOperation._authoritative_user_reconciliation_scope(pg_node, {"user-sync-id"}): + await pg_node.reconcile_users(users) + + assert len(epochs) == 2 + assert epochs[1] > epochs[0] + + +@pytest.mark.asyncio +async def test_failed_reconcile_scope_does_not_leak_unlocked_authorization(): + from app.nats.kv_cas import MemoryCasKv + from app.node.nats_memory import NatsUserSyncStore, UserSyncLeaseLostError + + store = NatsUserSyncStore(MemoryCasKv()) + pg_node = SimpleNamespace( + node_id="node-namespace", + worker_id="worker-a", + _user_sync_store=store, + ) + + with pytest.raises(RuntimeError, match="before acquire"): + async with NodeOperation._authoritative_user_reconciliation_scope(pg_node, {"user-sync-id"}): + raise RuntimeError("failed before acquire") + + with pytest.raises(UserSyncLeaseLostError, match="was not supplied"): + await store.acquire_user_sync_reconciliation_lease( + "node-namespace", "worker-a", ["user-sync-id"], 30 + ) + + +@pytest.mark.asyncio +async def test_cancelled_reconcile_scope_does_not_leak_unlocked_authorization(): + from app.nats.kv_cas import MemoryCasKv + from app.node.nats_memory import NatsUserSyncStore, UserSyncLeaseLostError + + store = NatsUserSyncStore(MemoryCasKv()) + pg_node = SimpleNamespace( + node_id="node-namespace", + worker_id="worker-a", + _user_sync_store=store, + ) + current = asyncio.current_task() + assert current is not None + + with pytest.raises(asyncio.CancelledError): + async with NodeOperation._authoritative_user_reconciliation_scope(pg_node, {"user-sync-id"}): + asyncio.get_running_loop().call_soon(current.cancel) + await asyncio.Future() + + with pytest.raises(UserSyncLeaseLostError, match="was not supplied"): + await store.acquire_user_sync_reconciliation_lease( + "node-namespace", "worker-a", ["user-sync-id"], 30 + ) + + +@pytest.mark.asyncio +async def test_ambiguous_commit_fresh_membership_splits_abort_and_finalize(monkeypatch): + removal = (ProtoUser(email="1"), ProtoUser(email="2")) + originals = (ProtoUser(email="1", inbounds=["in"]), ProtoUser(email="2", inbounds=["in"])) + revocation = node_sync_module.UserRevocation("delete-1-2", removal, originals) + + class Result: + @staticmethod + def scalars(): + return SimpleNamespace(all=lambda: ["1"]) + + class DB: + async def execute(self, _query): + return Result() + + class DBContext: + async def __aenter__(self): + return DB() + + async def __aexit__(self, *_args): + return False + + abort = AsyncMock() + finalize = AsyncMock() + monkeypatch.setattr(node_sync_module, "GetDB", DBContext) + monkeypatch.setattr(node_sync_module, "_dispatch_abort_with_topology_retry", abort) + monkeypatch.setattr(node_sync_module, "_dispatch_finalize_with_topology_retry", finalize) + failed_db = SimpleNamespace(rollback=AsyncMock()) + + await node_sync_module.resolve_user_removal_after_db_error(revocation, failed_db) + + failed_db.rollback.assert_awaited_once() + assert [user.email for user in abort.await_args.args[0].removal_users] == ["1"] + assert [user.email for user in finalize.await_args.args[0].removal_users] == ["2"] + + +@pytest.mark.asyncio +async def test_ambiguous_commit_unknown_membership_stays_fenced_for_retry(monkeypatch): + revocation = node_sync_module.UserRevocation( + "delete-1", + (ProtoUser(email="1"),), + (ProtoUser(email="1", inbounds=["in"]),), + ) + + class FailingDBContext: + async def __aenter__(self): + raise RuntimeError("database unavailable") + + async def __aexit__(self, *_args): + return False + + scheduled = [] + monkeypatch.setattr(node_sync_module, "GetDB", FailingDBContext) + monkeypatch.setattr(node_sync_module, "_schedule_resolution_retry", scheduled.append) + abort = AsyncMock() + finalize = AsyncMock() + monkeypatch.setattr(node_sync_module, "_dispatch_abort_with_topology_retry", abort) + monkeypatch.setattr(node_sync_module, "_dispatch_finalize_with_topology_retry", finalize) + + await node_sync_module.resolve_user_removal_after_db_error(revocation) + + assert scheduled == [revocation] + abort.assert_not_awaited() + finalize.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_explicit_lifecycle_recovery_requires_probe_match_and_expired_lease(monkeypatch): + operation = NodeOperation(OperatorType.API) + pg_node = AsyncMock() + pg_node.get_lifecycle_state.return_value = NodeLifecycleState( + operation=LifecycleOperation.START, + observed=LifecycleStatus.STARTING, + ) + pg_node.info.return_value = service.BaseInfoResponse(started=True, node_version="0.4.0", core_version="1.0.0") + pg_node.reconcile_lifecycle.side_effect = NodeAPIError(409, "still active") + monkeypatch.setattr(node_module.node_manager, "get_node", AsyncMock(return_value=pg_node)) + recovery = NodeLifecycleRecovery(observed=LifecycleStatus.HEALTHY, acknowledge_expired_operation=True) + + with pytest.raises(NodeAPIError, match="still active"): + await operation._recover_node_lifecycle_local(1, recovery) + + pg_node.reconcile_lifecycle.assert_awaited_once_with(LifecycleStatus.HEALTHY) + + +@pytest.mark.asyncio +async def test_explicit_lifecycle_recovery_rejects_operator_state_mismatch(monkeypatch): + operation = NodeOperation(OperatorType.API) + pg_node = AsyncMock() + pg_node.get_lifecycle_state.return_value = NodeLifecycleState( + operation=LifecycleOperation.STOP, + observed=LifecycleStatus.STOPPING, + ) + pg_node.info.return_value = service.BaseInfoResponse(started=True, node_version="0.4.0", core_version="1.0.0") + monkeypatch.setattr(node_module.node_manager, "get_node", AsyncMock(return_value=pg_node)) + recovery = NodeLifecycleRecovery(observed=LifecycleStatus.STOPPED, acknowledge_expired_operation=True) + + with pytest.raises(Exception, match="Observed node state is healthy, not stopped"): + await operation._recover_node_lifecycle_local(1, recovery) + + pg_node.reconcile_lifecycle.assert_not_awaited() + + +@pytest.fixture(autouse=True) +def _use_local_revocation_store_by_default(monkeypatch: pytest.MonkeyPatch): + """Keep local NodeManager tests independent from the deployment environment.""" + monkeypatch.setattr(node_module, "needs_shared_bridge_memory", lambda: False) + + +@pytest.mark.asyncio +async def test_required_shared_store_fails_closed_instead_of_creating_hybrid_manager( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr(node_module, "needs_shared_bridge_memory", lambda: True) + manager = NodeManager() + db_node = SimpleNamespace( + id=7, + connection_type=NodeConnectionType.rest, + address="127.0.0.1", + port=62050, + api_port=62051, + server_ca="ca", + api_key="key", + name="node-7", + default_timeout=10, + internal_timeout=10, + proxy_url=None, + usage_coefficient=1.0, + ) + monkeypatch.setattr(node_module, "get_bridge_memory", lambda: (None, None, "worker-a")) + monkeypatch.setattr(node_module, "ensure_bridge_memory", AsyncMock(return_value=(None, None))) + + with pytest.raises(NodeAPIError, match="shared node bridge memory is unavailable"): + manager._create_node_kwargs(db_node) + + old_runtime = _healthy_runtime_node() + manager._nodes = {7: old_runtime} + with pytest.raises(NodeAPIError, match="shared node bridge memory is unavailable"): + await manager.update_node(db_node) + assert manager._nodes == {7: old_runtime} + assert manager._retiring_nodes == {} + old_runtime.stop.assert_not_awaited() + + store = object() + coordinator = object() + monkeypatch.setattr( + node_module, + "get_bridge_memory", + lambda: (store, coordinator, "worker-a"), + ) + kwargs = manager._create_node_kwargs(db_node) + assert kwargs["user_sync_store"] is store + assert kwargs["lifecycle_coordinator"] is coordinator + assert manager.uses_shared_revocation_store is True + + db_node.bridge_id = "9d7eb13c-a227-4a0c-a94c-76f15bcb624a" + first_worker = manager._create_node_kwargs(db_node) + second_worker = NodeManager()._create_node_kwargs(db_node) + assert first_worker["node_id"] == db_node.bridge_id + assert second_worker["node_id"] == db_node.bridge_id + + +@pytest.mark.asyncio +async def test_revocation_uses_current_locking_reads_for_users_and_node_topology( + monkeypatch: pytest.MonkeyPatch, +): + statements = [] + result = SimpleNamespace(scalars=lambda: SimpleNamespace(all=lambda: [7])) + session = AsyncMock() + session.execute.side_effect = lambda statement: (statements.append(statement), result)[1] + user = SimpleNamespace(id=73) + monkeypatch.setattr(node_sync_module, "async_object_session", lambda _user: session) + + expected_node_ids = await node_sync_module._lock_users_for_revocation([user]) + + compiled = [str(statement.compile(dialect=mysql.dialect())).upper() for statement in statements] + assert expected_node_ids == {7} + assert len(compiled) == 2 + assert all(" FOR UPDATE" in statement for statement in compiled) def test_node_update_users_nats_chunks_respect_payload_limit(monkeypatch: pytest.MonkeyPatch): @@ -15,3 +436,1685 @@ def test_node_update_users_nats_chunks_respect_payload_limit(monkeypatch: pytest assert [len(chunk) for chunk in chunks] == [2, 2, 1] assert all(len(encode_node_command("update_users", {"users": chunk})) <= max_payload for chunk in chunks) + + +def test_revocation_chunks_measure_removal_and_original_payload(monkeypatch: pytest.MonkeyPatch): + removals = [{"email": f"user-{index}"} for index in range(3)] + originals = [{"email": f"user-{index}", "inbounds": ["x" * 1200]} for index in range(3)] + revocation_id = "r" * 32 + max_payload = len( + encode_node_command( + "revoke_users", + { + "users": removals[:1], + "original_users": originals[:1], + "revocation_id": revocation_id, + }, + ) + ) + monkeypatch.setattr(node_sync_module.nats_settings, "node_update_users_batch_size", 100) + monkeypatch.setattr(node_sync_module.nats_settings, "node_command_max_payload_bytes", max_payload) + + chunks = node_sync_module._chunk_serialized_revocations_for_nats(removals, originals, revocation_id) + + assert [len(users) for users, _ in chunks] == [1, 1, 1] + for users, original_users in chunks: + payload = {"users": users, "original_users": original_users, "revocation_id": revocation_id} + assert len(encode_node_command("revoke_users", payload)) <= max_payload + + +@pytest.mark.asyncio +async def test_unhealthy_node_blocks_revocation_before_any_node_mutation(monkeypatch: pytest.MonkeyPatch): + manager = NodeManager() + healthy = _healthy_runtime_node() + broken = _healthy_runtime_node() + broken.get_health.return_value = Health.BROKEN + manager._nodes = {1: healthy, 2: broken} + sync_batch = AsyncMock(return_value=0) + monkeypatch.setattr(manager, "_sync_user_batch_to_node", sync_batch) + + with pytest.raises(node_sync_module.NodeRevocationError, match="node ids: 2"): + await manager.revoke_users_and_wait( + [ProtoUser(email="error-node-user")], + "error-node-operation", + [ProtoUser(email="error-node-user", inbounds=["vless-in"])], + ) + + healthy.begin_user_revocation.assert_not_awaited() + broken.begin_user_revocation.assert_not_awaited() + sync_batch.assert_not_awaited() + assert manager._deleted_user_keys == set() + + +@pytest.mark.asyncio +async def test_partial_revocation_restores_exact_original_before_releasing_fence(monkeypatch: pytest.MonkeyPatch): + manager = NodeManager() + first = _healthy_runtime_node() + second = _healthy_runtime_node() + manager._nodes = {1: first, 2: second} + calls: list[tuple[int, tuple[str, ...], tuple[str, ...]]] = [] + + async def sync_batch(node, users, *, revocation_id=None): + node_id = 1 if node is first else 2 + calls.append((node_id, tuple(user.email for user in users), tuple(users[0].inbounds))) + assert revocation_id == "partial-operation" + if node is second and not users[0].inbounds: + return len(users) + return 0 + + monkeypatch.setattr(manager, "_sync_user_batch_to_node", sync_batch) + removal = ProtoUser(email="partial-user") + original = ProtoUser(email="partial-user", inbounds=["vless-in"]) + + with pytest.raises(node_sync_module.NodeRevocationError, match="failed to sync users to 1/2 nodes"): + await manager.revoke_users_and_wait([removal], "partial-operation", [original]) + + assert (1, ("partial-user",), ("vless-in",)) in calls + assert (2, ("partial-user",), ("vless-in",)) in calls + first.abort_user_revocation.assert_awaited_once_with(["partial-user"], "partial-operation") + second.abort_user_revocation.assert_awaited_once_with(["partial-user"], "partial-operation") + first.update_users.assert_awaited_once() + second.update_users.assert_awaited_once() + assert first.update_users.await_args.args[0][0].inbounds == ["vless-in"] + assert manager._deleted_user_keys == set() + assert manager._deletion_fence_owners == {} + + +@pytest.mark.asyncio +async def test_revocation_skips_node_keys_already_finalized_by_an_earlier_attempt( + monkeypatch: pytest.MonkeyPatch, +): + manager = NodeManager() + active_node = _healthy_runtime_node() + finalized_node = _healthy_runtime_node() + finalized_node.begin_user_revocation.side_effect = lambda user_keys, _revocation_id: SimpleNamespace( + active_user_keys=(), + finalized_user_keys=tuple(user_keys), + ) + manager._nodes = {1: active_node, 2: finalized_node} + sync_batch = AsyncMock(return_value=0) + monkeypatch.setattr(manager, "_sync_user_batch_to_node", sync_batch) + removal = ProtoUser(email="already-finalized") + + await manager.revoke_users_and_wait([removal], "retry-operation", [removal]) + await manager.finalize_user_revocations([removal], "retry-operation") + + sync_batch.assert_awaited_once_with( + active_node, + [removal], + revocation_id="retry-operation", + ) + active_node.finalize_user_revocation.assert_awaited_once_with( + ["already-finalized"], + "retry-operation", + ) + finalized_node.finalize_user_revocation.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_separate_rpc_chunks_with_same_revocation_id_close_only_their_own_keys( + monkeypatch: pytest.MonkeyPatch, +): + manager = NodeManager() + node = _healthy_runtime_node() + manager._nodes = {1: node} + sync_batch = AsyncMock(return_value=0) + monkeypatch.setattr(manager, "_sync_user_batch_to_node", sync_batch) + first_removal = ProtoUser(email="chunk-a") + first_original = ProtoUser(email="chunk-a", inbounds=["in-a"]) + second_removal = ProtoUser(email="chunk-b") + second_original = ProtoUser(email="chunk-b", inbounds=["in-b"]) + + await manager.revoke_users_and_wait( + [first_removal], + "shared-operation", + [first_original], + ) + await manager.revoke_users_and_wait( + [second_removal], + "shared-operation", + [second_original], + ) + + await manager.abort_user_revocations( + [first_removal], + "shared-operation", + [first_original], + ) + assert manager._revocation_nodes["shared-operation"][0][2] == frozenset({"chunk-b"}) + node.abort_user_revocation.assert_awaited_once_with(["chunk-a"], "shared-operation") + + await manager.finalize_user_revocations([second_removal], "shared-operation") + assert "shared-operation" not in manager._revocation_nodes + node.finalize_user_revocation.assert_awaited_once_with(["chunk-b"], "shared-operation") + + +@pytest.mark.asyncio +async def test_close_chunk_falls_back_when_same_operation_record_contains_only_other_keys(): + manager = NodeManager() + node = _healthy_runtime_node() + manager._nodes = {7: node} + manager._revocation_nodes = {"shared-operation": [(7, node, frozenset({"chunk-a"}))]} + + await manager.finalize_user_revocations( + [ProtoUser(email="chunk-b")], + "shared-operation", + expected_node_ids={7}, + ) + + node.finalize_user_revocation.assert_awaited_once_with(["chunk-b"], "shared-operation") + assert manager._revocation_nodes == {"shared-operation": [(7, node, frozenset({"chunk-a"}))]} + + +@pytest.mark.asyncio +async def test_transient_finalize_failure_retains_operation_state_for_retry(monkeypatch: pytest.MonkeyPatch): + manager = NodeManager() + node = _healthy_runtime_node() + node.finalize_user_revocation.side_effect = [RuntimeError("temporary close failure"), None] + manager._nodes = {7: node} + monkeypatch.setattr(manager, "_sync_user_batch_to_node", AsyncMock(return_value=0)) + user = ProtoUser(email="73") + + await manager.revoke_users_and_wait([user], "delete-73", [user], expected_node_ids={7}) + with pytest.raises(node_sync_module.NodeRevocationError, match="failed to finalize"): + await manager.finalize_user_revocations([user], "delete-73", expected_node_ids={7}) + + assert manager._revocation_nodes["delete-73"][0][2] == frozenset({"73"}) + assert manager._deletion_fence_owners == {"73": {"delete-73"}} + assert not manager._revocations_idle.is_set() + + await manager.finalize_user_revocations([user], "delete-73", expected_node_ids={7}) + + assert manager._revocation_nodes == {} + assert manager._deletion_fence_owners == {} + assert manager._deleted_user_keys == {"73"} + assert manager._revocations_idle.is_set() + + +def test_deletion_fence_uses_the_real_serialized_panel_user_key(): + proto_user = _serialize_user_for_node(73, {}) + + assert "id" not in proto_user.DESCRIPTOR.fields_by_name + assert NodeManager._user_key(proto_user) == "73" + + +@pytest.mark.asyncio +async def test_node_startup_waits_until_provisional_revocation_is_resolved(): + manager = NodeManager() + manager._acquire_deletion_fences({"73"}, "delete-73") + + startup = asyncio.create_task(manager.wait_for_user_revocations()) + await asyncio.sleep(0) + assert not startup.done() + + manager._release_deletion_fences({"73"}, "delete-73") + await asyncio.wait_for(startup, timeout=1) + + +def test_node_startup_filters_permanent_tombstones_from_old_db_snapshot(): + manager = NodeManager() + manager._acquire_deletion_fences({"73"}, "delete-73") + manager._finalize_deletion_fences({"73"}, "delete-73") + users = [ProtoUser(email="73", inbounds=["stale"]), ProtoUser(email="74", inbounds=["active"])] + + assert manager.filter_permanently_deleted_users(users) == [users[1]] + + +@pytest.mark.asyncio +async def test_revocation_never_uses_full_replacement_sync_on_legacy_node(): + manager = NodeManager() + node = _healthy_runtime_node() + node._supports_chunked_sync.return_value = (False, "legacy") + + with pytest.raises(node_sync_module.NodeRevocationError, match="partial chunked sync"): + await manager._sync_user_batch_to_node( + node, + [ProtoUser(email="73")], + revocation_id="delete-73", + ) + + node.sync_users.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_failed_old_runtime_stop_remains_in_revocation_preflight(): + manager = NodeManager() + active = _healthy_runtime_node() + retiring = _healthy_runtime_node() + retiring.stop.side_effect = RuntimeError("stop failed") + retiring.get_health.return_value = Health.INVALID + manager._nodes = {1: active} + manager._retiring_nodes = {1: [retiring]} + + await manager._finish_retiring_node(1, retiring) + + assert manager._retiring_nodes == {1: [retiring]} + with pytest.raises(node_sync_module.NodeRevocationError, match="node ids: 1"): + await manager.revoke_users_and_wait( + [ProtoUser(email="73")], + "delete-73", + [ProtoUser(email="73", inbounds=["active"])], + ) + active.begin_user_revocation.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_confirmed_old_runtime_stop_removes_it_from_revocation_topology(): + manager = NodeManager() + retiring = _healthy_runtime_node() + manager._retiring_nodes = {1: [retiring]} + + await manager._finish_retiring_node(1, retiring) + + assert manager._retiring_nodes == {} + retiring.stop.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_remove_node_waits_for_confirmed_stop_before_forgetting_runtime(): + manager = NodeManager() + node = _healthy_runtime_node() + manager._nodes = {1: node} + manager._user_sync_locks = {1: asyncio.Lock()} + + await manager.remove_node(1) + + node.stop.assert_awaited_once() + assert manager._nodes == {} + assert manager._retiring_nodes == {} + assert manager._user_sync_locks == {} + + +@pytest.mark.asyncio +async def test_stale_remove_namespace_cannot_remove_recreated_numeric_id(): + manager = NodeManager() + replacement = _healthy_runtime_node() + replacement.node_id = "new-bridge-id" + manager._nodes = {1: replacement} + + await manager.remove_node(1, expected_bridge_namespace="old-bridge-id") + + assert manager._nodes == {1: replacement} + assert manager._retiring_nodes == {} + replacement.stop.assert_not_awaited() + replacement.disconnect.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_periodic_orphan_recovery_runs_inside_fresh_db_transaction(monkeypatch): + from app.jobs import node_checker + + transaction_active = False + recovered: list[int] = [] + + class DBContext: + async def __aenter__(self): + nonlocal transaction_active + transaction_active = True + return self + + async def __aexit__(self, *_args): + nonlocal transaction_active + transaction_active = False + return False + + async def _recover(_db, db_node): + assert transaction_active + recovered.append(db_node.id) + + monkeypatch.setattr(node_checker.runtime_settings, "role", Role.NODE) + monkeypatch.setattr(node_checker.node_manager, "get_nodes", AsyncMock(return_value={7: object()})) + monkeypatch.setattr(node_checker, "GetDB", DBContext) + monkeypatch.setattr(node_checker, "get_node_by_id", AsyncMock(return_value=SimpleNamespace(id=7))) + monkeypatch.setattr(node_checker.node_operator, "reconcile_orphaned_user_sync", _recover) + + await node_checker.reconcile_orphaned_user_sync() + + assert recovered == [7] + assert transaction_active is False + + +@pytest.mark.asyncio +async def test_remove_node_failed_stop_retains_runtime_and_fails_closed(): + manager = NodeManager() + node = _healthy_runtime_node() + node.stop.side_effect = RuntimeError("unknown stop outcome") + manager._nodes = {1: node} + manager._user_sync_locks = {1: asyncio.Lock()} + + with pytest.raises(NodeAPIError, match="cannot confirm node 1 runtime shutdown"): + await manager.remove_node(1) + + assert manager._retiring_nodes == {1: [node]} + assert 1 in manager._user_sync_locks + assert manager._removing_node_ids == {1} + + +@pytest.mark.asyncio +async def test_remove_node_stop_failure_does_not_purge_shared_memory(monkeypatch): + operation = NodeOperation(OperatorType.API) + remove = AsyncMock(side_effect=NodeAPIError(503, "ambiguous stop")) + clear = AsyncMock() + monkeypatch.setattr(node_module.node_manager, "remove_node", remove) + monkeypatch.setattr("app.operation.node.clear_bridge_memory_for_node", clear) + + with pytest.raises(NodeAPIError, match="ambiguous stop"): + await operation._remove_node_local(7) + + clear.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_lost_delete_broadcast_still_fences_and_quiesces_sibling(monkeypatch): + from app.nats.kv_cas import MemoryCasKv + from app.node.nats_memory import NatsNodeLifecycleCoordinator + + coordinator = NatsNodeLifecycleCoordinator(MemoryCasKv()) + monkeypatch.setattr(node_module, "ensure_bridge_memory", AsyncMock()) + monkeypatch.setattr( + node_module, + "get_bridge_memory", + lambda: (object(), coordinator, "worker"), + ) + origin = NodeManager() + sibling = NodeManager() + origin._uses_shared_revocation_store = True + sibling._uses_shared_revocation_store = True + origin_runtime = _healthy_runtime_node() + origin_runtime.node_id = "bridge-old" + sibling_runtime = _healthy_runtime_node() + sibling_runtime.node_id = "bridge-old" + origin._nodes = {1: origin_runtime} + sibling._nodes = {1: sibling_runtime} + + await origin.remove_node( + 1, + expected_bridge_namespace="bridge-old", + permanent_delete=True, + ) + + # The sibling deliberately receives no broadcast. Its next DB-driven + # registration observes durable shared state and only disconnects locally. + with pytest.raises(NodeAPIError, match="permanently deleted"): + await sibling.update_node(_db_node(1, "bridge-old")) + sibling_runtime.disconnect.assert_awaited_once() + sibling_runtime.stop.assert_not_awaited() + assert sibling._nodes == {} + + # Reuse of the public numeric id is safe because the new row has a new + # stable Bridge namespace. + replacement = _healthy_runtime_node() + replacement.node_id = "bridge-new" + monkeypatch.setattr(node_module, "create_node", lambda **_kwargs: replacement) + assert await sibling.update_node(_db_node(1, "bridge-new")) is replacement + + +@pytest.mark.asyncio +async def test_temporary_disconnect_allows_same_incarnation_to_restart(monkeypatch): + from app.nats.kv_cas import MemoryCasKv + from app.node.nats_memory import NatsNodeLifecycleCoordinator + + coordinator = NatsNodeLifecycleCoordinator(MemoryCasKv()) + monkeypatch.setattr(node_module, "ensure_bridge_memory", AsyncMock()) + monkeypatch.setattr( + node_module, + "get_bridge_memory", + lambda: (object(), coordinator, "worker"), + ) + manager = NodeManager() + manager._uses_shared_revocation_store = True + old = _healthy_runtime_node() + old.node_id = "bridge-enabled" + manager._nodes = {1: old} + + await manager.remove_node(1, expected_bridge_namespace="bridge-enabled") + + assert await coordinator.is_deleted("bridge-enabled") is False + restarted = _healthy_runtime_node() + restarted.node_id = "bridge-enabled" + monkeypatch.setattr(node_module, "create_node", lambda **_kwargs: restarted) + assert await manager.update_node(_db_node(1, "bridge-enabled")) is restarted + + +@pytest.mark.asyncio +async def test_default_delete_failure_retains_row_runtime_fence_and_tombstone(monkeypatch): + from app.nats.kv_cas import MemoryCasKv + from app.node.nats_memory import NatsNodeLifecycleCoordinator + + coordinator = NatsNodeLifecycleCoordinator(MemoryCasKv()) + monkeypatch.setattr(node_module, "ensure_bridge_memory", AsyncMock()) + monkeypatch.setattr( + node_module, + "get_bridge_memory", + lambda: (object(), coordinator, "worker"), + ) + manager = NodeManager() + manager._uses_shared_revocation_store = True + runtime = _healthy_runtime_node() + runtime.node_id = "bridge-offline" + runtime.stop.side_effect = RuntimeError("offline") + manager._nodes = {1: runtime} + + with pytest.raises(NodeAPIError, match="cannot confirm"): + await manager.remove_node( + 1, + expected_bridge_namespace="bridge-offline", + permanent_delete=True, + ) + + assert await coordinator.is_deleted("bridge-offline") is True + assert manager._retiring_nodes == {1: [runtime]} + assert manager._removing_node_ids == {1} + + +@pytest.mark.asyncio +async def test_force_delete_tombstones_but_only_disconnects_local_runtime(monkeypatch): + from app.nats.kv_cas import MemoryCasKv + from app.node.nats_memory import NatsNodeLifecycleCoordinator + + coordinator = NatsNodeLifecycleCoordinator(MemoryCasKv()) + monkeypatch.setattr(node_module, "ensure_bridge_memory", AsyncMock()) + monkeypatch.setattr( + node_module, + "get_bridge_memory", + lambda: (object(), coordinator, "worker"), + ) + manager = NodeManager() + manager._uses_shared_revocation_store = True + runtime = _healthy_runtime_node() + runtime.node_id = "bridge-force" + manager._nodes = {1: runtime} + + await manager.remove_node( + 1, + remote_stop=False, + expected_bridge_namespace="bridge-force", + permanent_delete=True, + ) + + assert await coordinator.is_deleted("bridge-force") is True + runtime.disconnect.assert_awaited_once() + runtime.stop.assert_not_awaited() + assert manager._nodes == {} + + +@pytest.mark.asyncio +async def test_bulk_delete_reports_partial_failure_and_commits_each_success(monkeypatch): + operation = NodeOperation(OperatorType.API) + nodes = { + 1: SimpleNamespace(id=1, bridge_id="bridge-1", name="node-1"), + 2: SimpleNamespace(id=2, bridge_id="bridge-2", name="node-2"), + } + monkeypatch.setattr( + operation, + "get_validated_node", + AsyncMock(side_effect=lambda _db, node_id: nodes[node_id]), + ) + operation._remove_node_impl = AsyncMock( + side_effect=[None, NodeAPIError(503, "offline stop outcome")] + ) + committed_remove = AsyncMock() + monkeypatch.setattr("app.operation.node.remove_node", committed_remove) + monkeypatch.setattr( + "app.operation.node.NodeResponse.model_validate", + lambda node: SimpleNamespace(id=node.id, name=node.name), + ) + monkeypatch.setattr("app.operation.node.notification.remove_node", AsyncMock()) + admin = SimpleNamespace(username="operator") + + result = await operation.bulk_remove_nodes( + object(), + BulkNodeSelection(ids={1, 2}), + admin, + ) + + assert result.nodes == ["node-1"] + assert result.count == 1 + assert result.failed == {2: "offline stop outcome"} + committed_remove.assert_awaited_once_with(ANY, nodes[1]) + + +@pytest.mark.asyncio +async def test_single_node_delete_commit_ack_loss_is_resolved_from_fresh_db(monkeypatch): + operation = NodeOperation(OperatorType.API) + db_node = SimpleNamespace(id=1, bridge_id="bridge-1", name="node-1") + failed_db = SimpleNamespace(rollback=AsyncMock()) + monkeypatch.setattr(operation, "get_validated_node", AsyncMock(return_value=db_node)) + operation._remove_node_impl = AsyncMock() + monkeypatch.setattr( + "app.operation.node.NodeResponse.model_validate", + lambda node: SimpleNamespace(id=node.id, name=node.name), + ) + monkeypatch.setattr( + "app.operation.node.remove_node", + AsyncMock(side_effect=RuntimeError("commit acknowledgement lost")), + ) + monkeypatch.setattr("app.operation.node.notification.remove_node", AsyncMock()) + + class _FreshDB: + async def __aenter__(self): + return self + + async def __aexit__(self, *_args): + return False + + async def scalar(self, _query): + return None + + monkeypatch.setattr("app.operation.node.GetDB", _FreshDB) + + await operation.remove_node( + failed_db, + 1, + SimpleNamespace(username="operator"), + ) + + failed_db.rollback.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_bulk_node_delete_commit_ack_loss_counts_committed_row(monkeypatch): + operation = NodeOperation(OperatorType.API) + db_node = SimpleNamespace(id=1, bridge_id="bridge-1", name="node-1") + failed_db = SimpleNamespace(rollback=AsyncMock()) + monkeypatch.setattr(operation, "get_validated_node", AsyncMock(return_value=db_node)) + operation._remove_node_impl = AsyncMock() + monkeypatch.setattr( + "app.operation.node.NodeResponse.model_validate", + lambda node: SimpleNamespace(id=node.id, name=node.name), + ) + monkeypatch.setattr( + "app.operation.node.remove_node", + AsyncMock(side_effect=RuntimeError("commit acknowledgement lost")), + ) + monkeypatch.setattr("app.operation.node.notification.remove_node", AsyncMock()) + + class _FreshDB: + async def __aenter__(self): + return self + + async def __aexit__(self, *_args): + return False + + async def scalar(self, _query): + return None + + monkeypatch.setattr("app.operation.node.GetDB", _FreshDB) + + result = await operation.bulk_remove_nodes( + failed_db, + BulkNodeSelection(ids={1}), + SimpleNamespace(username="operator"), + ) + + assert result.nodes == ["node-1"] + assert result.count == 1 + assert result.failed == {} + failed_db.rollback.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_start_completing_after_delete_tombstone_is_stopped(monkeypatch): + from app.nats.kv_cas import MemoryCasKv + from app.node.nats_memory import NatsNodeLifecycleCoordinator + + coordinator = NatsNodeLifecycleCoordinator(MemoryCasKv()) + pg_node = AsyncMock() + pg_node.node_id = "bridge-race" + pg_node.worker_id = "worker-a" + pg_node._user_sync_store = None + pg_node._lifecycle_coordinator = coordinator + pg_node.get_lifecycle_state.return_value = None + pg_node.info.return_value = service.BaseInfoResponse(user_sync_epoch_supported=True) + + async def _start(**_kwargs): + await coordinator.mark_deleted("bridge-race") + return service.BaseInfoResponse(started=True, node_version="0.4.0", core_version="1.0.0") + + pg_node.start.side_effect = _start + core = SimpleNamespace(type=CoreType.wg, to_str=lambda: "{}") + + with pytest.raises(NodeAPIError, match="permanently deleted"): + await NodeOperation._start_or_attach_node( + pg_node, + SimpleNamespace(name="node-race", keep_alive=30), + core, + [], + service.BackendType.WIREGUARD, + ) + pg_node.stop.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_replacement_paused_during_delete_cannot_install_after_tombstone(monkeypatch): + from app.nats.kv_cas import MemoryCasKv + from app.node.nats_memory import NatsNodeLifecycleCoordinator + + coordinator = NatsNodeLifecycleCoordinator(MemoryCasKv()) + monkeypatch.setattr(node_module, "ensure_bridge_memory", AsyncMock()) + monkeypatch.setattr( + node_module, + "get_bridge_memory", + lambda: (object(), coordinator, "worker"), + ) + manager = NodeManager() + manager._uses_shared_revocation_store = True + old = _healthy_runtime_node() + old.node_id = "bridge-old" + stop_started = asyncio.Event() + allow_stop = asyncio.Event() + + async def _stop(): + stop_started.set() + await allow_stop.wait() + + old.stop.side_effect = _stop + manager._nodes = {1: old} + replacement = _healthy_runtime_node() + replacement.node_id = "bridge-old" + monkeypatch.setattr(node_module, "create_node", lambda **_kwargs: replacement) + update = asyncio.create_task(manager.update_node(_db_node(1, "bridge-old"))) + await stop_started.wait() + await coordinator.mark_deleted("bridge-old") + allow_stop.set() + + with pytest.raises(NodeAPIError, match="permanently deleted"): + await update + assert manager._nodes == {} + assert manager._removing_node_ids == set() + assert manager._replacing_node_ids == set() + replacement.disconnect.assert_awaited_once() + + new_incarnation = _healthy_runtime_node() + new_incarnation.node_id = "bridge-new" + monkeypatch.setattr(node_module, "create_node", lambda **_kwargs: new_incarnation) + assert await manager.update_node(_db_node(1, "bridge-new")) is new_incarnation + + +@pytest.mark.asyncio +async def test_local_replacement_paused_during_delete_cannot_install_after_tombstone(monkeypatch): + manager = NodeManager() + old = _healthy_runtime_node() + old.node_id = "bridge-local-old" + first_stop_started = asyncio.Event() + allow_first_stop = asyncio.Event() + stop_calls = 0 + + async def _stop(): + nonlocal stop_calls + stop_calls += 1 + if stop_calls == 1: + first_stop_started.set() + await allow_first_stop.wait() + + old.stop.side_effect = _stop + manager._nodes = {1: old} + replacement = _healthy_runtime_node() + replacement.node_id = "bridge-local-old" + monkeypatch.setattr(node_module, "ensure_bridge_memory", AsyncMock()) + monkeypatch.setattr(node_module, "create_node", lambda **_kwargs: replacement) + + update = asyncio.create_task(manager.update_node(_db_node(1, "bridge-local-old"))) + await first_stop_started.wait() + await manager.remove_node( + 1, + expected_bridge_namespace="bridge-local-old", + permanent_delete=True, + ) + allow_first_stop.set() + + with pytest.raises(NodeAPIError, match="permanently deleted"): + await update + assert stop_calls == 2 + assert manager._nodes == {} + replacement.disconnect.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_replacement_waits_for_all_old_runtimes_and_retries_failure(monkeypatch): + manager = NodeManager() + old = _healthy_runtime_node() + old.stop.side_effect = [RuntimeError("ambiguous stop"), None] + manager._nodes = {1: old} + first_new = _healthy_runtime_node() + second_new = _healthy_runtime_node() + created = iter((first_new, second_new)) + monkeypatch.setattr(node_module, "ensure_bridge_memory", AsyncMock()) + monkeypatch.setattr(node_module, "create_node", lambda **_kwargs: next(created)) + + with pytest.raises(NodeAPIError, match="cannot confirm old node"): + await manager.update_node(_db_node()) + assert manager._nodes == {} + assert manager._retiring_nodes == {1: [old]} + assert manager._replacing_node_ids == {1} + + assert await manager.update_node(_db_node()) is second_new + assert manager._nodes == {1: second_new} + assert manager._retiring_nodes == {} + assert old.stop.await_count == 2 + + +@pytest.mark.asyncio +async def test_cancelled_replacement_keeps_shielded_retiree_cleanup(monkeypatch): + manager = NodeManager() + old = _healthy_runtime_node() + stop_started = asyncio.Event() + allow_stop = asyncio.Event() + + async def _stop(): + stop_started.set() + await allow_stop.wait() + + old.stop.side_effect = _stop + manager._nodes = {1: old} + replacement = _healthy_runtime_node() + monkeypatch.setattr(node_module, "ensure_bridge_memory", AsyncMock()) + monkeypatch.setattr(node_module, "create_node", lambda **_kwargs: replacement) + task = asyncio.create_task(manager.update_node(_db_node())) + await stop_started.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + assert manager._nodes == {} + allow_stop.set() + for _ in range(20): + if manager._nodes.get(1) is replacement: + break + await asyncio.sleep(0) + assert manager._nodes == {1: replacement} + assert manager._retiring_nodes == {} + + +@pytest.mark.asyncio +async def test_health_pass_repairs_missed_runtime_configuration_upsert(monkeypatch): + from app.jobs import node_checker + + db_node = _db_node() + runtime = _healthy_runtime_node() + runtime.node_id = db_node.bridge_id + runtime._extra = {"config_signature": "stale-signature"} + update = AsyncMock() + monkeypatch.setattr(node_checker, "get_bridge_memory", lambda: (None, None, "worker")) + monkeypatch.setattr(node_checker.node_manager, "update_node", update) + + await node_checker.process_node_health_check(db_node, runtime) + + update.assert_awaited_once_with(db_node) + + +@pytest.mark.asyncio +async def test_health_pass_repairs_missing_runtime(monkeypatch): + from app.jobs import node_checker + + db_node = _db_node() + connect = AsyncMock() + + class _DB: + async def __aenter__(self): + return self + + async def __aexit__(self, *_args): + return False + + monkeypatch.setattr(node_checker, "get_bridge_memory", lambda: (None, None, "worker")) + monkeypatch.setattr(node_checker, "GetDB", lambda: _DB()) + monkeypatch.setattr(node_checker.node_operator, "connect_single_node", connect) + + await node_checker.process_node_health_check(db_node, None) + + connect.assert_awaited_once_with(ANY, db_node.id) + + +@pytest.mark.asyncio +async def test_health_pass_uses_uuid_namespace_for_active_lifecycle_lease(monkeypatch): + from app.jobs import node_checker + + db_node = _db_node(4, "uuid-bridge-4") + db_node.status = NodeStatus.connected + runtime = _healthy_runtime_node() + runtime.node_id = db_node.bridge_id + runtime._extra = {} + runtime.requires_hard_reset = lambda: False + runtime.get_lifecycle_state.return_value = NodeLifecycleState( + observed=LifecycleStatus.HEALTHY + ) + seen: list[str] = [] + + class _Coordinator: + async def is_deleted(self, _node_id): + return False + + async def has_active_lease(self, node_id): + seen.append(node_id) + return True + + monkeypatch.setattr( + node_checker, + "get_bridge_memory", + lambda: (None, _Coordinator(), "worker"), + ) + monkeypatch.setattr( + node_checker, + "verify_node_backend_health", + AsyncMock(return_value=(Health.NOT_CONNECTED, None, None)), + ) + monkeypatch.setattr(NodeOperation, "_attach_if_running", AsyncMock(return_value=None)) + reconnect = AsyncMock() + monkeypatch.setattr(node_checker.node_operator, "connect_single_node", reconnect) + + await node_checker.process_node_health_check(db_node, runtime) + + assert seen == ["uuid-bridge-4"] + reconnect.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_worker_with_incomplete_runtime_topology_cannot_commit_revocation(): + authoritative_worker = NodeManager() + node = _healthy_runtime_node() + authoritative_worker._nodes = {7: node} + lagging_worker = NodeManager() + + with pytest.raises(node_sync_module.NodeRevocationError, match="missing node ids: 7"): + await lagging_worker.revoke_users_and_wait( + [ProtoUser(email="73")], + "delete-73", + [ProtoUser(email="73", inbounds=["active"])], + expected_node_ids={7}, + ) + + node.begin_user_revocation.assert_not_awaited() + node.sync_users_chunked.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_shared_store_close_on_another_worker_has_no_process_local_gate(monkeypatch: pytest.MonkeyPatch): + begin_worker = NodeManager() + close_worker = NodeManager() + begin_worker._uses_shared_revocation_store = True + close_worker._uses_shared_revocation_store = True + begin_node = _healthy_runtime_node() + close_node = _healthy_runtime_node() + begin_worker._nodes = {7: begin_node} + close_worker._nodes = {7: close_node} + monkeypatch.setattr(begin_worker, "_sync_user_batch_to_node", AsyncMock(return_value=0)) + user = ProtoUser(email="73") + + await begin_worker.revoke_users_and_wait([user], "delete-73", [user], expected_node_ids={7}) + await close_worker.finalize_user_revocations([user], "delete-73", expected_node_ids={7}) + + close_node.finalize_user_revocation.assert_awaited_once_with(["73"], "delete-73") + assert begin_worker._revocation_nodes == {} + assert begin_worker._deletion_fence_owners == {} + assert begin_worker._deleted_user_keys == set() + assert begin_worker._revocations_idle.is_set() + + +@pytest.mark.asyncio +async def test_shared_store_abort_on_another_worker_restores_payload_without_local_owner( + monkeypatch: pytest.MonkeyPatch, +): + close_worker = NodeManager() + close_worker._uses_shared_revocation_store = True + close_node = _healthy_runtime_node() + close_worker._nodes = {7: close_node} + restored = [] + + async def sync_batch(_node, users, *, revocation_id=None): + restored.extend(users) + assert revocation_id == "delete-73" + return 0 + + monkeypatch.setattr(close_worker, "_sync_user_batch_to_node", sync_batch) + removal = ProtoUser(email="73") + original = ProtoUser(email="73", inbounds=["active"]) + + await close_worker.abort_user_revocations( + [removal], + "delete-73", + [original], + expected_node_ids={7}, + ) + + assert restored == [original] + close_node.abort_user_revocation.assert_awaited_once_with(["73"], "delete-73") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("close_action", ["abort", "finalize"]) +async def test_close_uses_recorded_runtime_after_confirmed_node_removal( + monkeypatch: pytest.MonkeyPatch, + close_action: str, +): + manager = NodeManager() + node = _healthy_runtime_node() + manager._nodes = {7: node} + monkeypatch.setattr(manager, "_sync_user_batch_to_node", AsyncMock(return_value=0)) + removal = ProtoUser(email="73") + original = ProtoUser(email="73", inbounds=["active"]) + + await manager.revoke_users_and_wait( + [removal], + "delete-73", + [original], + expected_node_ids={7}, + ) + manager._nodes.pop(7) + + if close_action == "abort": + await manager.abort_user_revocations( + [removal], + "delete-73", + [original], + expected_node_ids={7}, + ) + node.abort_user_revocation.assert_awaited_once_with(["73"], "delete-73") + assert manager._deleted_user_keys == set() + else: + await manager.finalize_user_revocations( + [removal], + "delete-73", + expected_node_ids={7}, + ) + node.finalize_user_revocation.assert_awaited_once_with(["73"], "delete-73") + assert manager._deleted_user_keys == {"73"} + + assert manager._deletion_fence_owners == {} + assert manager._revocations_idle.is_set() + + +@pytest.mark.asyncio +async def test_removal_waits_for_node_worker_rpc_ack(monkeypatch: pytest.MonkeyPatch): + request = AsyncMock() + monkeypatch.setattr(node_sync_module.runtime_settings, "role", Role.BACKEND) + monkeypatch.setattr(node_sync_module.node_nats_client, "request", request) + removal = ProtoUser(email="user-1") + original = ProtoUser(email="user-1", inbounds=["vless-in"]) + + await node_sync_module._dispatch_users_removal([removal], [original], {7}) + + payload = request.await_args.args[1] + assert request.await_args.args[0] == "revoke_users" + assert payload["revocation_id"] is not None + assert payload["users"][0]["email"] == "user-1" + assert payload["original_users"][0]["email"] == "user-1" + assert payload["original_users"][0]["inbounds"] == ["vless-in"] + assert payload["expected_node_ids"] == [7] + + +@pytest.mark.asyncio +async def test_single_removal_waits_for_node_worker_rpc_ack(monkeypatch: pytest.MonkeyPatch): + request = AsyncMock() + monkeypatch.setattr(node_sync_module.runtime_settings, "role", Role.BACKEND) + monkeypatch.setattr(node_sync_module.node_nats_client, "request", request) + + await node_sync_module._dispatch_user_removal( + ProtoUser(email="user-1"), ProtoUser(email="user-1", inbounds=["vless-in"]) + ) + + payload = request.await_args.args[1] + assert request.await_args.args[0] == "revoke_user" + assert payload["user"]["email"] == "user-1" + assert payload["original_user"]["inbounds"] == ["vless-in"] + assert payload["revocation_id"] is not None + + +@pytest.mark.asyncio +async def test_single_local_removal_waits_for_runtime_ack(monkeypatch: pytest.MonkeyPatch): + revoke_users_and_wait = AsyncMock() + monkeypatch.setattr(node_sync_module.runtime_settings, "role", Role.ALL_IN_ONE) + monkeypatch.setattr(node_sync_module.node_manager, "revoke_users_and_wait", revoke_users_and_wait) + proto_user = object() + original_user = object() + + await node_sync_module._dispatch_user_removal(proto_user, original_user) + + revoke_users_and_wait.assert_awaited_once_with( + [proto_user], + ANY, + [original_user], + expected_node_ids=None, + ) + + +@pytest.mark.asyncio +async def test_remote_removal_failure_is_retryable(monkeypatch: pytest.MonkeyPatch): + request = AsyncMock(side_effect=RuntimeError("NATS is not available")) + monkeypatch.setattr(node_sync_module.runtime_settings, "role", Role.BACKEND) + monkeypatch.setattr(node_sync_module.node_nats_client, "request", request) + + with pytest.raises(node_sync_module.NodeRevocationError, match="cannot confirm user revocation"): + await node_sync_module._dispatch_user_removal( + ProtoUser(email="user-1"), ProtoUser(email="user-1", inbounds=["vless-in"]) + ) + + +@pytest.mark.asyncio +async def test_remote_removal_abort_waits_for_node_worker_ack(monkeypatch: pytest.MonkeyPatch): + request = AsyncMock() + monkeypatch.setattr(node_sync_module.runtime_settings, "role", Role.BACKEND) + monkeypatch.setattr(node_sync_module.node_nats_client, "request", request) + + await node_sync_module._dispatch_users_removal_abort( + [ProtoUser(email="7")], + [ProtoUser(email="7", inbounds=["vless-in"])], + "operation-7", + frozenset({7}), + ) + + payload = request.await_args.args[1] + assert request.await_args.args[0] == "abort_revoke_users" + assert payload["users"][0]["email"] == "7" + assert payload["original_users"][0]["inbounds"] == ["vless-in"] + assert payload["revocation_id"] == "operation-7" + assert payload["expected_node_ids"] == [7] + + +@pytest.mark.asyncio +async def test_abort_refreshes_topology_only_after_stale_snapshot_is_rejected(monkeypatch: pytest.MonkeyPatch): + calls = [] + + async def dispatch(_users, _originals, _revocation_id, expected_node_ids): + calls.append(expected_node_ids) + if len(calls) == 1: + raise node_sync_module.NodeRevocationError( + "runtime topology is incomplete for user revocation (missing node ids: 7)" + ) + + monkeypatch.setattr(node_sync_module, "_dispatch_users_removal_abort", dispatch) + monkeypatch.setattr(node_sync_module, "_refresh_expected_node_ids", AsyncMock(return_value=frozenset())) + revocation = node_sync_module.UserRevocation( + "operation-7", + (ProtoUser(email="7"),), + (ProtoUser(email="7", inbounds=["active"]),), + frozenset({7}), + ) + + await node_sync_module.abort_user_removal(revocation) + + assert calls == [frozenset({7}), frozenset()] + node_sync_module._refresh_expected_node_ids.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_remote_abort_refreshes_topology_through_nested_rpc_error(monkeypatch: pytest.MonkeyPatch): + request = AsyncMock( + side_effect=[ + RuntimeError("runtime topology is incomplete for user revocation (missing node ids: 7)"), + {}, + ] + ) + monkeypatch.setattr(node_sync_module.runtime_settings, "role", Role.BACKEND) + monkeypatch.setattr(node_sync_module.node_nats_client, "request", request) + monkeypatch.setattr(node_sync_module, "_refresh_expected_node_ids", AsyncMock(return_value=frozenset())) + revocation = node_sync_module.UserRevocation( + "operation-remote-refresh", + (ProtoUser(email="7"),), + (ProtoUser(email="7", inbounds=["active"]),), + frozenset({7}), + ) + + await node_sync_module.abort_user_removal(revocation) + + assert request.await_count == 2 + assert request.await_args_list[0].args[1]["expected_node_ids"] == [7] + assert "expected_node_ids" in request.await_args_list[1].args[1] + assert request.await_args_list[1].args[1]["expected_node_ids"] == [] + node_sync_module._refresh_expected_node_ids.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_finalize_transient_failure_is_retried_until_acknowledged(monkeypatch: pytest.MonkeyPatch): + completed = asyncio.Event() + attempts = 0 + + async def finalize(_users, _revocation_id, _expected_node_ids): + nonlocal attempts + attempts += 1 + if attempts == 1: + raise node_sync_module.NodeRevocationError("temporary finalize failure") + completed.set() + + monkeypatch.setattr(node_sync_module, "_dispatch_users_removal_finalize", finalize) + revocation = node_sync_module.UserRevocation( + "operation-retry", + (ProtoUser(email="7"),), + (ProtoUser(email="7", inbounds=["active"]),), + frozenset({7}), + ) + + await node_sync_module.finalize_user_removal(revocation) + retry_task = node_sync_module._finalize_retry_tasks["operation-retry"] + await asyncio.wait_for(completed.wait(), timeout=1) + await asyncio.wait_for(retry_task, timeout=1) + + assert attempts == 2 + assert "operation-retry" not in node_sync_module._finalize_retry_tasks + + +@pytest.mark.asyncio +async def test_abort_transient_failure_keeps_retrying_after_api_failure(monkeypatch: pytest.MonkeyPatch): + completed = asyncio.Event() + attempts = 0 + + async def abort(_users, _originals, _revocation_id, _expected_node_ids): + nonlocal attempts + attempts += 1 + if attempts == 1: + raise node_sync_module.NodeRevocationError("temporary abort failure") + completed.set() + + monkeypatch.setattr(node_sync_module, "_dispatch_users_removal_abort", abort) + revocation = node_sync_module.UserRevocation( + "operation-abort-retry", + (ProtoUser(email="7"),), + (ProtoUser(email="7", inbounds=["active"]),), + frozenset({7}), + ) + + with pytest.raises(node_sync_module.NodeRevocationError, match="temporary abort failure"): + await node_sync_module.abort_user_removal(revocation) + retry_task = node_sync_module._abort_retry_tasks["operation-abort-retry"] + await asyncio.wait_for(completed.wait(), timeout=1) + await asyncio.wait_for(retry_task, timeout=1) + + assert attempts == 2 + assert "operation-abort-retry" not in node_sync_module._abort_retry_tasks + + +@pytest.mark.asyncio +async def test_finalize_cancellation_after_commit_schedules_retry(monkeypatch: pytest.MonkeyPatch): + first_started = asyncio.Event() + retry_completed = asyncio.Event() + attempts = 0 + + async def finalize(_users, _revocation_id, _expected_node_ids): + nonlocal attempts + attempts += 1 + if attempts == 1: + first_started.set() + await asyncio.Future() + retry_completed.set() + + monkeypatch.setattr(node_sync_module, "_dispatch_users_removal_finalize", finalize) + monkeypatch.setattr(node_sync_module.nats_settings, "node_rpc_timeout", 0.01) + revocation = node_sync_module.UserRevocation( + "operation-cancelled-finalize", + (ProtoUser(email="7"),), + (ProtoUser(email="7", inbounds=["active"]),), + frozenset({7}), + ) + + task = asyncio.create_task(node_sync_module.finalize_user_removal(revocation)) + await first_started.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + retry_task = node_sync_module._finalize_retry_tasks["operation-cancelled-finalize"] + await asyncio.wait_for(retry_completed.wait(), timeout=1) + await asyncio.wait_for(retry_task, timeout=1) + assert attempts == 2 + + +@pytest.mark.asyncio +async def test_remote_multi_chunk_abort_has_a_fresh_timeout_for_every_chunk(monkeypatch: pytest.MonkeyPatch): + completed: list[str] = [] + + async def request(_action, payload): + await asyncio.sleep(0.03) + completed.append(payload["users"][0]["email"]) + return {} + + monkeypatch.setattr(node_sync_module.runtime_settings, "role", Role.BACKEND) + monkeypatch.setattr(node_sync_module.nats_settings, "node_update_users_batch_size", 1) + monkeypatch.setattr(node_sync_module.nats_settings, "node_rpc_timeout", 0.05) + monkeypatch.setattr(node_sync_module.node_nats_client, "request", request) + + await node_sync_module._dispatch_users_removal_abort( + [ProtoUser(email="1"), ProtoUser(email="2")], + [ProtoUser(email="1", inbounds=["a"]), ProtoUser(email="2", inbounds=["b"])], + "operation-7", + ) + + assert completed == ["1", "2"] + + +@pytest.mark.asyncio +async def test_remote_chunk_failure_aborts_every_possibly_applied_chunk(monkeypatch: pytest.MonkeyPatch): + request = AsyncMock(side_effect=[{}, RuntimeError("second chunk failed"), {}, {}]) + monkeypatch.setattr(node_sync_module.runtime_settings, "role", Role.BACKEND) + monkeypatch.setattr(node_sync_module.nats_settings, "node_update_users_batch_size", 1) + monkeypatch.setattr(node_sync_module.node_nats_client, "request", request) + removals = [ProtoUser(email="1"), ProtoUser(email="2")] + originals = [ProtoUser(email="1", inbounds=["a"]), ProtoUser(email="2", inbounds=["b"])] + + with pytest.raises(node_sync_module.NodeRevocationError, match="second chunk failed"): + await node_sync_module._dispatch_users_removal(removals, originals, {7}) + + revocation_id = request.await_args_list[0].args[1]["revocation_id"] + assert [item.args[0] for item in request.await_args_list] == [ + "revoke_users", + "revoke_users", + "abort_revoke_users", + "abort_revoke_users", + ] + assert [item.args[1]["users"][0]["email"] for item in request.await_args_list] == ["1", "2", "1", "2"] + assert all(item.args[1]["revocation_id"] == revocation_id for item in request.await_args_list) + assert all(item.args[1]["expected_node_ids"] == [7] for item in request.await_args_list) + assert [item.args[1]["original_users"][0]["email"] for item in request.await_args_list] == [ + "1", + "2", + "1", + "2", + ] + + +@pytest.mark.asyncio +async def test_remote_single_cancellation_aborts_ambiguous_applied_revoke(monkeypatch: pytest.MonkeyPatch): + manager = NodeManager() + user = ProtoUser(email="single-cancel") + revoke_started = asyncio.Event() + calls = [] + + async def request(action, payload): + calls.append((action, payload)) + users = [user] + if action == "revoke_user": + manager._acquire_deletion_fences({user.email}, payload["revocation_id"]) + revoke_started.set() + await asyncio.Future() + else: + await manager.abort_user_revocations( + users, + payload["revocation_id"], + [ProtoUser(email=item["email"], inbounds=["active"]) for item in payload["original_users"]], + ) + return {} + + monkeypatch.setattr(node_sync_module.runtime_settings, "role", Role.BACKEND) + monkeypatch.setattr(node_sync_module.node_nats_client, "request", request) + + task = asyncio.create_task( + node_sync_module._dispatch_user_removal(user, ProtoUser(email=user.email, inbounds=["active"])) + ) + await revoke_started.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + assert [action for action, _ in calls] == ["revoke_user", "abort_revoke_users"] + assert manager._deleted_user_keys == set() + assert manager._deletion_fence_owners == {} + + +@pytest.mark.asyncio +async def test_remote_bulk_cancellation_aborts_acknowledged_and_inflight_chunks(monkeypatch: pytest.MonkeyPatch): + manager = NodeManager() + second_started = asyncio.Event() + calls = [] + + async def request(action, payload): + calls.append((action, payload)) + users = [ProtoUser(email=item["email"]) for item in payload["users"]] + if action == "revoke_users": + manager._acquire_deletion_fences({user.email for user in users}, payload["revocation_id"]) + if users[0].email == "bulk-2": + second_started.set() + await asyncio.Future() + else: + await manager.abort_user_revocations( + users, + payload["revocation_id"], + [ProtoUser(email=item["email"], inbounds=["active"]) for item in payload["original_users"]], + ) + return {} + + monkeypatch.setattr(node_sync_module.runtime_settings, "role", Role.BACKEND) + monkeypatch.setattr(node_sync_module.nats_settings, "node_update_users_batch_size", 1) + monkeypatch.setattr(node_sync_module.node_nats_client, "request", request) + + task = asyncio.create_task( + node_sync_module._dispatch_users_removal( + [ProtoUser(email="bulk-1"), ProtoUser(email="bulk-2")], + [ + ProtoUser(email="bulk-1", inbounds=["active"]), + ProtoUser(email="bulk-2", inbounds=["active"]), + ], + ) + ) + await second_started.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + assert [action for action, _ in calls] == [ + "revoke_users", + "revoke_users", + "abort_revoke_users", + "abort_revoke_users", + ] + assert manager._deleted_user_keys == set() + assert manager._deletion_fence_owners == {} + + +@pytest.mark.asyncio +async def test_remote_bulk_lost_reply_aborts_applied_inflight_chunk(monkeypatch: pytest.MonkeyPatch): + manager = NodeManager() + + async def request(action, payload): + users = [ProtoUser(email=item["email"]) for item in payload["users"]] + if action == "revoke_users": + manager._acquire_deletion_fences({user.email for user in users}, payload["revocation_id"]) + if users[0].email == "lost-reply-2": + raise TimeoutError("reply was lost") + else: + await manager.abort_user_revocations( + users, + payload["revocation_id"], + [ProtoUser(email=item["email"], inbounds=["active"]) for item in payload["original_users"]], + ) + return {} + + monkeypatch.setattr(node_sync_module.runtime_settings, "role", Role.BACKEND) + monkeypatch.setattr(node_sync_module.nats_settings, "node_update_users_batch_size", 1) + monkeypatch.setattr(node_sync_module.node_nats_client, "request", request) + + with pytest.raises(node_sync_module.NodeRevocationError, match="reply was lost"): + await node_sync_module._dispatch_users_removal( + [ProtoUser(email="lost-reply-1"), ProtoUser(email="lost-reply-2")], + [ + ProtoUser(email="lost-reply-1", inbounds=["active"]), + ProtoUser(email="lost-reply-2", inbounds=["active"]), + ], + ) + + assert manager._deleted_user_keys == set() + assert manager._deletion_fence_owners == {} + + +@pytest.mark.asyncio +async def test_removal_without_an_active_runtime_node_is_a_noop(): + manager = NodeManager() + + await manager.revoke_users_and_wait([ProtoUser(email="1")]) + + +@pytest.mark.asyncio +async def test_failed_permanent_removal_releases_the_update_fence(monkeypatch: pytest.MonkeyPatch): + manager = NodeManager() + failed_update = AsyncMock(side_effect=RuntimeError("node unavailable")) + monkeypatch.setattr(manager, "_update_users", failed_update) + + with pytest.raises(RuntimeError, match="node unavailable"): + await manager.revoke_users_and_wait([ProtoUser(email="7")]) + + assert "7" not in manager._deleted_user_keys + + +@pytest.mark.asyncio +async def test_child_cancelled_node_revocation_is_failure_and_releases_fences(monkeypatch: pytest.MonkeyPatch): + manager = NodeManager() + manager._nodes = {1: _healthy_runtime_node(), 2: _healthy_runtime_node()} + visited_nodes: set[int] = set() + + async def sync_node(node_id, _node, _users, **_kwargs): + visited_nodes.add(node_id) + if node_id == 2: + raise asyncio.CancelledError + + monkeypatch.setattr(manager, "_sync_users_to_node", sync_node) + + with pytest.raises(node_sync_module.NodeRevocationError, match="failed to sync users to 1/2 nodes"): + await manager.revoke_users_and_wait([ProtoUser(email="child-cancel")], "child-cancel-operation") + + assert visited_nodes == {1, 2} + assert manager._deleted_user_keys == set() + assert manager._deletion_fence_owners == {} + + +@pytest.mark.asyncio +async def test_caller_cancelled_two_node_revocation_cancels_children_and_releases_fences( + monkeypatch: pytest.MonkeyPatch, +): + manager = NodeManager() + manager._nodes = {1: _healthy_runtime_node(), 2: _healthy_runtime_node()} + both_started = asyncio.Event() + started_nodes: set[int] = set() + cancelled_nodes: set[int] = set() + + async def sync_node(node_id, _node, _users, **_kwargs): + started_nodes.add(node_id) + if len(started_nodes) == 2: + both_started.set() + try: + await asyncio.Future() + finally: + cancelled_nodes.add(node_id) + + monkeypatch.setattr(manager, "_sync_users_to_node", sync_node) + task = asyncio.create_task( + manager.revoke_users_and_wait([ProtoUser(email="caller-cancel")], "caller-cancel-operation") + ) + await both_started.wait() + task.cancel() + + with pytest.raises(asyncio.CancelledError): + await task + + assert cancelled_nodes == {1, 2} + assert manager._deleted_user_keys == set() + assert manager._deletion_fence_owners == {} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("user_count", [1, 2]) +async def test_cancelled_local_revocation_releases_single_and_bulk_fences(monkeypatch, user_count): + manager = NodeManager() + revocation_started = asyncio.Event() + + async def wait_forever(*_args, **_kwargs): + revocation_started.set() + await asyncio.Future() + + monkeypatch.setattr(manager, "_update_users", wait_forever) + users = [ProtoUser(email=f"local-cancel-{index}") for index in range(user_count)] + task = asyncio.create_task(manager.revoke_users_and_wait(users, "cancelled-operation")) + await revocation_started.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + assert manager._deleted_user_keys == set() + assert manager._deletion_fence_owners == {} + + +@pytest.mark.asyncio +async def test_explicit_abort_releases_fence_and_allows_subsequent_sync(): + manager = NodeManager() + removed_user = ProtoUser(email="8") + revocation_id = await manager.revoke_users_and_wait([removed_user]) + assert "8" in manager._deleted_user_keys + + await manager.abort_user_revocations([removed_user], revocation_id, [removed_user]) + node = AsyncMock() + manager._nodes[1] = node + await manager.update_user(ProtoUser(email="8", inbounds=["active-inbound"])) + + assert "8" not in manager._deleted_user_keys + node.update_user.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_failed_concurrent_revocation_releases_only_its_own_fence(monkeypatch: pytest.MonkeyPatch): + manager = NodeManager() + successful_revocation_started = asyncio.Event() + allow_success = asyncio.Event() + calls = 0 + + async def interleaved_update(*_args, **_kwargs): + nonlocal calls + calls += 1 + if calls == 1: + successful_revocation_started.set() + await allow_success.wait() + return + raise RuntimeError("concurrent node failure") + + monkeypatch.setattr(manager, "_update_users", interleaved_update) + user = ProtoUser(email="23") + successful = asyncio.create_task(manager.revoke_users_and_wait([user], "successful-operation")) + await successful_revocation_started.wait() + + with pytest.raises(RuntimeError, match="concurrent node failure"): + await manager.revoke_users_and_wait([user], "failed-operation") + allow_success.set() + await successful + + assert manager._deletion_fence_owners == {"23": {"successful-operation"}} + assert manager._deleted_user_keys == {"23"} + await manager.abort_user_revocations([user], "failed-operation", [user]) + assert manager._deleted_user_keys == {"23"} + + +@pytest.mark.asyncio +async def test_concurrent_database_abort_releases_only_matching_revocation(): + manager = NodeManager() + user = ProtoUser(email="24") + + await asyncio.gather( + manager.revoke_users_and_wait([user], "committed-operation"), + manager.revoke_users_and_wait([user], "rolled-back-operation"), + ) + await manager.abort_user_revocations([user], "rolled-back-operation", [user]) + + assert manager._deletion_fence_owners == {"24": {"committed-operation"}} + assert manager._deleted_user_keys == {"24"} + + +@pytest.mark.asyncio +async def test_successful_finalize_clears_all_operation_owners_but_keeps_tombstone(): + manager = NodeManager() + user = ProtoUser(email="25") + + await asyncio.gather( + manager.revoke_users_and_wait([user], "committed-operation"), + manager.revoke_users_and_wait([user], "rolled-back-operation"), + ) + await manager.finalize_user_revocations([user], "committed-operation") + await manager.abort_user_revocations([user], "rolled-back-operation", [user]) + + assert manager._deleted_user_keys == {"25"} + assert manager._deletion_fence_owners == {} + + await manager.revoke_users_and_wait([user], "stale-duplicate") + await manager.abort_user_revocations([user], "stale-duplicate", [user]) + assert manager._deleted_user_keys == {"25"} + assert manager._deletion_fence_owners == {} + + +@pytest.mark.asyncio +async def test_legacy_revoke_and_abort_without_id_use_same_deterministic_owner(): + manager = NodeManager() + users = [ProtoUser(email="legacy-2"), ProtoUser(email="legacy-1")] + + revocation_id = await manager.revoke_users_and_wait(users) + assert revocation_id == manager._resolve_revocation_id(list(reversed(users)), None) + assert revocation_id.startswith("legacy:") + + await manager.abort_user_revocations(list(reversed(users)), restore_users=list(reversed(users))) + assert manager._deleted_user_keys == set() + assert manager._deletion_fence_owners == {} + + +@pytest.mark.asyncio +async def test_legacy_worker_revoke_and_abort_without_id_pair(monkeypatch): + manager = NodeManager() + service = node_worker_module.NodeWorkerService.__new__(node_worker_module.NodeWorkerService) + monkeypatch.setattr(node_worker_module, "node_manager", manager) + payload = { + "users": [{"email": "legacy-worker"}], + "original_users": [{"email": "legacy-worker", "inbounds": ["active"]}], + } + + await service._rpc_revoke_users(payload) + assert manager._deleted_user_keys == {"legacy-worker"} + await service._rpc_abort_revoke_users(payload) + + assert manager._deleted_user_keys == set() + assert manager._deletion_fence_owners == {} + + +@pytest.mark.asyncio +async def test_permanent_removal_fences_a_concurrent_single_update(monkeypatch: pytest.MonkeyPatch): + manager = NodeManager() + node = _healthy_runtime_node() + stale_update_reached_snapshot = asyncio.Event() + allow_stale_update_to_continue = asyncio.Event() + snapshot_calls = 0 + + async def snapshot_nodes(): + nonlocal snapshot_calls + snapshot_calls += 1 + if snapshot_calls == 1: + stale_update_reached_snapshot.set() + await allow_stale_update_to_continue.wait() + return [(1, node)] + + removal_batches = AsyncMock(return_value=0) + monkeypatch.setattr(manager, "_snapshot_node_items", snapshot_nodes) + monkeypatch.setattr(manager, "_sync_user_batch_to_node", removal_batches) + active_user = ProtoUser(email="41", inbounds=["active-inbound"]) + removed_user = ProtoUser(email="41") + revocation_id = manager._resolve_revocation_id([removed_user], None) + + stale_update = asyncio.create_task(manager.update_user(active_user)) + await stale_update_reached_snapshot.wait() + await manager.revoke_users_and_wait([removed_user]) + assert "41" in manager._deleted_user_keys + allow_stale_update_to_continue.set() + await stale_update + + node.update_user.assert_not_awaited() + removal_batches.assert_awaited_once_with( + node, + [removed_user], + revocation_id=revocation_id, + ) + + +@pytest.mark.asyncio +async def test_permanent_removal_fences_a_concurrent_bulk_update(monkeypatch: pytest.MonkeyPatch): + manager = NodeManager() + node = _healthy_runtime_node() + stale_update_reached_snapshot = asyncio.Event() + allow_stale_update_to_continue = asyncio.Event() + snapshot_calls = 0 + + async def snapshot_nodes(): + nonlocal snapshot_calls + snapshot_calls += 1 + if snapshot_calls == 1: + stale_update_reached_snapshot.set() + await allow_stale_update_to_continue.wait() + return [(1, node)] + + sync_batches = AsyncMock(return_value=0) + monkeypatch.setattr(manager, "_snapshot_node_items", snapshot_nodes) + monkeypatch.setattr(manager, "_sync_user_batch_to_node", sync_batches) + active_user = ProtoUser(email="52", inbounds=["active-inbound"]) + removed_user = ProtoUser(email="52") + revocation_id = manager._resolve_revocation_id([removed_user], None) + + stale_update = asyncio.create_task(manager._update_users([active_user])) + await stale_update_reached_snapshot.wait() + await manager.revoke_users_and_wait([removed_user]) + allow_stale_update_to_continue.set() + await stale_update + + sync_batches.assert_awaited_once_with( + node, + [removed_user], + revocation_id=revocation_id, + ) diff --git a/tests/test_record_usages.py b/tests/test_record_usages.py index b2151bc43..762d55655 100644 --- a/tests/test_record_usages.py +++ b/tests/test_record_usages.py @@ -6,7 +6,7 @@ from unittest.mock import AsyncMock import pytest -from sqlalchemy import select +from sqlalchemy import delete, select from sqlalchemy.exc import SQLAlchemyError from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine from sqlalchemy.pool import NullPool, StaticPool @@ -115,6 +115,7 @@ async def test_record_user_usages_updates_users_and_admins(monkeypatch: pytest.M session.add_all([user_one, user_two]) await session.flush() user_one_id, user_two_id = user_one.id, user_two.id + user_one_sync_id, user_two_sync_id = user_one.sync_id, user_two.sync_id node_one = Node( name="node-1", @@ -146,8 +147,12 @@ async def test_record_user_usages_updates_users_and_admins(monkeypatch: pytest.M monkeypatch.setattr(record_usages.node_manager, "get_healthy_nodes", AsyncMock(return_value=nodes)) stats_map = { - node_one_id: [{"uid": str(user_one_id), "value": 100}, {"uid": str(user_two_id), "value": 50}], - node_two_id: [{"uid": str(user_one_id), "value": 75}], + node_one_id: [ + {"uid": user_one_sync_id, "value": 100}, + {"uid": user_two_sync_id, "value": 50}, + {"uid": "deleted-user-incarnation", "value": 999}, + ], + node_two_id: [{"uid": user_one_sync_id, "value": 75}], } async def fake_get_users_stats(node: DummyNode): @@ -183,6 +188,7 @@ async def fake_get_users_stats(node: DummyNode): (node_one_id, user_two_id), (node_two_id, user_one_id), } + assert len(node_usage_records) == 3 aggregated_usage = defaultdict(int) for record in node_usage_records: @@ -193,6 +199,98 @@ async def fake_get_users_stats(node: DummyNode): assert aggregated_usage[user_id] == total_usage +@pytest.mark.asyncio +async def test_numeric_cutover_stats_are_lossless_but_cannot_hit_reused_user_id( + monkeypatch: pytest.MonkeyPatch, session_factory +): + async with session_factory() as session: + admin = Admin(username="cutover-admin", hashed_password="secret", role_id=3) + session.add(admin) + await session.flush() + user = User( + username="legacy-user", + admin_id=admin.id, + proxy_settings=ProxyTable().dict(no_obj=True), + ) + node = Node( + name="cutover-node", + address="10.0.0.9", + port=1000, + api_port=1001, + server_ca="ca", + api_key="key", + core_config_id=None, + ) + session.add_all([user, node]) + await session.flush() + legacy_user_id = user.id + node_id = node.id + # d12 backfills already-running users to their existing Xray stat key. + user.sync_id = str(legacy_user_id) + await session.commit() + + monkeypatch.setattr( + record_usages.node_manager, + "get_healthy_nodes", + AsyncMock(return_value=[(node_id, DummyNode(node_id))]), + ) + monkeypatch.setattr( + record_usages, + "get_users_stats", + AsyncMock(return_value=[{"uid": str(legacy_user_id), "value": 125}]), + ) + monkeypatch.setattr(record_usages.usage_settings, "disable_recording_node_usage", False) + + # Counters drained from a core that still uses numeric UIDs are credited + # before the first full UUID snapshot reaches that core. + await record_usages.record_user_usages() + async with session_factory() as session: + assert await session.scalar(select(User.used_traffic).where(User.id == legacy_user_id)) == 125 + assert ( + await session.scalar( + select(NodeUserUsage.used_traffic).where( + NodeUserUsage.node_id == node_id, + NodeUserUsage.user_id == legacy_user_id, + ) + ) + == 125 + ) + await session.execute(delete(NodeUserUsage).where(NodeUserUsage.user_id == legacy_user_id)) + await session.execute(delete(User).where(User.id == legacy_user_id)) + await session.commit() + + replacement = User( + username="replacement-user", + admin_id=admin.id, + proxy_settings=ProxyTable().dict(no_obj=True), + ) + # Reproduce public-ID reuse on every supported dialect. SQLite may + # reuse the deleted highest row automatically, while MySQL and + # PostgreSQL sequences intentionally keep advancing. + replacement.id = legacy_user_id + session.add(replacement) + await session.flush() + assert replacement.id == legacy_user_id + assert replacement.sync_id != str(legacy_user_id) + replacement_sync_id = replacement.sync_id + await session.commit() + + # A late numeric counter from the deleted incarnation no longer matches + # the replacement even though the public numeric id was reused. + await record_usages.record_user_usages() + async with session_factory() as session: + replacement = await session.scalar(select(User).where(User.sync_id == replacement_sync_id)) + assert replacement is not None + assert replacement.used_traffic == 0 + usage = await session.scalar( + select(NodeUserUsage).where( + NodeUserUsage.node_id == node_id, + NodeUserUsage.user_id == legacy_user_id, + ) + ) + assert usage is None + + @pytest.mark.asyncio async def test_record_user_usages_limits_overused_admin(monkeypatch: pytest.MonkeyPatch, session_factory): async with session_factory() as session: @@ -213,7 +311,7 @@ async def test_record_user_usages_limits_overused_admin(monkeypatch: pytest.Monk ) session.add_all([user, node]) await session.flush() - user_id, node_id = user.id, node.id + user_sync_id, node_id = user.sync_id, node.id await session.commit() monkeypatch.setattr( @@ -221,7 +319,7 @@ async def test_record_user_usages_limits_overused_admin(monkeypatch: pytest.Monk ) async def fake_get_users_stats(_: DummyNode): - return [{"uid": str(user_id), "value": 150}] + return [{"uid": user_sync_id, "value": 150}] remove_users = AsyncMock() monkeypatch.setattr(record_usages, "get_users_stats", fake_get_users_stats) diff --git a/tests/test_rpc_service.py b/tests/test_rpc_service.py new file mode 100644 index 000000000..2287ecb14 --- /dev/null +++ b/tests/test_rpc_service.py @@ -0,0 +1,122 @@ +import json +from collections import defaultdict +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from app.nats import rpc_service as rpc_service_module +from app.nats.rpc_service import BaseRpcService +from app.node.errors import NodeRevocationError +from app.node.worker import NODE_RPC_QUEUE_GROUP, NodeWorkerService +from config import nats_settings + + +class _FakeSubscription: + async def unsubscribe(self): + return None + + +class _FakeNats: + def __init__(self): + self.is_closed = False + self.subscriptions = [] + + async def subscribe(self, subject, *, queue=None, cb): + self.subscriptions.append((subject, queue, cb)) + return _FakeSubscription() + + async def close(self): + self.is_closed = True + + async def deliver(self, subject, message): + grouped = defaultdict(list) + ungrouped = [] + for registered_subject, queue, callback in self.subscriptions: + if registered_subject != subject: + continue + if queue is None: + ungrouped.append(callback) + else: + grouped[queue].append(callback) + + for callback in ungrouped: + await callback(message) + for callbacks in grouped.values(): + await callbacks[0](message) + + +@pytest.mark.asyncio +async def test_rpc_service_preserves_retryable_error_code(): + service = BaseRpcService("test.rpc", MagicMock(), lambda: True) + service._dispatch_rpc = AsyncMock(side_effect=NodeRevocationError("node unavailable")) + message = MagicMock(respond=AsyncMock()) + + await service._run_rpc(message, "revoke_users", {}) + + assert json.loads(message.respond.await_args.args[0]) == { + "ok": False, + "error": "node unavailable", + "code": 503, + } + + +@pytest.mark.asyncio +async def test_node_worker_rpc_subscription_uses_stable_queue_group(monkeypatch): + connection = _FakeNats() + monkeypatch.setattr(rpc_service_module, "create_nats_client", AsyncMock(return_value=connection)) + monkeypatch.setattr(rpc_service_module, "is_nats_enabled", lambda: True) + service = NodeWorkerService() + service._role_check = lambda: True + + await service.start() + + assert connection.subscriptions[0][:2] == (service._rpc_subject, NODE_RPC_QUEUE_GROUP) + assert connection.subscriptions[1][:2] == (nats_settings.node_command_subject, None) + + await BaseRpcService.stop(service) + + +@pytest.mark.asyncio +async def test_rpc_queue_group_delivers_request_to_only_one_service(monkeypatch): + connection = _FakeNats() + monkeypatch.setattr(rpc_service_module, "create_nats_client", AsyncMock(return_value=connection)) + monkeypatch.setattr(rpc_service_module, "is_nats_enabled", lambda: True) + calls = [] + + async def handle(instance): + calls.append(instance) + + services = [BaseRpcService("test.rpc", MagicMock(), lambda: True, queue_group="test.workers") for _ in range(2)] + for service in services: + service._handle_rpc = lambda message, service=service: handle(service) + await service.start() + + await connection.deliver("test.rpc", MagicMock()) + + assert calls == [services[0]] + + for service in services: + await service.stop() + + +@pytest.mark.asyncio +async def test_rpc_without_queue_group_preserves_broadcast_delivery(monkeypatch): + connection = _FakeNats() + monkeypatch.setattr(rpc_service_module, "create_nats_client", AsyncMock(return_value=connection)) + monkeypatch.setattr(rpc_service_module, "is_nats_enabled", lambda: True) + calls = [] + + async def handle(instance): + calls.append(instance) + + services = [BaseRpcService("test.rpc", MagicMock(), lambda: True) for _ in range(2)] + for service in services: + service._handle_rpc = lambda message, service=service: handle(service) + await service.start() + + await connection.deliver("test.rpc", MagicMock()) + + assert calls == services + + for service in services: + await service.stop() diff --git a/tests/test_security_hardening.py b/tests/test_security_hardening.py new file mode 100644 index 000000000..c4e94b9a3 --- /dev/null +++ b/tests/test_security_hardening.py @@ -0,0 +1,557 @@ +import asyncio +from datetime import UTC, datetime, timedelta +from types import SimpleNamespace +from unittest.mock import ANY, AsyncMock + +import jwt +import pytest +from fastapi import HTTPException +from PasarGuardNodeBridge import Health +from PasarGuardNodeBridge.common.service_pb2 import User as ProtoUser +from pydantic import ValidationError + +import app.node as node_module +from app.db.models import UserStatus +from app.jobs import remove_expired_users as remove_expired_users_job +from app.models.subscription import SubscriptionUsageQuery +from app.models.user import BulkUsersSelection, ExpiredUsersQuery, UserCreate, UserNotificationResponse +from app.node import NodeManager, sync as node_sync_module +from app.notification import webhook as webhook_notification +from app.operation import OperatorType, admin as admin_operation_module, user as user_operation_module +from app.operation.admin import AdminOperation +from app.operation.subscription import SubscriptionOperation +from app.operation.user import UserOperation +from app.utils import jwt as jwt_utils +from role import Role + + +def _healthy_runtime_node() -> AsyncMock: + node = AsyncMock() + node.get_health.return_value = Health.HEALTHY + node._supports_chunked_sync.return_value = (True, "0.2.0") + node.sync_users_chunked.return_value = [] + node.begin_user_revocation = AsyncMock( + side_effect=lambda user_keys, _revocation_id: SimpleNamespace( + active_user_keys=tuple(user_keys), + finalized_user_keys=(), + ) + ) + node.abort_user_revocation = AsyncMock() + node.finalize_user_revocation = AsyncMock() + return node + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status", [UserStatus.disabled, UserStatus.expired, UserStatus.limited]) +async def test_subscription_config_requires_eligible_status(status): + operation = SubscriptionOperation(operator_type=OperatorType.API) + + with pytest.raises(HTTPException) as exc_info: + await operation.require_config_eligible(SimpleNamespace(status=status)) + + assert exc_info.value.status_code == 403 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("status", [UserStatus.active, UserStatus.on_hold]) +async def test_subscription_config_allows_runtime_eligible_statuses(status): + operation = SubscriptionOperation(operator_type=OperatorType.API) + + await operation.require_config_eligible(SimpleNamespace(status=status)) + + +@pytest.mark.asyncio +async def test_subscription_usage_rejects_ranges_over_31_days(): + operation = SubscriptionOperation(operator_type=OperatorType.API) + start = datetime(2026, 1, 1, tzinfo=UTC) + + with pytest.raises(HTTPException) as exc_info: + await operation.get_user_usage( + db=None, + token="unused", + query=SubscriptionUsageQuery(start=start, end=start + timedelta(days=32)), + ) + + assert exc_info.value.status_code == 400 + + +def test_user_auto_delete_days_are_bounded(): + with pytest.raises(ValidationError): + UserCreate(username="overflow", auto_delete_in_days=2_147_483_647) + + +@pytest.mark.asyncio +async def test_admin_tokens_require_expiry(monkeypatch): + secret = "test-secret" + monkeypatch.setattr(jwt_utils, "get_secret_key", AsyncMock(return_value=secret)) + permanent_token = jwt.encode( + {"sub": "admin", "access": "admin", "iat": datetime.now(UTC)}, + secret, + algorithm="HS256", + ) + + assert await jwt_utils.get_admin_payload(permanent_token) is None + issued_token = await jwt_utils.create_admin_token(1, "admin") + assert "exp" in jwt.decode(issued_token, secret, algorithms=["HS256"]) + + +@pytest.mark.asyncio +async def test_webhook_notification_redacts_subscription_credentials(monkeypatch): + enqueue = AsyncMock() + monkeypatch.setattr( + webhook_notification, + "webhook_settings", + AsyncMock(return_value=SimpleNamespace(enable=True)), + ) + monkeypatch.setattr(webhook_notification, "enqueue_webhook", enqueue) + user = UserNotificationResponse( + id=1, + username="subscriber", + status=UserStatus.active, + used_traffic=0, + created_at=datetime.now(UTC), + subscription_url="https://example.test/sub/bearer-token", + proxy_settings={"vless": {"id": "00000000-0000-4000-8000-000000000001"}}, + ) + + await webhook_notification.notify( + webhook_notification.ReachedUsagePercent( + username=user.username, + user=user, + used_percent=80, + ) + ) + + payload = enqueue.await_args.args[0] + assert "subscription_url" not in payload["user"] + assert "proxy_settings" not in payload["user"] + + +@pytest.mark.asyncio +async def test_scheduled_cleanup_revokes_deleted_user_on_nodes(monkeypatch): + user = SimpleNamespace(username="expired") + sync_remove_users = AsyncMock(return_value="scheduled-operation") + finalize = AsyncMock() + + class FakeDBContext: + def __init__(self): + self.db = SimpleNamespace(commit=AsyncMock(), rollback=AsyncMock()) + + async def __aenter__(self): + return self.db + + async def __aexit__(self, exc_type, exc, traceback): + return None + + db_context = FakeDBContext() + monkeypatch.setattr(remove_expired_users_job, "GetDB", lambda: db_context) + remove_users = AsyncMock() + monkeypatch.setattr( + remove_expired_users_job, "get_autodelete_expired_users", AsyncMock(return_value=([user], [user])) + ) + monkeypatch.setattr(remove_expired_users_job, "remove_users", remove_users) + monkeypatch.setattr(remove_expired_users_job, "remove_users_and_wait", sync_remove_users) + monkeypatch.setattr(remove_expired_users_job, "finalize_users_removal", finalize) + monkeypatch.setattr(remove_expired_users_job.notification, "remove_user", AsyncMock()) + + await remove_expired_users_job.remove_expired_users() + + sync_remove_users.assert_awaited_once_with([user]) + finalize.assert_awaited_once_with("scheduled-operation") + remove_users.assert_awaited_once_with(db_context.db, [user]) + + +@pytest.mark.asyncio +async def test_scheduled_cleanup_does_not_report_success_when_node_publish_fails(monkeypatch): + user = SimpleNamespace(username="expired") + + class FakeDBContext: + def __init__(self): + self.db = SimpleNamespace(commit=AsyncMock(), rollback=AsyncMock()) + + async def __aenter__(self): + return self.db + + async def __aexit__(self, exc_type, exc, traceback): + return None + + db_context = FakeDBContext() + monkeypatch.setattr(remove_expired_users_job, "GetDB", lambda: db_context) + monkeypatch.setattr( + remove_expired_users_job, "get_autodelete_expired_users", AsyncMock(return_value=([user], [user])) + ) + monkeypatch.setattr( + remove_expired_users_job, + "remove_users_and_wait", + AsyncMock(side_effect=RuntimeError("node unavailable")), + ) + notify = AsyncMock() + monkeypatch.setattr(remove_expired_users_job.notification, "remove_user", notify) + + with pytest.raises(RuntimeError, match="node unavailable"): + await remove_expired_users_job.remove_expired_users() + + notify.assert_not_called() + db_context.db.rollback.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_manual_cleanup_revokes_deleted_user_on_nodes(monkeypatch): + db_user = SimpleNamespace(username="expired") + notification_user = SimpleNamespace(username="expired") + operation = UserOperation(operator_type=OperatorType.API) + operation.validate_dates = AsyncMock(return_value=(None, None)) + operation.validate_user = AsyncMock(return_value=notification_user) + sync_remove_users = AsyncMock(return_value="manual-operation") + finalize = AsyncMock() + monkeypatch.setattr(user_operation_module, "get_expired_users", AsyncMock(return_value=[db_user])) + monkeypatch.setattr(user_operation_module, "remove_users", AsyncMock()) + monkeypatch.setattr(user_operation_module, "remove_users_and_wait", sync_remove_users) + monkeypatch.setattr(user_operation_module, "finalize_users_removal", finalize) + + db = SimpleNamespace() + response = await operation.delete_expired_users( + db=db, + admin=SimpleNamespace(username="admin"), + query=ExpiredUsersQuery(), + ) + + assert response.users == ["expired"] + sync_remove_users.assert_awaited_once_with([db_user]) + finalize.assert_awaited_once_with("manual-operation") + + +@pytest.mark.asyncio +async def test_bulk_delete_db_failure_resolves_ambiguous_commit(monkeypatch): + db_user = SimpleNamespace(id=81, username="bulk-survivor") + notification_user = SimpleNamespace(id=81, username="bulk-survivor") + operation = UserOperation(operator_type=OperatorType.API) + operation._get_validated_users_by_ids = AsyncMock(return_value=[db_user]) + operation.validate_user = AsyncMock(return_value=notification_user) + revoke = AsyncMock(return_value="bulk-operation") + resolve = AsyncMock() + monkeypatch.setattr(user_operation_module, "remove_users_and_wait", revoke) + monkeypatch.setattr(user_operation_module, "resolve_user_removal_after_db_error", resolve) + monkeypatch.setattr( + user_operation_module, + "remove_users", + AsyncMock(side_effect=RuntimeError("database commit failed")), + ) + + with pytest.raises(RuntimeError, match="database commit failed"): + await operation.bulk_remove_users( + SimpleNamespace(), + BulkUsersSelection(ids={81}), + SimpleNamespace(username="admin"), + ) + + revoke.assert_awaited_once_with([db_user]) + resolve.assert_awaited_once_with("bulk-operation", ANY) + + +@pytest.mark.asyncio +async def test_bulk_delete_success_finalizes_operation_metadata(monkeypatch): + db_user = SimpleNamespace(id=83, username="bulk-deleted") + notification_user = SimpleNamespace(id=83, username="bulk-deleted") + operation = UserOperation(operator_type=OperatorType.API) + operation._get_validated_users_by_ids = AsyncMock(return_value=[db_user]) + operation.validate_user = AsyncMock(return_value=notification_user) + finalize = AsyncMock() + monkeypatch.setattr(user_operation_module, "remove_users_and_wait", AsyncMock(return_value="bulk-commit")) + monkeypatch.setattr(user_operation_module, "remove_users", AsyncMock()) + monkeypatch.setattr(user_operation_module, "finalize_users_removal", finalize) + + await operation.bulk_remove_users( + SimpleNamespace(), BulkUsersSelection(ids={83}), SimpleNamespace(username="admin") + ) + + finalize.assert_awaited_once_with("bulk-commit") + + +@pytest.mark.asyncio +async def test_admin_delete_success_finalizes_operation_metadata(monkeypatch): + db_user = SimpleNamespace(id=84, username="admin-owned") + notification_user = SimpleNamespace(id=84, username="admin-owned") + operation = AdminOperation(operator_type=OperatorType.API) + finalize = AsyncMock() + monkeypatch.setattr(admin_operation_module, "get_users", AsyncMock(return_value=[db_user])) + monkeypatch.setattr( + admin_operation_module.UserOperation, "validate_user", AsyncMock(return_value=notification_user) + ) + monkeypatch.setattr(admin_operation_module, "remove_users_and_wait", AsyncMock(return_value="admin-commit")) + monkeypatch.setattr(admin_operation_module, "remove_users", AsyncMock()) + monkeypatch.setattr(admin_operation_module, "finalize_users_removal", finalize) + + await operation._remove_all_users_for_admin( + SimpleNamespace(), SimpleNamespace(username="owner"), SimpleNamespace(username="admin") + ) + + finalize.assert_awaited_once_with("admin-commit") + + +@pytest.mark.asyncio +async def test_bulk_delete_late_remote_chunk_failure_aborts_fences_before_database_delete(monkeypatch): + manager = NodeManager() + users = [ + UserNotificationResponse( + id=user_id, + username=f"user-{user_id}", + status=UserStatus.active, + used_traffic=0, + created_at=datetime.now(UTC), + proxy_settings={"vless": {"id": f"00000000-0000-4000-8000-{user_id:012d}"}}, + ) + for user_id in (101, 102) + ] + operation = UserOperation(operator_type=OperatorType.API) + operation._get_validated_users_by_ids = AsyncMock( + return_value=[ + SimpleNamespace( + id=user.id, + sync_id=f"sync-{user.id}", + username=user.username, + proxy_settings=user.proxy_settings.dict(), + status=UserStatus.active, + groups=[], + ) + for user in users + ] + ) + operation.validate_user = AsyncMock(side_effect=users) + database_delete = AsyncMock() + + async def remote_request(action, payload): + user_keys = {user["email"] for user in payload["users"]} + revocation_id = payload["revocation_id"] + if action == "revoke_users": + if user_keys == {"sync-102"}: + raise RuntimeError("second chunk failed") + manager._acquire_deletion_fences(user_keys, revocation_id) + return {} + assert action == "abort_revoke_users" + manager._release_deletion_fences(user_keys, revocation_id) + return {} + + monkeypatch.setattr(node_sync_module.runtime_settings, "role", Role.BACKEND) + monkeypatch.setattr(node_sync_module.nats_settings, "node_update_users_batch_size", 1) + monkeypatch.setattr(node_sync_module.node_nats_client, "request", remote_request) + monkeypatch.setattr(user_operation_module, "remove_users", database_delete) + + with pytest.raises(node_sync_module.NodeRevocationError, match="second chunk failed"): + await operation.bulk_remove_users( + SimpleNamespace(), + BulkUsersSelection(ids={101, 102}), + SimpleNamespace(username="admin"), + ) + + database_delete.assert_not_awaited() + assert manager._deleted_user_keys == set() + assert manager._deletion_fence_owners == {} + + +@pytest.mark.asyncio +async def test_bulk_delete_child_cancelled_node_does_not_delete_database_rows(monkeypatch): + manager = NodeManager() + manager._nodes = {1: _healthy_runtime_node(), 2: _healthy_runtime_node()} + user = UserNotificationResponse( + id=103, + username="child-cancelled-user", + status=UserStatus.active, + used_traffic=0, + created_at=datetime.now(UTC), + proxy_settings={"vless": {"id": "00000000-0000-4000-8000-000000000103"}}, + ) + operation = UserOperation(operator_type=OperatorType.API) + operation._get_validated_users_by_ids = AsyncMock( + return_value=[SimpleNamespace(id=user.id, username=user.username)] + ) + operation.validate_user = AsyncMock(return_value=user) + database_delete = AsyncMock() + + async def sync_node(node_id, _node, _users, **_kwargs): + if node_id == 2: + raise asyncio.CancelledError + + async def revoke_users(users): + proto_users = [ProtoUser(email=str(item.id)) for item in users] + return await manager.revoke_users_and_wait(proto_users, "bulk-child-cancel") + + monkeypatch.setattr(manager, "_sync_users_to_node", sync_node) + monkeypatch.setattr(user_operation_module, "remove_users_and_wait", revoke_users) + monkeypatch.setattr(user_operation_module, "remove_users", database_delete) + + with pytest.raises(node_sync_module.NodeRevocationError, match="failed to sync users to 1/2 nodes"): + await operation.bulk_remove_users( + SimpleNamespace(), + BulkUsersSelection(ids={103}), + SimpleNamespace(username="admin"), + ) + + database_delete.assert_not_awaited() + assert manager._deleted_user_keys == set() + assert manager._deletion_fence_owners == {} + + +@pytest.mark.asyncio +async def test_manual_expired_delete_db_failure_resolves_ambiguous_commit(monkeypatch): + db_user = SimpleNamespace(id=82, username="expired-survivor") + notification_user = SimpleNamespace(id=82, username="expired-survivor") + operation = UserOperation(operator_type=OperatorType.API) + operation.validate_dates = AsyncMock(return_value=(None, None)) + operation.validate_user = AsyncMock(return_value=notification_user) + resolve = AsyncMock() + monkeypatch.setattr(user_operation_module, "get_expired_users", AsyncMock(return_value=[db_user])) + monkeypatch.setattr(user_operation_module, "remove_users_and_wait", AsyncMock(return_value="expired-operation")) + monkeypatch.setattr(user_operation_module, "resolve_user_removal_after_db_error", resolve) + monkeypatch.setattr( + user_operation_module, + "remove_users", + AsyncMock(side_effect=RuntimeError("database commit failed")), + ) + + with pytest.raises(RuntimeError, match="database commit failed"): + await operation.delete_expired_users( + db=SimpleNamespace(), + admin=SimpleNamespace(username="admin"), + query=ExpiredUsersQuery(), + ) + + resolve.assert_awaited_once_with("expired-operation", ANY) + + +@pytest.mark.asyncio +async def test_single_delete_db_failure_aborts_fence_and_allows_future_sync(monkeypatch): + manager = NodeManager() + monkeypatch.setattr(node_sync_module.runtime_settings, "role", Role.ALL_IN_ONE) + monkeypatch.setattr(node_sync_module, "node_manager", manager) + db_user = SimpleNamespace( + id=91, + sync_id="sync-91", + username="surviving-user", + proxy_settings={"vless": {"id": "00000000-0000-4000-8000-000000000091"}}, + status=UserStatus.active, + groups=[], + ) + notification_user = UserNotificationResponse( + id=91, + username="surviving-user", + status=UserStatus.active, + used_traffic=0, + created_at=datetime.now(UTC), + proxy_settings={"vless": {"id": "00000000-0000-4000-8000-000000000091"}}, + ) + operation = UserOperation(operator_type=OperatorType.API) + operation.validate_user = AsyncMock(return_value=notification_user) + row_exists = True + + async def fail_database_delete(*_args, **_kwargs): + raise RuntimeError("database commit failed") + + monkeypatch.setattr(user_operation_module, "remove_user", fail_database_delete) + async def resolve_present(revocation, _db): + await node_sync_module.abort_user_removal(revocation) + + monkeypatch.setattr(user_operation_module, "resolve_user_removal_after_db_error", resolve_present) + + with pytest.raises(RuntimeError, match="database commit failed"): + await operation._remove_user(SimpleNamespace(), db_user, SimpleNamespace(username="admin")) + + assert row_exists is True + assert "sync-91" not in manager._deleted_user_keys + node = AsyncMock() + manager._nodes[1] = node + await manager.update_user( + node_sync_module._serialize_user_for_node("sync-91", notification_user.proxy_settings.dict()) + ) + node.update_user.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_single_delete_db_failure_restores_exact_original_node_state(monkeypatch): + manager = NodeManager() + node = _healthy_runtime_node() + manager._nodes = {1: node} + monkeypatch.setattr(node_sync_module.runtime_settings, "role", Role.ALL_IN_ONE) + monkeypatch.setattr(node_sync_module, "node_manager", manager) + group = SimpleNamespace( + is_disabled=False, + inbounds=[SimpleNamespace(tag="vless-in")], + ) + db_user = SimpleNamespace( + id=191, + sync_id="sync-191", + username="surviving-user", + proxy_settings={"vless": {"id": "00000000-0000-4000-8000-000000000191"}}, + status=UserStatus.active, + groups=[group], + ) + notification_user = UserNotificationResponse( + id=191, + username="surviving-user", + status=UserStatus.active, + used_traffic=0, + created_at=datetime.now(UTC), + proxy_settings=db_user.proxy_settings, + ) + operation = UserOperation(operator_type=OperatorType.API) + operation.validate_user = AsyncMock(return_value=notification_user) + applied: list[tuple[str, ...]] = [] + + async def sync_batch(_node, users, *, revocation_id=None): + assert revocation_id + applied.append(tuple(users[0].inbounds)) + return 0 + + monkeypatch.setattr(manager, "_sync_user_batch_to_node", sync_batch) + monkeypatch.setattr( + user_operation_module, + "remove_user", + AsyncMock(side_effect=RuntimeError("database commit failed")), + ) + async def resolve_present(revocation, _db): + await node_sync_module.abort_user_removal(revocation) + + monkeypatch.setattr(user_operation_module, "resolve_user_removal_after_db_error", resolve_present) + + with pytest.raises(RuntimeError, match="database commit failed"): + await operation._remove_user(SimpleNamespace(), db_user, SimpleNamespace(username="admin")) + + assert applied == [(), ("vless-in",)] + node.abort_user_revocation.assert_awaited_once_with(["sync-191"], ANY) + node.update_users.assert_awaited_once() + assert node.update_users.await_args.args[0][0].inbounds == ["vless-in"] + assert manager._deleted_user_keys == set() + assert manager._deletion_fence_owners == {} + + +@pytest.mark.asyncio +async def test_single_delete_success_keeps_permanent_fence(monkeypatch): + monkeypatch.setattr(node_module, "needs_shared_bridge_memory", lambda: False) + manager = NodeManager() + monkeypatch.setattr(node_sync_module.runtime_settings, "role", Role.ALL_IN_ONE) + monkeypatch.setattr(node_sync_module, "node_manager", manager) + db_user = SimpleNamespace( + id=92, + sync_id="sync-92", + username="deleted-user", + proxy_settings={"vless": {"id": "00000000-0000-4000-8000-000000000092"}}, + status=UserStatus.active, + groups=[], + ) + notification_user = UserNotificationResponse( + id=92, + username="deleted-user", + status=UserStatus.active, + used_traffic=0, + created_at=datetime.now(UTC), + proxy_settings={"vless": {"id": "00000000-0000-4000-8000-000000000092"}}, + ) + operation = UserOperation(operator_type=OperatorType.API) + operation.validate_user = AsyncMock(return_value=notification_user) + monkeypatch.setattr(user_operation_module, "remove_user", AsyncMock()) + monkeypatch.setattr(user_operation_module.notification, "remove_user", AsyncMock()) + + await operation._remove_user(SimpleNamespace(), db_user, SimpleNamespace(username="admin")) + + assert "sync-92" in manager._deleted_user_keys + assert manager._deletion_fence_owners == {} diff --git a/tests/test_sync_namespace_migration.py b/tests/test_sync_namespace_migration.py new file mode 100644 index 000000000..37aaa2cf9 --- /dev/null +++ b/tests/test_sync_namespace_migration.py @@ -0,0 +1,56 @@ +import importlib.util +from pathlib import Path + +import sqlalchemy as sa +from alembic.migration import MigrationContext +from alembic.operations import Operations +from sqlalchemy.dialects import mysql, postgresql + + +def _load_migration_module(): + path = Path("app/db/migrations/versions/d12f6a8b9c30_add_bridge_sync_namespaces.py") + spec = importlib.util.spec_from_file_location("bridge_sync_namespace_migration", path) + assert spec and spec.loader + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def test_namespace_migration_backfills_existing_node_and_user(monkeypatch): + module = _load_migration_module() + engine = sa.create_engine("sqlite://") + metadata = sa.MetaData() + nodes = sa.Table("nodes", metadata, sa.Column("id", sa.Integer, primary_key=True)) + users = sa.Table("users", metadata, sa.Column("id", sa.Integer, primary_key=True)) + metadata.create_all(engine) + + with engine.begin() as connection: + connection.execute(nodes.insert(), [{"id": 1}]) + connection.execute(users.insert(), [{"id": 7}]) + monkeypatch.setattr(module, "op", Operations(MigrationContext.configure(connection))) + + module.upgrade() + + bridge_id = connection.execute(sa.text("SELECT bridge_id FROM nodes WHERE id = 1")).scalar_one() + sync_id = connection.execute(sa.text("SELECT sync_id FROM users WHERE id = 7")).scalar_one() + node_constraints = sa.inspect(connection).get_unique_constraints("nodes") + user_constraints = sa.inspect(connection).get_unique_constraints("users") + + assert bridge_id == "1" + assert sync_id == "7" + assert {item["name"] for item in node_constraints} == {"uq_nodes_bridge_id"} + assert {item["name"] for item in user_constraints} == {"uq_users_sync_id"} + + +def test_namespace_backfill_uses_dialect_safe_cast(): + module = _load_migration_module() + statements = [] + + class CaptureConnection: + def execute(self, statement): + statements.append(statement) + + module._backfill_legacy_namespace(CaptureConnection(), "nodes", "bridge_id") + + assert "AS CHAR(36)" in str(statements[0].compile(dialect=mysql.dialect())) + assert "AS VARCHAR(36)" in str(statements[0].compile(dialect=postgresql.dialect())) diff --git a/tests/test_xray_loopback_migration.py b/tests/test_xray_loopback_migration.py new file mode 100644 index 000000000..7f6f71d02 --- /dev/null +++ b/tests/test_xray_loopback_migration.py @@ -0,0 +1,60 @@ +import importlib.util +from pathlib import Path + +import sqlalchemy as sa + + +def _load_migration_module(): + path = Path("app/db/migrations/versions/a8c2d491e705_bind_xray_client_inbounds_to_loopback.py") + spec = importlib.util.spec_from_file_location("xray_loopback_migration", path) + assert spec and spec.loader + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def test_xray_loopback_migration_updates_edited_system_template_without_reformatting(monkeypatch): + module = _load_migration_module() + canonical = '{"inbounds": [{"listen": "0.0.0.0"}]}' + customized = ( + '{"inbounds":[{"listen":"0.0.0.0", "custom": true}], ' + '"outbounds": [{"listen": "0.0.0.0"}], "label": "keep 0.0.0.0"}' + ) + already_safe = '{"inbounds": [{"listen": "127.0.0.1", "custom": true}]}' + + engine = sa.create_engine("sqlite://") + metadata = sa.MetaData() + templates = sa.Table( + "client_templates", + metadata, + sa.Column("id", sa.Integer, primary_key=True), + sa.Column("template_type", sa.String), + sa.Column("content", sa.Text), + sa.Column("is_system", sa.Boolean), + ) + metadata.create_all(engine) + with engine.begin() as connection: + connection.execute( + templates.insert(), + [ + {"id": 1, "template_type": "xray_subscription", "content": canonical, "is_system": True}, + {"id": 2, "template_type": "xray_subscription", "content": customized, "is_system": True}, + {"id": 3, "template_type": "xray_subscription", "content": customized, "is_system": False}, + {"id": 4, "template_type": "singbox_subscription", "content": customized, "is_system": True}, + {"id": 5, "template_type": "xray_subscription", "content": already_safe, "is_system": True}, + ], + ) + monkeypatch.setattr(module.op, "get_bind", lambda: connection) + + module.upgrade() + + contents = dict(connection.execute(sa.select(templates.c.id, templates.c.content)).all()) + + assert contents[1] == '{"inbounds": [{"listen": "127.0.0.1"}]}' + assert contents[2] == ( + '{"inbounds":[{"listen":"127.0.0.1", "custom": true}], ' + '"outbounds": [{"listen": "0.0.0.0"}], "label": "keep 0.0.0.0"}' + ) + assert contents[3] == customized + assert contents[4] == customized + assert contents[5] == already_safe diff --git a/uv.lock b/uv.lock index ae88fa869..fa317567f 100644 --- a/uv.lock +++ b/uv.lock @@ -868,7 +868,7 @@ requires-dist = [ { name = "jinja2", specifier = ">=3.1.6" }, { name = "nats-py", specifier = ">=2.15.0" }, { name = "packaging", specifier = ">=26.2" }, - { name = "pasarguard-node-bridge", specifier = ">=0.9.0" }, + { name = "pasarguard-node-bridge", git = "https://github.com/PasarGuard/node_bridge_py.git?rev=bb503222e373135e8166ddd025fc51348f0806b6" }, { name = "pip-system-certs", specifier = ">=5.3" }, { name = "psutil", specifier = ">=7.2.2" }, { name = "pydantic", specifier = ">=2.13.4" }, @@ -901,8 +901,8 @@ dev = [ [[package]] name = "pasarguard-node-bridge" -version = "0.9.0" -source = { registry = "https://pypi.org/simple" } +version = "0.10.0" +source = { git = "https://github.com/PasarGuard/node_bridge_py.git?rev=bb503222e373135e8166ddd025fc51348f0806b6#bb503222e373135e8166ddd025fc51348f0806b6" } dependencies = [ { name = "aiohttp" }, { name = "aiohttp-socks" }, @@ -911,10 +911,6 @@ dependencies = [ { name = "protobuf" }, { name = "python-socks" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/23/ac/557cc78c98e744b140d8c84cb8b48b08e539037c365608853a0453cd82ee/pasarguard_node_bridge-0.9.0.tar.gz", hash = "sha256:fe3dc842722a0b93b242acfcccdec3add33edbc4937e92ca659d33f76a3e84c6", size = 104313, upload-time = "2026-08-04T08:03:46.222Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/f4/68/76a9cfb929685881363844c49bfabe22d7af5f9e40cd5e2c5fbf9ce19cb5/pasarguard_node_bridge-0.9.0-py3-none-any.whl", hash = "sha256:62f4b49730c9102507c3978c91c9a635a69589efc54978f05063b35e815d3420", size = 56011, upload-time = "2026-08-04T08:03:45.092Z" }, -] [[package]] name = "pip" From d9d2cfbeaa3e6c36877adee834bbb2a0e835a533 Mon Sep 17 00:00:00 2001 From: Rerowros Date: Mon, 10 Aug 2026 01:20:33 +0400 Subject: [PATCH 2/2] Make subscription token revocation fail closed --- ...ve_subscription_revocation_microseconds.py | 56 ++++++++ app/db/models.py | 12 +- app/operation/__init__.py | 33 ++++- app/operation/user.py | 6 +- app/utils/jwt.py | 83 +++++++++-- tests/api/test_user.py | 30 ++++ .../test_subscription_revocation_precision.py | 65 +++++++++ tests/test_subscription_token_revocation.py | 135 ++++++++++++++++++ 8 files changed, 403 insertions(+), 17 deletions(-) create mode 100644 app/db/migrations/versions/9e0d7a1c4b52_preserve_subscription_revocation_microseconds.py create mode 100644 tests/test_subscription_revocation_precision.py create mode 100644 tests/test_subscription_token_revocation.py diff --git a/app/db/migrations/versions/9e0d7a1c4b52_preserve_subscription_revocation_microseconds.py b/app/db/migrations/versions/9e0d7a1c4b52_preserve_subscription_revocation_microseconds.py new file mode 100644 index 000000000..0b96c2b65 --- /dev/null +++ b/app/db/migrations/versions/9e0d7a1c4b52_preserve_subscription_revocation_microseconds.py @@ -0,0 +1,56 @@ +"""preserve subscription token timestamp microseconds + +Revision ID: 9e0d7a1c4b52 +Revises: d12f6a8b9c30 +Create Date: 2026-08-09 + +""" + +from alembic import op +from sqlalchemy.dialects import mysql + +# revision identifiers, used by Alembic. +revision = "9e0d7a1c4b52" +down_revision = "d12f6a8b9c30" +branch_labels = None +depends_on = None + + +def _is_mysql_family() -> bool: + return op.get_bind().dialect.name in {"mysql", "mariadb"} + + +def upgrade() -> None: + if _is_mysql_family(): + op.alter_column( + "users", + "created_at", + existing_type=mysql.DATETIME(fsp=0), + type_=mysql.DATETIME(fsp=6), + existing_nullable=False, + ) + op.alter_column( + "users", + "sub_revoked_at", + existing_type=mysql.DATETIME(fsp=0), + type_=mysql.DATETIME(fsp=6), + existing_nullable=True, + ) + + +def downgrade() -> None: + if _is_mysql_family(): + op.alter_column( + "users", + "sub_revoked_at", + existing_type=mysql.DATETIME(fsp=6), + type_=mysql.DATETIME(fsp=0), + existing_nullable=True, + ) + op.alter_column( + "users", + "created_at", + existing_type=mysql.DATETIME(fsp=6), + type_=mysql.DATETIME(fsp=0), + existing_nullable=False, + ) diff --git a/app/db/models.py b/app/db/models.py index 2d16aae10..b4d218479 100644 --- a/app/db/models.py +++ b/app/db/models.py @@ -23,6 +23,7 @@ func, or_, ) +from sqlalchemy.dialects import mysql from sqlalchemy.dialects.postgresql import JSONB from sqlalchemy.ext.asyncio import async_object_session from sqlalchemy.ext.hybrid import hybrid_property @@ -33,6 +34,7 @@ from app.db.compiles_types import CaseSensitiveString, DaysDiff, EnumArray, SqliteCompatibleBigInteger, StringArray PostgresJSONB = JSON().with_variant(JSONB(none_as_null=True), "postgresql") +SubscriptionTokenDateTime = DateTime(timezone=True).with_variant(mysql.DATETIME(fsp=6), "mysql", "mariadb") def fk_id_column(target: str, **column_kwargs: Any): @@ -202,6 +204,14 @@ class User(Base, CreatedAtUTCMixin): sync_id: Mapped[str] = mapped_column( String(36), default_factory=lambda: str(uuid4()), nullable=False, init=False ) + # Subscription tokens bind to this exact timestamp. MySQL/MariaDB default + # DATETIME precision is zero, so use microseconds explicitly just like the + # revocation timestamp below. + created_at: Mapped[dt] = mapped_column( + SubscriptionTokenDateTime, + default_factory=lambda: dt.now(UTC), + init=False, + ) username: Mapped[str] = mapped_column(CaseSensitiveString(128), unique=True, index=True) node_usages: Mapped[list[NodeUserUsage]] = relationship( back_populates="user", @@ -233,7 +243,7 @@ class User(Base, CreatedAtUTCMixin): ) _expire: Mapped[dt | None] = mapped_column("expire", DateTime(timezone=True), default=None, init=False) admin_id: Mapped[int | None] = fk_id_column("admins.id", default=None) - sub_revoked_at: Mapped[dt | None] = mapped_column(DateTime(timezone=True), default=None) + sub_revoked_at: Mapped[dt | None] = mapped_column(SubscriptionTokenDateTime, default=None) note: Mapped[str | None] = mapped_column(String(500), default=None) online_at: Mapped[dt | None] = mapped_column(DateTime(timezone=True), default=None) on_hold_expire_duration: Mapped[int | None] = mapped_column(BigInteger, default=None) diff --git a/app/operation/__init__.py b/app/operation/__init__.py index 7fab14bc3..db9a38dd3 100644 --- a/app/operation/__init__.py +++ b/app/operation/__init__.py @@ -149,11 +149,34 @@ async def get_validated_sub(self, db: AsyncSession, token: str, *, load_admin_ro elif sub.get("username"): db_user = await get_user(db, sub["username"], load_admin_role=load_admin_role) - if ( - not db_user - or db_user.created_at.astimezone(UTC) > sub["created_at"] - or (db_user.sub_revoked_at and db_user.sub_revoked_at.astimezone(UTC) > sub["created_at"]) - ): + token_created_at = sub["created_at"] if sub else None + token_version = sub.get("token_version", "legacy") if sub else "legacy" + token_subject_created_at = sub.get("subject_created_at") if sub else None + # v2/v3 encoded a rounded-up second, so their value cannot reliably + # order an event within that same second. Preserve their established + # strict-after compatibility rule. v4 carries microsecond precision, + # letting a reset revoke a token issued in the same database tick. + revoked_after_issuance = bool( + db_user + and db_user.sub_revoked_at + and ( + ensure_datetime_timezone(db_user.sub_revoked_at).astimezone(UTC) >= token_created_at + if token_version in {"legacy", "v4", "v5"} + else ensure_datetime_timezone(db_user.sub_revoked_at).astimezone(UTC) > token_created_at + ) + ) + + created_after_token = bool( + db_user + and ( + ensure_datetime_timezone(db_user.created_at).astimezone(UTC) != token_subject_created_at + if token_version == "v5" + else ensure_datetime_timezone(db_user.created_at).astimezone(UTC) + > (token_created_at + td(seconds=1) if token_version == "v4" else token_created_at) + ) + ) + + if not db_user or created_after_token or revoked_after_issuance: await self.raise_error(message="Not Found", code=404) return db_user diff --git a/app/operation/user.py b/app/operation/user.py index 5bc6ed257..cebc42524 100644 --- a/app/operation/user.py +++ b/app/operation/user.py @@ -226,7 +226,7 @@ async def generate_subscription_url(user: UserNotificationResponse): if user.admin and user.admin.sub_domain else (settings.url_prefix).replace("*", salt) ) - token = await create_subscription_token(user.id) + token = await create_subscription_token(user.id, user_created_at=user.created_at) return f"{url_prefix}/{subscription_env_settings.path}/{token}" async def _generate_usernames( @@ -1811,7 +1811,7 @@ def builder(username: str): db_admin = await get_admin(db, admin.username, load_users=False, load_usage_logs=False) try: - subscription_urls = await self._persist_bulk_users( + await self._persist_bulk_users( db, admin, db_admin, @@ -1829,8 +1829,10 @@ def builder(username: str): created_users = await self._load_users_by_usernames(db, [user.username for user in users_to_create]) await sync_users(created_users) + subscription_urls: list[str] = [] for db_user in created_users: user = await self.validate_user(db_user) + subscription_urls.append(user.subscription_url) asyncio.create_task(notification.create_user(user, admin)) return BulkUsersCreateResponse(subscription_urls=subscription_urls, created=len(subscription_urls)) diff --git a/app/utils/jwt.py b/app/utils/jwt.py index 7af5f57e7..77c79c070 100644 --- a/app/utils/jwt.py +++ b/app/utils/jwt.py @@ -3,13 +3,13 @@ from base64 import b64decode, b64encode from datetime import UTC, datetime, timedelta from hashlib import sha256 -from math import ceil import jwt from aiocache import cached from app.db import GetDB from app.db.crud.general import get_jwt_secret_key +from app.utils.helpers import ensure_datetime_timezone from config import jwt_settings @@ -63,8 +63,40 @@ async def get_admin_payload(token: str) -> dict | None: return -async def create_subscription_token(user_id: int) -> str: - data = "v3," + str(user_id) + "," + str(ceil(time.time())) +def _datetime_to_epoch_nanoseconds(value: datetime) -> int: + value = ensure_datetime_timezone(value).astimezone(UTC) + epoch = datetime(1970, 1, 1, tzinfo=UTC) + delta = value - epoch + return (delta.days * 86_400 + delta.seconds) * 1_000_000_000 + delta.microseconds * 1_000 + + +def _datetime_from_epoch_seconds(value: float) -> datetime | None: + try: + return datetime.fromtimestamp(value, tz=UTC) + except OverflowError, OSError, TypeError, ValueError: + return None + + +def _datetime_from_epoch_nanoseconds(value: int) -> datetime | None: + seconds, nanoseconds = divmod(value, 1_000_000_000) + created_at = _datetime_from_epoch_seconds(seconds) + if created_at is None: + return None + return created_at.replace(microsecond=nanoseconds // 1_000) + + +async def create_subscription_token(user_id: int, *, user_created_at: datetime | None = None) -> str: + # Subscription-token revocation compares its issuance time with database + # timestamps at microsecond precision. A rounded-up epoch second can make + # a token appear to have been issued *after* a subsequent revocation, while + # a rounded-down second can predate a freshly-created user. v5 stores the + # actual UTC epoch and the user's creation timestamp in nanoseconds; + # v2/v3/v4 remain accepted for compatibility. + issued_at_ns = time.time_ns() + subject_created_at_ns = ( + _datetime_to_epoch_nanoseconds(user_created_at) if user_created_at is not None else issued_at_ns + ) + data = "v5," + str(user_id) + "," + str(issued_at_ns) + "," + str(subject_created_at_ns) data_b64_str = b64encode(data.encode("utf-8"), altchars=b"-_").decode("utf-8").rstrip("=") secret = await get_secret_key() # HMAC-SHA256 over the payload, url-safe base64, no truncation. @@ -84,16 +116,42 @@ async def create_subscription_token(user_id: int) -> str: def _parse_subscription_data(data_str: str) -> dict | None: """Parse the decoded subscription payload string into a result dict.""" parts = data_str.split(",") - if len(parts) == 3 and parts[0] in ("v2", "v3"): - _, u_user_id_str, u_created_at_str = parts + if len(parts) == 3 and parts[0] in ("v2", "v3", "v4"): + version, u_user_id_str, u_created_at_str = parts try: u_user_id = int(u_user_id_str) u_created_at = int(u_created_at_str) except ValueError: return + if version == "v4": + created_at = _datetime_from_epoch_nanoseconds(u_created_at) + else: + created_at = _datetime_from_epoch_seconds(u_created_at) + if created_at is None: + return return { "user_id": u_user_id, - "created_at": datetime.fromtimestamp(u_created_at, tz=UTC), + "created_at": created_at, + "token_version": version, + } + + if len(parts) == 4 and parts[0] == "v5": + _, u_user_id_str, u_created_at_str, u_subject_created_at_str = parts + try: + u_user_id = int(u_user_id_str) + u_created_at = int(u_created_at_str) + u_subject_created_at = int(u_subject_created_at_str) + except ValueError: + return + created_at = _datetime_from_epoch_nanoseconds(u_created_at) + subject_created_at = _datetime_from_epoch_nanoseconds(u_subject_created_at) + if created_at is None or subject_created_at is None: + return + return { + "user_id": u_user_id, + "created_at": created_at, + "subject_created_at": subject_created_at, + "token_version": "v5", } if len(parts) == 2: @@ -102,9 +160,13 @@ def _parse_subscription_data(data_str: str) -> dict | None: u_created_at = int(u_created_at_str) except ValueError: return + created_at = _datetime_from_epoch_seconds(u_created_at) + if created_at is None: + return return { "username": u_username, - "created_at": datetime.fromtimestamp(u_created_at, tz=UTC), + "created_at": created_at, + "token_version": "legacy", } return @@ -132,9 +194,12 @@ async def get_subscription_payload(token: str) -> dict | None: username = payload.get("sub") if not username: return + created_at = _datetime_from_epoch_seconds(payload.get("iat")) + if created_at is None: + return return { "username": username, - "created_at": datetime.fromtimestamp(payload["iat"], tz=UTC), + "created_at": created_at, } else: return @@ -176,5 +241,5 @@ async def get_subscription_payload(token: str) -> dict | None: if u_signature in (u_token_resign, u_token_hex_resign): return _parse_subscription_data(u_token_dec_str) return - except jwt.exceptions.PyJWTError: + except jwt.exceptions.PyJWTError, OverflowError, OSError, TypeError, ValueError: return diff --git a/tests/api/test_user.py b/tests/api/test_user.py index cb7532d76..6194f3380 100644 --- a/tests/api/test_user.py +++ b/tests/api/test_user.py @@ -184,6 +184,25 @@ async def fake_get_secret_key(): assert old_v2_payload["user_id"] == 456 +def test_subscription_token_uses_precise_non_future_issuance_time(monkeypatch): + secret = "test-secret" + issued_at_ns = 1_723_000_000_123_456_789 + + async def fake_get_secret_key(): + return secret + + monkeypatch.setattr(jwt_utils, "get_secret_key", fake_get_secret_key) + monkeypatch.setattr(jwt_utils.time, "time_ns", lambda: issued_at_ns) + + token = asyncio.run(create_subscription_token(123)) + payload = asyncio.run(get_subscription_payload(token)) + + assert payload["user_id"] == 123 + assert payload["token_version"] == "v5" + assert payload["created_at"] == datetime(2024, 8, 7, 3, 6, 40, 123456, tzinfo=UTC) + assert payload["subject_created_at"] == payload["created_at"] + + def test_user_create_active(access_token): """Test that the user create active route is accessible.""" core, groups = setup_groups(access_token, 2) @@ -1963,11 +1982,16 @@ def test_revoke_user_subscription(access_token): payload={"username": unique_name("test_user_revoke")}, ) try: + old_subscription_url = user["subscription_url"] response = client.post( f"/api/user/{user['username']}/revoke_sub", headers={"Authorization": f"Bearer {access_token}"}, ) assert response.status_code == status.HTTP_200_OK + # This request intentionally follows immediately: MySQL/MariaDB used + # to truncate sub_revoked_at to whole seconds, leaving a v5 token + # issued earlier in the same second valid. + assert client.get(old_subscription_url).status_code == status.HTTP_404_NOT_FOUND finally: delete_user(access_token, user["username"]) cleanup_groups(access_token, core, groups) @@ -2512,6 +2536,8 @@ def test_bulk_create_users_from_template_sequence(access_token): assert response.status_code == status.HTTP_201_CREATED assert response.json()["created"] == count assert len(response.json()["subscription_urls"]) == count + for subscription_url in response.json()["subscription_urls"]: + assert client.get(subscription_url).status_code == status.HTTP_200_OK expected_usernames = [f"{base_username}{start_number + idx}" for idx in range(count)] @@ -2558,6 +2584,8 @@ def test_bulk_create_users_from_template_sequence_with_template_affixes(access_t assert response.status_code == status.HTTP_201_CREATED assert response.json()["created"] == count assert len(response.json()["subscription_urls"]) == count + for subscription_url in response.json()["subscription_urls"]: + assert client.get(subscription_url).status_code == status.HTTP_200_OK expected_usernames = [f"{prefix}{base_username}{suffix}{start_number + idx}" for idx in range(count)] @@ -2593,6 +2621,8 @@ def test_bulk_create_users_from_template_random(access_token): assert response.status_code == status.HTTP_201_CREATED assert response.json()["created"] == count assert len(response.json()["subscription_urls"]) == count + for subscription_url in response.json()["subscription_urls"]: + assert client.get(subscription_url).status_code == status.HTTP_200_OK users_response = client.get( "/api/users", diff --git a/tests/test_subscription_revocation_precision.py b/tests/test_subscription_revocation_precision.py new file mode 100644 index 000000000..0f013a072 --- /dev/null +++ b/tests/test_subscription_revocation_precision.py @@ -0,0 +1,65 @@ +from datetime import UTC, datetime +from importlib import import_module + +import pytest +from sqlalchemy import select +from sqlalchemy.dialects import mysql +from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine + +from app.db.models import User + +migration = import_module( + "app.db.migrations.versions.9e0d7a1c4b52_preserve_subscription_revocation_microseconds" +) + + +@pytest.mark.parametrize("column_name", ["created_at", "sub_revoked_at"]) +def test_mysql_subscription_token_columns_use_microsecond_precision(column_name): + column_type = User.__table__.c[column_name].type.dialect_impl(mysql.dialect()) + + assert column_type.fsp == 6 + + +def test_mysql_migration_upgrades_both_subscription_token_timestamps(monkeypatch): + altered_columns: list[tuple[str, str, int | None]] = [] + + def capture_alter_column(table_name, column_name, **kwargs): + altered_columns.append((table_name, column_name, kwargs["type_"].fsp)) + + monkeypatch.setattr(migration, "_is_mysql_family", lambda: True) + monkeypatch.setattr(migration.op, "alter_column", capture_alter_column) + + migration.upgrade() + + assert altered_columns == [ + ("users", "created_at", 6), + ("users", "sub_revoked_at", 6), + ] + + +@pytest.mark.asyncio +async def test_subscription_revocation_microseconds_survive_database_round_trip(): + engine = create_async_engine("sqlite+aiosqlite:///:memory:") + session_factory = async_sessionmaker(engine, expire_on_commit=False) + timestamp = datetime(2026, 8, 9, 12, 0, 0, 123456, tzinfo=UTC) + + async with engine.begin() as connection: + await connection.run_sync(lambda sync: User.__table__.create(sync, checkfirst=True)) + + async with session_factory() as session: + await session.execute( + User.__table__.insert().values( + sync_id="revocation-precision-sync", + username="revocation-precision", + status="active", + proxy_settings={}, + used_traffic=0, + created_at=timestamp, + sub_revoked_at=timestamp, + ) + ) + await session.commit() + stored = (await session.execute(select(User.sub_revoked_at))).scalar_one() + + assert stored.microsecond == 123456 + await engine.dispose() diff --git a/tests/test_subscription_token_revocation.py b/tests/test_subscription_token_revocation.py new file mode 100644 index 000000000..fd66d9a13 --- /dev/null +++ b/tests/test_subscription_token_revocation.py @@ -0,0 +1,135 @@ +from datetime import UTC, datetime, timedelta +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest +from fastapi import HTTPException + +import app.operation as operation_module +from app.operation import BaseOperation, OperatorType +from app.utils import jwt as jwt_utils + + +async def _assert_token_is_rejected(monkeypatch, payload: dict, user: SimpleNamespace): + monkeypatch.setattr(operation_module, "get_subscription_payload", AsyncMock(return_value=payload)) + monkeypatch.setattr(operation_module, "get_user_by_id", AsyncMock(return_value=user)) + + with pytest.raises(HTTPException) as exc_info: + await BaseOperation(OperatorType.API).get_validated_sub(db=None, token="subscription-token") + + assert exc_info.value.status_code == 404 + + +@pytest.mark.asyncio +async def test_legacy_v3_token_keeps_existing_strict_after_semantics(monkeypatch): + # A v3 timestamp was rounded to seconds, therefore an earlier sub-second + # database value cannot be distinguished from the original issuance. + encoded_timestamp = datetime(2026, 1, 1, 12, 0, 1, tzinfo=UTC) + user = SimpleNamespace( + created_at=encoded_timestamp - timedelta(microseconds=800_000), + sub_revoked_at=encoded_timestamp - timedelta(microseconds=200_000), + ) + monkeypatch.setattr( + operation_module, + "get_subscription_payload", + AsyncMock(return_value={"user_id": 1, "created_at": encoded_timestamp, "token_version": "v3"}), + ) + monkeypatch.setattr(operation_module, "get_user_by_id", AsyncMock(return_value=user)) + + assert await BaseOperation(OperatorType.API).get_validated_sub(db=None, token="subscription-token") is user + + +@pytest.mark.asyncio +async def test_v4_token_is_revoked_on_the_same_database_microsecond(monkeypatch): + issued_at = datetime(2026, 1, 1, 12, 0, 0, 123456, tzinfo=UTC) + await _assert_token_is_rejected( + monkeypatch, + {"user_id": 1, "created_at": issued_at, "token_version": "v4"}, + SimpleNamespace(created_at=issued_at - timedelta(microseconds=1), sub_revoked_at=issued_at), + ) + + +@pytest.mark.asyncio +async def test_legacy_token_is_revoked_on_a_timestamp_tie(monkeypatch): + issued_at = datetime(2026, 1, 1, 12, 0, 0, tzinfo=UTC) + await _assert_token_is_rejected( + monkeypatch, + {"user_id": 1, "created_at": issued_at}, + SimpleNamespace(created_at=issued_at - timedelta(microseconds=1), sub_revoked_at=issued_at), + ) + + +@pytest.mark.asyncio +async def test_v5_token_rejects_a_recreated_user_even_one_microsecond_later(monkeypatch): + issued_at = datetime(2026, 1, 1, 12, 0, 0, 123456, tzinfo=UTC) + await _assert_token_is_rejected( + monkeypatch, + { + "user_id": 1, + "created_at": issued_at, + "subject_created_at": issued_at, + "token_version": "v5", + }, + SimpleNamespace(created_at=issued_at + timedelta(microseconds=1), sub_revoked_at=None), + ) + + +@pytest.mark.asyncio +async def test_v5_token_allows_the_original_user_at_issuance(monkeypatch): + issued_at = datetime(2026, 1, 1, 12, 0, 0, 123456, tzinfo=UTC) + user = SimpleNamespace(created_at=issued_at, sub_revoked_at=None) + monkeypatch.setattr( + operation_module, + "get_subscription_payload", + AsyncMock( + return_value={ + "user_id": 1, + "created_at": issued_at, + "subject_created_at": issued_at, + "token_version": "v5", + } + ), + ) + monkeypatch.setattr(operation_module, "get_user_by_id", AsyncMock(return_value=user)) + + assert await BaseOperation(OperatorType.API).get_validated_sub(db=None, token="subscription-token") is user + + +@pytest.mark.asyncio +async def test_v5_token_treats_naive_database_timestamp_as_utc(monkeypatch): + issued_at = datetime(2026, 1, 1, 12, 0, 0, 123456, tzinfo=UTC) + user = SimpleNamespace(created_at=issued_at.replace(tzinfo=None), sub_revoked_at=None) + monkeypatch.setattr( + operation_module, + "get_subscription_payload", + AsyncMock( + return_value={ + "user_id": 1, + "created_at": issued_at, + "subject_created_at": issued_at, + "token_version": "v5", + } + ), + ) + monkeypatch.setattr(operation_module, "get_user_by_id", AsyncMock(return_value=user)) + + assert await BaseOperation(OperatorType.API).get_validated_sub(db=None, token="subscription-token") is user + + +def test_naive_subscription_timestamp_is_encoded_as_utc(): + aware = datetime(2026, 1, 1, 12, 0, 0, 123456, tzinfo=UTC) + naive = aware.replace(tzinfo=None) + + assert jwt_utils._datetime_to_epoch_nanoseconds(naive) == jwt_utils._datetime_to_epoch_nanoseconds(aware) + + +@pytest.mark.parametrize( + "payload", + [ + "v4,1,999999999999999999999999999999999999999999999999", + "v5,1,999999999999999999999999999999999999999999999999,0", + "subscriber,999999999999999999999999999999999999999999999999", + ], +) +def test_out_of_range_subscription_timestamps_fail_closed(payload): + assert jwt_utils._parse_subscription_data(payload) is None