🐸 Frogue(Fed)

~ololduck/Frogue

Raw

//! Shared registration logic between web and SSH TUI.

use chrono::{Duration, Utc};
use sea_orm::{ActiveModelTrait, ColumnTrait, EntityTrait, PaginatorTrait, QueryFilter, Set};
use tracing::info;
use uuid::Uuid;

use crate::auth::hash_password;
use crate::config::*;
use crate::entities::{invites, users};

/// Errors that can occur during registration.
#[derive(Debug, thiserror::Error)]
pub enum RegisterError {
    #[error("Invalid: {0}")]
    Invalid(String),
    #[error("Username or email already taken")]
    Taken,
    #[error("Database error: {0}")]
    Database(String),
    #[error("Registration on this server is invite-only.")]
    InviteRequired,
    #[error("Invalid invite: {0}")]
    InviteInvalid(String),
}

/// Errors that can occur when creating an invite.
#[derive(Debug, thiserror::Error)]
pub enum InviteError {
    #[error("You have already created the maximum number of invites for today.")]
    QuotaExceeded,
    #[error("Database error: {0}")]
    Database(String),
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum RegistrationMode {
    #[default]
    Open,
    InviteOnly,
}

/// Returns the registration mode, read from the environment or the default.
pub fn registration_mode() -> RegistrationMode {
    match std::env::var(REGISTRATION_MODE_VAR_KEY)
        .ok()
        .map(|v| v.to_lowercase())
    {
        Some(mode) if mode == "invite_only" => RegistrationMode::InviteOnly,
        _ => RegistrationMode::Open,
    }
}

/// Returns the invite validity duration, read from the environment or the
/// default.
pub fn invites_validity() -> Duration {
    std::env::var(INVITES_VALIDITY_SECONDS_VAR_KEY)
        .ok()
        .and_then(|v| v.parse().ok())
        .map(Duration::seconds)
        .unwrap_or_else(|| Duration::seconds(DEFAULT_INVITES_VALIDITY_SECONDS))
}

/// Returns the number of invites allowed per day, read from the environment or
/// the default.
pub fn invites_allowed_per_day() -> u64 {
    std::env::var(INVITES_ALLOWED_PER_DAY_VAR_KEY)
        .ok()
        .and_then(|v| v.parse().ok())
        .unwrap_or(DEFAULT_INVITES_ALLOWED_PER_DAY)
}

/// User registration data.
#[derive(Debug, Clone)]
pub struct Registration {
    pub name: String,
    pub email: String,
    pub password: String,
    pub ssh_key: Option<String>,
    pub invite: Option<Uuid>,
}

/// Checks whether registration is allowed given the current mode and invite.
///
/// Returns Ok(()) if registration should proceed, or an error if it should be
/// blocked.
#[tracing::instrument(skip(db))]
pub async fn check_mode_allows_registration(
    db: &sea_orm::DatabaseConnection,
    invite: Option<Uuid>,
) -> Result<(), RegisterError> {
    match registration_mode() {
        RegistrationMode::Open => Ok(()),
        RegistrationMode::InviteOnly => {
            // Check if this is the bootstrap case (no users yet)
            let user_count = users::Entity::find()
                .count(db)
                .await
                .map_err(|e| RegisterError::Database(e.to_string()))?;

            if user_count == 0 {
                // Bootstrap: allow registration without invite
                Ok(())
            } else if let Some(invite_uuid) = invite {
                // Validate the invite
                let invite = find_valid_invite(db, invite_uuid).await;
                match invite {
                    Some(_) => Ok(()),
                    None => Err(RegisterError::InviteInvalid(
                        "This invite is not valid (any more).".to_string(),
                    )),
                }
            } else {
                Err(RegisterError::InviteRequired)
            }
        }
    }
}

/// Finds a valid invite by UUID (exists, unused, unexpired).
#[tracing::instrument(skip(db))]
pub async fn find_valid_invite(
    db: &sea_orm::DatabaseConnection,
    uuid: Uuid,
) -> Option<invites::Model> {
    let uuid_str = uuid.to_string();
    let now_timestamp = Utc::now().timestamp();

    let invite = invites::Entity::find()
        .filter(invites::Column::Uuid.eq(&uuid_str))
        .filter(invites::Column::UsedBy.is_null())
        .one(db)
        .await
        .ok()
        .flatten()?;

    if invite.expires_at.timestamp() > now_timestamp {
        Some(invite)
    } else {
        None
    }
}

/// Registers a new user.
pub async fn register(
    db: &sea_orm::DatabaseConnection,
    registration: Registration,
) -> Result<users::Model, RegisterError> {
    // Check registration mode before validation
    check_mode_allows_registration(db, registration.invite).await?;

    // Validate the registration data
    if registration.name.len() < 3 || registration.name.len() > 30 {
        return Err(RegisterError::Invalid(
            "Username must be between 3 and 30 characters".to_string(),
        ));
    }

    if !registration
        .name
        .chars()
        .all(|c| c.is_alphanumeric() || c == '-' || c == '_')
    {
        return Err(RegisterError::Invalid(
            "Username can only contain letters, numbers, hyphens, and underscores".to_string(),
        ));
    }

    if registration.email.is_empty() || !registration.email.contains('@') {
        return Err(RegisterError::Invalid("Invalid email address".to_string()));
    }

    if registration.password.len() < 8 {
        return Err(RegisterError::Invalid(
            "Password must be at least 8 characters".to_string(),
        ));
    }

    // Check if username or email already exists
    let existing = users::Entity::find()
        .filter(
            sea_orm::Condition::any()
                .add(users::Column::Name.eq(&registration.name))
                .add(users::Column::Email.eq(&registration.email)),
        )
        .one(db)
        .await
        .map_err(|e| RegisterError::Database(e.to_string()))?;

    if existing.is_some() {
        return Err(RegisterError::Taken);
    }

    // Hash the password
    let password_hash = hash_password(&registration.password)
        .map_err(|e| RegisterError::Database(e.to_string()))?;

    // Look up the invite to get the inviter's ID (if applicable)
    let invited_by = if let Some(invite_uuid) = registration.invite {
        let invite = find_valid_invite(db, invite_uuid).await;
        invite.map(|inv| inv.created_by)
    } else {
        None
    };

    // Create the user
    let user = users::ActiveModel {
        uuid: Set(Uuid::new_v4()),
        name: Set(registration.name),
        email: Set(registration.email),
        password_hash: Set(password_hash),
        invited_by: Set(invited_by),
        ..Default::default()
    };

    let user = user
        .insert(db)
        .await
        .map_err(|e| RegisterError::Database(e.to_string()))?;

    // Mark the invite as used (if applicable)
    if let Some(invite_uuid) = registration.invite {
        let uuid_str = invite_uuid.to_string();
        let invite = invites::Entity::find()
            .filter(invites::Column::Uuid.eq(&uuid_str))
            .one(db)
            .await
            .map_err(|e| RegisterError::Database(e.to_string()))?;

        if let Some(invite) = invite {
            let mut active: invites::ActiveModel = invite.into();
            active.used_by = Set(Some(user.id));
            active.used_at = Set(Some(Utc::now().into()));
            active
                .update(db)
                .await
                .map_err(|e| RegisterError::Database(e.to_string()))?;

            info!(
                user_id = user.id,
                invite_uuid = %uuid_str,
                "user registered via invite"
            );
        }
    }

    // Add SSH key if provided
    if let Some(key) = registration.ssh_key
        && !key.trim().is_empty()
    {
        let key_model = crate::entities::ssh_keys::ActiveModel {
            user_id: Set(user.id),
            public_key: Set(key),
            ..Default::default()
        };
        key_model
            .insert(db)
            .await
            .map_err(|e| RegisterError::Database(format!("Failed to add SSH key: {}", e)))?;
    }

    Ok(user)
}

/// Creates a new invite for a user, enforcing the per-day limit.
#[tracing::instrument(skip(db))]
pub async fn create_invite(
    db: &sea_orm::DatabaseConnection,
    user_id: i64,
) -> Result<invites::Model, InviteError> {
    // Check per-day quota
    let now = Utc::now();
    let day_ago = now - Duration::hours(24);

    let recent_count = invites::Entity::find()
        .filter(invites::Column::CreatedBy.eq(user_id))
        .filter(invites::Column::CreatedAt.gt(day_ago))
        .count(db)
        .await
        .map_err(|e| InviteError::Database(e.to_string()))?;

    if recent_count >= invites_allowed_per_day() {
        return Err(InviteError::QuotaExceeded);
    }

    // Generate the invite
    let uuid = Uuid::new_v4();
    let uuid_str = uuid.to_string();
    let expires_at = now + invites_validity();

    let invite = invites::ActiveModel {
        uuid: Set(uuid_str),
        created_by: Set(user_id),
        expires_at: Set(expires_at.into()),
        ..Default::default()
    };

    let invite = invite
        .insert(db)
        .await
        .map_err(|e| InviteError::Database(e.to_string()))?;

    info!(
        user_id,
        invite_uuid = %uuid,
        expires_at = %expires_at,
        "invite created"
    );

    Ok(invite)
}

/// Lists all invites created by a user.
#[tracing::instrument(skip(db))]
pub async fn list_invites(
    db: &sea_orm::DatabaseConnection,
    user_id: i64,
) -> Result<Vec<invites::Model>, sea_orm::DbErr> {
    invites::Entity::find()
        .filter(invites::Column::CreatedBy.eq(user_id))
        .all(db)
        .await
}