Something went wrong. Try again.
This repository has no description
Something went wrong. Try again.
54 kB · 1340 lines
Rust
at master
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314131513161317131813191320132113221323132413251326132713281329133013311332133313341335133613371338133913401341use crate::{assets::AssetRecord, error::ApiError, state::AppState};use axum::{Json, extract::State};use base64::{Engine, engine::general_purpose::STANDARD};use chrono::Utc;use serde_json::{Map, Value, json};use sha2::{Digest, Sha256};use sqlx::{Postgres, Transaction};use std::{ collections::{HashMap, HashSet}, path::{Component, Path, PathBuf},};use tokio::{fs, io::AsyncWriteExt};use uuid::Uuid;
// Ordering is also the restore order. Self references are checked at the end// of each INSERT, so a collection's rows can be restored in one statement.const TABLES: &[&str] = &[ "notes", "tags", "note_revisions", "note_tags", "note_links", "annotations", "problems", "shelf_collections", "shelf_items", "learning_spaces", "learning_lessons", "whiteboards", "spatial_documents", "workspace_states", "reader_materials", "reader_legacy_imports", "provider_settings", "agent_conversations", "agent_messages", "reader_chunks", "reader_highlights", "reader_notes", "reader_review_cards", "reader_exercise_attempts", "reader_material_revisions", "reader_rewrite_proposals", "agent_lesson_revisions",];const MAX_FILE_BYTES: usize = 256 * 1024 * 1024;const MAX_TOTAL_BYTES: usize = 512 * 1024 * 1024;const MAX_PAYLOAD_JSON: usize = 768 * 1024 * 1024;const MAX_FILES: usize = 20_000;
fn invalid(message: impl Into<String>) -> ApiError { ApiError::BadRequest(message.into())}
pub async fn export_data(State(state): State<AppState>) -> Result<Json<Value>, ApiError> { let mut tx = state.platform_db.begin().await?; sqlx::query("set transaction isolation level repeatable read, read only") .execute(&mut *tx) .await?; let mut data = Map::new(); for &table in TABLES { let selection = match table { "provider_settings" => "select provider, model, base_url, enabled, name, models, options, updated_at from provider_settings".to_owned(), "workspace_states" => "select tabs, active_tab_id, layout, updated_at from workspace_states".to_owned(), "reader_chunks" => "select id, material_id, page, heading, text from reader_chunks order by id".to_owned(), "note_revisions" => "select * from note_revisions order by note_id, revision_number".to_owned(), "agent_messages" => "select * from agent_messages order by conversation_id, created_at, id".to_owned(), _ => format!("select * from {table}"), }; let query = format!( "select coalesce(jsonb_agg(to_jsonb(export_row)), '[]'::jsonb) from ({selection}) export_row" ); let mut rows: Value = sqlx::query_scalar(&query).fetch_one(&mut *tx).await?; if table == "provider_settings" { sanitize_providers(&mut rows)?; } data.insert(table.into(), rows); } let root = fs::canonicalize(&state.config.storage_dir).await?; let records = sqlx::query_as::<_, AssetRecord>("select * from assets order by created_at") .fetch_all(&mut *tx) .await?; if records.len() > MAX_FILES { return Err(invalid("export contains too many files")); } let mut total = 0usize; let mut assets = Vec::with_capacity(records.len()); for record in records { validate_asset_key(record.id, &record.storage_key)?; let path = checked_path(&root, &record.storage_key, false).await?; let bytes = bounded_read(&path, &mut total).await?; verify_content(&bytes, record.byte_size, &record.checksum_sha256)?; assets.push(json!({"record":record,"content_base64":STANDARD.encode(bytes)})); } data.insert("assets".into(), Value::Array(assets)); let mut reader_files = Vec::new(); if let Some(materials) = data["reader_materials"].as_array() { for material in materials { let id = row_uuid(material, "id")?; let directory = format!("reader/{id}"); let path = checked_path(&root, &directory, true).await?; if !fs::try_exists(&path).await? { continue; } let mut pending = vec![path.clone()]; while let Some(directory) = pending.pop() { let mut entries = fs::read_dir(directory).await?; while let Some(entry) = entries.next_entry().await? { let kind = entry.file_type().await?; if kind.is_symlink() { return Err(invalid("reader storage contains a symbolic link")); } if kind.is_dir() { pending.push(entry.path()); continue; } if !kind.is_file() { return Err(invalid("reader storage contains a nonregular file")); } if reader_files.len() + assets_count(&data) >= MAX_FILES { return Err(invalid("export contains too many files")); } let relative = entry .path() .strip_prefix(&path) .map_err(anyhow::Error::from)? .to_str() .ok_or_else(|| invalid("reader asset path must be UTF-8"))? .replace('\\', "/"); validate_relative_path(&relative)?; let bytes = bounded_read(&entry.path(), &mut total).await?; reader_files.push(json!({"material_id":id,"path":relative,"byte_size":bytes.len(), "checksum_sha256":checksum(&bytes),"content_base64":STANDARD.encode(bytes)})); } } } } data.insert("reader_files".into(), Value::Array(reader_files)); tx.commit().await?; let payload = json!({"format":"odre-platform-export","version":2,"exported_at":Utc::now(), "data":data,"excluded_secrets":["provider API keys"]}); check_payload_size(&payload)?; Ok(Json(payload))}
fn assets_count(data: &Map<String, Value>) -> usize { data.get("assets") .and_then(Value::as_array) .map_or(0, Vec::len)}
pub async fn import_data( State(state): State<AppState>, Json(mut payload): Json<Value>,) -> Result<Json<Value>, ApiError> { check_payload_size(&payload)?; validate_payload(&payload)?; let Value::Object(mut data) = payload["data"].take() else { unreachable!("data was validated as an object") }; if !data.contains_key("agent_lesson_revisions") { if let Some(rows) = data.remove("learning_lesson_revisions") { data.insert("agent_lesson_revisions".into(), rows); } } normalize_import(&mut data)?; let mut files = stage_files(&state.config.storage_dir, &data).await?; let result = restore_data(&state, &data, &mut files).await; if result.is_err() { files.rollback().await; } files.cleanup_stage().await; result?; if let Err(error) = crate::reader_pdf::resume_pending(&state).await { tracing::warn!(?error, "could not resume imported PDF indexing"); } Ok(Json(json!({"ok":true,"imported_at":Utc::now(), "warnings":["Provider API keys are excluded. Existing keys were preserved; new providers require a key."]})))}
fn validate_payload(payload: &Value) -> Result<&Map<String, Value>, ApiError> { let supported = matches!( ( payload.get("format").and_then(Value::as_str), payload.get("version").and_then(Value::as_i64) ), (Some("forge-platform-export"), Some(1)) | (Some("odre-platform-export"), Some(2)) ); if !supported { return Err(invalid( "unsupported Odre export format (Forge v1 and Odre v2 are supported)", )); } let data = payload .get("data") .and_then(Value::as_object) .ok_or_else(|| invalid("export data is missing"))?; for key in TABLES .iter() .copied() .chain(["assets", "reader_files", "learning_lesson_revisions"]) { if let Some(rows) = data.get(key) { let rows = rows .as_array() .ok_or_else(|| invalid(format!("{key} must be an array")))?; if rows.iter().any(|row| !row.is_object()) { return Err(invalid(format!("{key} must contain objects"))); } } } Ok(data)}
fn normalize_import(data: &mut Map<String, Value>) -> Result<(), ApiError> { if let Some(workspaces) = data .get_mut("workspace_states") .and_then(Value::as_array_mut) { workspaces.truncate(1); if let Some(workspace) = workspaces.first_mut() { workspace["id"] = json!(true); } } if let Some(notes) = data.get_mut("notes").and_then(Value::as_array_mut) { for note in notes { note["status"] = json!(if note["status"] == "archived" { "archived" } else { "draft" }); } } if let Some(providers) = data.get_mut("provider_settings") { sanitize_providers(providers)?; } if let Some(messages) = data.get_mut("agent_messages").and_then(Value::as_array_mut) { let mut last_timestamps = HashMap::new(); for message in messages { if message["metadata"]["status"] == "streaming" { message["metadata"]["status"] = json!("cancelled"); message["metadata"]["error"] = json!("Imported partial response. Generation was interrupted by the export."); } // Old exports can contain several turns stamped by transaction // now(). Preserve their array order when the reader reloads them. if let (Some(conversation), Some(timestamp)) = ( message["conversation_id"].as_str(), message["created_at"].as_str(), ) { if let Ok(timestamp) = chrono::DateTime::parse_from_rfc3339(timestamp) { let adjusted = match last_timestamps.get(conversation) { Some(previous) if timestamp <= *previous => { *previous + chrono::Duration::microseconds(1) } _ => timestamp, }; last_timestamps.insert(conversation.to_owned(), adjusted); message["created_at"] = json!(adjusted.to_rfc3339()); } } } } if let Some(proposals) = data .get_mut("reader_rewrite_proposals") .and_then(Value::as_array_mut) { for proposal in proposals { if proposal["status"] == "streaming" { proposal["status"] = json!("cancelled"); proposal["error"] = json!("Imported partial proposal. Generation was interrupted by the export."); } } } Ok(())}
fn sanitize_providers(rows: &mut Value) -> Result<(), ApiError> { let rows = rows .as_array_mut() .ok_or_else(|| invalid("provider_settings must be an array"))?; for row in rows { let row = row .as_object_mut() .ok_or_else(|| invalid("provider_settings must contain objects"))?; row.retain(|key, _| { matches!( key.as_str(), "provider" | "model" | "base_url" | "enabled" | "name" | "models" | "options" | "updated_at" ) }); let provider = row .get("provider") .and_then(Value::as_str) .ok_or_else(|| invalid("provider ID is missing"))? .to_owned(); crate::domains::provider::valid_provider(&provider)?; if let Some(url) = row.get("base_url").and_then(Value::as_str) { let url = reqwest::Url::parse(url).map_err(|_| invalid("invalid provider base URL"))?; if !matches!(url.scheme(), "http" | "https") || !url.username().is_empty() || url.password().is_some() { return Err(invalid( "provider URL must use HTTP(S) without embedded credentials", )); } } if row .get("name") .is_none_or(|v| v.is_null() || v.as_str() == Some("")) { row.insert("name".into(), json!(provider)); } if row.get("models").is_none_or(Value::is_null) { row.insert( "models".into(), json!([row.get("model").and_then(Value::as_str).unwrap_or("")]), ); } if row.get("options").is_none_or(Value::is_null) { row.insert("options".into(), json!({"protocol":match provider.as_str(){"openai"=>"responses","anthropic"=>"anthropic",_=>"chat"}})); } if !row["options"].is_object() || !row["models"].is_array() { return Err(invalid( "provider options must be an object and models an array", )); } strip_secrets(row.get_mut("options").expect("options initialized")); } Ok(())}
fn strip_secrets(value: &mut Value) { match value { Value::Object(object) => { object.retain(|key, _| { let key: String = key .chars() .filter(|c| c.is_ascii_alphanumeric()) .flat_map(char::to_lowercase) .collect(); !key.contains("apikey") && !key.contains("secret") && !key.contains("password") && !key.ends_with("token") && key != "authorization" }); for value in object.values_mut() { strip_secrets(value); } } Value::Array(array) => { for value in array { strip_secrets(value); } } _ => {} }}
async fn restore_data( state: &AppState, data: &Map<String, Value>, files: &mut StagedFiles,) -> Result<(), ApiError> { let mut tx = state.platform_db.begin().await?; // Imports share the filesystem, so serialize their restore/commit window. sqlx::query("select pg_advisory_xact_lock(739271580119::bigint)") .execute(&mut *tx) .await?; sqlx::query("lock table reader_chunks in share row exclusive mode") .execute(&mut *tx) .await?; import_assets(&mut tx, data).await?; for &table in TABLES { let Some(rows) = data.get(table) else { continue; }; if rows.as_array().is_none_or(Vec::is_empty) { continue; } if table == "reader_chunks" { import_reader_chunks(&mut tx, rows).await?; continue; } import_table(&mut tx, table, rows).await?; } let imported_materials: Vec<Uuid> = data .get("reader_materials") .and_then(Value::as_array) .into_iter() .flatten() .map(|row| row_uuid(row, "id")) .collect::<Result<_, _>>()?; crate::reader_native::reconcile_import_tx(&mut tx, &imported_materials).await?; for file in &files.files { if let Some(id) = file.material_id { let exists: bool = sqlx::query_scalar("select exists(select 1 from reader_materials where id=$1)") .bind(id) .fetch_one(&mut *tx) .await?; if !exists { return Err(invalid("reader asset references a missing material")); } } } sqlx::query("select setval(pg_get_serial_sequence('reader_chunks','id'),greatest(coalesce((select max(id) from reader_chunks),1),(select last_value from reader_chunks_id_seq)),(select is_called from reader_chunks_id_seq) or (select exists(select 1 from reader_chunks)))") .execute(&mut *tx).await?; files.publish().await?; // On an uncertain commit outcome, retain files: deleting them could break // a commit which reached PostgreSQL before the connection failed. files.created.clear(); tx.commit().await?; Ok(())}
async fn import_reader_chunks( tx: &mut Transaction<'_, Postgres>, rows: &Value,) -> Result<(), ApiError> { let mut ids = HashSet::new(); for row in rows.as_array().expect("validated collection") { let id = row["id"] .as_i64() .filter(|id| *id > 0) .ok_or_else(|| invalid("reader chunk ID must be a positive integer"))?; if !ids.insert(id) { return Err(invalid("duplicate reader chunk ID in export")); } } sqlx::query("delete from reader_chunks where material_id in (select material_id from jsonb_populate_recordset(null::reader_chunks,$1))") .bind(rows).execute(&mut **tx).await.map_err(import_sql_error)?; sqlx::query("select setval(pg_get_serial_sequence('reader_chunks','id'),greatest(coalesce((select max(id) from reader_chunks),1),coalesce((select max(id) from jsonb_populate_recordset(null::reader_chunks,$1)),1),(select last_value from reader_chunks_id_seq)),true)") .bind(rows).execute(&mut **tx).await.map_err(import_sql_error)?; let collisions:Vec<i64>=sqlx::query_scalar("select current.id from reader_chunks current join jsonb_populate_recordset(null::reader_chunks,$1) incoming using(id)") .bind(rows).fetch_all(&mut **tx).await.map_err(import_sql_error)?; let collisions: HashSet<_> = collisions.into_iter().collect(); let mut imported = rows.clone(); for chunk in imported.as_array_mut().expect("validated collection") { if chunk["id"] .as_i64() .is_some_and(|id| collisions.contains(&id)) { let id: i64 = sqlx::query_scalar("select nextval(pg_get_serial_sequence('reader_chunks','id'))") .fetch_one(&mut **tx) .await?; chunk["id"] = json!(id); } } import_table(tx, "reader_chunks", &imported).await}
#[derive(sqlx::FromRow)]struct ImportColumn { name: String, default_expression: Option<String>, not_null: bool,}
async fn import_table( tx: &mut Transaction<'_, Postgres>, table: &str, rows: &Value,) -> Result<(), ApiError> { let columns = sqlx::query_as::<_,ImportColumn>( "select a.attname::text name, pg_get_expr(d.adbin,d.adrelid) default_expression,a.attnotnull not_null from pg_attribute a left join pg_attrdef d on a.attrelid=d.adrelid and a.attnum=d.adnum where a.attrelid=to_regclass($1) and a.attnum>0 and not a.attisdropped and a.attgenerated='' order by a.attnum") .bind(table).fetch_all(&mut **tx).await?; let columns: Vec<_> = columns .into_iter() .filter(|c| c.name != "api_key_ciphertext") .collect(); let keys: &[&str] = match table { "provider_settings" => &["provider"], "workspace_states" => &[], "note_tags" => &["note_id", "tag_id"], "note_links" => &["source_note_id", "target_note_id", "source_block_id"], "reader_exercise_attempts" => &["material_id", "exercise_key"], "reader_legacy_imports" => &["fingerprint"], _ => &["id"], }; for row in rows.as_array().expect("validated collection") { for &key in keys { if row.get(key).is_none_or(Value::is_null) { return Err(invalid(format!("{table}.{key} is required"))); } } } let names = columns .iter() .map(|c| format!("\"{}\"", c.name)) .collect::<Vec<_>>() .join(","); let values = columns .iter() .map(|c| { let column = format!("imported.\"{}\"", c.name); if c.not_null { if let Some(default) = &c.default_expression { return format!("coalesce({column},{default})"); } } column }) .collect::<Vec<_>>() .join(","); let conflict = if table == "workspace_states" { "id".to_owned() } else { keys.join(",") }; let updates = columns .iter() .filter(|c| !keys.contains(&c.name.as_str()) && c.name != "created_at" && c.name != "id") .map(|c| { if table == "reader_materials" && c.name == "note_id" { "note_id=coalesce(excluded.note_id,reader_materials.note_id)".into() } else { format!("\"{0}\"=excluded.\"{0}\"", c.name) } }) .collect::<Vec<_>>() .join(","); let on_conflict = if updates.is_empty() || matches!( table, "note_revisions" | "reader_material_revisions" | "agent_lesson_revisions" ) { "do nothing".to_owned() } else { format!("do update set {updates}") }; let query = format!( "insert into {table} ({names}) select {values} from jsonb_populate_recordset(null::{table},$1) imported on conflict ({conflict}) {on_conflict}" ); sqlx::query(&query) .bind(rows) .execute(&mut **tx) .await .map_err(import_sql_error)?; Ok(())}
fn import_sql_error(error: sqlx::Error) -> ApiError { if let sqlx::Error::Database(ref db) = error { if db.code().is_some_and(|code| { matches!( code.as_ref(), "23502" | "23503" | "23505" | "23514" | "22P02" | "22007" | "21000" ) }) { return invalid(format!( "export contains invalid or conflicting data: {}", db.message() )); } } error.into()}
async fn import_assets( tx: &mut Transaction<'_, Postgres>, data: &Map<String, Value>,) -> Result<(), ApiError> { let Some(assets) = data.get("assets").and_then(Value::as_array) else { return Ok(()); }; let mut rows = Vec::with_capacity(assets.len()); for item in assets { let record = &item["record"]; let id = row_uuid(record, "id")?; let incoming_key = record["storage_key"] .as_str() .ok_or_else(|| invalid("asset storage key is missing"))?; let incoming_hash = record["checksum_sha256"] .as_str() .ok_or_else(|| invalid("asset checksum is missing"))?; let incoming_size = record["byte_size"] .as_i64() .ok_or_else(|| invalid("asset size is missing"))?; let existing = sqlx::query_as::<_, AssetRecord>( "select * from assets where id=$1 or storage_key=$2 for update", ) .bind(id) .bind(incoming_key) .fetch_all(&mut **tx) .await?; if existing.iter().any(|row| { row.id != id || row.storage_key != incoming_key || row.byte_size != incoming_size || !row.checksum_sha256.eq_ignore_ascii_case(incoming_hash) }) { return Err(ApiError::Conflict("An existing asset has a different identity or content. The import did not overwrite it.".into())); } rows.push(record.clone()); } import_table(tx, "assets", &Value::Array(rows)).await}
struct StagedFile { stage: PathBuf, relative: String, hash: String, size: i64, material_id: Option<Uuid>,}struct StagedFiles { root: PathBuf, stage_dir: PathBuf, files: Vec<StagedFile>, created: Vec<PathBuf>,}
impl StagedFiles { async fn publish(&mut self) -> Result<(), ApiError> { for file in &self.files { let destination = checked_path(&self.root, &file.relative, true).await?; create_safe_parents(&self.root, &destination).await?; // Hard links provide atomic no-clobber publication on the same volume. match fs::hard_link(&file.stage, &destination).await { Ok(()) => self.created.push(destination), Err(error) if error.kind() == std::io::ErrorKind::AlreadyExists => { let destination = checked_path(&self.root, &file.relative, false).await?; verify_existing(&destination, file.size, &file.hash).await?; } Err(error) => return Err(error.into()), } } Ok(()) } async fn rollback(&mut self) { for path in self.created.drain(..).rev() { let _ = fs::remove_file(path).await; } } async fn cleanup_stage(&self) { let _ = fs::remove_dir_all(&self.stage_dir).await; }}
impl Drop for StagedFiles { fn drop(&mut self) { let paths = std::mem::take(&mut self.created); let directory = self.stage_dir.clone(); // Request cancellation must not leave partly published files behind. if let Ok(runtime) = tokio::runtime::Handle::try_current() { runtime.spawn(async move { for path in paths.into_iter().rev() { let _ = fs::remove_file(path).await; } let _ = fs::remove_dir_all(directory).await; }); } }}
async fn stage_files(storage: &Path, data: &Map<String, Value>) -> Result<StagedFiles, ApiError> { fs::create_dir_all(storage).await?; let root = fs::canonicalize(storage).await?; let stage_dir = root.join(format!(".transfer-{}", Uuid::new_v4())); fs::create_dir(&stage_dir).await?; let mut staged = StagedFiles { root, stage_dir, files: Vec::new(), created: Vec::new(), }; let mut paths = HashSet::new(); let mut identities = HashSet::new(); let mut total = 0usize; for item in data .get("assets") .and_then(Value::as_array) .into_iter() .flatten() { let record = &item["record"]; let id = row_uuid(record, "id")?; if !identities.insert(id) { return Err(invalid("duplicate asset ID in export")); } let relative = record["storage_key"] .as_str() .ok_or_else(|| invalid("asset storage key is missing"))?; validate_asset_key(id, relative)?; stage_one( &mut staged, &mut paths, &mut total, relative, None, record, item, ) .await?; } for item in data .get("reader_files") .and_then(Value::as_array) .into_iter() .flatten() { let id = row_uuid(item, "material_id")?; let relative = item["path"] .as_str() .ok_or_else(|| invalid("reader asset path is missing"))?; validate_relative_path(relative)?; stage_one( &mut staged, &mut paths, &mut total, &format!("reader/{id}/{relative}"), Some(id), item, item, ) .await?; } Ok(staged)}
async fn stage_one( staged: &mut StagedFiles, paths: &mut HashSet<String>, total: &mut usize, relative: &str, material_id: Option<Uuid>, record: &Value, item: &Value,) -> Result<(), ApiError> { if !paths.insert(relative.to_owned()) { return Err(invalid("duplicate storage path in export")); } if paths.len() > MAX_FILES { return Err(invalid("export contains too many files")); } let size = record["byte_size"] .as_i64() .ok_or_else(|| invalid("asset size is missing"))?; let hash = record["checksum_sha256"] .as_str() .ok_or_else(|| invalid("asset checksum is missing"))?; let encoded = item["content_base64"] .as_str() .ok_or_else(|| invalid("asset content is missing"))?; if size < 0 || size as u64 > MAX_FILE_BYTES as u64 || encoded.len() > MAX_FILE_BYTES.div_ceil(3) * 4 + 4 { return Err(invalid("asset exceeds the 256 MiB restore limit")); } let bytes = STANDARD .decode(encoded) .map_err(|_| invalid("invalid base64 asset content"))?; account_bytes(bytes.len(), total)?; verify_content(&bytes, size, hash)?; let destination = checked_path(&staged.root, relative, true).await?; if fs::try_exists(&destination).await? { verify_existing(&destination, size, hash).await?; } let stage = staged.stage_dir.join(staged.files.len().to_string()); let mut file = fs::OpenOptions::new() .write(true) .create_new(true) .open(&stage) .await?; file.write_all(&bytes).await?; file.sync_all().await?; staged.files.push(StagedFile { stage, relative: relative.to_owned(), hash: hash.to_owned(), size, material_id, }); Ok(())}
fn check_payload_size(payload: &Value) -> Result<(), ApiError> { struct Counter { size: usize, } impl std::io::Write for Counter { fn write(&mut self, bytes: &[u8]) -> std::io::Result<usize> { self.size = self .size .checked_add(bytes.len()) .ok_or_else(|| std::io::Error::other("export is too large"))?; if self.size > MAX_PAYLOAD_JSON { return Err(std::io::Error::other("export is too large")); } Ok(bytes.len()) } fn flush(&mut self) -> std::io::Result<()> { Ok(()) } } serde_json::to_writer(Counter { size: 0 }, payload).map_err(|_| { invalid("export exceeds the 768 MiB JSON limit; use a database and storage backup for larger libraries") })}
fn row_uuid(row: &Value, field: &str) -> Result<Uuid, ApiError> { row.get(field) .and_then(Value::as_str) .and_then(|id| Uuid::parse_str(id).ok()) .ok_or_else(|| invalid(format!("{field} must be a UUID")))}fn checksum(bytes: &[u8]) -> String { format!("{:x}", Sha256::digest(bytes))}fn verify_content(bytes: &[u8], size: i64, hash: &str) -> Result<(), ApiError> { if size < 0 || size as u64 != bytes.len() as u64 || hash.len() != 64 || !checksum(bytes).eq_ignore_ascii_case(hash) { return Err(invalid( "asset size or SHA-256 checksum does not match its content", )); } Ok(())}fn account_bytes(size: usize, total: &mut usize) -> Result<(), ApiError> { *total = total .checked_add(size) .ok_or_else(|| invalid("export is too large"))?; if size > MAX_FILE_BYTES || *total > MAX_TOTAL_BYTES { return Err(invalid( "export exceeds 256 MiB per file or 512 MiB of total file contents", )); } Ok(())}async fn bounded_read(path: &Path, total: &mut usize) -> Result<Vec<u8>, ApiError> { let size = fs::metadata(path).await?.len(); if size > MAX_FILE_BYTES as u64 { return Err(invalid("asset exceeds the 256 MiB export limit")); } let bytes = fs::read(path).await?; account_bytes(bytes.len(), total)?; Ok(bytes)}async fn verify_existing(path: &Path, size: i64, hash: &str) -> Result<(), ApiError> { let metadata = fs::symlink_metadata(path).await?; if !metadata.is_file() || metadata.file_type().is_symlink() || metadata.len() != size as u64 { return Err(ApiError::Conflict( "Existing storage content differs from the export; it was preserved.".into(), )); } let bytes = fs::read(path).await?; if !checksum(&bytes).eq_ignore_ascii_case(hash) { return Err(ApiError::Conflict( "Existing storage content differs from the export; it was preserved.".into(), )); } Ok(())}fn validate_relative_path(key: &str) -> Result<(), ApiError> { if key.is_empty() || key.len() > 2048 || key.contains(['\\', '\0']) || (key.as_bytes().get(1) == Some(&b':') && key.as_bytes()[0].is_ascii_alphabetic()) || (cfg!(windows) && key.contains(':')) || key.starts_with('/') || key.ends_with('/') || key .split('/') .any(|part| part.is_empty() || part == "." || part == "..") || Path::new(key) .components() .any(|part| !matches!(part, Component::Normal(_))) { return Err(invalid("unsafe storage path in export")); } Ok(())}fn validate_asset_key(id: Uuid, key: &str) -> Result<(), ApiError> { validate_relative_path(key)?; let (owner, _) = key .split_once('/') .ok_or_else(|| invalid("asset path must be owned by its UUID"))?; if owner != id.to_string() { return Err(invalid("asset storage path belongs to a different UUID")); } Ok(())}async fn checked_path(root: &Path, key: &str, allow_missing: bool) -> Result<PathBuf, ApiError> { validate_relative_path(key)?; let mut path = root.to_path_buf(); let parts: Vec<_> = key.split('/').collect(); for (index, part) in parts.iter().enumerate() { path.push(part); match fs::symlink_metadata(&path).await { Ok(metadata) if metadata.file_type().is_symlink() => { return Err(invalid("symbolic links are not allowed in storage paths")); } Ok(metadata) if index + 1 < parts.len() && !metadata.is_dir() => { return Err(invalid("storage parent must be a directory")); } Ok(_) => {} Err(error) if allow_missing && error.kind() == std::io::ErrorKind::NotFound => {} Err(error) => return Err(error.into()), } } Ok(path)}async fn create_safe_parents(root: &Path, destination: &Path) -> Result<(), ApiError> { let parent = destination .parent() .ok_or_else(|| invalid("asset parent is missing"))?; let relative = parent .strip_prefix(root) .map_err(|_| invalid("asset escapes storage"))?; let mut path = root.to_path_buf(); for part in relative.components() { path.push(part); match fs::create_dir(&path).await { Ok(()) => {} Err(error) if error.kind() == std::io::ErrorKind::AlreadyExists => {} Err(error) => return Err(error.into()), } let metadata = fs::symlink_metadata(&path).await?; if !metadata.is_dir() || metadata.file_type().is_symlink() { return Err(invalid("unsafe asset parent directory")); } } Ok(())}
#[cfg(test)]mod tests { use super::*;
#[test] fn accepts_legacy_and_unified_exports_but_rejects_malformed_collections() { for (format, version) in [("forge-platform-export", 1), ("odre-platform-export", 2)] { assert!( validate_payload(&json!({"format":format,"version":version,"data":{"notes":[]}})) .is_ok() ); } assert!( validate_payload( &json!({"format":"odre-platform-export","version":2,"data":{"notes":{}}}) ) .is_err() ); assert!( validate_payload(&json!({"format":"odre-platform-export","version":1,"data":{}})) .is_err() ); } #[test] fn rejects_traversal_and_cross_asset_identity() { for path in [ "../escape", "/tmp/file", "foo/../file", "foo//file", "foo/./file", "C:/file", "foo\\file", "foo/", "foo\0file", ] { assert!(validate_relative_path(path).is_err(), "{path}"); } let id = Uuid::new_v4(); assert!(validate_asset_key(id, &format!("{id}/book.pdf")).is_ok()); assert!(validate_asset_key(id, &format!("{}/book.pdf", Uuid::new_v4())).is_err()); } #[test] fn checks_bytes_against_export_metadata() { let content = b"known test content"; assert!(verify_content(content, content.len() as i64, &checksum(content)).is_ok()); assert!(verify_content(content, 0, &checksum(content)).is_err()); assert!(verify_content(content, content.len() as i64, &"0".repeat(64)).is_err()); } #[test] fn file_limits_are_consistent_without_allocating_large_payloads() { let mut total = 0; assert!(account_bytes(MAX_FILE_BYTES, &mut total).is_ok()); assert!(account_bytes(MAX_FILE_BYTES, &mut total).is_ok()); assert_eq!(total, MAX_TOTAL_BYTES); assert!(account_bytes(1, &mut total).is_err()); assert!(account_bytes(MAX_FILE_BYTES + 1, &mut 0).is_err()); } #[test] fn imported_tied_timestamps_keep_original_turn_order() { let mut data = json!({"agent_messages":[ {"conversation_id":"qa","created_at":"2026-10-07T10:00:00Z","content":"first"}, {"conversation_id":"qa","created_at":"2026-10-07T10:00:00Z","content":"second"}, {"conversation_id":"qa","created_at":"2026-10-07T10:00:01Z","content":"third"} ]}) .as_object() .unwrap() .clone(); normalize_import(&mut data).unwrap(); let times: Vec<_> = data["agent_messages"] .as_array() .unwrap() .iter() .map(|row| { chrono::DateTime::parse_from_rfc3339(row["created_at"].as_str().unwrap()).unwrap() }) .collect(); assert!(times[0] < times[1] && times[1] < times[2]); assert_eq!( times[0], chrono::DateTime::parse_from_rfc3339("2026-10-07T10:00:00Z").unwrap() ); assert_eq!( times[2], chrono::DateTime::parse_from_rfc3339("2026-10-07T10:00:01Z").unwrap() ); } #[test] fn strips_credentials_and_preserves_partial_content() { let mut data=json!({"provider_settings":[{"provider":"openai","model":"test","api_key":"private","api_key_ciphertext":"encrypted","options":{"protocol":"responses","nested":{"access_token":"private","max_output_tokens":42}}}], "agent_messages":[{"content":"partial","metadata":{"status":"streaming","citations":[1]}}], "reader_rewrite_proposals":[{"replacement":"partial","status":"streaming"}],"notes":[{"status":"published"}]}).as_object().unwrap().clone(); normalize_import(&mut data).unwrap(); assert!( data["provider_settings"][0] .get("api_key_ciphertext") .is_none() ); assert!( data["provider_settings"][0]["options"]["nested"] .get("access_token") .is_none() ); assert_eq!( data["provider_settings"][0]["options"]["nested"]["max_output_tokens"], 42 ); assert_eq!(data["agent_messages"][0]["metadata"]["status"], "cancelled"); assert_eq!(data["agent_messages"][0]["content"], "partial"); assert_eq!(data["reader_rewrite_proposals"][0]["status"], "cancelled"); assert_eq!(data["notes"][0]["status"], "draft"); } #[tokio::test] async fn staged_restore_preserves_existing_files_and_removes_new_files_on_rollback() { let root = std::env::temp_dir().join(format!("odre-transfer-test-{}", Uuid::new_v4())); let id = Uuid::new_v4(); let bytes = b"asset bytes"; let data = json!({"assets":[{"record":{"id":id,"storage_key":format!("{id}/book.pdf"),"byte_size":bytes.len(),"checksum_sha256":checksum(bytes)},"content_base64":STANDARD.encode(bytes)}]}); let mut staged = stage_files(&root, data.as_object().unwrap()).await.unwrap(); staged.publish().await.unwrap(); let destination = root.join(format!("{id}/book.pdf")); assert!(destination.exists()); staged.rollback().await; assert!(!destination.exists()); fs::write(&destination, b"original unrelated bytes") .await .unwrap(); assert!(stage_files(&root, data.as_object().unwrap()).await.is_err()); assert_eq!( fs::read(&destination).await.unwrap(), b"original unrelated bytes" ); staged.cleanup_stage().await; fs::remove_dir_all(root).await.unwrap(); } #[tokio::test] async fn partial_publish_failure_rolls_back_only_new_files() { let root = std::env::temp_dir().join(format!("odre-transfer-partial-{}", Uuid::new_v4())); let id = Uuid::new_v4(); let bytes = b"fixture"; let files:Vec<_>=["first.txt","second.txt"].into_iter().map(|path|json!({"material_id":id,"path":path, "byte_size":bytes.len(),"checksum_sha256":checksum(bytes),"content_base64":STANDARD.encode(bytes)})).collect(); let data = json!({"reader_files":files}); let mut staged = stage_files(&root, data.as_object().unwrap()).await.unwrap(); let existing = root.join(format!("reader/{id}/second.txt")); fs::create_dir_all(&existing).await.unwrap(); assert!(staged.publish().await.is_err()); let first = root.join(format!("reader/{id}/first.txt")); assert!(first.exists()); staged.rollback().await; assert!(!first.exists()); assert!(existing.is_dir()); staged.cleanup_stage().await; fs::remove_dir_all(root).await.unwrap(); } #[cfg(unix)] #[tokio::test] async fn rejects_symlinked_parents() { let root = std::env::temp_dir().join(format!("odre-transfer-link-{}", Uuid::new_v4())); fs::create_dir(&root).await.unwrap(); std::os::unix::fs::symlink(std::env::temp_dir(), root.join("linked")).unwrap(); assert!(checked_path(&root, "linked/escape", true).await.is_err()); fs::remove_dir_all(root).await.unwrap(); }
#[tokio::test] #[ignore = "requires disposable PostgreSQL via ODRE_TRANSFER_TEST_DATABASE_URL"] async fn postgres_roundtrip_preserves_all_domains_defaults_search_and_local_keys() { use sqlx::postgres::PgPoolOptions; let url = std::env::var("ODRE_TRANSFER_TEST_DATABASE_URL") .expect("use a disposable test database"); let admin = PgPoolOptions::new() .max_connections(1) .connect(&url) .await .unwrap(); let schema = format!("transfer_test_{}", Uuid::new_v4().simple()); sqlx::query(&format!("create schema {schema}")) .execute(&admin) .await .unwrap(); let scoped_schema = schema.clone(); let pool = PgPoolOptions::new() .max_connections(2) .after_connect(move |connection, _| { let command = format!("set search_path to {scoped_schema},public"); Box::pin(async move { sqlx::query(&command).execute(connection).await?; Ok(()) }) }) .connect(&url) .await .unwrap(); sqlx::migrate!("../../migrations/platform") .run(&pool) .await .unwrap(); let storage = std::env::temp_dir().join(format!("odre-transfer-db-test-{}", Uuid::new_v4())); let state = AppState { platform_db: pool.clone(), config: crate::config::Config { platform_database_url: url, bind: "127.0.0.1:0".parse().unwrap(), web_origin: "http://localhost".into(), storage_dir: storage.clone(), api_key_encryption_secret: "qa-only-no-real-credentials".into(), }, }; let note = Uuid::new_v4(); let revision = Uuid::new_v4(); let tag = Uuid::new_v4(); let problem = Uuid::new_v4(); let collection = Uuid::new_v4(); let shelf = Uuid::new_v4(); let space = Uuid::new_v4(); let lesson = Uuid::new_v4(); let article = Uuid::new_v4(); let book = Uuid::new_v4(); let asset = Uuid::new_v4(); let chat = Uuid::new_v4(); let message = Uuid::new_v4(); let learning_note = Uuid::new_v4(); let bytes = b"QA PDF fixture bytes"; let image = b"<svg xmlns=\"http://www.w3.org/2000/svg\"></svg>"; let payload = json!({"format":"odre-platform-export","version":2,"data":{ "notes":[{"id":note,"slug":"qa-note","title":"QA note","kind":"card","status":"published"}], "tags":[{"id":tag,"name":"qa","display_name":"QA"}], "note_revisions":[{"id":revision,"note_id":note,"revision_number":1,"title":"QA note","excerpt":"","document":{}}], "note_tags":[{"note_id":note,"tag_id":tag}], "annotations":[{"id":Uuid::new_v4(),"note_id":note,"block_id":"qa-block"}], "problems":[{"id":problem,"title":"QA problem"}], "shelf_collections":[{"id":collection,"name":"QA collection"}], "shelf_items":[{"id":shelf,"collection_id":collection,"kind":"book","title":"QA book","asset_id":asset,"metadata":{"reader_material_id":book}}], "learning_spaces":[{"id":space,"kind":"course","title":"QA course"}], "learning_lessons":[{"id":lesson,"space_id":space,"title":"QA lesson","document":{"reader_material_id":article}}], "whiteboards":[{"id":Uuid::new_v4(),"title":"QA board"}], "spatial_documents":[{"id":Uuid::new_v4(),"title":"QA spatial"}], "workspace_states":[{"tabs":[],"layout":{"test":true}}], "assets":[{"record":{"id":asset,"kind":"book","filename":"book.pdf","mime_type":"application/pdf","storage_key":format!("{asset}/book.pdf"),"byte_size":bytes.len(),"checksum_sha256":checksum(bytes)},"content_base64":STANDARD.encode(bytes)}], "reader_materials":[{"id":article,"kind":"article","title":"QA article","body":"# searchable lesson"},{"id":book,"kind":"book","title":"QA book","asset_id":asset}], "provider_settings":[{"provider":"openai","model":"qa-model","enabled":false}], "agent_conversations":[{"id":chat,"title":"QA chat","context_type":"reader","context_id":article,"provider":"openai","model":"qa-model"}], "agent_messages":[{"id":message,"conversation_id":chat,"role":"assistant","content":"Partial response","metadata":{"status":"streaming"}}], "reader_chunks":[{"id":71,"material_id":article,"text":"searchable lesson"}], "reader_highlights":[{"id":Uuid::new_v4(),"material_id":article,"quote":"lesson"}], "reader_notes":[{"id":learning_note,"material_id":article,"title":"QA captured","body":"Learning body","chat_id":chat,"message_id":message,"captured_note_id":note}], "reader_review_cards":[{"id":Uuid::new_v4(),"material_id":article,"question":"Question","answer":"Answer","note_id":learning_note}], "reader_exercise_attempts":[{"material_id":article,"exercise_key":"exercise-1","answer":"42"}], "reader_material_revisions":[{"id":Uuid::new_v4(),"material_id":article,"version":1,"title":"QA previous","subject":"","body":"Old body","sources":[]}], "reader_rewrite_proposals":[{"id":Uuid::new_v4(),"material_id":article,"operation":"rewrite","instruction":"Improve","title":"QA article","replacement":"Better body","base_version":1,"status":"complete","provider":"openai","model":"qa-model"}], "agent_lesson_revisions":[{"id":Uuid::new_v4(),"lesson_id":lesson,"version":1,"title":"QA previous lesson","document":{}}], "reader_files":[{"material_id":article,"path":"assets/diagram.svg","byte_size":image.len(),"checksum_sha256":checksum(image),"content_base64":STANDARD.encode(image)}] }}); let _ = import_data(State(state.clone()), Json(payload)) .await .unwrap(); sqlx::query("update provider_settings set api_key_ciphertext='qa-encrypted-sentinel'") .execute(&pool) .await .unwrap(); let exported = export_data(State(state.clone())).await.unwrap().0; assert_eq!(exported["format"], "odre-platform-export"); assert!( exported["data"]["provider_settings"][0] .get("api_key_ciphertext") .is_none() ); assert_eq!(exported["data"]["notes"][0]["status"], "draft"); assert_eq!( exported["data"]["agent_messages"][0]["metadata"]["status"], "cancelled" ); assert!(exported["data"]["whiteboards"][0]["canvas"].is_object()); assert!( !serde_json::to_string(&exported) .unwrap() .contains("qa-encrypted-sentinel") ); assert!( exported["data"]["reader_chunks"][0] .get("search_vector") .is_none() ); for &table in TABLES { assert!(exported["data"][table].is_array(), "{table}"); } sqlx::query("truncate notes,tags,problems,shelf_collections,learning_spaces,whiteboards,spatial_documents,reader_materials,agent_conversations,assets cascade").execute(&pool).await.unwrap(); let _ = import_data(State(state.clone()), Json(exported.clone())) .await .unwrap(); let roundtrip = export_data(State(state.clone())).await.unwrap().0; for &table in TABLES { assert_eq!( exported["data"][table].as_array().unwrap().len(), roundtrip["data"][table].as_array().unwrap().len(), "{table}" ); } let key_preserved:bool=sqlx::query_scalar("select api_key_ciphertext='qa-encrypted-sentinel' from provider_settings where provider='openai'").fetch_one(&pool).await.unwrap(); assert!(key_preserved); let search_count:i64=sqlx::query_scalar("select count(*) from reader_chunks where search_vector @@ plainto_tsquery('simple','searchable')").fetch_one(&pool).await.unwrap(); assert_eq!(search_count, 1); let next: i64 = sqlx::query_scalar("select nextval(pg_get_serial_sequence('reader_chunks','id'))") .fetch_one(&pool) .await .unwrap(); assert_eq!(next, 72); let independent_article = Uuid::new_v4(); let independent = json!({"format":"odre-platform-export","version":2,"data":{ "reader_materials":[{"id":independent_article,"kind":"article","title":"Independent source"}], "reader_chunks":[{"id":71,"material_id":independent_article,"text":"Independent source index"}] }}); let _ = import_data(State(state.clone()), Json(independent.clone())) .await .unwrap(); let _ = import_data(State(state.clone()), Json(independent)) .await .unwrap(); let imported_id: i64 = sqlx::query_scalar("select id from reader_chunks where material_id=$1") .bind(independent_article) .fetch_one(&pool) .await .unwrap(); assert_ne!(imported_id, 71); let retained_text: String = sqlx::query_scalar("select text from reader_chunks where id=71") .fetch_one(&pool) .await .unwrap(); assert_eq!(retained_text, "searchable lesson"); let index_count: i64 = sqlx::query_scalar("select count(*) from reader_chunks") .fetch_one(&pool) .await .unwrap(); assert_eq!(index_count, 2); assert_eq!( fs::read(storage.join(format!("reader/{article}/assets/diagram.svg"))) .await .unwrap(), image ); let legacy = json!({"format":"forge-platform-export","version":1,"data":{"whiteboards":[{"id":Uuid::new_v4(),"title":"Legacy without canvas"}]}}); let _ = import_data(State(state.clone()), Json(legacy)) .await .unwrap(); let mut conflict = exported.clone(); conflict["data"]["assets"][0]["record"]["storage_key"] = json!(format!("{asset}/different.pdf")); assert!(matches!( import_data(State(state.clone()), Json(conflict)).await, Err(ApiError::Conflict(_)) )); assert_eq!( fs::read(storage.join(format!("{asset}/book.pdf"))) .await .unwrap(), bytes ); assert!(!storage.join(format!("{asset}/different.pdf")).exists()); let failed_asset = Uuid::new_v4(); let mut invalid_restore = exported.clone(); invalid_restore["data"]["notes"][0]["title"] = json!("Must roll back"); invalid_restore["data"]["reader_notes"][0]["chat_id"] = json!(Uuid::new_v4()); invalid_restore["data"]["assets"].as_array_mut().unwrap().push(json!({ "record":{"id":failed_asset,"kind":"file","filename":"new.txt","mime_type":"text/plain", "storage_key":format!("{failed_asset}/new.txt"),"byte_size":bytes.len(),"checksum_sha256":checksum(bytes)}, "content_base64":STANDARD.encode(bytes) })); assert!(matches!( import_data(State(state.clone()), Json(invalid_restore)).await, Err(ApiError::BadRequest(_)) )); let unchanged_title: String = sqlx::query_scalar("select title from notes where id=$1") .bind(note) .fetch_one(&pool) .await .unwrap(); assert_eq!(unchanged_title, "QA note"); let retained: bool = sqlx::query_scalar("select exists(select 1 from assets where id=$1)") .bind(failed_asset) .fetch_one(&pool) .await .unwrap(); assert!(!retained); assert!(!storage.join(format!("{failed_asset}/new.txt")).exists()); pool.close().await; sqlx::query(&format!("drop schema {schema} cascade")) .execute(&admin) .await .unwrap(); fs::remove_dir_all(storage).await.unwrap(); }}