passki 0.2.0

// Copyright 2026 Grzegorz Blach
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! # Passkeys Demo Server (Warp)
//!
//! This example demonstrates WebAuthn/Passkey authentication using the Passki library
//! with the Warp web framework, including optional PRF extension support for key derivation.
//!
//! ## Authentication Flows
//!
//! ### Registration (creating a new passkey)
//! 1. Client sends username to `/register/start`
//! 2. Server generates a challenge and returns WebAuthn options (with PRF probe)
//! 3. Client calls `navigator.credentials.create()` with these options
//! 4. User authenticates with their device (fingerprint, face, PIN, etc.)
//! 5. Client sends the credential to `/register/finish`
//! 6. Server verifies, stores the passkey, and reports PRF support
//!
//! ### Authentication (using an existing passkey)
//! Two modes are supported:
//!
//! **Passwordless** (username provided):
//! - Server returns only the credentials registered to that user
//! - Browser shows only matching passkeys
//!
//! **Usernameless** (no username):
//! - Server returns empty credential list
//! - Browser shows all available passkeys (discoverable credentials)
//! - Server identifies the user by the credential used
//!
//! ### PRF key derivation (optional)
//! If the client sends a `prf_salt` with the authentication request, the server
//! includes it as `extensions.prf.eval.first` in the challenge. The authenticator
//! derives a deterministic 32-byte key returned hex-encoded in `prf_output`.
//!
//! ## Running
//! ```sh
//! cargo run --example warp
//! ```
//! Then open http://localhost:3000 in your browser.

use passki::{
    AttestationConveyancePreference, AuthenticationCredential, AuthenticationExtensions,
    AuthenticationState, ClientData, ClientExtensionResults, Passki, PrfEval, PrfInput,
    RegistrationCredential, RegistrationExtensions, RegistrationState, ResidentKeyRequirement,
    StoredPasskey, UserVerificationRequirement,
};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::convert::Infallible;
use std::sync::{Arc, Mutex};
use uuid::Uuid;
use warp::{Filter, Reply, http::StatusCode, reject::Reject, reply};

// =============================================================================
// Error handling
// =============================================================================

#[derive(Debug)]
struct AppError(String);

impl Reject for AppError {}

/// Converts rejections into HTTP 400 responses with the error message.
async fn handle_rejection(err: warp::Rejection) -> Result<impl Reply, Infallible> {
    let message = if let Some(e) = err.find::<AppError>() {
        e.0.clone()
    } else {
        "Internal server error".to_string()
    };

    Ok(reply::with_status(message, StatusCode::BAD_REQUEST))
}

// =============================================================================
// Storage
// =============================================================================

/// In-memory storage for users and pending WebAuthn ceremonies.
///
/// In production, you would use a database for users and a cache (e.g., Redis)
/// for pending states with appropriate expiration.
#[derive(Clone, Default)]
struct Store {
    /// Maps username -> User data
    users: Arc<Mutex<HashMap<String, User>>>,

    /// Pending registration ceremonies, keyed by challenge.
    /// The state must be kept between start and finish calls.
    pending_registrations: Arc<Mutex<HashMap<String, RegistrationState>>>,

    /// Pending authentication ceremonies, keyed by challenge.
    /// The state must be kept between start and finish calls.
    pending_authentications: Arc<Mutex<HashMap<String, AuthenticationState>>>,
}

/// User record containing their profile and registered passkeys.
#[derive(Clone)]
#[allow(unused)]
struct User {
    /// Unique user identifier.
    id: Uuid,
    /// Username or account identifier.
    username: String,
    /// Human-readable display name.
    display_name: String,
    /// All passkeys registered by this user. A user can have multiple passkeys
    /// (e.g., one on their phone, one on their laptop, one security key).
    passkeys: Vec<StoredPasskey>,
    /// Whether any of this user's passkeys reported PRF support during registration.
    prf_supported: bool,
}

// =============================================================================
// Request/Response types
// =============================================================================

#[derive(Deserialize)]
struct RegisterStartRequest {
    username: String,
}

/// Data sent by the client after WebAuthn credential creation.
#[derive(Deserialize)]
struct RegisterFinishRequest {
    /// Base64url-encoded credential ID from the authenticator
    credential_id: String,
    /// Base64url-encoded attestation object containing the public key
    public_key: String,
    /// Base64url-encoded client data JSON
    client_data_json: String,
    /// Extension results from the browser (e.g., PRF support flag)
    client_extension_results: Option<ClientExtensionResults>,
}

/// Request to start authentication. Username and PRF salt are both optional.
#[derive(Deserialize, Default)]
struct AuthStartRequest {
    /// If provided: passwordless flow (server specifies allowed credentials)
    /// If omitted: usernameless flow (browser shows all available passkeys)
    #[serde(default)]
    username: Option<String>,
    /// Base64url-encoded PRF context string. When present, the server requests
    /// a PRF derivation for this salt.
    #[serde(default)]
    prf_salt: Option<String>,
}

/// Data sent by the client after WebAuthn authentication.
#[derive(Deserialize)]
struct AuthFinishRequest {
    /// Base64url-encoded credential ID identifying which passkey was used
    credential_id: String,
    /// Base64url-encoded authenticator data (contains flags and counter)
    authenticator_data: String,
    /// Base64url-encoded client data JSON
    client_data_json: String,
    /// Base64url-encoded signature over authenticator_data + hash(client_data_json)
    signature: String,
    /// Extension results from the browser (e.g., PRF outputs)
    client_extension_results: Option<ClientExtensionResults>,
}

#[derive(Serialize)]
struct ApiResponse {
    success: bool,
    message: String,
    /// In usernameless flow, returns the identified username
    #[serde(skip_serializing_if = "Option::is_none")]
    username: Option<String>,
    /// Registration only: whether a residental key was created
    #[serde(skip_serializing_if = "Option::is_none")]
    residental_key: Option<bool>,
    /// Registration only: whether this passkey supports PRF
    #[serde(skip_serializing_if = "Option::is_none")]
    prf_supported: Option<bool>,
    /// Authentication only: hex-encoded 32-byte derived key, when prf_salt was provided
    #[serde(skip_serializing_if = "Option::is_none")]
    prf_output: Option<String>,
}

// =============================================================================
// Application state
// =============================================================================

#[derive(Clone)]
struct AppState {
    passki: Arc<Passki>,
    store: Store,
}

// =============================================================================
// Filters
// =============================================================================

/// Extracts AppState for injection into handlers.
fn with_state(state: AppState) -> impl Filter<Extract = (AppState,), Error = Infallible> + Clone {
    warp::any().map(move || state.clone())
}

// =============================================================================
// Handlers
// =============================================================================

async fn index() -> Result<impl Reply, warp::Rejection> {
    Ok(reply::html(include_str!("index.html")))
}

/// POST /register/start - Begin passkey registration
///
/// Generates a challenge and WebAuthn options for the client to create a credential.
/// The challenge is random bytes that the authenticator must sign, preventing replay attacks.
async fn register_start(
    state: AppState,
    req: RegisterStartRequest,
) -> Result<impl Reply, warp::Rejection> {
    // Generate a unique user ID. This should be random and opaque (not the username)
    // to prevent tracking users across sites.
    let user_id = Uuid::new_v4().as_bytes().to_vec();

    // If user exists, get their existing passkeys to exclude them from re-registration
    let existing = state
        .store
        .users
        .lock()
        .unwrap()
        .get(&req.username)
        .map(|u| u.passkeys.clone());

    let (challenge, reg_state) = state
        .passki
        .start_passkey_registration(
            &user_id,
            &req.username, // User handle (displayed by authenticator)
            &req.username, // Display name
            60000,         // Timeout in milliseconds
            AttestationConveyancePreference::None, // Don't request attestation
            ResidentKeyRequirement::Preferred, // Request discoverable credential if possible
            UserVerificationRequirement::Preferred, // Request user verification if available
            existing.as_deref(), // Exclude existing credentials
            Some(RegistrationExtensions {
                cred_props: Some(true),
                prf: Some(PrfInput { eval: None }),
            }), // Probe PRF support
        )
        .map_err(|e| warp::reject::custom(AppError(e.to_string())))?;

    // Store state for verification in finish step, keyed by the challenge
    state
        .store
        .pending_registrations
        .lock()
        .unwrap()
        .insert(challenge.challenge.clone(), reg_state);

    // Return challenge to client (will be passed to navigator.credentials.create())
    Ok(reply::json(&challenge))
}

/// POST /register/finish - Complete passkey registration
///
/// Verifies the credential created by the authenticator, stores it, and reports
/// whether this passkey supports the PRF extension.
async fn register_finish(
    state: AppState,
    req: RegisterFinishRequest,
) -> Result<impl Reply, warp::Rejection> {
    // Parse client data to extract challenge
    let client_data = ClientData::from_base64(&req.client_data_json)
        .map_err(|e| warp::reject::custom(AppError(e.to_string())))?;

    // Retrieve and remove the pending registration state
    let reg_state = state
        .store
        .pending_registrations
        .lock()
        .unwrap()
        .remove(&client_data.challenge)
        .ok_or_else(|| warp::reject::custom(AppError("No pending registration".into())))?;

    let prf_supported = req
        .client_extension_results
        .as_ref()
        .and_then(|ext| ext.prf.as_ref())
        .and_then(|prf| prf.enabled)
        .unwrap_or(false);

    // Package the credential data from the client
    let credential = RegistrationCredential {
        credential_id: req.credential_id,
        public_key: req.public_key,
        client_data_json: req.client_data_json,
        client_extension_results: req.client_extension_results,
    };

    // Verify the credential (checks origin, challenge, parses public key)
    let passkey = state
        .passki
        .finish_passkey_registration(&credential, &reg_state)
        .map_err(|e| warp::reject::custom(AppError(e.to_string())))?;
    let residental_key = passkey.rk;

    // Decode user ID from base64url to UUID
    let user_id_bytes = Passki::base64_decode(&reg_state.user.id)
        .map_err(|e| warp::reject::custom(AppError(e.to_string())))?;
    let user_id = Uuid::from_slice(&user_id_bytes)
        .map_err(|e| warp::reject::custom(AppError(e.to_string())))?;

    // Store the passkey for future authentication.
    let mut users = state.store.users.lock().unwrap();
    users
        .entry(reg_state.user.name.clone())
        // If user exists, add passkey to their list.
        .and_modify(|user| {
            user.passkeys.push(passkey.clone());
            user.prf_supported |= prf_supported;
        })
        // If new user, create user record with their info.
        .or_insert(User {
            id: user_id,
            username: reg_state.user.name,
            display_name: reg_state.user.display_name,
            passkeys: vec![passkey],
            prf_supported,
        });

    Ok(reply::json(&ApiResponse {
        success: true,
        message: "Registration successful".into(),
        username: None,
        residental_key,
        prf_supported: Some(prf_supported),
        prf_output: None,
    }))
}

/// POST /auth/start - Begin passkey authentication
///
/// Two modes based on whether username is provided:
/// - Passwordless: returns challenge with user's credential IDs (browser filters to these)
/// - Usernameless: returns challenge with empty credential list (browser shows all passkeys)
///
/// If prf_salt is provided, the challenge includes a PRF eval so the authenticator
/// will derive a key for that salt.
async fn auth_start(state: AppState, req: AuthStartRequest) -> Result<impl Reply, warp::Rejection> {
    let passkeys = if let Some(ref username) = req.username {
        // Passwordless flow: get user's passkeys to include in allowCredentials
        let users = state.store.users.lock().unwrap();
        let user = users
            .get(username)
            .ok_or_else(|| warp::reject::custom(AppError("User not found".into())))?;
        user.passkeys.clone()
    } else {
        // Usernameless flow: empty allowCredentials lets browser show all passkeys
        vec![]
    };

    let extensions = req.prf_salt.map(|salt| AuthenticationExtensions {
        prf: PrfInput {
            eval: Some(PrfEval {
                first: salt,
                second: None,
            }),
        },
    });

    let (challenge, auth_state) = state.passki.start_passkey_authentication(
        &passkeys,
        60000,                                  // Timeout in milliseconds
        UserVerificationRequirement::Preferred, // Request user verification if available
        extensions,
    );

    // Store state for verification in finish step, keyed by the challenge
    state
        .store
        .pending_authentications
        .lock()
        .unwrap()
        .insert(challenge.challenge.clone(), auth_state);

    // Return challenge to client
    Ok(reply::json(&challenge))
}

/// POST /auth/finish - Complete passkey authentication
///
/// Verifies the signature from the authenticator and, if a PRF salt was provided,
/// returns the derived key hex-encoded in prf_output.
async fn auth_finish(
    state: AppState,
    req: AuthFinishRequest,
) -> Result<impl Reply, warp::Rejection> {
    // Parse client data to extract challenge
    let client_data = ClientData::from_base64(&req.client_data_json)
        .map_err(|e| warp::reject::custom(AppError(e.to_string())))?;

    // Retrieve pending state using challenge
    let auth_state = state
        .store
        .pending_authentications
        .lock()
        .unwrap()
        .remove(&client_data.challenge)
        .ok_or_else(|| warp::reject::custom(AppError("No pending authentication".into())))?;

    // Decode credential ID to find the matching passkey
    let credential_id = Passki::base64_decode(&req.credential_id)
        .map_err(|e| warp::reject::custom(AppError(e.to_string())))?;

    // Find user by credential_id
    let mut users = state.store.users.lock().unwrap();
    let (username, passkey) = users
        .iter_mut()
        .find_map(|(name, user)| {
            user.passkeys
                .iter_mut()
                .find(|pk| pk.credential_id == credential_id)
                .map(|pk| (name.clone(), pk))
        })
        .ok_or_else(|| warp::reject::custom(AppError("Unknown credential".into())))?;

    // Package the authentication response from the client
    let credential = AuthenticationCredential {
        credential_id: req.credential_id,
        authenticator_data: req.authenticator_data,
        client_data_json: req.client_data_json,
        signature: req.signature,
        client_extension_results: req.client_extension_results,
    };

    // Verify the signature (checks origin, challenge, signature, counter)
    let result = state
        .passki
        .finish_passkey_authentication(&credential, &auth_state, passkey)
        .map_err(|e| warp::reject::custom(AppError(e.to_string())))?;

    // Update the counter to detect cloned authenticators.
    // If counter goes backwards, it may indicate the credential was cloned.
    passkey.counter = result.counter;

    let prf_output = result
        .prf_first
        .map(|bytes| bytes.iter().map(|b| format!("{b:02x}")).collect());

    Ok(reply::json(&ApiResponse {
        success: true,
        message: format!("Welcome back, {}!", username),
        username: Some(username),
        prf_supported: None,
        residental_key: None,
        prf_output,
    }))
}

// =============================================================================
// Main
// =============================================================================

#[tokio::main]
async fn main() {
    // Initialize Passki with relying party information.
    // - rp_id: The domain name (no protocol or port). Credentials are bound to this.
    // - origin: The full origin URL. Must match what the browser sends.
    // - rp_name: Human-readable name shown by authenticators.
    let state = AppState {
        passki: Arc::new(Passki::new(
            "localhost",
            "http://localhost:3000",
            "Passkeys Demo",
        )),
        store: Store::default(),
    };

    let index = warp::path::end().and(warp::get()).and_then(index);

    let register_start = warp::path!("register" / "start")
        .and(warp::post())
        .and(with_state(state.clone()))
        .and(warp::body::json())
        .and_then(register_start);

    let register_finish = warp::path!("register" / "finish")
        .and(warp::post())
        .and(with_state(state.clone()))
        .and(warp::body::json())
        .and_then(register_finish);

    let auth_start = warp::path!("auth" / "start")
        .and(warp::post())
        .and(with_state(state.clone()))
        .and(warp::body::json())
        .and_then(auth_start);

    let auth_finish = warp::path!("auth" / "finish")
        .and(warp::post())
        .and(with_state(state.clone()))
        .and(warp::body::json())
        .and_then(auth_finish);

    let routes = index
        .or(register_start)
        .or(register_finish)
        .or(auth_start)
        .or(auth_finish)
        .recover(handle_rejection);

    println!("Server starting on http://localhost:3000");
    warp::serve(routes).run(([0, 0, 0, 0], 3000)).await;
}