systemprompt-oauth 0.44.0

//! Linking new passkeys to existing user accounts.
//!
//! Copyright (c) systemprompt.io — Business Source License 1.1.
//! See <https://systemprompt.io> for licensing details.

use super::WebAuthnService;
use crate::error::{OauthError, OauthResult as Result};
use crate::repository::{StoreChallengeParams, TokenValidationResult, WebAuthnChallengeKind};
use crate::services::webauthn::token::hash_token;
use serde::{Deserialize, Serialize};
use std::time::Duration;
use systemprompt_identifiers::{TokenId, UserId};
use tracing::instrument;
use uuid::Uuid;
use webauthn_rs::prelude::*;

const LINK_CHALLENGE_TTL: Duration = Duration::from_secs(300);

#[derive(Debug, Serialize, Deserialize)]
struct LinkRegistrationState {
    reg_state: PasskeyRegistration,
    token_id: TokenId,
}

#[derive(Debug, Clone)]
pub struct LinkUserInfo {
    pub id: UserId,
    pub email: String,
    pub name: String,
}

impl WebAuthnService {
    #[instrument(skip(self, setup_token))]
    pub async fn start_registration_with_token(
        &self,
        setup_token: &str,
    ) -> Result<(CreationChallengeResponse, String, LinkUserInfo)> {
        let token_hash = hash_token(setup_token);
        let validation = self.oauth_repo.validate_setup_token(&token_hash).await?;

        let token_record = match validation {
            TokenValidationResult::Valid(record) => record,
            TokenValidationResult::Expired => {
                return Err(OauthError::Internal("Setup token has expired".to_owned()));
            },
            TokenValidationResult::AlreadyUsed => {
                return Err(OauthError::Internal(
                    "Setup token has already been used".to_owned(),
                ));
            },
            TokenValidationResult::NotFound => {
                return Err(OauthError::Internal("Invalid setup token".to_owned()));
            },
        };

        let user = self
            .oauth_repo
            .get_authenticated_user(&token_record.user_id)
            .await?;

        let existing_creds = self.get_user_credentials(&token_record.user_id).await?;
        let exclude_credentials: Vec<CredentialID> =
            existing_creds.iter().map(|c| c.cred_id().clone()).collect();

        let user_unique_id = Uuid::parse_str(token_record.user_id.as_str()).map_err(|e| {
            OauthError::Internal(format!(
                "user_id {:?} is not a valid UUID: {e}",
                token_record.user_id.as_str()
            ))
        })?;

        let (challenge, reg_state) = self.webauthn.start_passkey_registration(
            user_unique_id,
            &user.username,
            &user.username,
            if exclude_credentials.is_empty() {
                None
            } else {
                Some(exclude_credentials)
            },
        )?;

        let challenge_id = Uuid::new_v4().to_string();
        let state = serde_json::to_value(LinkRegistrationState {
            reg_state,
            token_id: token_record.id.clone(),
        })?;
        self.oauth_repo
            .store_webauthn_challenge(StoreChallengeParams {
                challenge: &challenge_id,
                kind: WebAuthnChallengeKind::Link,
                user_id: Some(&token_record.user_id),
                state: &state,
                oauth_state: None,
                ttl: LINK_CHALLENGE_TTL,
            })
            .await?;

        let user_info = LinkUserInfo {
            id: token_record.user_id.clone(),
            email: user.email,
            name: user.username,
        };

        tracing::info!(
            user_id = %user_info.id,
            challenge_id = %challenge_id,
            "Link registration ceremony initiated"
        );

        Ok((challenge, challenge_id, user_info))
    }

    #[instrument(skip(self, setup_token, credential))]
    pub async fn finish_registration_with_token(
        &self,
        challenge_id: &str,
        setup_token: &str,
        credential: &RegisterPublicKeyCredential,
    ) -> Result<UserId> {
        let token_hash = hash_token(setup_token);
        let validation = self.oauth_repo.validate_setup_token(&token_hash).await?;

        let TokenValidationResult::Valid(token_record) = validation else {
            return Err(OauthError::Internal(
                "Invalid or expired setup token".to_owned(),
            ));
        };

        let consumed = self
            .oauth_repo
            .consume_webauthn_challenge(challenge_id, WebAuthnChallengeKind::Link)
            .await?
            .ok_or_else(|| {
                OauthError::Internal("Registration session not found or expired".to_owned())
            })?;
        let user_id = consumed.user_id.ok_or_else(|| {
            OauthError::Internal("Registration session has no user id".to_owned())
        })?;
        let state: LinkRegistrationState = serde_json::from_value(consumed.state)?;

        if state.token_id != token_record.id {
            return Err(OauthError::Internal("Token mismatch".to_owned()));
        }

        let passkey = self
            .webauthn
            .finish_passkey_registration(credential, &state.reg_state)?;

        self.store_credential(&user_id, &passkey, "Linked Passkey")
            .await?;

        self.oauth_repo
            .consume_setup_token(&token_record.id)
            .await?;

        tracing::info!(
            user_id = %user_id,
            "WebAuthn credential linked to existing user"
        );

        Ok(user_id)
    }
}