danvics/warden-worker
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 1

feat: enhance import functionality to handle existing folders and improve cipher insertion logic

Browse source
This commit is contained in:
qaz741wsd856 2025-12-01 03:32:38 +00:00
parent eb64650594
commit 388469a84c
2 changed files with 133 additions and 50 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

146
src/handlers/import.rs
View file

@ -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<Arc<Env>>,
Json(mut payload): Json<ImportRequest>,
Json(data): Json<ImportRequest>,
) -> Result<Json<()>, 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<D1PreparedStatement> = 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::<FolderIdRow>()?;
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<String> = existing_folder_rows
.into_iter()
.map(|row| row.id)
.collect();
// Process folders and build the folder_id list
let mut folder_statements: Vec<D1PreparedStatement> = Vec::new();
let mut folders: Vec<String> = 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<usize, usize> = 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<D1PreparedStatement> = Vec::with_capacity(payload.ciphers.len());
let mut cipher_statements: Vec<D1PreparedStatement> = 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(()))
}
}
/// Helper struct for querying existing folder IDs
#[derive(serde::Deserialize)]
struct FolderIdRow {
id: String,
}

37
src/models/import.rs
View file

@ -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<String>,
/// Folder ID is not included in import - determined by folder_relationships instead
#[allow(dead_code)]
pub folder_id: Option<String>,
#[serde(alias = "organizationID")]
pub organization_id: Option<String>,
#[serde(rename = "type")]
pub r#type: i32,
pub folder_id: Option<String>,
pub organization_id: Option<String>,
pub name: String,
pub notes: Option<String>,
pub favorite: bool,
pub fields: Option<Value>,
// Only one of these should exist, depending on type
pub login: Option<Value>,
pub secure_note: Option<Value>,
pub card: Option<Value>,
pub identity: Option<Value>,
pub secure_note: Option<Value>,
pub fields: Option<Value>,
pub password_history: Option<Value>,
#[serde(default)]
pub favorite: Option<bool>,
pub reprompt: Option<i32>,
pub password_history: Option<Value>,
/// The revision datetime (in ISO 8601 format) of the client's local copy
#[allow(dead_code)]
pub last_known_revision_date: Option<String>,
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<String>,
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 {