added more support for supabase-js

auth/src/mfa.rs

@@ -0,0 +1,205 @@
+use axum::{
+    extract::State,
+    http::StatusCode,
+    response::{IntoResponse, Json},
+    Extension,
+};
+use common::ProjectContext;
+use serde::{Deserialize, Serialize};
+use sqlx::{PgPool, Row};
+use totp_rs::{Algorithm, Secret, TOTP};
+use uuid::Uuid;
+use crate::middleware::AuthContext;
+use crate::handlers::AuthState;
+
+#[derive(Serialize)]
+pub struct EnrollResponse {
+    pub id: Uuid,
+    pub type_: String,
+    pub totp: TotpResponse,
+}
+
+#[derive(Serialize)]
+pub struct TotpResponse {
+    pub qr_code: String, // SVG or PNG base64
+    pub secret: String,
+    pub uri: String,
+}
+
+#[derive(Deserialize)]
+pub struct VerifyRequest {
+    pub factor_id: Uuid,
+    pub code: String,
+    pub challenge_id: Option<Uuid>, // For future use
+}
+
+#[derive(Serialize)]
+pub struct VerifyResponse {
+    pub access_token: String, // Potentially upgraded token
+    pub token_type: String,
+    pub expires_in: usize,
+    pub refresh_token: String,
+    pub user: serde_json::Value,
+}
+
+// Enroll MFA (Generate Secret & QR)
+pub async fn enroll(
+    State(state): State<AuthState>,
+    Extension(auth_ctx): Extension<AuthContext>,
+    Extension(project_ctx): Extension<ProjectContext>,
+) -> Result<impl IntoResponse, (StatusCode, String)> {
+    let user_id = auth_ctx.claims.as_ref()
+        .and_then(|c| Uuid::parse_str(&c.sub).ok())
+        .ok_or((StatusCode::UNAUTHORIZED, "Invalid user".to_string()))?;
+
+    // 1. Generate TOTP Secret
+    let secret = Secret::generate_secret();
+    let totp = TOTP::new(
+        Algorithm::SHA1,
+        6,
+        1,
+        30,
+        secret.to_bytes().unwrap(),
+        Some(project_ctx.project_ref.clone()), // Issuer
+        auth_ctx.claims.as_ref().and_then(|c| c.email.clone()).unwrap_or("user".to_string()), // Account Name
+    ).unwrap();
+
+    let secret_str = totp.get_secret_base32();
+    let qr_code = totp.get_qr_base64().map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e))?;
+    let uri = totp.get_url();
+
+    // 2. Store in DB (Unverified)
+    let row = sqlx::query(
+        "INSERT INTO auth.mfa_factors (user_id, factor_type, secret, status) VALUES ($1, 'totp', $2, 'unverified') RETURNING id"
+    )
+    .bind(user_id)
+    .bind(&secret_str)
+    .fetch_one(&state.db)
+    .await
+    .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
+
+    let factor_id: Uuid = row.get("id");
+
+    Ok(Json(EnrollResponse {
+        id: factor_id,
+        type_: "totp".to_string(),
+        totp: TotpResponse {
+            qr_code,
+            secret: secret_str,
+            uri,
+        }
+    }))
+}
+
+// Verify MFA (Activate Factor)
+pub async fn verify(
+    State(state): State<AuthState>,
+    Extension(auth_ctx): Extension<AuthContext>,
+    Extension(_project_ctx): Extension<ProjectContext>,
+    Json(payload): Json<VerifyRequest>,
+) -> Result<impl IntoResponse, (StatusCode, String)> {
+    let user_id = auth_ctx.claims.as_ref()
+        .and_then(|c| Uuid::parse_str(&c.sub).ok())
+        .ok_or((StatusCode::UNAUTHORIZED, "Invalid user".to_string()))?;
+
+    // 1. Fetch Factor
+    let row = sqlx::query(
+        "SELECT secret, status FROM auth.mfa_factors WHERE id = $1 AND user_id = $2"
+    )
+    .bind(payload.factor_id)
+    .bind(user_id)
+    .fetch_optional(&state.db)
+    .await
+    .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?
+    .ok_or((StatusCode::NOT_FOUND, "Factor not found".to_string()))?;
+
+    let secret_str: String = row.get("secret");
+    let status: String = row.get("status");
+
+    // 2. Validate Code
+    let secret_bytes = base32::decode(base32::Alphabet::RFC4648 { padding: false }, &secret_str)
+        .ok_or((StatusCode::INTERNAL_SERVER_ERROR, "Invalid secret format".to_string()))?;
+
+    let totp = TOTP::new(
+        Algorithm::SHA1,
+        6,
+        1,
+        30,
+        secret_bytes,
+        None,
+        "".to_string(),
+    ).unwrap();
+
+    let is_valid = totp.check_current(&payload.code).unwrap_or(false);
+
+    if !is_valid {
+        return Err((StatusCode::BAD_REQUEST, "Invalid code".to_string()));
+    }
+
+    // 3. Update Status if Unverified
+    if status == "unverified" {
+        sqlx::query("UPDATE auth.mfa_factors SET status = 'verified', updated_at = now() WHERE id = $1")
+            .bind(payload.factor_id)
+            .execute(&state.db)
+            .await
+            .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
+    }
+
+    // 4. Return Success (In a real scenario, this might return an upgraded JWT with `aal: 2`)
+    // For now, we just confirm verification.
+    
+    Ok(Json(serde_json::json!({
+        "status": "verified",
+        "factor_id": payload.factor_id
+    })))
+}
+
+// Challenge (Login with MFA)
+pub async fn challenge(
+    State(state): State<AuthState>,
+    Extension(auth_ctx): Extension<AuthContext>,
+    Json(payload): Json<VerifyRequest>,
+) -> Result<impl IntoResponse, (StatusCode, String)> {
+    // This is essentially the same as verify for now, but semantically distinct.
+    // It implies checking a code against an ALREADY verified factor to allow login proceed.
+    
+    let user_id = auth_ctx.claims.as_ref()
+        .and_then(|c| Uuid::parse_str(&c.sub).ok())
+        .ok_or((StatusCode::UNAUTHORIZED, "Invalid user".to_string()))?;
+
+    let row = sqlx::query(
+        "SELECT secret FROM auth.mfa_factors WHERE id = $1 AND user_id = $2 AND status = 'verified'"
+    )
+    .bind(payload.factor_id)
+    .bind(user_id)
+    .fetch_optional(&state.db)
+    .await
+    .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?
+    .ok_or((StatusCode::BAD_REQUEST, "Factor not found or not verified".to_string()))?;
+
+    let secret_str: String = row.get("secret");
+    
+    let secret_bytes = base32::decode(base32::Alphabet::RFC4648 { padding: false }, &secret_str)
+        .ok_or((StatusCode::INTERNAL_SERVER_ERROR, "Invalid secret format".to_string()))?;
+
+    let totp = TOTP::new(
+        Algorithm::SHA1,
+        6,
+        1,
+        30,
+        secret_bytes,
+        None,
+        "".to_string(),
+    ).unwrap();
+
+    let is_valid = totp.check_current(&payload.code).unwrap_or(false);
+
+    if !is_valid {
+        return Err((StatusCode::BAD_REQUEST, "Invalid code".to_string()));
+    }
+
+    Ok(Json(serde_json::json!({
+        "status": "success",
+        "factor_id": payload.factor_id
+    })))
+}