bitwarden_organization_invite_link/

invite_link_user_client.rs

use std::sync::Arc;

use bitwarden_api_api::models::{
    AcceptOrganizationInviteLinkRequestModel, ConfirmOrganizationInviteLinkRequestModel,
    GetOrganizationInviteLinkStatusRequestModel, GetOrganizationInviteRequestModel,
    OrganizationInviteLinkValidateEmailDomainRequestModel,
};
use bitwarden_core::{
    FromClient, OrganizationId,
    client::ApiConfigurations,
    key_management::{KeySlotIds, PrivateKeySlotId, SymmetricKeySlotId},
    require,
};
use bitwarden_crypto::{KeyStore, PrimitiveEncryptable, PublicKey, SpkiPublicKeyBytes};
use bitwarden_encoding::B64;
use bitwarden_organization_crypto::{
    account_recovery::{AccountRecovery, AccountRecoveryMakeError},
    invite::{Invite, InviteSecret},
    organization_shared_key::OrganizationSharedKey,
    trust_proof::{CryptographicProof, TrustProof},
};

use crate::{AcceptInviteLinkError, InviteLinkError, OrganizationInviteLinkStatusView};

/// Client for organization invite link invitee (user) operations: checking link status, validating
/// email eligibility, and accepting or self-confirming an invite.
#[bitwarden_ffi::wasm_object]
#[derive(FromClient)]
pub struct InviteLinkUserClient {
    pub(crate) key_store: KeyStore<KeySlotIds>,
    pub(crate) api_configurations: Arc<ApiConfigurations>,
}

#[bitwarden_ffi::wasm_export]
impl InviteLinkUserClient {
    /// Retrieves the status of an invite link.
    /// Used to verify basic availability before attempting to accept.
    pub async fn get_status(
        &self,
        organization_id: OrganizationId,
        code: String,
    ) -> Result<OrganizationInviteLinkStatusView, InviteLinkError> {
        let code =
            uuid::Uuid::parse_str(&code).map_err(|_| InviteLinkError::ParseFailure("code"))?;

        let response = self
            .api_configurations
            .api_client
            .organization_invite_links_api()
            .get_status(Some(GetOrganizationInviteLinkStatusRequestModel {
                organization_id: organization_id.into(),
                code,
            }))
            .await?;

        OrganizationInviteLinkStatusView::try_from(response)
    }

    /// Returns whether the given email address is in the allowed domains for an invite link.
    pub async fn is_email_allowed(
        &self,
        organization_id: OrganizationId,
        code: String,
        email: String,
    ) -> Result<bool, InviteLinkError> {
        let code =
            uuid::Uuid::parse_str(&code).map_err(|_| InviteLinkError::ParseFailure("code"))?;

        let response = self
            .api_configurations
            .api_client
            .organization_invite_links_api()
            .validate_email_domain(Some(
                OrganizationInviteLinkValidateEmailDomainRequestModel {
                    organization_id: organization_id.into(),
                    code,
                    email,
                },
            ))
            .await?;

        Ok(require!(response.is_allowed))
    }

    /// Accepts an organization invite for the current user, optionally enrolling into account
    /// recovery (when `enroll_into_account_recovery` is set) and — when the invite supports
    /// confirmation — self-confirming.
    ///
    /// # Errors
    /// Server-reported failures are mapped onto typed variants: a `404` from any of the acceptance
    /// endpoints becomes [`AcceptInviteLinkError::LinkNotFound`], and a `400` validation
    /// problem becomes the variant matching its error code (or
    /// [`AcceptInviteLinkError::Unknown`], carrying the server's English description, for an
    /// unrecognized code). Responses in the legacy error format carry no code and remain
    /// [`AcceptInviteLinkError::Api`].
    pub async fn accept_and_optionally_confirm(
        &self,
        organization_id: OrganizationId,
        code: String,
        invite_secret: InviteSecret,
        default_collection_name: String,
        enroll_into_account_recovery: bool,
    ) -> Result<(), AcceptInviteLinkError> {
        let code = uuid::Uuid::parse_str(&code)
            .map_err(|_| AcceptInviteLinkError::ParseFailure("code"))?;

        // When enrolling into account recovery, fetch the organization's public key (which is the
        // account-recovery public key) from the server.
        let recovery_public_key = if enroll_into_account_recovery {
            let response = self
                .api_configurations
                .api_client
                .organizations_api()
                .get_public_key(&organization_id.to_string())
                .await?;
            Some(
                require!(response.public_key)
                    .parse::<B64>()
                    .map_err(|_| AcceptInviteLinkError::ParseFailure("public_key"))?,
            )
        } else {
            None
        };

        let invite_response = self
            .api_configurations
            .api_client
            .organization_users_api()
            .get_invite(Some(GetOrganizationInviteRequestModel {
                organization_id: organization_id.into(),
                code,
            }))
            .await
            .map_err(AcceptInviteLinkError::from_api_error)?;

        let invite: Invite = require!(invite_response.invite).parse()?;

        // Confine the (non-Send) key store context to a synchronous scope; it produces the owned
        // request payload consumed after the `.await`s below.
        let request = {
            let mut ctx = self.key_store.context();

            // Recover the invite key from the invite secret the invitee holds.
            let invite_key =
                invite.unseal_invite_key_with_invite_secret(&invite_secret, &mut ctx)?;

            // Enroll into account recovery when requested. The account-recovery public key must
            // match the organization public-key thumbprint bound into the invite: a substituted
            // recovery key would not match, so the user key cannot be captured by an
            // attacker-supplied key.
            let reset_password_key = recovery_public_key
                .as_ref()
                .map(
                    |recovery_public_key| -> Result<String, AcceptInviteLinkError> {
                        let recovery_public_key =
                            PublicKey::from_der(&SpkiPublicKeyBytes::from(recovery_public_key))?;
                        let bound_thumbprint =
                            invite.get_public_key_thumbprint(invite_key, &mut ctx)?;
                        let account_recovery = AccountRecovery::make(
                            SymmetricKeySlotId::User,
                            &recovery_public_key,
                            &TrustProof::ManuallyTrustedThumbprint(bound_thumbprint),
                            &ctx,
                        )
                        .map_err(|error| match error {
                            AccountRecoveryMakeError::UntrustedPublicKey => {
                                AcceptInviteLinkError::RecoveryKeyMismatch
                            }
                            error @ (AccountRecoveryMakeError::PrivateKeyNotInContext
                            | AccountRecoveryMakeError::EncapsulationFailed) => error.into(),
                        })?;
                        Ok(String::from(&account_recovery))
                    },
                )
                .transpose()?;

            if invite.supports_confirmation() {
                // Self-confirm: recover the organization key and encapsulate it to the user.
                let org_key = invite.unseal_organization_key(invite_key, &mut ctx)?;
                let user_public_key = ctx.get_public_key(PrivateKeySlotId::UserPrivateKey)?;
                let organization_shared_key = OrganizationSharedKey::make(
                    org_key,
                    &user_public_key,
                    &TrustProof::CryptographicallyTrusted(CryptographicProof::OwnPrivateKey(
                        PrivateKeySlotId::UserPrivateKey,
                    )),
                    &ctx,
                )?;
                let org_user_key = String::from(&organization_shared_key);
                let default_user_collection_name = default_collection_name
                    .encrypt(&mut ctx, org_key)?
                    .to_string();
                PendingPost::Confirm(ConfirmOrganizationInviteLinkRequestModel {
                    organization_id: organization_id.into(),
                    code,
                    org_user_key,
                    reset_password_key,
                    default_user_collection_name,
                })
            } else {
                PendingPost::Accept(AcceptOrganizationInviteLinkRequestModel {
                    organization_id: organization_id.into(),
                    code,
                    reset_password_key,
                })
            }
        };

        let organization_users_api = self.api_configurations.api_client.organization_users_api();
        match request {
            PendingPost::Confirm(model) => organization_users_api
                .confirm_invite_link(Some(model))
                .await
                .map_err(AcceptInviteLinkError::from_api_error)?,
            PendingPost::Accept(model) => organization_users_api
                .accept_invite_link(Some(model))
                .await
                .map_err(AcceptInviteLinkError::from_api_error)?,
        }

        Ok(())
    }
}

/// A prepared invite acceptance request, built while the key store context is held and posted once
/// it has been dropped.
enum PendingPost {
    Confirm(ConfirmOrganizationInviteLinkRequestModel),
    Accept(AcceptOrganizationInviteLinkRequestModel),
}

#[cfg(test)]
mod tests {
    use bitwarden_api_api::{
        apis::ApiClient,
        models::{
            OrganizationInviteLinkSsoResponseModel, OrganizationInviteLinkStatusResponseModel,
            OrganizationInviteLinkValidateEmailDomainResponseModel,
            OrganizationInviteResponseModel, OrganizationPublicKeyResponseModel,
        },
    };
    use bitwarden_core::{
        client::ApiConfigurations, key_management::create_test_crypto_with_user_and_org_key,
    };
    use bitwarden_crypto::{
        PublicKeyEncryptionAlgorithm, SymmetricCryptoKey, SymmetricKeyAlgorithm,
    };

    use super::*;
    use crate::validation_problem::tests::response_error;

    fn make_client(org_id: OrganizationId, api_client: ApiClient) -> InviteLinkUserClient {
        let user_key = SymmetricCryptoKey::make(SymmetricKeyAlgorithm::Aes256CbcHmac);
        let org_key = SymmetricCryptoKey::make(SymmetricKeyAlgorithm::Aes256CbcHmac);
        let key_store = create_test_crypto_with_user_and_org_key(user_key, org_id, org_key);
        // Give the store a user private key so the confirmation branch can derive a user public
        // key.
        {
            let mut ctx = key_store.context_mut();
            let local = ctx.make_private_key(PublicKeyEncryptionAlgorithm::RsaOaepSha1);
            ctx.persist_private_key(local, PrivateKeySlotId::UserPrivateKey)
                .expect("persisting the user private key should work");
        }
        InviteLinkUserClient {
            key_store,
            api_configurations: Arc::new(ApiConfigurations::from_api_client(api_client)),
        }
    }

    /// Builds an invite + its secret and the organization public key it binds, all consistent with
    /// the client's org key.
    fn build_invite(
        client: &InviteLinkUserClient,
        org_id: OrganizationId,
    ) -> (InviteSecret, Invite, B64) {
        let mut ctx = client.key_store.context();
        let org_key = SymmetricKeySlotId::Organization(org_id);
        let private_key = ctx.make_private_key(PublicKeyEncryptionAlgorithm::RsaOaepSha1);
        let org_public_key = B64::from(
            ctx.get_public_key(private_key)
                .unwrap()
                .to_der()
                .unwrap()
                .as_ref(),
        );
        let wrapped = ctx.wrap_private_key(org_key, private_key).unwrap();
        let (secret, invite) = Invite::make_for_private_key(org_key, &wrapped, &mut ctx).unwrap();
        (secret, invite, org_public_key)
    }

    #[tokio::test]
    async fn accept_and_confirm_succeeds_for_confirmable_invite() {
        let org_id = OrganizationId::new_v4();
        // `get_public_key` returns the base64 key held in this cell, and `get_invite` returns the
        // serialized invite; both are filled after the invite is generated below.
        let recovery = Arc::new(std::sync::Mutex::new(None::<String>));
        let invite_cell = Arc::new(std::sync::Mutex::new(None::<String>));
        let recovery_mock = recovery.clone();
        let invite_mock = invite_cell.clone();
        let client = make_client(
            org_id,
            ApiClient::new_mocked(move |mock| {
                mock.organizations_api
                    .expect_get_public_key()
                    .returning(move |_id| {
                        Ok(OrganizationPublicKeyResponseModel {
                            object: None,
                            public_key: recovery_mock.lock().unwrap().clone(),
                        })
                    })
                    .once();
                mock.organization_users_api
                    .expect_get_invite()
                    .returning(move |_model| {
                        Ok(OrganizationInviteResponseModel {
                            invite: invite_mock.lock().unwrap().clone(),
                        })
                    })
                    .once();
                mock.organization_users_api
                    .expect_confirm_invite_link()
                    .returning(|_model| Ok(()))
                    .once();
            }),
        );

        let (secret, invite, org_public_key) = build_invite(&client, org_id);
        assert!(invite.supports_confirmation());
        // The recovery public key returned by the "server" matches the invite's bound org key.
        *recovery.lock().unwrap() = Some(String::from(&org_public_key));
        *invite_cell.lock().unwrap() = Some(String::from(&invite));

        client
            .accept_and_optionally_confirm(
                org_id,
                uuid::Uuid::new_v4().to_string(),
                secret,
                "Default".to_string(),
                true,
            )
            .await
            .unwrap();
    }

    #[tokio::test]
    async fn accept_without_enrollment_confirms_without_recovery_key() {
        let org_id = OrganizationId::new_v4();
        // Without enrollment the recovery key is never fetched, so only `get_invite` and
        // `confirm_invite_link` run.
        let invite_cell = Arc::new(std::sync::Mutex::new(None::<String>));
        let invite_mock = invite_cell.clone();
        let client = make_client(
            org_id,
            ApiClient::new_mocked(move |mock| {
                mock.organization_users_api
                    .expect_get_invite()
                    .returning(move |_model| {
                        Ok(OrganizationInviteResponseModel {
                            invite: invite_mock.lock().unwrap().clone(),
                        })
                    })
                    .once();
                mock.organization_users_api
                    .expect_confirm_invite_link()
                    .returning(|_model| Ok(()))
                    .once();
            }),
        );

        let (secret, invite, _org_public_key) = build_invite(&client, org_id);
        *invite_cell.lock().unwrap() = Some(String::from(&invite));
        client
            .accept_and_optionally_confirm(
                org_id,
                uuid::Uuid::new_v4().to_string(),
                secret,
                "Default".to_string(),
                false,
            )
            .await
            .unwrap();
    }

    #[tokio::test]
    async fn accept_without_confirmation_posts_acceptance() {
        let org_id = OrganizationId::new_v4();
        let invite_cell = Arc::new(std::sync::Mutex::new(None::<String>));
        let invite_mock = invite_cell.clone();
        let client = make_client(
            org_id,
            ApiClient::new_mocked(move |mock| {
                mock.organization_users_api
                    .expect_get_invite()
                    .returning(move |_model| {
                        Ok(OrganizationInviteResponseModel {
                            invite: invite_mock.lock().unwrap().clone(),
                        })
                    })
                    .once();
                mock.organization_users_api
                    .expect_accept_invite_link()
                    .returning(|_model| Ok(()))
                    .once();
            }),
        );

        // An invite with confirmation disabled routes to the acceptance branch.
        let (secret, mut invite, _org_public_key) = build_invite(&client, org_id);
        invite.disable_confirmation();
        assert!(!invite.supports_confirmation());
        *invite_cell.lock().unwrap() = Some(String::from(&invite));

        client
            .accept_and_optionally_confirm(
                org_id,
                uuid::Uuid::new_v4().to_string(),
                secret,
                "Default".to_string(),
                false,
            )
            .await
            .unwrap();
    }

    #[tokio::test]
    async fn accept_with_mismatched_recovery_key_fails() {
        let org_id = OrganizationId::new_v4();
        let recovery = Arc::new(std::sync::Mutex::new(None::<String>));
        let invite_cell = Arc::new(std::sync::Mutex::new(None::<String>));
        let recovery_mock = recovery.clone();
        let invite_mock = invite_cell.clone();
        let client = make_client(
            org_id,
            ApiClient::new_mocked(move |mock| {
                mock.organizations_api
                    .expect_get_public_key()
                    .returning(move |_id| {
                        Ok(OrganizationPublicKeyResponseModel {
                            object: None,
                            public_key: recovery_mock.lock().unwrap().clone(),
                        })
                    })
                    .once();
                mock.organization_users_api
                    .expect_get_invite()
                    .returning(move |_model| {
                        Ok(OrganizationInviteResponseModel {
                            invite: invite_mock.lock().unwrap().clone(),
                        })
                    })
                    .once();
            }),
        );

        let (secret, invite, _org_public_key) = build_invite(&client, org_id);
        *invite_cell.lock().unwrap() = Some(String::from(&invite));
        // The "server" returns an unrelated public key that must not match the invite's bound
        // thumbprint.
        let unrelated_public_key = {
            let mut ctx = client.key_store.context();
            let private_key = ctx.make_private_key(PublicKeyEncryptionAlgorithm::RsaOaepSha1);
            B64::from(
                ctx.get_public_key(private_key)
                    .unwrap()
                    .to_der()
                    .unwrap()
                    .as_ref(),
            )
        };
        *recovery.lock().unwrap() = Some(String::from(&unrelated_public_key));

        let result = client
            .accept_and_optionally_confirm(
                org_id,
                uuid::Uuid::new_v4().to_string(),
                secret,
                "Default".to_string(),
                true,
            )
            .await;

        assert!(matches!(
            result,
            Err(AcceptInviteLinkError::RecoveryKeyMismatch)
        ));
    }

    #[tokio::test]
    async fn get_status_returns_mapped_view() {
        let org_id = OrganizationId::new_v4();
        let client = make_client(
            org_id,
            ApiClient::new_mocked(|mock| {
                mock.organization_invite_links_api
                    .expect_get_status()
                    .returning(|_model| {
                        Ok(OrganizationInviteLinkStatusResponseModel {
                            object: None,
                            organization_name: Some("Test Org".to_string()),
                            links_enabled: Some(true),
                            seats_available: Some(true),
                            supports_confirmation: Some(true),
                            sso: Some(Box::new(OrganizationInviteLinkSsoResponseModel {
                                object: None,
                                org_sso_id: Some("sso-id".to_string()),
                                required: Some(true),
                            })),
                        })
                    })
                    .once();
            }),
        );

        let status = client
            .get_status(org_id, uuid::Uuid::new_v4().to_string())
            .await
            .unwrap();

        assert_eq!(status.organization_name, "Test Org");
        assert!(status.links_enabled);
        assert!(status.seats_available);
        assert!(status.supports_confirmation);
        let sso = status.sso.expect("sso should be present");
        assert_eq!(sso.org_sso_id.as_deref(), Some("sso-id"));
        assert!(sso.required);
    }

    #[tokio::test]
    async fn is_email_allowed_returns_is_allowed() {
        let org_id = OrganizationId::new_v4();
        let client = make_client(
            org_id,
            ApiClient::new_mocked(|mock| {
                mock.organization_invite_links_api
                    .expect_validate_email_domain()
                    .returning(|_model| {
                        Ok(OrganizationInviteLinkValidateEmailDomainResponseModel {
                            is_allowed: Some(true),
                        })
                    })
                    .once();
            }),
        );

        let allowed = client
            .is_email_allowed(
                org_id,
                uuid::Uuid::new_v4().to_string(),
                "[email protected]".to_string(),
            )
            .await
            .unwrap();

        assert!(allowed);
    }

    /// Runs `accept_and_optionally_confirm` against a confirmable invite whose confirm request
    /// fails with the given status and body.
    async fn accept_with_confirm_failure(status: u16, body: &'static str) -> AcceptInviteLinkError {
        let org_id = OrganizationId::new_v4();
        let invite_cell = Arc::new(std::sync::Mutex::new(None::<String>));
        let invite_mock = invite_cell.clone();
        let client = make_client(
            org_id,
            ApiClient::new_mocked(move |mock| {
                mock.organization_users_api
                    .expect_get_invite()
                    .returning(move |_model| {
                        Ok(OrganizationInviteResponseModel {
                            invite: invite_mock.lock().unwrap().clone(),
                        })
                    })
                    .once();
                mock.organization_users_api
                    .expect_confirm_invite_link()
                    .returning(move |_model| Err(response_error(status, body)))
                    .once();
            }),
        );

        let (secret, invite, _org_public_key) = build_invite(&client, org_id);
        *invite_cell.lock().unwrap() = Some(String::from(&invite));

        client
            .accept_and_optionally_confirm(
                org_id,
                uuid::Uuid::new_v4().to_string(),
                secret,
                "Default".to_string(),
                false,
            )
            .await
            .unwrap_err()
    }

    #[tokio::test]
    async fn confirm_maps_validation_code_to_variant() {
        let error = accept_with_confirm_failure(
            400,
            r#"{"type":"validation_error","status":400,"errors":{"code":[{"type":"already_organization_member","detail":"You're already a member of Acme."}]}}"#,
        )
        .await;

        assert!(matches!(
            error,
            AcceptInviteLinkError::AlreadyOrganizationMember
        ));
    }

    #[tokio::test]
    async fn confirm_maps_not_found_to_link_not_found() {
        let error = accept_with_confirm_failure(
            404,
            r#"{"message":"Invite link not found.","object":"error"}"#,
        )
        .await;

        assert!(matches!(error, AcceptInviteLinkError::LinkNotFound));
    }

    #[tokio::test]
    async fn confirm_maps_unmapped_code_to_unknown() {
        let error = accept_with_confirm_failure(
            400,
            r#"{"type":"validation_error","status":400,"errors":{"code":[{"type":"some_future_code","detail":"Something new."}]}}"#,
        )
        .await;

        assert!(
            matches!(error, AcceptInviteLinkError::Unknown(detail) if detail == "Something new.")
        );
    }

    #[tokio::test]
    async fn confirm_keeps_legacy_error_as_api_error() {
        let error = accept_with_confirm_failure(
            400,
            r#"{"message":"You're already a member of Acme.","object":"error"}"#,
        )
        .await;

        assert!(matches!(error, AcceptInviteLinkError::Api(_)));
    }

    #[tokio::test]
    async fn get_invite_not_found_maps_to_link_not_found() {
        let org_id = OrganizationId::new_v4();
        let client = make_client(
            org_id,
            ApiClient::new_mocked(|mock| {
                mock.organization_users_api
                    .expect_get_invite()
                    .returning(|_model| Err(response_error(404, "")))
                    .once();
            }),
        );
        let (secret, _invite, _org_public_key) = build_invite(&client, org_id);

        let result = client
            .accept_and_optionally_confirm(
                org_id,
                uuid::Uuid::new_v4().to_string(),
                secret,
                "Default".to_string(),
                false,
            )
            .await;

        assert!(matches!(result, Err(AcceptInviteLinkError::LinkNotFound)));
    }
}