1use crate::at_rest::{
26 MigrationGalaxyState, MigrationLedger, read_migration_ledger, write_migration_ledger,
27};
28use crate::store::MemoryStore;
29use lmdb::{Cursor, Database, Transaction, WriteFlags};
30use wm_core::{CoreError, Galaxy, Result};
31
32pub const DEFAULT_MIGRATION_BATCH: usize = 256;
34
35pub const RECORD_GALAXIES: [Galaxy; 12] = [
41 Galaxy::Aria,
42 Galaxy::Citta,
43 Galaxy::Codex,
44 Galaxy::Journals,
45 Galaxy::Dreams,
46 Galaxy::Research,
47 Galaxy::Sessions,
48 Galaxy::Substrate,
49 Galaxy::Tutorial,
50 Galaxy::Universal,
51 Galaxy::Valkyrie,
52 Galaxy::Telemetry,
53];
54
55#[derive(Debug, Clone, PartialEq, Eq)]
57pub struct GalaxyMigrationReport {
58 pub galaxy: String,
60 pub scanned: u64,
62 pub encrypted: u64,
64 pub already_sealed: u64,
66 pub skipped_non_record: u64,
68 pub undecodable: u64,
70 pub done: bool,
72}
73
74#[derive(Debug, Clone, PartialEq, Eq)]
76pub struct AtRestMigrationReport {
77 pub galaxies: Vec<GalaxyMigrationReport>,
79 pub no_keyring: bool,
81 pub total_encrypted: u64,
83 pub total_already_sealed: u64,
85 pub total_undecodable: u64,
87}
88
89impl AtRestMigrationReport {
90 const fn empty_no_keyring() -> Self {
91 Self {
92 galaxies: Vec::new(),
93 no_keyring: true,
94 total_encrypted: 0,
95 total_already_sealed: 0,
96 total_undecodable: 0,
97 }
98 }
99
100 #[must_use]
102 pub fn all_done(&self) -> bool {
103 self.no_keyring || self.galaxies.iter().all(|g| g.done)
104 }
105}
106
107#[derive(Debug, Clone, PartialEq, Eq, Default)]
110pub struct GalaxyAtRestCounts {
111 pub galaxy: String,
113 pub sealed: u64,
115 pub plaintext: u64,
117 pub non_record: u64,
119}
120
121pub fn at_rest_record_counts(store: &MemoryStore) -> Result<Vec<GalaxyAtRestCounts>> {
125 let mut out = Vec::with_capacity(RECORD_GALAXIES.len());
126 for galaxy in RECORD_GALAXIES {
127 let db = store.galaxy_db(galaxy)?;
128 let tx = store
129 .env()
130 .begin_ro_txn()
131 .map_err(|e| CoreError::Memory(format!("LMDB ro_txn failed (at-rest counts): {e}")))?;
132 let mut cursor = tx
133 .open_ro_cursor(db)
134 .map_err(|e| CoreError::Memory(format!("LMDB cursor failed (at-rest counts): {e}")))?;
135 let mut counts = GalaxyAtRestCounts {
136 galaxy: galaxy.db_name().to_string(),
137 ..GalaxyAtRestCounts::default()
138 };
139 for (key, value) in cursor.iter() {
140 if !crate::at_rest::migration_candidate_key(key) {
141 counts.non_record += 1;
142 } else if crate::codec::is_sealed_record(value) {
143 counts.sealed += 1;
144 } else {
145 counts.plaintext += 1;
146 }
147 }
148 drop(cursor);
149 tx.commit()
150 .map_err(|e| CoreError::Memory(format!("LMDB commit failed (at-rest counts): {e}")))?;
151 out.push(counts);
152 }
153 Ok(out)
154}
155
156pub fn migration_ledger(store: &MemoryStore) -> Result<Option<MigrationLedger>> {
161 let Some(db) = store.keyring_db() else {
162 return Ok(None);
163 };
164 read_migration_ledger(store.env(), db).map(Some)
165}
166
167pub fn migrate_at_rest_records(
181 store: &MemoryStore,
182 galaxies: &[Galaxy],
183 batch: usize,
184) -> Result<AtRestMigrationReport> {
185 let Some(keyring_db) = store.keyring_db() else {
186 return Ok(AtRestMigrationReport::empty_no_keyring());
187 };
188 let batch = if batch == 0 {
189 DEFAULT_MIGRATION_BATCH
190 } else {
191 batch
192 };
193
194 let mut ledger = read_migration_ledger(store.env(), keyring_db)?;
195 let mut report = AtRestMigrationReport {
196 galaxies: Vec::new(),
197 no_keyring: false,
198 total_encrypted: 0,
199 total_already_sealed: 0,
200 total_undecodable: 0,
201 };
202
203 for &galaxy in galaxies {
204 let name = galaxy.db_name().to_string();
205 let resume = ledger.galaxies.get(&name).cloned().unwrap_or_default();
206 let galaxy_report = migrate_galaxy(store, galaxy, keyring_db, batch, &resume, &mut ledger)?;
207
208 report.total_encrypted += galaxy_report.encrypted;
209 report.total_already_sealed += galaxy_report.already_sealed;
210 report.total_undecodable += galaxy_report.undecodable;
211 report.galaxies.push(galaxy_report);
212 }
213
214 Ok(report)
215}
216
217fn migrate_galaxy(
221 store: &MemoryStore,
222 galaxy: Galaxy,
223 keyring_db: Database,
224 batch: usize,
225 resume: &MigrationGalaxyState,
226 ledger: &mut MigrationLedger,
227) -> Result<GalaxyMigrationReport> {
228 let Some(dek) = store.record_cipher(galaxy) else {
229 return Err(CoreError::Memory(format!(
230 "at-rest migration requested for {} but its DEK is not loaded",
231 galaxy.db_name()
232 )));
233 };
234 let db = store.galaxy_db(galaxy)?;
235 let resume_key = if resume.done {
239 Vec::new()
240 } else {
241 decode_cursor(&resume.cursor_hex)?
242 };
243 let Some(mut cursor_key) = first_key_at_or_after(store, db, &resume_key)? else {
249 let entry = ledger
250 .galaxies
251 .entry(galaxy.db_name().to_string())
252 .or_default();
253 entry.encrypted = resume.encrypted;
254 entry.cursor_hex = String::new();
255 entry.done = true;
256 ledger.updated_at = chrono::Utc::now().to_rfc3339();
257 write_migration_ledger(store.env(), keyring_db, ledger)?;
258 return Ok(GalaxyMigrationReport {
259 galaxy: galaxy.db_name().to_string(),
260 scanned: 0,
261 encrypted: 0,
262 already_sealed: 0,
263 skipped_non_record: 0,
264 undecodable: 0,
265 done: true,
266 });
267 };
268 let base_encrypted = resume.encrypted;
271 let mut anchor_examined = !resume_key.is_empty() && cursor_key == resume_key;
274
275 let mut report = GalaxyMigrationReport {
276 galaxy: galaxy.db_name().to_string(),
277 scanned: 0,
278 encrypted: 0,
279 already_sealed: 0,
280 skipped_non_record: 0,
281 undecodable: 0,
282 done: false,
283 };
284
285 loop {
286 let mut candidates: Vec<(Vec<u8>, Vec<u8>, u64)> = Vec::new();
290 let mut last_examined: Option<Vec<u8>> = None;
291 {
292 let tx = store
293 .env()
294 .begin_ro_txn()
295 .map_err(|e| CoreError::Memory(format!("LMDB ro_txn failed (migration): {e}")))?;
296 let mut cursor = tx
297 .open_ro_cursor(db)
298 .map_err(|e| CoreError::Memory(format!("LMDB cursor failed (migration): {e}")))?;
299 for (key, value) in cursor.iter_from(cursor_key.as_slice()) {
300 if anchor_examined && key == cursor_key.as_slice() {
301 continue;
302 }
303 if candidates.len() >= batch {
304 break;
305 }
306 last_examined = Some(key.to_vec());
307 report.scanned += 1;
308 if !crate::at_rest::migration_candidate_key(key) {
309 report.skipped_non_record += 1;
310 continue;
311 }
312 if crate::codec::is_sealed_record(value) {
313 report.already_sealed += 1;
314 continue;
315 }
316 match crate::at_rest::decode_plaintext_for_migration(value) {
317 Some(memory) => {
318 candidates.push((key.to_vec(), value.to_vec(), memory.metadata.version));
319 }
320 None => report.undecodable += 1,
321 }
322 }
323 drop(cursor);
324 tx.commit()
325 .map_err(|e| CoreError::Memory(format!("LMDB commit failed (migration): {e}")))?;
326 }
327
328 let Some(last_key) = last_examined else {
329 report.done = true;
331 break;
332 };
333
334 if !candidates.is_empty() {
336 let mut tx = store
337 .env()
338 .begin_rw_txn()
339 .map_err(|e| CoreError::Memory(format!("LMDB rw_txn failed (migration): {e}")))?;
340 for (key, plaintext, version) in &candidates {
341 let record_id: [u8; 16] = key
342 .as_slice()
343 .try_into()
344 .map_err(|_| CoreError::Memory("migration key is not a 16-byte id".into()))?;
345 let sealed = crate::at_rest::seal_migrated_record(
346 plaintext,
347 dek,
348 galaxy.db_name(),
349 &record_id,
350 *version,
351 )?;
352 tx.put(db, key, &sealed, WriteFlags::default())
353 .map_err(|e| CoreError::Memory(format!("LMDB put failed (migration): {e}")))?;
354 report.encrypted += 1;
355 }
356 tx.commit()
357 .map_err(|e| CoreError::Memory(format!("LMDB commit failed (migration): {e}")))?;
358 }
359
360 cursor_key = last_key;
361 anchor_examined = true;
362 let entry = ledger.galaxies.entry(report.galaxy.clone()).or_default();
363 entry.encrypted = base_encrypted + report.encrypted;
364 entry.cursor_hex = encode_cursor(&cursor_key);
365 entry.done = false;
366 ledger.updated_at = chrono::Utc::now().to_rfc3339();
367 write_migration_ledger(store.env(), keyring_db, ledger)?;
368 }
369
370 let entry = ledger.galaxies.entry(report.galaxy.clone()).or_default();
371 entry.encrypted = base_encrypted + report.encrypted;
372 entry.cursor_hex = String::new();
373 entry.done = true;
374 ledger.updated_at = chrono::Utc::now().to_rfc3339();
375 write_migration_ledger(store.env(), keyring_db, ledger)?;
376
377 Ok(report)
378}
379
380fn first_key_at_or_after(
384 store: &MemoryStore,
385 db: Database,
386 resume_key: &[u8],
387) -> Result<Option<Vec<u8>>> {
388 let tx = store
389 .env()
390 .begin_ro_txn()
391 .map_err(|e| CoreError::Memory(format!("LMDB ro_txn failed (migration): {e}")))?;
392 let mut cursor = tx
393 .open_ro_cursor(db)
394 .map_err(|e| CoreError::Memory(format!("LMDB cursor failed (migration): {e}")))?;
395 const MDB_SET_RANGE: u32 = 17;
401 let found = if resume_key.is_empty() {
402 cursor.iter().next().map(|(key, _)| key.to_vec())
403 } else {
404 match cursor.get(Some(resume_key), None, MDB_SET_RANGE) {
405 Ok((key, _)) => key.map(<[u8]>::to_vec),
406 Err(lmdb::Error::NotFound) => None,
407 Err(e) => {
408 return Err(CoreError::Memory(format!(
409 "LMDB cursor seek failed (migration): {e}"
410 )));
411 }
412 }
413 };
414 drop(cursor);
415 tx.commit()
416 .map_err(|e| CoreError::Memory(format!("LMDB commit failed (migration): {e}")))?;
417 Ok(found)
418}
419
420fn encode_cursor(key: &[u8]) -> String {
421 const HEX: &[u8; 16] = b"0123456789abcdef";
422 let mut out = String::with_capacity(key.len() * 2);
423 for b in key {
424 out.push(HEX[(b >> 4) as usize] as char);
425 out.push(HEX[(b & 0x0f) as usize] as char);
426 }
427 out
428}
429
430fn decode_cursor(hex: &str) -> Result<Vec<u8>> {
431 if hex.is_empty() {
432 return Ok(Vec::new());
433 }
434 if hex.len() % 2 != 0 {
435 return Err(CoreError::Memory(
436 "migration ledger cursor is not valid hex".into(),
437 ));
438 }
439 let mut out = Vec::with_capacity(hex.len() / 2);
440 let bytes = hex.as_bytes();
441 for chunk in bytes.chunks_exact(2) {
442 let hi = hex_val(chunk[0])
443 .ok_or_else(|| CoreError::Memory("migration ledger cursor is not valid hex".into()))?;
444 let lo = hex_val(chunk[1])
445 .ok_or_else(|| CoreError::Memory("migration ledger cursor is not valid hex".into()))?;
446 out.push((hi << 4) | lo);
447 }
448 Ok(out)
449}
450
451const fn hex_val(b: u8) -> Option<u8> {
452 match b {
453 b'0'..=b'9' => Some(b - b'0'),
454 b'a'..=b'f' => Some(b - b'a' + 10),
455 b'A'..=b'F' => Some(b - b'A' + 10),
456 _ => None,
457 }
458}
459
460#[cfg(test)]
461mod tests {
462 use super::*;
463 use crate::at_rest::AtRestConfig;
464 use crate::memory::Memory;
465 use crate::store::MemoryStore;
466
467 const TEST_MAP: usize = 16 * 1024 * 1024;
468
469 fn open_plain(store_dir: &std::path::Path) -> MemoryStore {
470 MemoryStore::open_with_at_rest(store_dir, TEST_MAP, &AtRestConfig::off()).unwrap()
471 }
472
473 fn open_keyfile(store_dir: &std::path::Path) -> MemoryStore {
474 MemoryStore::open_with_at_rest(store_dir, TEST_MAP, &AtRestConfig::keyfile()).unwrap()
475 }
476
477 fn is_sealed(store: &MemoryStore, galaxy: Galaxy, id: uuid::Uuid) -> bool {
478 store
479 .get_raw(galaxy, id.as_bytes())
480 .unwrap()
481 .is_some_and(|raw| crate::codec::is_sealed_record(&raw))
482 }
483
484 #[test]
485 fn plaintext_store_is_a_provable_noop() {
486 let tmp = tempfile::tempdir().unwrap();
487 let store = open_plain(tmp.path());
488 let mem = Memory::new(Galaxy::Codex, "still plaintext".into());
489 let id = mem.metadata.id;
490 store.put(Galaxy::Codex, &mem).unwrap();
491
492 let report =
493 migrate_at_rest_records(&store, &RECORD_GALAXIES, DEFAULT_MIGRATION_BATCH).unwrap();
494 assert!(report.no_keyring);
495 assert!(report.all_done());
496 assert_eq!(report.total_encrypted, 0);
497 assert!(migration_ledger(&store).unwrap().is_none());
498 assert!(
499 !is_sealed(&store, Galaxy::Codex, id),
500 "plaintext store must stay byte-identical"
501 );
502 }
503
504 #[test]
505 fn mixed_store_migrates_all_records_and_is_idempotent() {
506 let tmp = tempfile::tempdir().unwrap();
507 let mut ids = Vec::new();
508 {
509 let store = open_plain(tmp.path());
510 for i in 0..5 {
511 let mem = Memory::new(Galaxy::Codex, format!("legacy record {i}"));
512 ids.push(mem.metadata.id);
513 store.put(Galaxy::Codex, &mem).unwrap();
514 }
515 }
516
517 let store = open_keyfile(tmp.path());
518 let report =
519 migrate_at_rest_records(&store, &RECORD_GALAXIES, DEFAULT_MIGRATION_BATCH).unwrap();
520 assert!(!report.no_keyring);
521 assert!(report.all_done(), "{report:?}");
522 assert_eq!(report.total_encrypted, 5, "{report:?}");
523 assert!(ids.iter().all(|id| is_sealed(&store, Galaxy::Codex, *id)));
524 assert_eq!(store.scan_all(Galaxy::Codex).unwrap().len(), 5);
525
526 let ledger = migration_ledger(&store).unwrap().unwrap();
527 assert_eq!(ledger.galaxies["codex"].encrypted, 5);
528 assert!(ledger.galaxies["codex"].done);
529
530 let second =
532 migrate_at_rest_records(&store, &RECORD_GALAXIES, DEFAULT_MIGRATION_BATCH).unwrap();
533 assert_eq!(second.total_encrypted, 0, "{second:?}");
534 assert!(second.all_done(), "{second:?}");
535 assert_eq!(store.scan_all(Galaxy::Codex).unwrap().len(), 5);
536 }
537
538 #[test]
539 fn resume_from_a_mid_way_ledger_completes_the_remaining_records() {
540 let tmp = tempfile::tempdir().unwrap();
541 let mut ids = Vec::new();
542 {
543 let store = open_plain(tmp.path());
544 for i in 0..10 {
545 let mem = Memory::new(Galaxy::Codex, format!("resume record {i}"));
546 ids.push(mem.metadata.id);
547 store.put(Galaxy::Codex, &mem).unwrap();
548 }
549 }
550 ids.sort_by_key(|id| *id.as_bytes());
554
555 let store = open_keyfile(tmp.path());
556 let keyring_db = store.keyring_db().expect("keyring");
557 {
558 let dek = *store
559 .at_rest_state()
560 .unwrap()
561 .galaxy_dek(Galaxy::Codex.db_name())
562 .unwrap();
563 let db = store.galaxy_db(Galaxy::Codex).unwrap();
564 let mut tx = store.env().begin_rw_txn().unwrap();
565 for id in &ids[..4] {
566 let memory = store.get(Galaxy::Codex, *id).unwrap().unwrap();
567 let sealed = crate::codec::seal_record(
568 &rmp_serde::to_vec_named(&memory).unwrap(),
569 &dek,
570 Galaxy::Codex.db_name(),
571 id.as_bytes(),
572 memory.metadata.version,
573 )
574 .unwrap();
575 tx.put(db, id.as_bytes(), &sealed, WriteFlags::default())
576 .unwrap();
577 }
578 tx.commit().unwrap();
579
580 let mut ledger = MigrationLedger::default();
581 ledger.galaxies.insert(
582 Galaxy::Codex.db_name().to_string(),
583 MigrationGalaxyState {
584 encrypted: 4,
585 cursor_hex: encode_cursor(ids[3].as_bytes()),
586 done: false,
587 },
588 );
589 write_migration_ledger(store.env(), keyring_db, &ledger).unwrap();
590 }
591
592 let report = migrate_at_rest_records(&store, &[Galaxy::Codex], 4).unwrap();
593 assert!(report.all_done(), "{report:?}");
594 assert_eq!(report.total_encrypted, 6, "{report:?}");
595 assert_eq!(report.total_already_sealed, 0, "{report:?}");
596 assert!(ids.iter().all(|id| is_sealed(&store, Galaxy::Codex, *id)));
597 assert_eq!(store.scan_all(Galaxy::Codex).unwrap().len(), 10);
598
599 let ledger = migration_ledger(&store).unwrap().unwrap();
600 assert!(ledger.galaxies["codex"].done);
601 assert_eq!(ledger.galaxies["codex"].encrypted, 10);
602 }
603
604 #[test]
605 fn a_deleted_resume_cursor_resumes_at_its_successor() {
606 let tmp = tempfile::tempdir().unwrap();
607 let mut ids = Vec::new();
608 {
609 let store = open_plain(tmp.path());
610 for i in 0..3 {
611 let mem = Memory::new(Galaxy::Codex, format!("shrunk store {i}"));
612 ids.push(mem.metadata.id);
613 store.put(Galaxy::Codex, &mem).unwrap();
614 }
615 }
616 ids.sort_by_key(|id| *id.as_bytes());
617
618 let store = open_keyfile(tmp.path());
619 let keyring_db = store.keyring_db().expect("keyring");
620 store.delete(Galaxy::Codex, ids[0]).unwrap();
624 let mut ledger = MigrationLedger::default();
625 ledger.galaxies.insert(
626 Galaxy::Codex.db_name().to_string(),
627 MigrationGalaxyState {
628 encrypted: 0,
629 cursor_hex: encode_cursor(ids[0].as_bytes()),
630 done: false,
631 },
632 );
633 write_migration_ledger(store.env(), keyring_db, &ledger).unwrap();
634
635 let report = migrate_at_rest_records(&store, &[Galaxy::Codex], 0).unwrap();
636 assert!(report.all_done(), "{report:?}");
637 assert_eq!(report.total_encrypted, 2, "{report:?}");
638 assert!(is_sealed(&store, Galaxy::Codex, ids[1]));
639 assert!(is_sealed(&store, Galaxy::Codex, ids[2]));
640 }
641
642 #[test]
643 fn undecodable_plaintext_is_counted_skipped_and_left_untouched() {
644 let tmp = tempfile::tempdir().unwrap();
645 let id = uuid::Uuid::new_v4();
646 {
647 let store = open_plain(tmp.path());
648 store
649 .put_raw(Galaxy::Codex, id.as_bytes(), b"not a memory record")
650 .unwrap();
651 let mem = Memory::new(Galaxy::Codex, "good record".into());
652 store.put(Galaxy::Codex, &mem).unwrap();
653 }
654
655 let store = open_keyfile(tmp.path());
656 let report = migrate_at_rest_records(&store, &[Galaxy::Codex], 0).unwrap();
657 assert_eq!(report.total_encrypted, 1, "{report:?}");
658 assert_eq!(report.total_undecodable, 1, "{report:?}");
659 assert_eq!(
660 store
661 .get_raw(Galaxy::Codex, id.as_bytes())
662 .unwrap()
663 .unwrap(),
664 b"not a memory record",
665 "undecodable rows are never rewritten"
666 );
667 }
668
669 #[test]
670 fn non_record_galaxies_are_never_touched() {
671 let tmp = tempfile::tempdir().unwrap();
672 {
673 let store = open_plain(tmp.path());
674 store
675 .put_raw(Galaxy::Karma, b"karma:entry:1", b"chain-payload")
676 .unwrap();
677 store
678 .put_raw(Galaxy::Associations, &[0u8; 32], b"assoc-payload")
679 .unwrap();
680 }
681
682 let store = open_keyfile(tmp.path());
683 let report =
684 migrate_at_rest_records(&store, &RECORD_GALAXIES, DEFAULT_MIGRATION_BATCH).unwrap();
685 assert_eq!(report.total_encrypted, 0, "{report:?}");
686 assert_eq!(report.total_undecodable, 0, "{report:?}");
687 assert_eq!(
688 store
689 .get_raw(Galaxy::Karma, b"karma:entry:1")
690 .unwrap()
691 .unwrap(),
692 b"chain-payload"
693 );
694 assert_eq!(
695 store
696 .get_raw(Galaxy::Associations, &[0u8; 32])
697 .unwrap()
698 .unwrap(),
699 b"assoc-payload"
700 );
701 }
702
703 #[test]
704 fn a_plaintext_record_added_after_done_is_resealed_on_the_next_run() {
705 let tmp = tempfile::tempdir().unwrap();
706 let store = open_keyfile(tmp.path());
707 let report =
708 migrate_at_rest_records(&store, &RECORD_GALAXIES, DEFAULT_MIGRATION_BATCH).unwrap();
709 assert!(report.all_done(), "{report:?}");
710
711 let legacy = Memory::new(Galaxy::Codex, "restored plaintext".into());
713 let id = legacy.metadata.id;
714 let plaintext = {
715 let off = open_plain(&tmp.path().join("scratch"));
716 off.put(Galaxy::Codex, &legacy).unwrap();
717 off.get_raw(Galaxy::Codex, id.as_bytes()).unwrap().unwrap()
718 };
719 store
720 .put_raw(Galaxy::Codex, id.as_bytes(), &plaintext)
721 .unwrap();
722 assert!(!is_sealed(&store, Galaxy::Codex, id));
723
724 let second =
725 migrate_at_rest_records(&store, &RECORD_GALAXIES, DEFAULT_MIGRATION_BATCH).unwrap();
726 assert_eq!(second.total_encrypted, 1, "{second:?}");
727 assert!(is_sealed(&store, Galaxy::Codex, id));
728 }
729}