From 388469a84c25e3fdb7bdd391784d2013cde932f3 Mon Sep 17 00:00:00 2001 From: qaz741wsd856 Date: Mon, 1 Dec 2025 03:32:38 +0000 Subject: [PATCH] feat: enhance import functionality to handle existing folders and improve cipher insertion logic --- src/handlers/import.rs | 146 +++++++++++++++++++++++++++++------------ src/models/import.rs | 37 ++++++++--- 2 files changed, 133 insertions(+), 50 deletions(-) diff --git a/src/handlers/import.rs b/src/handlers/import.rs index b88bb41..77016da 100644 --- a/src/handlers/import.rs +++ b/src/handlers/import.rs @@ -1,5 +1,6 @@ use axum::{extract::State, Json}; use chrono::Utc; +use std::collections::{HashMap, HashSet}; use std::sync::Arc; use uuid::Uuid; use worker::{query, D1PreparedStatement, Env}; @@ -13,62 +14,117 @@ use crate::models::import::ImportRequest; use super::get_batch_size; +/// Import ciphers and folders. +/// Aligned with vaultwarden's POST /ciphers/import implementation. #[worker::send] pub async fn import_data( claims: Claims, State(env): State>, - Json(mut payload): Json, + Json(data): Json, ) -> Result, AppError> { let db = db::get_db(&env)?; let now = Utc::now(); let now = now.format("%Y-%m-%dT%H:%M:%S%.3fZ").to_string(); let batch_size = get_batch_size(&env); - // Prepare all folder insert statements - let mut folder_statements: Vec = Vec::with_capacity(payload.folders.len()); + // Get existing folders for this user + let existing_folder_rows = query!( + &db, + "SELECT id FROM folders WHERE user_id = ?1", + &claims.sub + ) + .map_err(|_| AppError::Database)? + .all() + .await? + .results::()?; - for import_folder in &payload.folders { - let folder = Folder { - id: import_folder.id.clone(), - user_id: claims.sub.clone(), - name: import_folder.name.clone(), - created_at: now.clone(), - updated_at: now.clone(), + let existing_folders: HashSet = existing_folder_rows + .into_iter() + .map(|row| row.id) + .collect(); + + // Process folders and build the folder_id list + let mut folder_statements: Vec = Vec::new(); + let mut folders: Vec = Vec::with_capacity(data.folders.len()); + + for import_folder in data.folders { + let folder_id = if let Some(ref id) = import_folder.id { + if existing_folders.contains(id) { + // Folder already exists, use existing ID + id.clone() + } else { + // Folder doesn't exist, create new one with provided ID + let folder = Folder { + id: id.clone(), + user_id: claims.sub.clone(), + name: import_folder.name.clone(), + created_at: now.clone(), + updated_at: now.clone(), + }; + + let stmt = query!( + &db, + "INSERT INTO folders (id, user_id, name, created_at, updated_at) VALUES (?1, ?2, ?3, ?4, ?5)", + folder.id, + folder.user_id, + folder.name, + folder.created_at, + folder.updated_at + ) + .map_err(|_| AppError::Database)?; + + folder_statements.push(stmt); + id.clone() + } + } else { + // No ID provided, create new folder with generated UUID + let new_id = Uuid::new_v4().to_string(); + let folder = Folder { + id: new_id.clone(), + user_id: claims.sub.clone(), + name: import_folder.name.clone(), + created_at: now.clone(), + updated_at: now.clone(), + }; + + let stmt = query!( + &db, + "INSERT INTO folders (id, user_id, name, created_at, updated_at) VALUES (?1, ?2, ?3, ?4, ?5)", + folder.id, + folder.user_id, + folder.name, + folder.created_at, + folder.updated_at + ) + .map_err(|_| AppError::Database)?; + + folder_statements.push(stmt); + new_id }; - let stmt = query!( - &db, - "INSERT OR IGNORE INTO folders (id, user_id, name, created_at, updated_at) VALUES (?1, ?2, ?3, ?4, ?5)", - folder.id, - folder.user_id, - folder.name, - folder.created_at, - folder.updated_at - ) - .map_err(|_| AppError::Database)?; - - folder_statements.push(stmt); + folders.push(folder_id); } // Execute folder inserts in batches - db::execute_in_batches(&db, folder_statements, batch_size).await?; + if !folder_statements.is_empty() { + db::execute_in_batches(&db, folder_statements, batch_size).await?; + } - // Process folder relationships - 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()); - } - } + // Build the relations map: cipher_index -> folder_index + // Each cipher can only be in one folder at a time + let mut relations_map: HashMap = HashMap::with_capacity(data.folder_relationships.len()); + for relation in data.folder_relationships { + relations_map.insert(relation.key, relation.value); } // Prepare all cipher insert statements - let mut cipher_statements: Vec = Vec::with_capacity(payload.ciphers.len()); + let mut cipher_statements: Vec = Vec::with_capacity(data.ciphers.len()); - for import_cipher in payload.ciphers { - if import_cipher.encrypted_for != claims.sub { - return Err(AppError::BadRequest("Cipher encrypted for wrong user".to_string())); - } + for (index, import_cipher) in data.ciphers.into_iter().enumerate() { + // Determine folder_id from folder_relationships + let folder_id = relations_map + .get(&index) + .and_then(|folder_idx| folders.get(*folder_idx).cloned()); let cipher_data = CipherData { name: import_cipher.name, @@ -87,11 +143,11 @@ pub async fn import_data( let cipher = Cipher { id: Uuid::new_v4().to_string(), user_id: Some(claims.sub.clone()), - organization_id: import_cipher.organization_id.clone(), + organization_id: import_cipher.organization_id, r#type: import_cipher.r#type, data: data_value, - favorite: import_cipher.favorite, - folder_id: import_cipher.folder_id.clone(), + favorite: import_cipher.favorite.unwrap_or(false), + folder_id, deleted_at: None, created_at: now.clone(), updated_at: now.clone(), @@ -106,7 +162,7 @@ pub async fn import_data( let stmt = query!( &db, - "INSERT OR IGNORE INTO ciphers (id, user_id, organization_id, type, data, favorite, folder_id, created_at, updated_at) + "INSERT INTO ciphers (id, user_id, organization_id, type, data, favorite, folder_id, created_at, updated_at) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9)", cipher.id, cipher.user_id, @@ -123,9 +179,17 @@ pub async fn import_data( } // Execute cipher inserts in batches - db::execute_in_batches(&db, cipher_statements, batch_size).await?; + if !cipher_statements.is_empty() { + db::execute_in_batches(&db, cipher_statements, batch_size).await?; + } touch_user_updated_at(&db, &claims.sub).await?; Ok(Json(())) -} \ No newline at end of file +} + +/// Helper struct for querying existing folder IDs +#[derive(serde::Deserialize)] +struct FolderIdRow { + id: String, +} diff --git a/src/models/import.rs b/src/models/import.rs index 0391a8e..bd62eee 100644 --- a/src/models/import.rs +++ b/src/models/import.rs @@ -1,42 +1,61 @@ use serde::Deserialize; use serde_json::Value; +/// Cipher data structure for import requests. +/// Aligned with vaultwarden's CipherData used in /ciphers/import endpoint. +/// Note: The `key` field is omitted as this service's compatibility version doesn't support per-cipher encryption keys. #[derive(Deserialize, Debug)] #[serde(rename_all = "camelCase")] pub struct ImportCipher { + /// Optional cipher ID (only used in bulk share scenarios, not import) + #[allow(dead_code)] + pub id: Option, + /// Folder ID is not included in import - determined by folder_relationships instead + #[allow(dead_code)] + pub folder_id: Option, + #[serde(alias = "organizationID")] + pub organization_id: Option, #[serde(rename = "type")] pub r#type: i32, - pub folder_id: Option, - pub organization_id: Option, pub name: String, pub notes: Option, - pub favorite: bool, + pub fields: Option, + // Only one of these should exist, depending on type pub login: Option, + pub secure_note: Option, pub card: Option, pub identity: Option, - pub secure_note: Option, - pub fields: Option, - pub password_history: Option, + #[serde(default)] + pub favorite: Option, pub reprompt: Option, + pub password_history: Option, + /// The revision datetime (in ISO 8601 format) of the client's local copy + #[allow(dead_code)] pub last_known_revision_date: Option, - pub encrypted_for: String, } - +/// Folder data structure for import requests. +/// Aligned with vaultwarden's FolderData. #[derive(Deserialize, Debug)] #[serde(rename_all = "camelCase")] pub struct ImportFolder { - pub id: String, + /// Optional folder ID - if provided and exists, the existing folder is used + pub id: Option, pub name: String, } +/// Relationship between cipher index and folder index in the import arrays. #[derive(Deserialize, Debug)] #[serde(rename_all = "camelCase")] pub struct FolderRelationship { + /// Cipher index in the ciphers array pub key: usize, + /// Folder index in the folders array pub value: usize, } +/// Import request payload structure. +/// Aligned with vaultwarden's ImportData used in POST /ciphers/import. #[derive(Deserialize, Debug)] #[serde(rename_all = "camelCase")] pub struct ImportRequest {