diff --git a/README.md b/README.md index dc6cc6b2..8330d953 100644 --- a/README.md +++ b/README.md @@ -52,7 +52,7 @@ wrangler secret put TWO_FACTOR_ENC_KEY - JWT_SECRET：访问令牌签名密钥 - JWT_REFRESH_SECRET：刷新令牌签名密钥 -- ALLOWED_EMAILS：首个账号注册白名单（仅在“数据库还没有任何用户”时启用），多个邮箱用英文逗号分隔 +- ALLOWED_EMAILS：注册白名单，多个邮箱用英文逗号分隔；可使用 `*` 表示允许所有邮箱注册 - TWO_FACTOR_ENC_KEY：可选，Base64 的 32 字节密钥；用于加密存储 TOTP 秘钥（不设置则以 `plain:` 形式存储） ### 4. 部署 diff --git a/sql/migrations/20260129_add_kdf_fields.sql b/sql/migrations/20260129_add_kdf_fields.sql new file mode 100644 index 00000000..46a2f8e0 --- /dev/null +++ b/sql/migrations/20260129_add_kdf_fields.sql @@ -0,0 +1,5 @@ +-- Add KDF memory and parallelism columns to users table +-- These are required for Argon2id support + +ALTER TABLE users ADD COLUMN kdf_memory INTEGER; +ALTER TABLE users ADD COLUMN kdf_parallelism INTEGER; diff --git a/sql/schema_full.sql b/sql/schema_full.sql index f8a01293..7d428aa4 100644 --- a/sql/schema_full.sql +++ b/sql/schema_full.sql @@ -72,6 +72,12 @@ CREATE TABLE IF NOT EXISTS devices ( FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE ); +CREATE TABLE IF NOT EXISTS login_rate_limits ( + key TEXT PRIMARY KEY NOT NULL, + count INTEGER NOT NULL, + reset_at INTEGER NOT NULL +); + CREATE INDEX IF NOT EXISTS idx_ciphers_user_id ON ciphers(user_id); CREATE INDEX IF NOT EXISTS idx_ciphers_folder_id ON ciphers(folder_id); CREATE INDEX IF NOT EXISTS idx_folders_user_id ON folders(user_id); diff --git a/src/auth.rs b/src/auth.rs index 77458ab3..fa329c99 100644 --- a/src/auth.rs +++ b/src/auth.rs @@ -2,7 +2,7 @@ use axum::{ extract::FromRequestParts, http::{header, request::Parts}, }; -use jsonwebtoken::{decode, DecodingKey, Validation}; +use jsonwebtoken::{decode, Algorithm, DecodingKey, Validation}; use serde::{Deserialize, Serialize}; use std::sync::Arc; use worker::Env; @@ -45,7 +45,11 @@ impl FromRequestParts> for Claims // Decode and validate the token let decoding_key = DecodingKey::from_secret(secret.to_string().as_ref()); - let token_data = decode::(&token, &decoding_key, &Validation::default()) + let mut validation = Validation::new(Algorithm::HS256); + validation.validate_nbf = true; + validation.leeway = 60; + + let token_data = decode::(&token, &decoding_key, &validation) .map_err(|_| AppError::Unauthorized("Invalid token".to_string()))?; Ok(token_data.claims) diff --git a/src/db.rs b/src/db.rs index 7416df11..812c10f8 100644 --- a/src/db.rs +++ b/src/db.rs @@ -3,5 +3,37 @@ use std::sync::Arc; use worker::{D1Database, Env}; pub fn get_db(env: &Arc) -> Result { - env.d1("vault1").map_err(AppError::Worker) + // 尝试获取名为 "warden-mima" 的数据库绑定 (根据 wrangler.jsonc) + // 如果失败,尝试旧的 "vault1" 以保持兼容性,或者报错 + env.d1("warden-mima") + .or_else(|_| env.d1("vault1")) + .map_err(|e| { + log::error!("❌ Failed to get database binding: {}", e); + AppError::Worker(e) + }) +} + +/// 将 worker::Error 转换为更有意义的 AppError +pub fn handle_db_error(error: worker::Error) -> AppError { + let error_str = error.to_string(); + + // 记录原始错误以便调试 + log::error!("🗄️ Raw Database Error: {}", error_str); + + if error_str.contains("UNIQUE constraint failed") { + if error_str.contains("email") { + return AppError::DatabaseConstraint("Email already registered".to_string()); + } + return AppError::DatabaseConstraint(format!("Record already exists: {}", error_str)); + } + + if error_str.contains("NOT NULL constraint failed") { + return AppError::BadRequest(format!("Missing required field: {}", error_str)); + } + + if error_str.contains("FOREIGN KEY constraint failed") { + return AppError::BadRequest("Invalid reference".to_string()); + } + + AppError::Database(error_str) } diff --git a/src/error.rs b/src/error.rs index 56a819ec..c99dd6f5 100644 --- a/src/error.rs +++ b/src/error.rs @@ -9,8 +9,11 @@ pub enum AppError { #[error("Worker error: {0}")] Worker(#[from] worker::Error), - #[error("Database query failed")] - Database, + #[error("Database query failed: {0}")] + Database(String), + + #[error("Database constraint violation: {0}")] + DatabaseConstraint(String), #[error("Not found: {0}")] NotFound(String), @@ -21,39 +24,59 @@ pub enum AppError { #[error("Unauthorized: {0}")] Unauthorized(String), + #[error("Too many requests: {0}")] + TooManyRequests(String), + #[error("Cryptography error: {0}")] Crypto(String), #[error(transparent)] JsonWebToken(#[from] jsonwebtoken::errors::Error), - #[error("Internal server error")] - Internal, + #[error("Internal server error: {0}")] + Internal(String), } impl IntoResponse for AppError { fn into_response(self) -> Response { - let (status, error_message) = match self { - AppError::Worker(e) => ( - StatusCode::INTERNAL_SERVER_ERROR, - format!("Worker error: {}", e), - ), - AppError::Database => ( - StatusCode::INTERNAL_SERVER_ERROR, - "Database error".to_string(), - ), - AppError::NotFound(msg) => (StatusCode::NOT_FOUND, msg), - AppError::BadRequest(msg) => (StatusCode::BAD_REQUEST, msg), - AppError::Unauthorized(msg) => (StatusCode::UNAUTHORIZED, msg), - AppError::Crypto(msg) => ( - StatusCode::INTERNAL_SERVER_ERROR, - format!("Crypto error: {}", msg), - ), + let (status, error_message) = match &self { + AppError::Worker(e) => { + log::error!("🔴 Worker error: {}", e); + ( + StatusCode::INTERNAL_SERVER_ERROR, + format!("Worker error: {}", e), + ) + } + AppError::Database(msg) => { + log::error!("💾 Database error: {}", msg); + ( + StatusCode::INTERNAL_SERVER_ERROR, + "Database error".to_string(), + ) + } + AppError::DatabaseConstraint(msg) => { + log::warn!("⚠️ Database constraint violation: {}", msg); + (StatusCode::BAD_REQUEST, msg.clone()) + } + AppError::NotFound(msg) => (StatusCode::NOT_FOUND, msg.clone()), + AppError::BadRequest(msg) => (StatusCode::BAD_REQUEST, msg.clone()), + AppError::Unauthorized(msg) => (StatusCode::UNAUTHORIZED, msg.clone()), + AppError::TooManyRequests(msg) => (StatusCode::TOO_MANY_REQUESTS, msg.clone()), + AppError::Crypto(msg) => { + log::error!("🔐 Crypto error: {}", msg); + ( + StatusCode::INTERNAL_SERVER_ERROR, + format!("Crypto error: {}", msg), + ) + } AppError::JsonWebToken(_) => (StatusCode::UNAUTHORIZED, "Invalid token".to_string()), - AppError::Internal => ( - StatusCode::INTERNAL_SERVER_ERROR, - "Internal server error".to_string(), - ), + AppError::Internal(msg) => { + log::error!("🔥 Internal error: {}", msg); + ( + StatusCode::INTERNAL_SERVER_ERROR, + "Internal server error".to_string(), + ) + } }; let body = Json(json!({ "error": error_message })); diff --git a/src/handlers/accounts.rs b/src/handlers/accounts.rs index b691345d..9b6eceb9 100644 --- a/src/handlers/accounts.rs +++ b/src/handlers/accounts.rs @@ -29,6 +29,10 @@ pub struct ChangeMasterPasswordRequest { pub kdf: Option, #[serde(default)] pub kdf_iterations: Option, + #[serde(default)] + pub kdf_memory: Option, + #[serde(default)] + pub kdf_parallelism: Option, } #[derive(Debug, Deserialize)] @@ -42,6 +46,10 @@ pub struct ChangeEmailRequest { pub kdf: Option, #[serde(default)] pub kdf_iterations: Option, + #[serde(default)] + pub kdf_memory: Option, + #[serde(default)] + pub kdf_parallelism: Option, } #[worker::send] @@ -56,7 +64,7 @@ pub async fn profile( "SELECT * FROM users WHERE id = ?1", claims.sub ) - .map_err(|_| AppError::Database)? + .map_err(|e| AppError::Database(e.to_string()))? .first(None) .await? .ok_or(AppError::NotFound("User not found".to_string()))?; @@ -96,19 +104,40 @@ pub async fn prelogin( .ok_or_else(|| AppError::BadRequest("Missing email".to_string()))?; let db = db::get_db(&env)?; - let stmt = db.prepare("SELECT kdf_iterations FROM users WHERE email = ?1"); + let stmt = db.prepare("SELECT kdf_type, kdf_iterations, kdf_memory, kdf_parallelism FROM users WHERE email = ?1"); let query = stmt.bind(&[email.into()])?; - let kdf_iterations: Option = query - .first(Some("kdf_iterations")) + let user_kdf: Option = query + .first(None) .await - .map_err(|_| AppError::Database)?; + .map_err(|e| AppError::Database(e.to_string()))?; + + if let Some(user) = user_kdf { + let kdf_type = user.get("kdf_type").and_then(|v| v.as_i64()).unwrap_or(0) as i32; + let kdf_iterations = user.get("kdf_iterations").and_then(|v| v.as_i64()).unwrap_or(600_000) as i32; + let mut kdf_memory = user.get("kdf_memory").and_then(|v| v.as_i64()).map(|v| v as i32); + let mut kdf_parallelism = user.get("kdf_parallelism").and_then(|v| v.as_i64()).map(|v| v as i32); + + // Fallback for Argon2id if values are missing (legacy users) + if kdf_type == 1 { + if kdf_memory.is_none() { kdf_memory = Some(64 * 1024); } + if kdf_parallelism.is_none() { kdf_parallelism = Some(4); } + } - Ok(Json(PreloginResponse { - kdf: 0, // PBKDF2 - kdf_iterations: kdf_iterations.unwrap_or(600_000), - kdf_memory: None, - kdf_parallelism: None, - })) + Ok(Json(PreloginResponse { + kdf: kdf_type, + kdf_iterations, + kdf_memory, + kdf_parallelism, + })) + } else { + // User not found, return default (PBKDF2) to prevent user enumeration + Ok(Json(PreloginResponse { + kdf: 0, + kdf_iterations: 600_000, + kdf_memory: None, + kdf_parallelism: None, + })) + } } #[worker::send] @@ -117,32 +146,48 @@ pub async fn register( Json(payload): Json, ) -> Result, AppError> { let db = db::get_db(&env)?; - let user_count: Option = db - .prepare("SELECT COUNT(1) AS user_count FROM users") - .first(Some("user_count")) + + // 1. 验证邮箱白名单 + let allowed_emails = env + .secret("ALLOWED_EMAILS") + .map_err(|e| { + log::error!("Failed to read ALLOWED_EMAILS secret: {:?}", e); + AppError::Internal("Configuration error".to_string()) + })?; + let allowed_emails = allowed_emails + .as_ref() + .as_string() + .ok_or_else(|| AppError::Internal("Invalid ALLOWED_EMAILS format".to_string()))?; + + let email = payload.email.to_lowercase(); + let allowed = allowed_emails + .split(',') + .map(|s| s.trim().to_lowercase()) + .any(|s| s == "*" || s == email); + + if !allowed { + log::warn!("📧 Registration attempt with unauthorized email: {}", email); + return Err(AppError::Unauthorized("Not allowed to signup".to_string())); + } + + // 2. 检查邮箱是否已存在 (主动检查,避免依赖数据库错误) + let existing: Option = db + .prepare("SELECT id FROM users WHERE email = ?1") + .bind(&[email.clone().into()])? + .first(Some("id")) .await - .map_err(|_| AppError::Database)?; - let user_count = user_count.unwrap_or(0); - if user_count == 0 { - let allowed_emails = env - .secret("ALLOWED_EMAILS") - .map_err(|_| AppError::Internal)?; - let allowed_emails = allowed_emails - .as_ref() - .as_string() - .ok_or_else(|| AppError::Internal)?; - if allowed_emails - .split(",") - .all(|email| email.trim() != payload.email) - { - return Err(AppError::Unauthorized("Not allowed to signup".to_string())); - } + .map_err(|e| db::handle_db_error(e))?; + + if existing.is_some() { + log::info!("📧 Email already registered: {}", email); + return Err(AppError::BadRequest("Email already registered".to_string())); } + let now = Utc::now().to_rfc3339(); let user = User { id: Uuid::new_v4().to_string(), name: payload.name, - email: payload.email.to_lowercase(), + email: email.clone(), email_verified: false, master_password_hash: payload.master_password_hash, master_password_hint: payload.master_password_hint, @@ -151,15 +196,17 @@ pub async fn register( public_key: payload.user_asymmetric_keys.public_key, kdf_type: payload.kdf, kdf_iterations: payload.kdf_iterations, + kdf_memory: payload.kdf_memory, + kdf_parallelism: payload.kdf_parallelism, security_stamp: Uuid::new_v4().to_string(), created_at: now.clone(), updated_at: now, }; - let query = query!( + query!( &db, - "INSERT INTO users (id, name, email, master_password_hash, key, private_key, public_key, kdf_iterations, security_stamp, created_at, updated_at) - VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11)", + "INSERT INTO users (id, name, email, master_password_hash, key, private_key, public_key, kdf_iterations, security_stamp, created_at, updated_at, kdf_type, kdf_memory, kdf_parallelism) + VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14)", user.id, user.name, user.email, @@ -170,16 +217,17 @@ pub async fn register( user.kdf_iterations, user.security_stamp, user.created_at, - user.updated_at - ).map_err(|error|{ - AppError::Database - })? + user.updated_at, + user.kdf_type, + user.kdf_memory, + user.kdf_parallelism + ) + .map_err(|e| db::handle_db_error(e))? .run() .await - .map_err(|error|{ - AppError::Database - })?; + .map_err(|e| db::handle_db_error(e))?; + log::info!("✅ User registered successfully: {} ({})", user.email, user.id); Ok(Json(json!({}))) } @@ -202,9 +250,9 @@ pub async fn change_master_password( .bind(&[claims.sub.clone().into()])? .first(None) .await - .map_err(|_| AppError::Database)? + .map_err(|e| AppError::Database(e.to_string()))? .ok_or_else(|| AppError::NotFound("User not found".to_string()))?; - let user: User = serde_json::from_value(user).map_err(|_| AppError::Internal)?; + let user: User = serde_json::from_value(user).map_err(|e| AppError::Internal(e.to_string()))?; if !constant_time_eq( user.master_password_hash.as_bytes(), @@ -228,9 +276,11 @@ pub async fn change_master_password( .unwrap_or_else(|| user.public_key.clone()); let kdf_type = payload.kdf.unwrap_or(user.kdf_type); let kdf_iterations = payload.kdf_iterations.unwrap_or(user.kdf_iterations); + let kdf_memory = payload.kdf_memory.or(user.kdf_memory); + let kdf_parallelism = payload.kdf_parallelism.or(user.kdf_parallelism); db.prepare( - "UPDATE users SET master_password_hash = ?1, master_password_hint = ?2, key = ?3, private_key = ?4, public_key = ?5, kdf_type = ?6, kdf_iterations = ?7, security_stamp = ?8, updated_at = ?9 WHERE id = ?10", + "UPDATE users SET master_password_hash = ?1, master_password_hint = ?2, key = ?3, private_key = ?4, public_key = ?5, kdf_type = ?6, kdf_iterations = ?7, security_stamp = ?8, updated_at = ?9, kdf_memory = ?10, kdf_parallelism = ?11 WHERE id = ?12", ) .bind(&[ payload.new_master_password_hash.into(), @@ -242,11 +292,13 @@ pub async fn change_master_password( kdf_iterations.into(), security_stamp.into(), now.into(), + to_js_val(kdf_memory), + to_js_val(kdf_parallelism), claims.sub.into(), ])? .run() .await - .map_err(|_| AppError::Database)?; + .map_err(|e| AppError::Database(e.to_string()))?; Ok(Json(json!({}))) } @@ -275,9 +327,9 @@ pub async fn change_email( .bind(&[claims.sub.clone().into()])? .first(None) .await - .map_err(|_| AppError::Database)? + .map_err(|e| AppError::Database(e.to_string()))? .ok_or_else(|| AppError::NotFound("User not found".to_string()))?; - let user: User = serde_json::from_value(user).map_err(|_| AppError::Internal)?; + let user: User = serde_json::from_value(user).map_err(|e| AppError::Internal(e.to_string()))?; if !constant_time_eq( user.master_password_hash.as_bytes(), @@ -290,9 +342,11 @@ pub async fn change_email( let security_stamp = Uuid::new_v4().to_string(); let kdf_type = payload.kdf.unwrap_or(user.kdf_type); let kdf_iterations = payload.kdf_iterations.unwrap_or(user.kdf_iterations); + let kdf_memory = payload.kdf_memory.or(user.kdf_memory); + let kdf_parallelism = payload.kdf_parallelism.or(user.kdf_parallelism); db.prepare( - "UPDATE users SET email = ?1, email_verified = ?2, master_password_hash = ?3, key = ?4, kdf_type = ?5, kdf_iterations = ?6, security_stamp = ?7, updated_at = ?8 WHERE id = ?9", + "UPDATE users SET email = ?1, email_verified = ?2, master_password_hash = ?3, key = ?4, kdf_type = ?5, kdf_iterations = ?6, security_stamp = ?7, updated_at = ?8, kdf_memory = ?9, kdf_parallelism = ?10 WHERE id = ?11", ) .bind(&[ new_email.into(), @@ -303,15 +357,18 @@ pub async fn change_email( kdf_iterations.into(), security_stamp.into(), now.into(), + to_js_val(kdf_memory), + to_js_val(kdf_parallelism), claims.sub.into(), ])? .run() .await .map_err(|e| { - if e.to_string().contains("UNIQUE") { + let error_str = e.to_string(); + if error_str.contains("UNIQUE") { AppError::BadRequest("Email already in use".to_string()) } else { - AppError::Database + AppError::Database(error_str) } })?; diff --git a/src/handlers/ciphers.rs b/src/handlers/ciphers.rs index 6b7d0e5d..fcaf28ef 100644 --- a/src/handlers/ciphers.rs +++ b/src/handlers/ciphers.rs @@ -10,6 +10,28 @@ use crate::error::AppError; use crate::models::cipher::{Cipher, CipherData, CipherRequestData, CreateCipherRequest, CipherRequestFlat}; use axum::extract::Path; +async fn ensure_folder_owned_by_user( + db: &worker::D1Database, + user_id: &str, + folder_id: &str, +) -> Result<(), AppError> { + let exists: Option = query!( + db, + "SELECT 1 AS ok FROM folders WHERE id = ?1 AND user_id = ?2 LIMIT 1", + folder_id, + user_id + ) + .map_err(|e| AppError::Database(e.to_string()))? + .first(Some("ok")) + .await + .map_err(|e| AppError::Database(e.to_string()))?; + + if exists.is_none() { + return Err(AppError::NotFound("Folder not found".to_string())); + } + Ok(()) +} + async fn create_cipher_inner( claims: Claims, env: &Arc, @@ -20,6 +42,10 @@ async fn create_cipher_inner( let now = Utc::now(); let now = now.format("%Y-%m-%dT%H:%M:%S%.3fZ").to_string(); + if let Some(folder_id) = cipher_data_req.folder_id.as_deref() { + ensure_folder_owned_by_user(&db, &claims.sub, folder_id).await?; + } + let cipher_data = CipherData { name: cipher_data_req.name, notes: cipher_data_req.notes, @@ -32,7 +58,7 @@ async fn create_cipher_inner( reprompt: cipher_data_req.reprompt, }; - let data_value = serde_json::to_value(&cipher_data).map_err(|_| AppError::Internal)?; + let data_value = serde_json::to_value(&cipher_data).map_err(|e| AppError::Internal(e.to_string()))?; let cipher = Cipher { id: Uuid::new_v4().to_string(), @@ -56,7 +82,7 @@ async fn create_cipher_inner( }, }; - let data = serde_json::to_string(&cipher.data).map_err(|_| AppError::Internal)?; + let data = serde_json::to_string(&cipher.data).map_err(|e| AppError::Internal(e.to_string()))?; query!( &db, @@ -71,7 +97,7 @@ async fn create_cipher_inner( cipher.folder_id, cipher.created_at, cipher.updated_at, - ).map_err(|_|AppError::Database)? + ).map_err(|e| AppError::Database(e.to_string()))? .run() .await?; @@ -113,13 +139,17 @@ pub async fn update_cipher( id, claims.sub ) - .map_err(|_| AppError::Database)? + .map_err(|e| AppError::Database(e.to_string()))? .first(None) .await? .ok_or(AppError::NotFound("Cipher not found".to_string()))?; let cipher_data_req = payload; + if let Some(folder_id) = cipher_data_req.folder_id.as_deref() { + ensure_folder_owned_by_user(&db, &claims.sub, folder_id).await?; + } + let cipher_data = CipherData { name: cipher_data_req.name, notes: cipher_data_req.notes, @@ -132,7 +162,7 @@ pub async fn update_cipher( reprompt: cipher_data_req.reprompt, }; - let data_value = serde_json::to_value(&cipher_data).map_err(|_| AppError::Internal)?; + let data_value = serde_json::to_value(&cipher_data).map_err(|e| AppError::Internal(e.to_string()))?; let cipher = Cipher { id: id.clone(), @@ -152,7 +182,7 @@ pub async fn update_cipher( collection_ids: None, }; - let data = serde_json::to_string(&cipher.data).map_err(|_| AppError::Internal)?; + let data = serde_json::to_string(&cipher.data).map_err(|e| AppError::Internal(e.to_string()))?; query!( &db, @@ -165,7 +195,7 @@ pub async fn update_cipher( cipher.updated_at, id, claims.sub, - ).map_err(|_|AppError::Database)? + ).map_err(|e| AppError::Database(e.to_string()))? .run() .await?; @@ -186,7 +216,7 @@ pub async fn delete_cipher( id, claims.sub ) - .map_err(|_| AppError::Database)? + .map_err(|e| AppError::Database(e.to_string()))? .run() .await?; diff --git a/src/handlers/devices.rs b/src/handlers/devices.rs index 44f3e747..7c735b18 100644 --- a/src/handlers/devices.rs +++ b/src/handlers/devices.rs @@ -23,7 +23,7 @@ async fn ensure_devices_table(db: &worker::D1Database) -> Result<(), AppError> { ) .run() .await - .map_err(|_| AppError::Database)?; + .map_err(|e| AppError::Database(e.to_string()))?; let _ = db .prepare("ALTER TABLE devices ADD COLUMN remember_token_hash TEXT") @@ -63,7 +63,7 @@ pub async fn knowndevice( .bind(&[email.into()])? .first(Some("id")) .await - .map_err(|_| AppError::Database)?; + .map_err(|e| AppError::Database(e.to_string()))?; let Some(user_id) = user_id else { return Ok(Json(json!(false))); @@ -74,7 +74,7 @@ pub async fn knowndevice( .bind(&[user_id.into(), device_identifier.into()])? .first(Some("ok")) .await - .map_err(|_| AppError::Database)?; + .map_err(|e| AppError::Database(e.to_string()))?; Ok(Json(json!(exists.is_some()))) } diff --git a/src/handlers/folders.rs b/src/handlers/folders.rs index d7e82c83..a18b59c8 100644 --- a/src/handlers/folders.rs +++ b/src/handlers/folders.rs @@ -37,7 +37,7 @@ pub async fn create_folder( folder.created_at, folder.updated_at ) - .map_err(|_| AppError::Database)? + .map_err(|e| AppError::Database(e.to_string()))? .run() .await?; @@ -65,7 +65,7 @@ pub async fn delete_folder( id, claims.sub ) - .map_err(|_| AppError::Database)? + .map_err(|e| AppError::Database(e.to_string()))? .run() .await?; @@ -88,7 +88,7 @@ pub async fn update_folder( id, claims.sub ) - .map_err(|_| AppError::Database)? + .map_err(|e| AppError::Database(e.to_string()))? .first(None) .await? .ok_or(AppError::NotFound("Folder not found".to_string()))?; @@ -109,7 +109,7 @@ pub async fn update_folder( folder.id, folder.user_id ) - .map_err(|_| AppError::Database)? + .map_err(|e| AppError::Database(e.to_string()))? .run() .await?; diff --git a/src/handlers/identity.rs b/src/handlers/identity.rs index 2fdc60bb..3bab7662 100644 --- a/src/handlers/identity.rs +++ b/src/handlers/identity.rs @@ -1,4 +1,4 @@ -use axum::{extract::State, response::IntoResponse, Form, Json}; +use axum::{extract::State, http::HeaderMap, response::IntoResponse, Form, Json}; use axum::http::StatusCode; use axum::response::Response; use chrono::{Duration, Utc}; @@ -103,14 +103,17 @@ fn generate_tokens_and_response( &EncodingKey::from_secret(jwt_refresh_secret.as_ref()), )?; - Ok(json!({ + let kdf_memory = user.kdf_memory.or_else(|| if user.kdf_type == 1 { Some(64 * 1024) } else { None }); + let kdf_parallelism = user.kdf_parallelism.or_else(|| if user.kdf_type == 1 { Some(4) } else { None }); + + let response = json!({ "ForcePasswordReset": false, "Kdf": user.kdf_type, "KdfIterations": user.kdf_iterations, - "KdfMemory": null, - "KdfParallelism": null, + "KdfMemory": kdf_memory, + "KdfParallelism": kdf_parallelism, "Key": user.key, - "MasterPasswordPolicy": { "Object": "masterPasswordPolicy" }, + "MasterPasswordPolicy": { "Object": "masterPasswordPolicy", "object": "masterPasswordPolicy" }, "PrivateKey": user.private_key, "ResetMasterPassword": false, "UserDecryptionOptions": { @@ -119,29 +122,35 @@ fn generate_tokens_and_response( "Kdf": { "KdfType": user.kdf_type, "Iterations": user.kdf_iterations, - "Memory": null, - "Parallelism": null + "Memory": kdf_memory, + "Parallelism": kdf_parallelism }, "MasterKeyEncryptedUserKey": user.key, "MasterKeyWrappedUserKey": user.key, "Salt": user.email }, - "Object": "userDecryptionOptions" + "Object": "userDecryptionOptions", + "object": "userDecryptionOptions" }, "AccountKeys": { "publicKeyEncryptionKeyPair": { "wrappedPrivateKey": user.private_key, "publicKey": user.public_key, - "Object": "publicKeyEncryptionKeyPair" + "Object": "publicKeyEncryptionKeyPair", + "object": "publicKeyEncryptionKeyPair" }, - "Object": "privateKeys" + "Object": "privateKeys", + "object": "privateKeys" }, "access_token": access_token, "expires_in": expires_in.num_seconds(), "refresh_token": refresh_token, "scope": "api offline_access", "token_type": "Bearer" - })) + }); + + log::info!("🔑 Generated token response for user {}: {}", user.id, response); + Ok(response) } async fn ensure_devices_table(db: &worker::D1Database) -> Result<(), AppError> { @@ -161,7 +170,7 @@ async fn ensure_devices_table(db: &worker::D1Database) -> Result<(), AppError> { ) .run() .await - .map_err(|_| AppError::Database)?; + .map_err(|e| AppError::Database(e.to_string()))?; let _ = db .prepare("ALTER TABLE devices ADD COLUMN remember_token_hash TEXT") @@ -171,6 +180,81 @@ async fn ensure_devices_table(db: &worker::D1Database) -> Result<(), AppError> { Ok(()) } +async fn ensure_login_rate_limits_table(db: &worker::D1Database) -> Result<(), AppError> { + db.prepare( + "CREATE TABLE IF NOT EXISTS login_rate_limits ( + key TEXT PRIMARY KEY NOT NULL, + count INTEGER NOT NULL, + reset_at INTEGER NOT NULL + )", + ) + .run() + .await + .map_err(|e| AppError::Database(e.to_string()))?; + Ok(()) +} + +fn client_ip(headers: &HeaderMap) -> String { + headers + .get("cf-connecting-ip") + .and_then(|v| v.to_str().ok()) + .map(|s| s.to_string()) + .or_else(|| { + headers + .get("x-forwarded-for") + .and_then(|v| v.to_str().ok()) + .and_then(|s| s.split(',').next()) + .map(|s| s.trim().to_string()) + }) + .unwrap_or_else(|| "unknown".to_string()) +} + +async fn enforce_login_rate_limit(db: &worker::D1Database, key: &str) -> Result<(), AppError> { + ensure_login_rate_limits_table(db).await?; + let now = Utc::now().timestamp_millis(); + let window_ms: i64 = 5 * 60 * 1000; + let max_attempts: i64 = 30; + + let row: Option = db + .prepare("SELECT count, reset_at FROM login_rate_limits WHERE key = ?1") + .bind(&[key.into()])? + .first(None) + .await + .map_err(|e| AppError::Database(e.to_string()))?; + + let (mut count, mut reset_at) = if let Some(row) = row { + let count = row.get("count").and_then(|v| v.as_i64()).unwrap_or(0); + let reset_at = row.get("reset_at").and_then(|v| v.as_i64()).unwrap_or(0); + (count, reset_at) + } else { + (0, 0) + }; + + if now >= reset_at { + count = 0; + reset_at = now + window_ms; + } + + count += 1; + + db.prepare( + "INSERT INTO login_rate_limits (key, count, reset_at) VALUES (?1, ?2, ?3) + ON CONFLICT(key) DO UPDATE SET count = excluded.count, reset_at = excluded.reset_at", + ) + .bind(&[key.into(), (count as f64).into(), (reset_at as f64).into()])? + .run() + .await + .map_err(|e| AppError::Database(e.to_string()))?; + + if count > max_attempts { + return Err(AppError::TooManyRequests( + "Too many login attempts. Please try again later.".to_string(), + )); + } + + Ok(()) +} + fn sha256_hex(input: &str) -> String { let mut hasher = Sha256::new(); hasher.update(input.as_bytes()); @@ -214,11 +298,16 @@ fn invalid_two_factor_response() -> Response { #[worker::send] pub async fn token( State(env): State>, + headers: HeaderMap, Form(payload): Form, ) -> Result { let db = db::get_db(&env)?; match payload.grant_type.as_str() { "password" => { + let ip = client_ip(&headers); + let rate_key = format!("login:{}", ip); + enforce_login_rate_limit(&db, &rate_key).await?; + let username = payload .username .ok_or_else(|| AppError::BadRequest("Missing username".to_string()))?; @@ -231,9 +320,9 @@ pub async fn token( .bind(&[username.to_lowercase().into()])? .first(None) .await - .map_err(|_| AppError::Unauthorized("Invalid credentials".to_string()))? + .map_err(|e| AppError::Database(e.to_string()))? .ok_or_else(|| AppError::Unauthorized("Invalid credentials".to_string()))?; - let user: User = serde_json::from_value(user).map_err(|_| AppError::Internal)?; + let user: User = serde_json::from_value(user).map_err(|e| AppError::Internal(e.to_string()))?; // Securely compare the provided hash with the stored hash if !constant_time_eq( user.master_password_hash.as_bytes(), @@ -264,7 +353,7 @@ pub async fn token( .bind(&[user.id.clone().into(), device_identifier.into()])? .first(None) .await - .map_err(|_| AppError::Database)?; + .map_err(|e| AppError::Database(e.to_string()))?; let stored_hash = row .and_then(|v| v.get("remember_token_hash").cloned()) .and_then(|v| v.as_str().map(|s| s.to_string())); @@ -282,7 +371,7 @@ pub async fn token( let secret_enc = two_factor::get_authenticator_secret_enc(&db, &user.id) .await? - .ok_or_else(|| AppError::Internal)?; + .ok_or_else(|| AppError::Internal("Missing authenticator secret".to_string()))?; let two_factor_key_b64 = env.secret("TWO_FACTOR_ENC_KEY").ok().map(|s| s.to_string()); let secret_encoded = two_factor::decrypt_secret_with_optional_key( @@ -332,7 +421,7 @@ pub async fn token( user_id.clone().into(), device_identifier.into(), device_name.clone().into(), - device_type.map(|v| v as i64).into(), + device_type.into(), remember_hash.clone().into(), now.clone().into(), now.into(), @@ -371,7 +460,7 @@ pub async fn token( .await .map_err(|_| AppError::Unauthorized("Invalid user".to_string()))? .ok_or_else(|| AppError::Unauthorized("Invalid user".to_string()))?; - let user: User = serde_json::from_value(user).map_err(|_| AppError::Internal)?; + let user: User = serde_json::from_value(user).map_err(|e| AppError::Internal(e.to_string()))?; let response = generate_tokens_and_response(user, &env)?; Ok(Json(response).into_response()) diff --git a/src/handlers/import.rs b/src/handlers/import.rs index d79313a3..5bf2812f 100644 --- a/src/handlers/import.rs +++ b/src/handlers/import.rs @@ -24,12 +24,15 @@ pub async fn import_data( let now = Utc::now(); let now = now.format("%Y-%m-%dT%H:%M:%S%.3fZ").to_string(); - let folder_query = "INSERT OR IGNORE INTO folders (id, user_id, name, created_at, updated_at) VALUES (?1, ?2, ?3, ?4, ?5)"; + let folder_query = "INSERT INTO folders (id, user_id, name, created_at, updated_at) VALUES (?1, ?2, ?3, ?4, ?5)"; + let mut folder_id_map: std::collections::HashMap = std::collections::HashMap::new(); let mut folder_stmts: Vec = Vec::new(); for import_folder in &payload.folders { + let new_id = Uuid::new_v4().to_string(); + folder_id_map.insert(import_folder.id.clone(), new_id.clone()); let folder = Folder { - id: import_folder.id.clone(), + id: new_id, user_id: claims.sub.clone(), name: import_folder.name.clone(), created_at: now.clone(), @@ -53,7 +56,7 @@ pub async fn import_data( for relationship in payload.folder_relationships { if let Some(cipher) = payload.ciphers.get_mut(relationship.key) { if let Some(folder) = payload.folders.get(relationship.value) { - cipher.folder_id = Some(folder.id.clone()); + cipher.folder_id = folder_id_map.get(&folder.id).cloned(); } } } @@ -80,7 +83,10 @@ pub async fn import_data( let id = Uuid::new_v4().to_string(); let user_id = claims.sub.clone(); - let data = serde_json::to_string(&cipher_data).map_err(|_| AppError::Internal)?; + let data = serde_json::to_string(&cipher_data).map_err(|e| AppError::Internal(e.to_string()))?; + let folder_id = import_cipher + .folder_id + .and_then(|id| folder_id_map.get(&id).cloned()); cipher_stmts.push(db.prepare(cipher_query).bind(&[ id.into(), @@ -89,7 +95,7 @@ pub async fn import_data( import_cipher.r#type.into(), data.into(), import_cipher.favorite.into(), - to_js_val(import_cipher.folder_id), + to_js_val(folder_id), now.clone().into(), now.clone().into(), ])?); @@ -113,6 +119,6 @@ async fn run_batch(db: &D1Database, stmts: &mut Vec) -> Res } let stmts = std::mem::take(stmts); - db.batch(stmts).await.map_err(|_| AppError::Database)?; + db.batch(stmts).await.map_err(|e| AppError::Database(e.to_string()))?; Ok(()) } diff --git a/src/handlers/sync.rs b/src/handlers/sync.rs index a58319ef..d5162de1 100644 --- a/src/handlers/sync.rs +++ b/src/handlers/sync.rs @@ -65,7 +65,7 @@ pub async fn get_sync_data( .collect::>(); let time = chrono::DateTime::parse_from_rfc3339(&user.created_at) - .map_err(|_| AppError::Internal)? + .map_err(|e| AppError::Internal(e.to_string()))? .to_rfc3339_opts(chrono::SecondsFormat::Micros, true); let profile = Profile { id: user.id, diff --git a/src/handlers/two_factor.rs b/src/handlers/two_factor.rs index 2c360195..27ac7740 100644 --- a/src/handlers/two_factor.rs +++ b/src/handlers/two_factor.rs @@ -54,7 +54,7 @@ pub async fn authenticator_request( .bind(&[claims.sub.clone().into()])? .first(Some("email")) .await - .map_err(|_| AppError::Database)?; + .map_err(|e| AppError::Database(e.to_string()))?; let user_email = user_email.ok_or_else(|| AppError::NotFound("User not found".to_string()))?; let secret = Secret::generate_secret(); @@ -82,13 +82,13 @@ pub async fn authenticator_request( 30, Secret::Encoded(secret_encoded.clone()) .to_bytes() - .map_err(|_| AppError::Internal)?, + .map_err(|e| AppError::Internal(e.to_string()))?, Some(issuer.clone()), account.clone(), ) - .map_err(|_| AppError::Internal)?; + .map_err(|e| AppError::Internal(e.to_string()))?; let otpauth = totp.get_url(); - let qr_base64 = totp.get_qr_base64().map_err(|_| AppError::Internal)?; + let qr_base64 = totp.get_qr_base64().map_err(|e| AppError::Internal(e.to_string()))?; Ok(Json(json!({ "secret": secret_encoded, diff --git a/src/models/user.rs b/src/models/user.rs index 58769532..872e5a62 100644 --- a/src/models/user.rs +++ b/src/models/user.rs @@ -14,6 +14,8 @@ pub struct User { pub public_key: String, pub kdf_type: i32, pub kdf_iterations: i32, + pub kdf_memory: Option, + pub kdf_parallelism: Option, pub security_stamp: String, pub created_at: String, pub updated_at: String, @@ -68,6 +70,8 @@ pub struct RegisterRequest { pub user_asymmetric_keys: KeyData, pub kdf: i32, pub kdf_iterations: i32, + pub kdf_memory: Option, + pub kdf_parallelism: Option, } #[derive(Debug, Deserialize)] diff --git a/src/two_factor.rs b/src/two_factor.rs index f3d2e74a..a6071189 100644 --- a/src/two_factor.rs +++ b/src/two_factor.rs @@ -24,7 +24,7 @@ pub async fn ensure_two_factor_authenticator_table(db: &D1Database) -> Result<() ) .run() .await - .map_err(|_| AppError::Database)?; + .map_err(|e| AppError::Database(e.to_string()))?; Ok(()) } @@ -35,7 +35,7 @@ pub async fn is_authenticator_enabled(db: &D1Database, user_id: &str) -> Result< .bind(&[user_id.into()])? .first(Some("enabled")) .await - .map_err(|_| AppError::Database)?; + .map_err(|e| AppError::Database(e.to_string()))?; Ok(matches!(enabled, Some(1))) } @@ -49,7 +49,7 @@ pub async fn get_authenticator_secret_enc( .bind(&[user_id.into()])? .first(Some("secret_enc")) .await - .map_err(|_| AppError::Database)?; + .map_err(|e| AppError::Database(e.to_string()))?; Ok(secret_enc) } @@ -75,7 +75,7 @@ pub async fn upsert_authenticator_secret( ])? .run() .await - .map_err(|_| AppError::Database)?; + .map_err(|e| AppError::Database(e.to_string()))?; Ok(()) } @@ -85,7 +85,7 @@ pub async fn disable_authenticator(db: &D1Database, user_id: &str) -> Result<(), .bind(&[user_id.into()])? .run() .await - .map_err(|_| AppError::Database)?; + .map_err(|e| AppError::Database(e.to_string()))?; Ok(()) } @@ -118,7 +118,7 @@ pub fn encrypt_secret_with_optional_key( aad: user_id.as_bytes(), }, ) - .map_err(|_| AppError::Internal)?; + .map_err(|e| AppError::Internal(e.to_string()))?; let mut blob = Vec::with_capacity(nonce_bytes.len() + ct.len()); blob.extend_from_slice(&nonce_bytes); @@ -135,7 +135,7 @@ pub fn decrypt_secret_with_optional_key( return Ok(rest.to_string()); } let Some(rest) = secret_enc.strip_prefix("gcm:") else { - return Err(AppError::Internal); + return Err(AppError::Internal("Invalid secret prefix".to_string())); }; let Some(key_b64) = two_factor_enc_key_b64 else { return Err(AppError::BadRequest("Missing TWO_FACTOR_ENC_KEY".to_string())); @@ -150,9 +150,9 @@ pub fn decrypt_secret_with_optional_key( let blob = general_purpose::STANDARD .decode(rest) - .map_err(|_| AppError::Internal)?; + .map_err(|e| AppError::Internal(e.to_string()))?; if blob.len() < 12 { - return Err(AppError::Internal); + return Err(AppError::Internal("Invalid secret blob length".to_string())); } let (nonce_bytes, ct) = blob.split_at(12); let cipher = Aes256Gcm::new(Key::::from_slice(&key_bytes)); @@ -171,7 +171,7 @@ pub fn decrypt_secret_with_optional_key( ) })?; - Ok(String::from_utf8(pt).map_err(|_| AppError::Internal)?) + Ok(String::from_utf8(pt).map_err(|e| AppError::Internal(e.to_string()))?) } pub fn verify_totp_code(secret_encoded: &str, token: &str) -> Result { @@ -186,11 +186,11 @@ pub fn verify_totp_code(secret_encoded: &str, token: &str) -> Result