diff --git a/api/extensions/application_services/account.py b/api/extensions/application_services/account.py index 5cf7ba6573c..7f74720fc17 100644 --- a/api/extensions/application_services/account.py +++ b/api/extensions/application_services/account.py @@ -196,22 +196,23 @@ def build_account_services( sessions=RedisAccountSessionGateway(redis=redis), now=naive_utc_now, ) + activation = AccountActivationService( + tokens=invitation_tokens, + accounts=activation_accounts, + workspace_policy=DeploymentWorkspaceInvitePolicy(), + eligibility=BillingAccountActivationEligibility( + enabled=deployment_edition == DeploymentEdition.CLOUD, + ), + membership_cache=BillingWorkspaceMembershipCache( + enabled=deployment_edition == DeploymentEdition.CLOUD, + ), + member_access_sync=RBACWorkspaceMemberAccessSync( + enabled=dify_config.RBAC_ENABLED, + ), + ) return AccountServices( lifecycle=lifecycle, - activation=AccountActivationService( - tokens=invitation_tokens, - accounts=activation_accounts, - workspace_policy=DeploymentWorkspaceInvitePolicy(), - eligibility=BillingAccountActivationEligibility( - enabled=deployment_edition == DeploymentEdition.CLOUD, - ), - membership_cache=BillingWorkspaceMembershipCache( - enabled=deployment_edition == DeploymentEdition.CLOUD, - ), - member_access_sync=RBACWorkspaceMemberAccessSync( - enabled=dify_config.RBAC_ENABLED, - ), - ), + activation=activation, identity=AccountIdentityGateway(accounts=accounts, redis=redis), access=AccountAccessService( accounts=accounts, @@ -226,6 +227,7 @@ def build_account_services( invitations=AccountActivationConsoleAuthInvitationGateway( tokens=invitation_tokens, accounts=activation_accounts, + activation=activation, ), policies=DeploymentConsoleAuthPolicyGateway( billing_enabled=deployment_edition == DeploymentEdition.CLOUD, diff --git a/api/extensions/ext_redis.py b/api/extensions/ext_redis.py index 870d265da73..8ff5db1643d 100644 --- a/api/extensions/ext_redis.py +++ b/api/extensions/ext_redis.py @@ -149,6 +149,9 @@ class RedisClientWrapper: def ttl(self, name: str | bytes) -> Any: return self._require_client().ttl(_serialize_redis_name_arg(name, self._get_prefix())) + def pttl(self, name: str | bytes) -> Any: + return self._require_client().pttl(_serialize_redis_name_arg(name, self._get_prefix())) + def getdel(self, name: str | bytes) -> Any: return self._require_client().getdel(_serialize_redis_name_arg(name, self._get_prefix())) diff --git a/api/repositories/account_activation_repository.py b/api/repositories/account_activation_repository.py index b8b30689335..0f00bc8c6dc 100644 --- a/api/repositories/account_activation_repository.py +++ b/api/repositories/account_activation_repository.py @@ -58,6 +58,7 @@ class SQLAlchemyAccountActivationRepository(AccountActivationRepository): *, role: str, setup: AccountSetup | None, + create_membership: bool = True, ) -> ActivationPersistenceResult | None: with self._session_factory.begin() as session: tenant_id = session.scalar( @@ -83,6 +84,8 @@ class SQLAlchemyAccountActivationRepository(AccountActivationRepository): TenantAccountJoin.account_id == account.id, ) ) + if membership is None and not create_membership: + return None membership_created = membership is None if membership is None: membership = TenantAccountJoin( diff --git a/api/services/account/adapters.py b/api/services/account/adapters.py index b72e244c00d..3cf720dfb94 100644 --- a/api/services/account/adapters.py +++ b/api/services/account/adapters.py @@ -185,6 +185,18 @@ class RedisInvitationTokenStore(InvitationTokenStore): return None return _invitation_token_adapter.validate_json(data) + @override + def consume_membership(self, invitation: InvitationLookup) -> bool: + key = self._invitation_token_key(invitation.token) + ttl = self._redis.pttl(key) + if ttl <= 0: + return False + return bool(self._redis.set(self._membership_used_key(invitation.token), "1", nx=True, px=ttl)) + + @override + def is_membership_consumed(self, invitation: InvitationLookup) -> bool: + return self._redis.get(self._membership_used_key(invitation.token)) == b"1" + @override def revoke(self, invitation: InvitationLookup) -> None: if invitation.workspace_id is not None and invitation.email is not None: @@ -196,6 +208,10 @@ class RedisInvitationTokenStore(InvitationTokenStore): def _invitation_token_key(token: str) -> str: return f"member_invite:token:{token}" + @classmethod + def _membership_used_key(cls, token: str) -> str: + return f"{cls._invitation_token_key(token)}:membership_used" + @staticmethod def _workspace_invitation_key(invitation: InvitationLookup) -> str: assert invitation.workspace_id is not None diff --git a/api/services/account/login_adapters.py b/api/services/account/login_adapters.py index 1d443dfa12b..9e64a10fb4b 100644 --- a/api/services/account/login_adapters.py +++ b/api/services/account/login_adapters.py @@ -27,7 +27,11 @@ from services.account.login_service import ( HumanVerificationGateway, ResetPasswordEmailGateway, ) -from services.account_activation_service import AccountActivationRepository, InvitationTokenStore +from services.account_activation_service import ( + AccountActivationRepository, + AccountActivationService, + InvitationTokenStore, +) from services.account_security_gateway import RedisAccountEmailSecurityGateway from services.billing_service import BillingService from services.email_code_login_challenge import EmailCodeLoginChallengeStore, EmailCodeLoginChallengeUnavailableError @@ -62,11 +66,18 @@ _RESET_PASSWORD_RATE_LIMIT_SECONDS = 60 class AccountActivationConsoleAuthInvitationGateway(ConsoleAuthInvitationGateway): - """Resolve login invitations through the shared activation contracts.""" + """Resolve login invitations and restore invited membership through activation.""" - def __init__(self, *, tokens: InvitationTokenStore, accounts: AccountActivationRepository) -> None: + def __init__( + self, + *, + tokens: InvitationTokenStore, + accounts: AccountActivationRepository, + activation: AccountActivationService, + ) -> None: self._tokens = tokens self._accounts = accounts + self._activation = activation @override def resolve(self, *, email: str, token: str) -> LoginInvitation | None: @@ -82,6 +93,10 @@ class AccountActivationConsoleAuthInvitationGateway(ConsoleAuthInvitationGateway return LoginInvitation(email=invitation.account_email) return None + @override + def ensure_membership(self, *, account_id: str, email: str, token: str) -> bool: + return self._activation.ensure_membership(account_id=account_id, email=email, token=token) + class DeploymentConsoleAuthPolicyGateway(ConsoleAuthPolicyGateway): def __init__(self, *, billing_enabled: bool) -> None: diff --git a/api/services/account/login_service.py b/api/services/account/login_service.py index 3e97df2f2ae..fec375b0eea 100644 --- a/api/services/account/login_service.py +++ b/api/services/account/login_service.py @@ -41,6 +41,8 @@ class ConsoleAuthWorkspaceQuery(Protocol): class ConsoleAuthInvitationGateway(Protocol): def resolve(self, *, email: str, token: str) -> LoginInvitation | None: ... + def ensure_membership(self, *, account_id: str, email: str, token: str) -> bool: ... + class ConsoleAuthPolicyGateway(Protocol): def get_email_freeze_type(self, email: str) -> str | None: ... @@ -200,9 +202,12 @@ class ConsoleAuthenticationService: self._invalid_credentials(normalized_email, command.ip_address) if not self._workspaces.has_active_for_account(account.id): - if self._policies.is_workspace_creation_allowed() and not self._policies.has_workspace_capacity(): - raise account_errors.LoginWorkspaceLimitError - return PasswordLoginResult(tokens=None, workspace_found=False) + if invitation_token is not None: + self._invitations.ensure_membership(account_id=account.id, email=command.email, token=invitation_token) + if not self._workspaces.has_active_for_account(account.id): + if self._policies.is_workspace_creation_allowed() and not self._policies.has_workspace_capacity(): + raise account_errors.LoginWorkspaceLimitError + return PasswordLoginResult(tokens=None, workspace_found=False) tokens = self._issue_session( account_id=account.id, diff --git a/api/services/account_activation_service.py b/api/services/account_activation_service.py index ebc1298dd48..a85a4f975d5 100644 --- a/api/services/account_activation_service.py +++ b/api/services/account_activation_service.py @@ -25,6 +25,10 @@ _PENDING_ACCOUNT_STATUS = "pending" class InvitationTokenStore(Protocol): + def consume_membership(self, invitation: InvitationLookup) -> bool: ... + + def is_membership_consumed(self, invitation: InvitationLookup) -> bool: ... + def find(self, invitation: InvitationLookup) -> InvitationToken | None: ... def revoke(self, invitation: InvitationLookup) -> None: ... @@ -39,6 +43,7 @@ class AccountActivationRepository(Protocol): *, role: str, setup: AccountSetup | None, + create_membership: bool = True, ) -> ActivationPersistenceResult | None: ... @@ -108,8 +113,9 @@ class AccountActivationService: raise FrozenAccountError setup = self._resolve_setup(invitation, command) - raw_role = invitation.role - role = raw_role if raw_role is not None and raw_role in _NON_OWNER_ROLES else _DEFAULT_ROLE + role = self._membership_role(invitation.role) + # An invitation that already restored membership at sign-in must not grant it again after removal. + membership_consumed = self._tokens.is_membership_consumed(command.invitation) normalized_email = command.invitation.email.lower() if command.invitation.email else None self._tokens.revoke( @@ -119,13 +125,29 @@ class AccountActivationService: token=command.invitation.token, ) ) - result = self._accounts.activate(invitation, role=role, setup=setup) + result = self._accounts.activate(invitation, role=role, setup=setup, create_membership=not membership_consumed) if result is None: raise InvalidInvitationError if result.membership_created: self._membership_cache.invalidate(invitation.workspace_id) self._member_access_sync.sync(invitation.workspace_id, invitation.account_id) + def ensure_membership(self, *, account_id: str, email: str, token: str) -> bool: + lookup = InvitationLookup(workspace_id=None, email=email, token=token) + invitation = self._resolve(lookup) + if invitation is None or invitation.account_id != account_id: + return False + if not self._tokens.consume_membership(lookup): + return False + + result = self._accounts.activate(invitation, role=self._membership_role(invitation.role), setup=None) + if result is None: + return False + if result.membership_created: + self._membership_cache.invalidate(invitation.workspace_id) + self._member_access_sync.sync(invitation.workspace_id, invitation.account_id) + return True + def _resolve(self, invitation: InvitationLookup) -> AccountInvitation | None: token = self._tokens.find(invitation) resolved = self._accounts.resolve(token) if token is not None else None @@ -146,6 +168,12 @@ class AccountActivationService: return None return self._accounts.resolve(token) + @staticmethod + def _membership_role(role: str | None) -> str: + if role is not None and role in _NON_OWNER_ROLES: + return role + return _DEFAULT_ROLE + @staticmethod def _requires_setup(invitation: AccountInvitation) -> bool: if invitation.requires_setup is not None: diff --git a/api/tests/unit_tests/repositories/test_account_activation_repository.py b/api/tests/unit_tests/repositories/test_account_activation_repository.py index 8c24093661c..0835cc245cf 100644 --- a/api/tests/unit_tests/repositories/test_account_activation_repository.py +++ b/api/tests/unit_tests/repositories/test_account_activation_repository.py @@ -176,3 +176,52 @@ class TestPersistActivation: assert membership is not None assert membership.role == TenantAccountRole.EDITOR assert membership.current is True + + def test_does_not_create_membership_when_creation_is_disallowed( + self, + sqlite_session: Session, + sqlite_session_factory: sessionmaker[Session], + ) -> None: + account, _ = _persist_invitation_state(sqlite_session) + repository = SQLAlchemyAccountActivationRepository(sqlite_session_factory) + + result = repository.activate( + _invitation(), + role="admin", + setup=AccountSetup(name="Name", interface_language="en-US", timezone="UTC"), + create_membership=False, + ) + + assert result is None + sqlite_session.expire_all() + assert sqlite_session.scalar(select(TenantAccountJoin)) is None + reloaded = sqlite_session.get(Account, account.id) + assert reloaded is not None + assert reloaded.status == AccountStatus.PENDING + + def test_keeps_existing_membership_when_creation_is_disallowed( + self, + sqlite_session: Session, + sqlite_session_factory: sessionmaker[Session], + ) -> None: + account, tenant = _persist_invitation_state(sqlite_session) + sqlite_session.add( + TenantAccountJoin( + tenant_id=tenant.id, + account_id=account.id, + role=TenantAccountRole.EDITOR, + ) + ) + account.status = AccountStatus.ACTIVE + sqlite_session.commit() + repository = SQLAlchemyAccountActivationRepository(sqlite_session_factory) + + result = repository.activate(_invitation(), role="admin", setup=None, create_membership=False) + + assert result is not None + assert result.membership_created is False + sqlite_session.expire_all() + membership = sqlite_session.scalar(select(TenantAccountJoin)) + assert membership is not None + assert membership.role == TenantAccountRole.EDITOR + assert membership.current is True diff --git a/api/tests/unit_tests/services/account/test_adapters.py b/api/tests/unit_tests/services/account/test_adapters.py index b92a2106da6..0e7be136fa8 100644 --- a/api/tests/unit_tests/services/account/test_adapters.py +++ b/api/tests/unit_tests/services/account/test_adapters.py @@ -84,6 +84,53 @@ def test_invitation_token_store_revokes_its_redis_key(redis_transport: tuple[Red assert commands.call_count == 1 +def test_invitation_token_store_marks_membership_use_with_token_ttl( + redis_transport: tuple[RedisClientWrapper, MagicMock], +) -> None: + redis, commands = redis_transport + commands.side_effect = [1200, b"OK"] + lookup = InvitationLookup(workspace_id=None, email="invitee@example.com", token="token-1") + + assert RedisInvitationTokenStore(redis=redis).consume_membership(lookup) is True + + assert commands.call_args_list[0].args == ("PTTL", "member_invite:token:token-1") + set_args = commands.call_args_list[1].args + assert set_args[:3] == ("SET", "member_invite:token:token-1:membership_used", "1") + assert "NX" in set_args + assert set_args[set_args.index("PX") + 1] == 1200 + + +@pytest.mark.parametrize("ttl", [-2, -1, 0]) +def test_invitation_token_store_refuses_membership_use_without_live_token( + redis_transport: tuple[RedisClientWrapper, MagicMock], ttl: int +) -> None: + redis, commands = redis_transport + commands.return_value = ttl + lookup = InvitationLookup(workspace_id=None, email="invitee@example.com", token="token-1") + + assert RedisInvitationTokenStore(redis=redis).consume_membership(lookup) is False + + assert commands.call_count == 1 + assert commands.call_args.args == ("PTTL", "member_invite:token:token-1") + + +@pytest.mark.parametrize(("stored", "expected"), [(b"1", True), (None, False)]) +def test_invitation_token_store_reports_membership_use( + redis_transport: tuple[RedisClientWrapper, MagicMock], stored: bytes | None, expected: bool +) -> None: + redis, commands = redis_transport + commands.return_value = stored + lookup = InvitationLookup(workspace_id="workspace-1", email="invitee@example.com", token="token-1") + + assert RedisInvitationTokenStore(redis=redis).is_membership_consumed(lookup) is expected + + commands.assert_called_once_with( + "GET", + "member_invite:token:token-1:membership_used", + keys=["member_invite:token:token-1:membership_used"], + ) + + def test_billing_eligibility_skips_gateway_when_disabled() -> None: with patch("services.account.adapters.BillingService.get_email_freeze_type") as get_freeze_type: result = BillingAccountActivationEligibility(enabled=False).get_freeze_type("invitee@example.com") diff --git a/api/tests/unit_tests/services/account/test_login_adapters.py b/api/tests/unit_tests/services/account/test_login_adapters.py index 170528ab5dd..eaac915d367 100644 --- a/api/tests/unit_tests/services/account/test_login_adapters.py +++ b/api/tests/unit_tests/services/account/test_login_adapters.py @@ -16,17 +16,33 @@ from extensions.ext_redis import RedisClientWrapper from libs.helper import RateLimiter from models.account import Account, AccountStatus, Tenant, TenantAccountJoin, TenantAccountRole from repositories.account.repository import SQLAlchemyAccountRepository +from repositories.account_activation_repository import SQLAlchemyAccountActivationRepository +from repositories.workspace.workspace_repository import WorkspaceRepository from services import account_errors from services.account import login_adapters as adapters +from services.account.adapters import ( + BillingAccountActivationEligibility, + BillingWorkspaceMembershipCache, + DeploymentWorkspaceInvitePolicy, + RBACWorkspaceMemberAccessSync, + RedisInvitationTokenStore, +) +from services.account_activation_service import AccountActivationService from services.email_code_login_challenge import ( EmailCodeLoginChallengeResult, EmailCodeLoginChallengeStatus, EmailCodeLoginChallengeUnavailableError, ) +from services.entities.account_activation_entities import InvitationToken from services.entities.account_entities import AccountSessionTokens -from services.entities.account_login_entities import EmailCodeChallengeStatus, RefreshAccountStatus +from services.entities.account_login_entities import ( + EmailCodeChallengeStatus, + PasswordLoginCommand, + RefreshAccountStatus, +) from services.turnstile_service import TurnstileChallengeRejectedError, TurnstileUpstreamError from tests.unit_tests.account_domain import AccountDomain +from tests.unit_tests.services.account.test_login_service import Dependencies class FakeRateLimiter(RateLimiter): @@ -67,6 +83,134 @@ def _persist_account(session: Session) -> Account: return account +def test_invitation_gateway_restores_editor_membership_without_revoking_token( + sqlite_session_factory: sessionmaker[Session], + redis_transport: tuple[RedisClientWrapper, MagicMock], +) -> None: + redis, commands = redis_transport + commands.return_value = ( + b'{"account_id":"account-1","email":"user@example.com","workspace_id":"workspace-1",' + b'"role":"editor","requires_setup":false}' + ) + with sqlite_session_factory.begin() as session: + account = Account(name="User", email="user@example.com", status=AccountStatus.ACTIVE) + account.id = "account-1" + tenant = Tenant(name="Workspace") + tenant.id = "workspace-1" + session.add_all([account, tenant]) + tokens = RedisInvitationTokenStore(redis=redis) + accounts = SQLAlchemyAccountActivationRepository(sqlite_session_factory) + activation = AccountActivationService( + tokens=tokens, + accounts=accounts, + workspace_policy=DeploymentWorkspaceInvitePolicy(), + eligibility=BillingAccountActivationEligibility(enabled=False), + membership_cache=BillingWorkspaceMembershipCache(enabled=False), + member_access_sync=RBACWorkspaceMemberAccessSync(enabled=False), + ) + gateway = adapters.AccountActivationConsoleAuthInvitationGateway( + tokens=tokens, + accounts=accounts, + activation=activation, + ) + + commands.side_effect = lambda command, *_args, **_kwargs: ( + 3600 if command == "PTTL" else b"OK" if command == "SET" else commands.return_value + ) + assert gateway.ensure_membership(account_id="different-account", email="user@example.com", token="token-1") is False + with sqlite_session_factory() as session: + assert session.scalar(select(TenantAccountJoin)) is None + assert gateway.ensure_membership(account_id="account-1", email="user@example.com", token="token-1") is True + + with sqlite_session_factory() as session: + persisted = session.get(Account, "account-1") + membership = session.scalar(select(TenantAccountJoin)) + assert persisted is not None + assert persisted.id == "account-1" + assert persisted.status == AccountStatus.ACTIVE + assert persisted.name == "User" + assert membership is not None + assert membership.role == TenantAccountRole.EDITOR + assert membership.current is True + assert not any(command.args and command.args[0] == "DEL" for command in commands.call_args_list) + + +def test_login_restores_membership_when_invitation_is_stored_under_prefixed_key( + sqlite_session_factory: sessionmaker[Session], + redis_transport: tuple[RedisClientWrapper, MagicMock], + config_overrides: Callable[..., None], +) -> None: + redis, commands = redis_transport + config_overrides(REDIS_KEY_PREFIX="enterprise-a") + stored: dict[str, object] = {} + + def execute(command: object, *args: object, **_kwargs: object) -> object: + operation, key = str(command).upper(), str(args[0]) + if operation == "SETEX": + stored[key] = args[2] + return True + if operation == "GET": + return stored.get(key) + if operation == "PTTL": + return 3600 if key in stored else -2 + if operation == "SET" and "NX" in args and key not in stored: + stored[key] = b"1" + return True + return None + + commands.side_effect = execute + with sqlite_session_factory.begin() as session: + account = Account( + name="User", + email="user@example.com", + status=AccountStatus.ACTIVE, + password="hash", + password_salt="salt", + ) + account.id = "account-1" + workspace = Tenant(name="Workspace") + workspace.id = "workspace-1" + session.add_all([account, workspace]) + tokens = RedisInvitationTokenStore(redis=redis) + token = tokens.create( + InvitationToken( + account_id="account-1", + email="user@example.com", + workspace_id="workspace-1", + role="editor", + requires_setup=False, + ) + ) + activation_accounts = SQLAlchemyAccountActivationRepository(sqlite_session_factory) + gateway = adapters.AccountActivationConsoleAuthInvitationGateway( + tokens=tokens, + accounts=activation_accounts, + activation=AccountActivationService( + tokens=tokens, + accounts=activation_accounts, + workspace_policy=DeploymentWorkspaceInvitePolicy(), + eligibility=BillingAccountActivationEligibility(enabled=False), + membership_cache=BillingWorkspaceMembershipCache(enabled=False), + member_access_sync=RBACWorkspaceMemberAccessSync(enabled=False), + ), + ) + dependencies = Dependencies() + dependencies.accounts = SQLAlchemyAccountRepository(sqlite_session_factory) # type: ignore[assignment] + dependencies.workspaces = WorkspaceRepository(sqlite_session_factory) # type: ignore[assignment] + dependencies.invitations = gateway # type: ignore[assignment] + result = dependencies.service().login_with_password( + PasswordLoginCommand(email="user@example.com", password="password", invite_token=token, ip_address="127.0.0.1") + ) + + assert f"enterprise-a:member_invite:token:{token}" in stored + assert result.workspace_found is True + assert result.tokens == AccountSessionTokens(access_token="access", refresh_token="refresh", csrf_token="csrf") + with sqlite_session_factory() as session: + membership = session.scalar(select(TenantAccountJoin)) + assert membership is not None + assert membership.role == TenantAccountRole.EDITOR + + def test_security_gateway_owns_login_failure_state(redis_transport: tuple[RedisClientWrapper, MagicMock]) -> None: redis, commands = redis_transport commands.side_effect = [b"6", b"2", True, 1] diff --git a/api/tests/unit_tests/services/account/test_login_service.py b/api/tests/unit_tests/services/account/test_login_service.py index 3c847c67549..30c0e0d4ea9 100644 --- a/api/tests/unit_tests/services/account/test_login_service.py +++ b/api/tests/unit_tests/services/account/test_login_service.py @@ -49,20 +49,30 @@ class FakeAccounts: @dataclass class FakeWorkspaces: active: bool = False + checks: list[str] = field(default_factory=list) def has_active_for_account(self, account_id: str) -> bool: - _ = account_id + self.checks.append(account_id) return self.active @dataclass class FakeInvitations: invitation: LoginInvitation | None = None + ensured: list[tuple[str, str, str]] = field(default_factory=list) + restored_workspace: FakeWorkspaces | None = None def resolve(self, *, email: str, token: str) -> LoginInvitation | None: _ = email, token return self.invitation + def ensure_membership(self, *, account_id: str, email: str, token: str) -> bool: + self.ensured.append((account_id, email, token)) + if self.restored_workspace is None: + return False + self.restored_workspace.active = True + return True + @dataclass class FakePolicies: @@ -410,6 +420,7 @@ def test_invitation_login_rejects_empty_password_without_completing_account() -> assert dependencies.passwords.hashed_passwords == [] assert dependencies.accounts.completions == [] assert dependencies.security.failures == ["user@example.com"] + assert dependencies.invitations.ensured == [] def test_invitation_login_does_not_initialize_password_when_token_is_unknown() -> None: @@ -421,6 +432,7 @@ def test_invitation_login_does_not_initialize_password_when_token_is_unknown() - assert dependencies.passwords.hashed_passwords == [] assert dependencies.accounts.completions == [] + assert dependencies.invitations.ensured == [] def test_invitation_login_applies_new_password_policy_before_completion() -> None: @@ -436,6 +448,7 @@ def test_invitation_login_applies_new_password_policy_before_completion() -> Non assert dependencies.passwords.hashed_passwords == ["letters-only"] assert dependencies.accounts.completions == [] + assert dependencies.invitations.ensured == [] def test_invalid_password_records_rate_limit_and_audit() -> None: @@ -448,6 +461,20 @@ def test_invalid_password_records_rate_limit_and_audit() -> None: assert dependencies.security.failures == ["user@example.com"] assert dependencies.audit.failures == [("user@example.com", LoginFailureReason.INVALID_CREDENTIALS, "127.0.0.1")] + assert dependencies.invitations.ensured == [] + + +def test_wrong_password_does_not_restore_invited_membership() -> None: + dependencies = Dependencies() + dependencies.accounts.candidates = [_account()] + dependencies.invitations.invitation = LoginInvitation(email="user@example.com") + dependencies.passwords.valid = False + + with pytest.raises(account_errors.InvalidLoginCredentialsError): + dependencies.service().login_with_password(_password_command(invite_token="invite")) + + assert dependencies.invitations.ensured == [] + assert dependencies.security.failures == ["user@example.com"] def test_password_login_rejects_locked_account_before_loading_candidates() -> None: @@ -504,6 +531,7 @@ def test_password_login_rejects_invitation_for_another_email() -> None: assert dependencies.audit.failures == [ ("user@example.com", LoginFailureReason.INVALID_INVITATION_EMAIL, "127.0.0.1") ] + assert dependencies.invitations.ensured == [] def test_password_login_tries_lowercase_candidate_after_exact_candidate_password_fails() -> None: @@ -532,6 +560,40 @@ def test_password_login_returns_no_workspace_without_creating_one() -> None: assert not result.workspace_found assert result.tokens is None assert dependencies.workspaces_provisioning.account_ids == [] + assert dependencies.invitations.ensured == [] + + +def test_password_login_restores_invited_membership_when_account_has_no_workspace() -> None: + dependencies = Dependencies() + dependencies.accounts.candidates = [_account()] + dependencies.invitations.invitation = LoginInvitation(email="user@example.com") + dependencies.invitations.restored_workspace = dependencies.workspaces + dependencies.policies.workspace_capacity = False + + result = dependencies.service().login_with_password(_password_command(invite_token="invite")) + + assert dependencies.invitations.ensured == [("account-1", "User@Example.com", "invite")] + assert dependencies.workspaces.checks == ["account-1", "account-1"] + assert result.workspace_found + assert result.tokens == TOKENS + assert dependencies.sessions.issued == ["account-1"] + assert dependencies.security.resets == ["user@example.com"] + + +@pytest.mark.parametrize("invite_token", [None, "invite"]) +def test_password_login_does_not_restore_membership_for_an_existing_workspace(invite_token: str | None) -> None: + dependencies = Dependencies() + dependencies.accounts.candidates = [_account()] + dependencies.workspaces.active = True + if invite_token is not None: + dependencies.invitations.invitation = LoginInvitation(email="user@example.com") + + result = dependencies.service().login_with_password(_password_command(invite_token=invite_token)) + + assert result.workspace_found + assert result.tokens == TOKENS + assert dependencies.invitations.ensured == [] + assert dependencies.workspaces.checks == ["account-1"] def test_password_login_enforces_workspace_capacity() -> None: @@ -542,6 +604,22 @@ def test_password_login_enforces_workspace_capacity() -> None: with pytest.raises(account_errors.LoginWorkspaceLimitError): dependencies.service().login_with_password(_password_command()) + assert dependencies.invitations.ensured == [] + + +def test_password_login_enforces_workspace_capacity_when_invitation_does_not_restore_membership() -> None: + dependencies = Dependencies() + dependencies.accounts.candidates = [_account()] + dependencies.invitations.invitation = LoginInvitation(email="user@example.com") + dependencies.policies.workspace_capacity = False + + with pytest.raises(account_errors.LoginWorkspaceLimitError): + dependencies.service().login_with_password(_password_command(invite_token="invite")) + + assert dependencies.invitations.ensured == [("account-1", "User@Example.com", "invite")] + assert dependencies.workspaces.checks == ["account-1", "account-1"] + assert dependencies.sessions.issued == [] + def test_send_email_code_normalizes_identity_and_preserves_account_recipient_address() -> None: dependencies = Dependencies() diff --git a/api/tests/unit_tests/services/test_account_activation_service.py b/api/tests/unit_tests/services/test_account_activation_service.py index 232a2928dfe..e4a83d7be40 100644 --- a/api/tests/unit_tests/services/test_account_activation_service.py +++ b/api/tests/unit_tests/services/test_account_activation_service.py @@ -220,6 +220,7 @@ class TestActivateInvitation: _invitation(role=role, requires_setup=True if role == "owner" else None), role="normal", setup=AccountSetup(name="John Doe", interface_language="en-US", timezone="UTC"), + create_membership=True, ) boundaries["cache"].assert_called_once_with("workspace-1") boundaries["sync"].assert_called_once_with("workspace-1", "account-1", operator_account_id=None) @@ -264,7 +265,10 @@ class TestActivateInvitation: ) boundaries["activate"].assert_called_once_with( - _invitation(account_status="active", role="editor", requires_setup=False), role="editor", setup=None + _invitation(account_status="active", role="editor", requires_setup=False), + role="editor", + setup=None, + create_membership=True, ) boundaries["cache"].assert_not_called() boundaries["sync"].assert_called_once_with("workspace-1", "account-1", operator_account_id=None) @@ -276,3 +280,227 @@ class TestActivateInvitation: assert len(memberships) == 1 assert memberships[0].role == TenantAccountRole.EDITOR assert memberships[0].current is True + + +class TestEnsureMembership: + def test_missing_token_returns_false_without_revoking( + self, + service: AccountActivationService, + boundaries: dict[str, MagicMock], + redis_transport: tuple[RedisClientWrapper, MagicMock], + ) -> None: + redis_transport[1].return_value = None + + assert service.ensure_membership(account_id="account-1", email="invitee@example.com", token="missing") is False + + boundaries["revoke"].assert_not_called() + boundaries["activate"].assert_not_called() + boundaries["cache"].assert_not_called() + boundaries["sync"].assert_not_called() + + @pytest.mark.parametrize( + ("role", "stored_role"), + [ + ("editor", TenantAccountRole.EDITOR), + (None, TenantAccountRole.NORMAL), + ("owner", TenantAccountRole.NORMAL), + ], + ) + def test_restores_membership_once_and_preserves_activation( + self, + service: AccountActivationService, + boundaries: dict[str, MagicMock], + sqlite_session_factory: sessionmaker[Session], + redis_transport: tuple[RedisClientWrapper, MagicMock], + role: str | None, + stored_role: TenantAccountRole, + ) -> None: + with sqlite_session_factory.begin() as session: + account = session.get(Account, "account-1") + assert account is not None + account.status = AccountStatus.ACTIVE + role_json = f',"role":"{role}"' if role is not None else "" + redis_transport[1].return_value = ( + '{"account_id":"account-1","email":"invitee@example.com","workspace_id":"workspace-1"' + f'{role_json},"requires_setup":false}}' + ).encode() + + used: set[object] = set() + + def execute(command: object, key: object, *args: object, **_kwargs: object) -> object: + if command == "PTTL": + return 3600 + if command == "SET": + assert "NX" in args + assert "PX" in args + if key in used: + return None + used.add(key) + return b"OK" + if command == "GET" and str(key).endswith(":membership_used"): + return b"1" if key in used else None + return redis_transport[1].return_value + + redis_transport[1].side_effect = execute + restored = service.ensure_membership(account_id="account-1", email="Invitee@Example.com", token="token-1") + + assert restored is True + boundaries["revoke"].assert_not_called() + boundaries["activate"].assert_called_once_with( + _invitation(account_status="active", role=role, requires_setup=False), + role=stored_role, + setup=None, + ) + boundaries["cache"].assert_called_once_with("workspace-1") + boundaries["sync"].assert_called_once_with("workspace-1", "account-1", operator_account_id=None) + with sqlite_session_factory() as session: + account = session.get(Account, "account-1") + membership = session.scalar(select(TenantAccountJoin)) + assert account is not None + assert account.id == "account-1" + assert account.status == AccountStatus.ACTIVE + assert account.name == "Test User" + assert membership is not None + assert membership.role == stored_role + assert membership.current is True + check = service.check(InvitationLookup(workspace_id=None, email=None, token="token-1")) + assert check.is_valid is True + boundaries["revoke"].assert_not_called() + + with sqlite_session_factory.begin() as session: + session.execute(delete(TenantAccountJoin)) + boundaries["activate"].reset_mock() + assert service.ensure_membership(account_id="account-1", email="invitee@example.com", token="token-1") is False + boundaries["activate"].assert_not_called() + with sqlite_session_factory() as session: + assert session.scalar(select(TenantAccountJoin)) is None + + with pytest.raises(InvalidInvitationError): + service.activate( + ActivationCommand(invitation=InvitationLookup(workspace_id=None, email=None, token="token-1")), + authenticated_account_id="account-1", + ) + boundaries["revoke"].assert_called_once() + boundaries["activate"].assert_called_once_with( + _invitation(account_status="active", role=role, requires_setup=False), + role=stored_role, + setup=None, + create_membership=False, + ) + with sqlite_session_factory() as session: + assert session.scalar(select(TenantAccountJoin)) is None + + def test_activation_after_restore_keeps_existing_membership( + self, + service: AccountActivationService, + boundaries: dict[str, MagicMock], + sqlite_session_factory: sessionmaker[Session], + redis_transport: tuple[RedisClientWrapper, MagicMock], + ) -> None: + with sqlite_session_factory.begin() as session: + account = session.get(Account, "account-1") + assert account is not None + account.status = AccountStatus.ACTIVE + redis_transport[1].return_value = ( + b'{"account_id":"account-1","email":"invitee@example.com","workspace_id":"workspace-1",' + b'"role":"editor","requires_setup":false}' + ) + used: set[object] = set() + + def execute(command: object, key: object, *_args: object, **_kwargs: object) -> object: + if command == "PTTL": + return 3600 + if command == "SET": + used.add(key) + return b"OK" + if command == "GET" and str(key).endswith(":membership_used"): + return b"1" if key in used else None + return redis_transport[1].return_value + + redis_transport[1].side_effect = execute + assert service.ensure_membership(account_id="account-1", email="invitee@example.com", token="token-1") is True + boundaries["cache"].reset_mock() + boundaries["sync"].reset_mock() + + service.activate( + ActivationCommand(invitation=InvitationLookup(workspace_id=None, email=None, token="token-1")), + authenticated_account_id="account-1", + ) + + boundaries["revoke"].assert_called_once() + boundaries["cache"].assert_not_called() + boundaries["sync"].assert_called_once_with("workspace-1", "account-1", operator_account_id=None) + with sqlite_session_factory() as session: + memberships = session.scalars(select(TenantAccountJoin)).all() + assert len(memberships) == 1 + assert memberships[0].role == TenantAccountRole.EDITOR + assert memberships[0].current is True + + def test_existing_membership_is_not_rewritten_or_resynced( + self, + service: AccountActivationService, + boundaries: dict[str, MagicMock], + sqlite_session_factory: sessionmaker[Session], + redis_transport: tuple[RedisClientWrapper, MagicMock], + ) -> None: + with sqlite_session_factory.begin() as session: + session.add( + TenantAccountJoin(tenant_id="workspace-1", account_id="account-1", role=TenantAccountRole.ADMIN) + ) + redis_transport[1].side_effect = lambda command, *_args, **_kwargs: ( + 3600 + if command == "PTTL" + else b"OK" + if command == "SET" + else b'{"account_id":"account-1","email":"invitee@example.com","workspace_id":"workspace-1"}' + ) + + assert service.ensure_membership(account_id="account-1", email="invitee@example.com", token="token-1") is True + + boundaries["activate"].assert_called_once() + boundaries["cache"].assert_not_called() + boundaries["sync"].assert_not_called() + with sqlite_session_factory() as session: + memberships = session.scalars(select(TenantAccountJoin)).all() + assert len(memberships) == 1 + assert memberships[0].role == TenantAccountRole.ADMIN + + def test_returns_false_when_activation_finds_no_workspace( + self, + service: AccountActivationService, + boundaries: dict[str, MagicMock], + mocker: MockerFixture, + redis_transport: tuple[RedisClientWrapper, MagicMock], + ) -> None: + redis_transport[1].side_effect = lambda command, *_args, **_kwargs: ( + 3600 + if command == "PTTL" + else b"OK" + if command == "SET" + else b'{"account_id":"account-1","email":"invitee@example.com","workspace_id":"workspace-1"}' + ) + mocker.patch.object(service._accounts, "activate", return_value=None) + + assert service.ensure_membership(account_id="account-1", email="invitee@example.com", token="token-1") is False + + boundaries["cache"].assert_not_called() + boundaries["sync"].assert_not_called() + + def test_account_mismatch_does_not_write_membership_or_consume_token( + self, + service: AccountActivationService, + boundaries: dict[str, MagicMock], + redis_transport: tuple[RedisClientWrapper, MagicMock], + ) -> None: + redis_transport[ + 1 + ].return_value = b'{"account_id":"account-1","email":"invitee@example.com","workspace_id":"workspace-1"}' + assert ( + service.ensure_membership(account_id="different-account", email="invitee@example.com", token="token-1") + is False + ) + boundaries["activate"].assert_not_called() + boundaries["revoke"].assert_not_called() + boundaries["cache"].assert_not_called() + boundaries["sync"].assert_not_called() + assert all(args[0] == "GET" for args, _ in redis_transport[1].call_args_list)