rustdesk-server/src/api/auth.rs

use crate::api::email;
use crate::api::error::ApiError;
use crate::api::middleware::{sha256_token, AuthedUser};
use crate::api::state::AppState;
use crate::api::users::{verify_password, UserPayload};
use crate::database::UserRow;
use axum::extract::Extension;
use axum::http::StatusCode;
use axum::Json;
use serde::Deserialize;
use serde_json::{json, Value};
use std::sync::Arc;
use totp_rs::{Algorithm, Secret, TOTP};

const EMAIL_CODE_TTL_SECS: i64 = 600;

/// `LoginRequest` mirrors the Flutter client at
/// flutter/lib/common/hbbs/hbbs.dart:133. M1 only consults `username`,
/// `password`, and `type`; the other fields are tolerated for forward-compat.
#[derive(Debug, Deserialize)]
pub struct LoginRequest {
    #[serde(default)]
    pub username: Option<String>,
    #[serde(default)]
    pub password: Option<String>,
    #[serde(default)]
    pub id: Option<String>,
    #[serde(default)]
    pub uuid: Option<String>,
    #[serde(default, rename = "type")]
    pub kind: Option<String>,
    #[serde(default, rename = "deviceInfo")]
    pub device_info: Option<Value>,
    // Tolerated, ignored in M1:
    #[serde(default)]
    pub auto_login: Option<bool>,
    #[serde(default, rename = "verificationCode")]
    pub verification_code: Option<String>,
    #[serde(default, rename = "tfaCode")]
    pub tfa_code: Option<String>,
    #[serde(default)]
    pub secret: Option<String>,
}

#[derive(Debug, Deserialize)]
pub struct IdUuidBody {
    #[serde(default)]
    pub id: Option<String>,
    #[serde(default)]
    pub uuid: Option<String>,
}

pub async fn login_options_head() -> StatusCode {
    StatusCode::OK
}

pub async fn login_options(Extension(state): Extension<Arc<AppState>>) -> Json<Vec<String>> {
    // Static base set from config (account / email_code), plus a dynamic
    // `oidc/<name>` entry per enabled provider in the DB. Recomputed per
    // request so adding a provider via SQL takes effect without a restart.
    let mut out = state.cfg.login_options.clone();
    if !state.cfg.public_base_url.is_empty() {
        if let Ok(providers) = state.db.oidc_provider_list_enabled().await {
            for p in providers {
                out.push(format!("oidc/{}", p.name));
            }
        }
    }
    Json(out)
}

const TFA_CHALLENGE_TTL_SECS: i64 = 300;

pub async fn login(
    Extension(state): Extension<Arc<AppState>>,
    Json(req): Json<LoginRequest>,
) -> Result<Json<Value>, ApiError> {
    // The desktop client reuses the email-code dialog for the TOTP second
    // leg: it POSTs `type: "email_code"` with `tfaCode` set (and the email
    // `verificationCode` field also set, but we ignore that when tfaCode is
    // present). Detect that shape up-front and route to the TOTP verifier;
    // otherwise dispatch on the declared `type`.
    let has_tfa = req.tfa_code.as_deref().is_some_and(|s| !s.is_empty())
        && req.secret.as_deref().is_some_and(|s| !s.is_empty());
    if has_tfa {
        return login_tfa_code(state, req).await;
    }
    let kind = req.kind.as_deref().unwrap_or("account");
    match kind {
        "account" | "" => login_account(state, req).await,
        "tfa_code" => login_tfa_code(state, req).await,
        "email_code" => login_email_code(state, req).await,
        other => Err(ApiError::BadRequest(format!(
            "unsupported login type: {}",
            other
        ))),
    }
}

/// Two-leg passwordless login by email. Leg 1 (no `verificationCode`) mints a
/// fresh 6-digit code and emails it to the user (or logs to stdout when SMTP
/// is unconfigured). Leg 2 (with `verificationCode`) verifies the code,
/// consumes it, and issues an access token.
async fn login_email_code(
    state: Arc<AppState>,
    req: LoginRequest,
) -> Result<Json<Value>, ApiError> {
    // The Flutter client passes the email/username in the `username` field;
    // accept it either as a literal email or as a username we can map to one.
    let identifier = req
        .username
        .as_deref()
        .filter(|s| !s.is_empty())
        .ok_or_else(|| ApiError::BadRequest("username (email) required".into()))?;
    let user = resolve_user_by_identifier(&state, identifier).await?;
    let email = if !user.email.is_empty() {
        user.email.clone()
    } else if user.username.contains('@') {
        // Operator bootstraps users with email-as-username — accept that.
        user.username.clone()
    } else {
        return Err(ApiError::BadRequest(
            "user has no email address on file".into(),
        ));
    };

    if let Some(code) = req
        .verification_code
        .as_deref()
        .filter(|s| !s.is_empty())
    {
        // Leg 2: verify.
        let supplied_hash = sodiumoxide::crypto::hash::sha256::hash(code.as_bytes())
            .as_ref()
            .to_vec();
        let ok = state
            .db
            .email_code_verify(&email, &supplied_hash)
            .await
            .map_err(|e| ApiError::Internal(e.to_string()))?;
        if !ok {
            return Err(ApiError::BadCredentials);
        }
        if user.status == 0 {
            return Err(ApiError::AccountDisabled);
        }
        if user.status == -1 {
            return Err(ApiError::Unverified);
        }
        return issue_session(&state, &req, &user).await;
    }

    // Leg 1: mint + send a fresh code.
    let (code, code_hash) = email::mint_code();
    state
        .db
        .email_code_create(&email, &code_hash, EMAIL_CODE_TTL_SECS)
        .await
        .map_err(|e| ApiError::Internal(e.to_string()))?;
    if let Err(e) = email::send_login_code(state.cfg.email.as_ref(), &email, &code).await {
        return Err(ApiError::Internal(e));
    }
    Ok(Json(json!({ "type": "email_check" })))
}

async fn resolve_user_by_identifier(
    state: &AppState,
    identifier: &str,
) -> Result<UserRow, ApiError> {
    if identifier.contains('@') {
        if let Some(u) = state
            .db
            .user_find_by_email(identifier)
            .await
            .map_err(|e| ApiError::Internal(e.to_string()))?
        {
            return Ok(u);
        }
    }
    state
        .db
        .user_find_by_username(identifier)
        .await
        .map_err(|e| ApiError::Internal(e.to_string()))?
        .ok_or(ApiError::BadCredentials)
}

async fn login_account(
    state: Arc<AppState>,
    req: LoginRequest,
) -> Result<Json<Value>, ApiError> {
    let username = req
        .username
        .as_deref()
        .filter(|s| !s.is_empty())
        .ok_or_else(|| ApiError::BadRequest("username required".into()))?;
    let password = req
        .password
        .as_deref()
        .filter(|s| !s.is_empty())
        .ok_or_else(|| ApiError::BadRequest("password required".into()))?;

    let user = state
        .db
        .user_find_by_username(username)
        .await?
        .ok_or(ApiError::BadCredentials)?;

    let ok = verify_password(user.password_hash.clone(), password.to_string()).await?;
    if !ok {
        return Err(ApiError::BadCredentials);
    }
    if user.status == 0 {
        return Err(ApiError::AccountDisabled);
    }
    if user.status == -1 {
        return Err(ApiError::Unverified);
    }

    // 2FA gate: if the user has TOTP enrolled, mint a short-lived nonce and
    // tell the client we want the TOTP code in a follow-up POST. The client
    // echoes the nonce back as `secret`.
    //
    // Wire shape matches the Flutter client's expectations
    // (flutter/lib/common/widgets/login.dart:485): the outer `type` is the
    // generic `email_check` envelope (the dialog the client opens for any
    // second-leg challenge), and `tfa_type` distinguishes TOTP (`tfa_check`)
    // from email (`email_check`). Returning `type:"tfa_check"` directly
    // would miss the switch's only branch and surface as the unhelpful
    // "bad response from server" toast.
    if state.db.totp_get_secret(user.id).await?.is_some() {
        let nonce = state
            .db
            .tfa_challenge_create(user.id, TFA_CHALLENGE_TTL_SECS)
            .await?;
        return Ok(Json(json!({
            "type":     "email_check",
            "tfa_type": "tfa_check",
            "secret":   nonce,
        })));
    }

    issue_session(&state, &req, &user).await
}

async fn login_tfa_code(
    state: Arc<AppState>,
    req: LoginRequest,
) -> Result<Json<Value>, ApiError> {
    let nonce = req
        .secret
        .as_deref()
        .filter(|s| !s.is_empty())
        .ok_or_else(|| ApiError::BadRequest("secret required".into()))?;
    let code = req
        .tfa_code
        .as_deref()
        .filter(|s| !s.is_empty())
        .ok_or_else(|| ApiError::BadRequest("tfaCode required".into()))?;

    let user_id = state
        .db
        .tfa_challenge_lookup(nonce)
        .await?
        .ok_or_else(|| ApiError::BadRequest("invalid or expired challenge".into()))?;
    let secret_b32 = state
        .db
        .totp_get_secret(user_id)
        .await?
        .ok_or_else(|| ApiError::BadRequest("TOTP not enrolled".into()))?;

    if !verify_totp(&secret_b32, code)? {
        // Leave the challenge row alive — operators may want short retries.
        return Err(ApiError::BadCredentials);
    }
    state.db.tfa_challenge_consume(nonce).await?;

    let user = state
        .db
        .user_find_by_id(user_id)
        .await?
        .ok_or(ApiError::Unauthorized)?;
    issue_session(&state, &req, &user).await
}

/// Build and persist a fresh access token, claim the calling device, and
/// return the standard logged-in response shape. Shared by the password,
/// post-TOTP, post-email-code, and (later) post-OIDC paths.
async fn issue_session(
    state: &AppState,
    req: &LoginRequest,
    user: &UserRow,
) -> Result<Json<Value>, ApiError> {
    let token = mint_token();
    let sha = sha256_token(&token);
    let device_info_json = req
        .device_info
        .as_ref()
        .map(|v| v.to_string())
        .unwrap_or_default();
    state
        .db
        .token_insert(
            user.id,
            &sha,
            req.id.as_deref().unwrap_or_default(),
            req.uuid.as_deref().unwrap_or_default(),
            &device_info_json,
            state.cfg.session_ttl_secs,
        )
        .await?;
    // Bind the calling device to this user so /api/peers shows it correctly.
    state
        .db
        .device_claim(
            user.id,
            req.id.as_deref().unwrap_or_default(),
            req.uuid.as_deref().unwrap_or_default(),
        )
        .await;

    Ok(Json(json!({
        "access_token": token,
        "type":         "access_token",
        "user":         UserPayload::from(user),
    })))
}

pub(crate) fn verify_totp(secret_b32: &str, code: &str) -> Result<bool, ApiError> {
    let secret = Secret::Encoded(secret_b32.to_string())
        .to_bytes()
        .map_err(|e| ApiError::Internal(format!("bad TOTP secret: {:?}", e)))?;
    let totp = TOTP::new(Algorithm::SHA1, 6, 1, 30, secret)
        .map_err(|e| ApiError::Internal(format!("TOTP init: {}", e)))?;
    totp.check_current(code)
        .map_err(|e| ApiError::Internal(format!("TOTP check: {}", e)))
}

pub async fn current_user(
    Extension(state): Extension<Arc<AppState>>,
    user: AuthedUser,
    // Body is required by the client but its content is purely advisory.
    Json(_body): Json<IdUuidBody>,
) -> Result<Json<UserPayload>, ApiError> {
    let row = state
        .db
        .user_find_by_id(user.user_id)
        .await?
        .ok_or(ApiError::Unauthorized)?;
    Ok(Json(UserPayload::from(&row)))
}

pub async fn logout(
    Extension(state): Extension<Arc<AppState>>,
    headers: axum::http::HeaderMap,
    Json(_body): Json<IdUuidBody>,
) -> StatusCode {
    // Best-effort: parse the bearer ourselves so a missing/invalid token still
    // returns 200 (matches the client's fire-and-forget logout flow).
    if let Some(auth) = headers.get(axum::http::header::AUTHORIZATION) {
        if let Ok(s) = auth.to_str() {
            if let Some(tok) = s.strip_prefix("Bearer ").map(str::trim) {
                if !tok.is_empty() {
                    let sha = sha256_token(tok);
                    let _ = state.db.token_delete(&sha).await;
                }
            }
        }
    }
    StatusCode::OK
}

pub(crate) fn mint_token() -> String {
    let bytes = sodiumoxide::randombytes::randombytes(32);
    base64::encode_config(bytes, base64::URL_SAFE_NO_PAD)
}