groupchat/server/src/handlers/nostr_auth.rs

use axum::{extract::State, http::StatusCode, Json};
use jsonwebtoken::{decode, encode, DecodingKey, EncodingKey, Header, Validation};
use nostr::prelude::*;
use std::sync::Arc;
use uuid::Uuid;

use crate::{
    middleware::auth::create_token,
    models::{AuthResponse, NostrChallengeResponse, NostrVerifyRequest, UserPublic},
    AppState,
};

#[derive(Debug, serde::Serialize, serde::Deserialize)]
struct ChallengeClaims {
    pub nonce: String,
    pub exp: usize,
}

/// GET /api/auth/nostr/challenge — return a short-lived JWT containing a random nonce
pub async fn challenge(
    State(state): State<Arc<AppState>>,
) -> Result<Json<NostrChallengeResponse>, (StatusCode, String)> {
    // Generate 32 random bytes as hex nonce
    let mut nonce_bytes = [0u8; 32];
    use rand::RngCore;
    rand::thread_rng().fill_bytes(&mut nonce_bytes);
    let nonce = hex::encode(&nonce_bytes);

    let exp = (chrono::Utc::now().timestamp() + 120) as usize; // 2 minutes

    let claims = ChallengeClaims {
        nonce,
        exp,
    };

    let token = encode(
        &Header::default(),
        &claims,
        &EncodingKey::from_secret(state.jwt_secret.as_bytes()),
    )
    .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;

    Ok(Json(NostrChallengeResponse { challenge: token }))
}

/// Simple hex encoder (avoid adding the `hex` crate just for this)
mod hex {
    pub fn encode(bytes: &[u8]) -> String {
        bytes.iter().map(|b| format!("{:02x}", b)).collect()
    }
}

/// POST /api/auth/nostr/verify — verify signed event, create/login user
pub async fn verify(
    State(state): State<Arc<AppState>>,
    Json(body): Json<NostrVerifyRequest>,
) -> Result<Json<AuthResponse>, (StatusCode, String)> {
    // 1. Decode challenge JWT, verify not expired, extract nonce
    let challenge_data = decode::<ChallengeClaims>(
        &body.challenge,
        &DecodingKey::from_secret(state.jwt_secret.as_bytes()),
        &Validation::default(),
    )
    .map_err(|_| (StatusCode::BAD_REQUEST, "Invalid or expired challenge".to_string()))?;

    let nonce = &challenge_data.claims.nonce;

    // 2. Deserialize signed_event as nostr::Event
    let event: Event = serde_json::from_str(&body.signed_event)
        .map_err(|e| (StatusCode::BAD_REQUEST, format!("Invalid event JSON: {}", e)))?;

    // 3. Verify Schnorr signature
    if !event.verify_signature() {
        return Err((StatusCode::UNAUTHORIZED, "Invalid event signature".to_string()));
    }

    // 4. Verify event.content == nonce
    if event.content.as_str() != nonce.as_str() {
        return Err((StatusCode::BAD_REQUEST, "Nonce mismatch".to_string()));
    }

    // 5. Verify event.created_at within 5 minutes
    let now = chrono::Utc::now().timestamp() as u64;
    let event_ts = event.created_at.as_secs();
    if now.abs_diff(event_ts) > 300 {
        return Err((StatusCode::BAD_REQUEST, "Event timestamp too far off".to_string()));
    }

    // 6. Extract pubkey hex
    let pubkey_hex = event.pubkey.to_hex();

    // 7. Lookup user by nostr_pubkey
    let existing = sqlx::query_as::<_, (String, String, String, Option<String>)>(
        "SELECT id, email, display_name, avatar_url FROM users WHERE nostr_pubkey = ?",
    )
    .bind(&pubkey_hex)
    .fetch_optional(&state.db)
    .await
    .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;

    let (user_id, email, display_name, avatar_url) = if let Some(user) = existing {
        // Update avatar if provided and user doesn't have a custom one
        if let Some(ref pic) = body.profile_picture {
            if user.3.is_none() || user.3.as_deref() == Some("") {
                let _ = sqlx::query("UPDATE users SET avatar_url = ? WHERE id = ?")
                    .bind(pic)
                    .bind(&user.0)
                    .execute(&state.db)
                    .await;
                (user.0, user.1, user.2, Some(pic.clone()))
            } else {
                user
            }
        } else {
            user
        }
    } else {
        // Create new user
        let user_id = Uuid::new_v4().to_string();
        let sentinel_email = format!("nostr:{}", &pubkey_hex[..16]);
        let display_name = body
            .profile_name
            .clone()
            .filter(|n| !n.trim().is_empty())
            .unwrap_or_else(|| {
                let npub = PublicKey::from_hex(&pubkey_hex)
                    .map(|pk| pk.to_bech32().unwrap_or_default())
                    .unwrap_or_default();
                if npub.len() > 8 {
                    format!("npub...{}", &npub[npub.len() - 8..])
                } else {
                    format!("nostr-{}", &pubkey_hex[..8])
                }
            });

        let avatar_url = body.profile_picture.clone();

        sqlx::query(
            "INSERT INTO users (id, email, display_name, password_hash, nostr_pubkey, avatar_url) VALUES (?, ?, ?, ?, ?, ?)",
        )
        .bind(&user_id)
        .bind(&sentinel_email)
        .bind(&display_name)
        .bind("")
        .bind(&pubkey_hex)
        .bind(&avatar_url)
        .execute(&state.db)
        .await
        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;

        (user_id, sentinel_email, display_name, avatar_url)
    };

    // 8. Issue JWT, return AuthResponse
    let token = create_token(&user_id, &email, &display_name, &state.jwt_secret)
        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;

    Ok(Json(AuthResponse {
        token,
        user: UserPublic {
            id: user_id,
            email: crate::models::public_email(&email),
            display_name,
            avatar_url,
        },
    }))
}