1use super::{
2 ledger,
3 manager::MEMORY_MANAGER,
4 registry::{
5 MemoryRange, MemoryRangeAuthority, MemoryRegistry, MemoryRegistryError, defer_register,
6 defer_register_with_key_metadata,
7 },
8 runtime::{MemoryRuntimeApi, registry::MemoryRegistryRuntime},
9};
10use crate::cdk::structures::{
11 DefaultMemoryImpl,
12 memory::{MemoryId, VirtualMemory},
13};
14use ic_memory::{
15 AllocationSession, AllocationSessionError, AllocationSlotDescriptor, MemoryManagerSlotError,
16 StableKey, StorageSubstrate,
17};
18
19pub struct MemoryApi;
26
27#[derive(Clone, Debug, Eq, PartialEq)]
33pub struct MemoryInspection {
34 pub id: u8,
36 pub owner: String,
38 pub range: MemoryRange,
40 pub label: Option<String>,
42 pub stable_key: Option<String>,
44 pub schema_version: Option<u32>,
46 pub schema_fingerprint: Option<String>,
48}
49
50#[derive(Clone, Debug, Eq, PartialEq)]
56pub struct RegisteredMemory {
57 pub id: u8,
59 pub owner: String,
61 pub range: MemoryRange,
63 pub label: String,
65 pub stable_key: String,
67 pub schema_version: Option<u32>,
69 pub schema_fingerprint: Option<String>,
71}
72
73#[derive(Clone, Debug, Eq, PartialEq)]
79pub struct LedgerSnapshot {
80 pub magic: u64,
82 pub format_id: u32,
84 pub schema_version: u32,
86 pub layout_epoch: u32,
88 pub header_len: u32,
90 pub header_checksum: u64,
92 pub current_generation: u64,
94 pub authorities: Vec<MemoryRangeAuthority>,
96 pub ranges: Vec<(String, MemoryRange)>,
98 pub entries: Vec<(u8, super::registry::MemoryRegistryEntry)>,
100}
101
102impl MemoryApi {
103 pub fn bootstrap_owner_range(
105 crate_name: &'static str,
106 start: u8,
107 end: u8,
108 ) -> Result<(), MemoryRegistryError> {
109 let _ = MemoryRuntimeApi::bootstrap_registry(crate_name, start, end)?;
110 Ok(())
111 }
112
113 pub fn bootstrap_pending() -> Result<(), MemoryRegistryError> {
116 let _ = MemoryRuntimeApi::bootstrap_registry_without_range()?;
117 Ok(())
118 }
119
120 pub fn declare(id: u8, crate_name: &str, label: &str) -> Result<(), MemoryRegistryError> {
124 if MemoryRegistryRuntime::is_initialized() {
125 return Err(MemoryRegistryError::RegistrationAfterBootstrap {
126 ranges: 0,
127 registrations: 1,
128 });
129 }
130
131 defer_register(id, crate_name, label)
132 }
133
134 pub fn declare_with_key(
138 id: u8,
139 crate_name: &str,
140 label: &str,
141 stable_key: &str,
142 ) -> Result<(), MemoryRegistryError> {
143 Self::declare_with_key_metadata(id, crate_name, label, stable_key, None, None)
144 }
145
146 pub fn declare_with_key_metadata(
151 id: u8,
152 crate_name: &str,
153 label: &str,
154 stable_key: &str,
155 schema_version: Option<u32>,
156 schema_fingerprint: Option<&str>,
157 ) -> Result<(), MemoryRegistryError> {
158 if MemoryRegistryRuntime::is_initialized() {
159 return Err(MemoryRegistryError::RegistrationAfterBootstrap {
160 ranges: 0,
161 registrations: 1,
162 });
163 }
164
165 defer_register_with_key_metadata(
166 id,
167 crate_name,
168 label,
169 stable_key,
170 schema_version,
171 schema_fingerprint,
172 )
173 }
174
175 pub fn register(
180 id: u8,
181 crate_name: &str,
182 label: &str,
183 ) -> Result<VirtualMemory<DefaultMemoryImpl>, MemoryRegistryError> {
184 if !MemoryRegistryRuntime::is_initialized() {
185 return Err(MemoryRegistryError::RegistryNotBootstrapped);
186 }
187
188 if let Some(entry) = MemoryRegistry::get(id)
189 && entry.crate_name == crate_name
190 && entry.label == label
191 {
192 return Self::register_with_key(id, crate_name, label, &entry.stable_key);
193 }
194
195 Err(MemoryRegistryError::RegistrationAfterBootstrap {
196 ranges: 0,
197 registrations: 1,
198 })
199 }
200
201 pub fn register_with_key(
203 id: u8,
204 _crate_name: &str,
205 _label: &str,
206 stable_key: &str,
207 ) -> Result<VirtualMemory<DefaultMemoryImpl>, MemoryRegistryError> {
208 if !MemoryRegistryRuntime::is_initialized() {
209 return Err(MemoryRegistryError::RegistryNotBootstrapped);
210 }
211
212 let key =
213 StableKey::parse(stable_key).map_err(|err| MemoryRegistryError::InvalidStableKey {
214 stable_key: err.stable_key,
215 reason: err.reason,
216 })?;
217 let validated = MemoryRegistryRuntime::validated_allocations()?;
218 let slot =
219 validated
220 .slot_for(&key)
221 .ok_or(MemoryRegistryError::RegistrationAfterBootstrap {
222 ranges: 0,
223 registrations: 1,
224 })?;
225 let slot_id = memory_manager_id_from_slot(slot)?;
226 if slot_id != id {
227 return Err(MemoryRegistryError::RegistrationAfterBootstrap {
228 ranges: 0,
229 registrations: 1,
230 });
231 }
232
233 let session = AllocationSession::new(MemoryManagerSubstrate, validated);
234 session
235 .open(&key)
236 .map_err(memory_registry_error_from_session_error)
237 }
238
239 #[must_use]
241 pub fn inspect(id: u8) -> Option<MemoryInspection> {
242 let range = MemoryRegistry::export_range_entries()
243 .into_iter()
244 .find(|entry| entry.range.contains(id))?;
245 let entry = MemoryRegistry::get(id);
246 let label = entry.as_ref().map(|entry| entry.label.clone());
247 let stable_key = entry.as_ref().map(|entry| entry.stable_key.clone());
248 let schema_version = entry.as_ref().and_then(|entry| entry.schema_version);
249 let schema_fingerprint = entry.and_then(|entry| entry.schema_fingerprint);
250
251 Some(MemoryInspection {
252 id,
253 owner: range.owner,
254 range: range.range,
255 label,
256 stable_key,
257 schema_version,
258 schema_fingerprint,
259 })
260 }
261
262 #[must_use]
264 pub fn registered() -> Vec<RegisteredMemory> {
265 MemoryRegistry::export_ids_by_range()
266 .into_iter()
267 .flat_map(|snapshot| {
268 snapshot
269 .entries
270 .into_iter()
271 .map(move |(id, entry)| RegisteredMemory {
272 id,
273 owner: snapshot.owner.clone(),
274 range: snapshot.range,
275 label: entry.label,
276 stable_key: entry.stable_key,
277 schema_version: entry.schema_version,
278 schema_fingerprint: entry.schema_fingerprint,
279 })
280 })
281 .collect()
282 }
283
284 #[must_use]
286 pub fn registered_for_owner(owner: &str) -> Vec<RegisteredMemory> {
287 Self::registered()
288 .into_iter()
289 .filter(|entry| entry.owner == owner)
290 .collect()
291 }
292
293 #[must_use]
295 pub fn find(owner: &str, label: &str) -> Option<RegisteredMemory> {
296 Self::registered()
297 .into_iter()
298 .find(|entry| entry.owner == owner && entry.label == label)
299 }
300
301 pub fn ledger_snapshot() -> Result<LedgerSnapshot, MemoryRegistryError> {
303 #[cfg(target_arch = "wasm32")]
304 {
305 let snapshot = ledger::try_diagnostic_snapshot()?;
306 Ok(LedgerSnapshot::from(snapshot))
307 }
308
309 #[cfg(not(target_arch = "wasm32"))]
310 {
311 let snapshot = ledger::try_snapshot()?;
312 Ok(LedgerSnapshot::from(snapshot))
313 }
314 }
315}
316
317impl From<ledger::MemoryLayoutLedgerSnapshot> for LedgerSnapshot {
318 fn from(snapshot: ledger::MemoryLayoutLedgerSnapshot) -> Self {
319 Self {
320 magic: snapshot.magic,
321 format_id: snapshot.format_id,
322 schema_version: snapshot.schema_version,
323 layout_epoch: snapshot.layout_epoch,
324 header_len: snapshot.header_len,
325 header_checksum: snapshot.header_checksum,
326 current_generation: snapshot.current_generation,
327 authorities: snapshot.authorities,
328 ranges: snapshot.ranges,
329 entries: snapshot.entries,
330 }
331 }
332}
333
334fn open_memory(id: u8) -> VirtualMemory<DefaultMemoryImpl> {
336 MEMORY_MANAGER.with_borrow_mut(|mgr| mgr.get(MemoryId::new(id)))
337}
338
339#[derive(Clone, Copy, Debug, Eq, PartialEq)]
340struct MemoryManagerSubstrate;
341
342impl StorageSubstrate for MemoryManagerSubstrate {
343 type Slot = u8;
344 type LedgerMemory = VirtualMemory<DefaultMemoryImpl>;
345 type MemoryHandle = VirtualMemory<DefaultMemoryImpl>;
346 type Error = MemoryRegistryError;
347
348 fn open_ledger(&self) -> Result<Self::LedgerMemory, Self::Error> {
349 Ok(open_memory(ledger::MEMORY_LAYOUT_LEDGER_ID))
350 }
351
352 fn open_slot(
353 &self,
354 slot: &AllocationSlotDescriptor,
355 ) -> Result<Self::MemoryHandle, Self::Error> {
356 let id = memory_manager_id_from_slot(slot)?;
357 Ok(open_memory(id))
358 }
359
360 fn describe_slot(&self, slot: &Self::Slot) -> AllocationSlotDescriptor {
361 AllocationSlotDescriptor::memory_manager(*slot)
362 }
363}
364
365fn memory_registry_error_from_session_error(
366 err: AllocationSessionError<MemoryRegistryError>,
367) -> MemoryRegistryError {
368 match err {
369 AllocationSessionError::UnknownStableKey(_) => {
370 MemoryRegistryError::RegistrationAfterBootstrap {
371 ranges: 0,
372 registrations: 1,
373 }
374 }
375 AllocationSessionError::Substrate(err) => err,
376 }
377}
378
379fn memory_manager_id_from_slot(slot: &AllocationSlotDescriptor) -> Result<u8, MemoryRegistryError> {
380 slot.memory_manager_id()
381 .map_err(memory_registry_error_from_slot_error)
382}
383
384fn memory_registry_error_from_slot_error(err: MemoryManagerSlotError) -> MemoryRegistryError {
385 match err {
386 MemoryManagerSlotError::InvalidMemoryManagerId { id } => {
387 MemoryRegistryError::ReservedInternalId { id }
388 }
389 MemoryManagerSlotError::UnsupportedSlot
390 | MemoryManagerSlotError::UnsupportedSubstrate { .. }
391 | MemoryManagerSlotError::UnsupportedDescriptorVersion { .. } => {
392 MemoryRegistryError::LedgerCorrupt {
393 reason: "unsupported MemoryManager allocation slot descriptor",
394 }
395 }
396 }
397}
398
399#[cfg(test)]
404mod tests {
405 use super::super::registry::{
406 MemoryRegistryError, defer_register, defer_reserve_range, reset_for_tests,
407 };
408 use super::*;
409
410 #[test]
411 fn register_memory_opens_validated_memory_for_reserved_slot() {
412 reset_for_tests();
413 defer_reserve_range("crate_a", 100, 102).expect("defer range");
414 defer_register(101, "crate_a", "slot").expect("defer register");
415 MemoryApi::bootstrap_pending().expect("bootstrap registry");
416
417 let _memory = MemoryApi::register(101, "crate_a", "slot").expect("open memory");
418 }
419
420 #[test]
421 fn register_memory_is_idempotent_for_same_entry() {
422 reset_for_tests();
423 defer_reserve_range("crate_a", 100, 102).expect("defer range");
424 defer_register(101, "crate_a", "slot").expect("defer register");
425 MemoryApi::bootstrap_pending().expect("bootstrap registry");
426 let _ = MemoryApi::register(101, "crate_a", "slot").expect("first open succeeds");
427
428 let _ = MemoryApi::register(101, "crate_a", "slot").expect("second open succeeds");
429 }
430
431 #[test]
432 fn register_with_key_opens_validated_explicit_key() {
433 reset_for_tests();
434 defer_reserve_range("crate_a", 100, 102).expect("defer range");
435 MemoryApi::declare_with_key(101, "crate_a", "slot", "app.crate_a.slot.v1")
436 .expect("defer register");
437 MemoryApi::bootstrap_pending().expect("bootstrap registry");
438
439 let _memory = MemoryApi::register_with_key(101, "crate_a", "slot", "app.crate_a.slot.v1")
440 .expect("open memory");
441 }
442
443 #[test]
444 fn declare_with_key_metadata_records_schema_metadata() {
445 reset_for_tests();
446 defer_reserve_range("crate_a", 100, 102).expect("defer range");
447 MemoryApi::declare_with_key_metadata(
448 101,
449 "crate_a",
450 "slot",
451 "app.crate_a.slot.v1",
452 Some(3),
453 Some("sha256:abc123"),
454 )
455 .expect("defer register");
456 MemoryApi::bootstrap_pending().expect("bootstrap registry");
457
458 let registered = MemoryApi::find("crate_a", "slot").expect("registered memory");
459 assert_eq!(registered.schema_version, Some(3));
460 assert_eq!(
461 registered.schema_fingerprint.as_deref(),
462 Some("sha256:abc123")
463 );
464
465 let snapshot = MemoryApi::ledger_snapshot().expect("ledger snapshot");
466 assert_eq!(snapshot.format_id, 1);
467 assert_eq!(snapshot.schema_version, 1);
468 assert_eq!(snapshot.layout_epoch, 1);
469 assert!(snapshot.current_generation > 0);
470 let (_, entry) = snapshot
471 .entries
472 .into_iter()
473 .find(|(id, _)| *id == 101)
474 .expect("ledger entry");
475 assert_eq!(entry.schema_version, Some(3));
476 assert_eq!(entry.schema_fingerprint.as_deref(), Some("sha256:abc123"));
477 }
478
479 #[test]
480 fn declare_memory_does_not_open_before_bootstrap() {
481 reset_for_tests();
482
483 MemoryApi::declare_with_key(101, "crate_a", "slot", "app.crate_a.slot.v1")
484 .expect("declare memory");
485
486 assert!(MemoryRegistry::get(101).is_none());
487 }
488
489 #[test]
490 fn declare_memory_rejects_after_bootstrap_seal() {
491 reset_for_tests();
492 MemoryApi::bootstrap_owner_range("crate_a", 100, 102).expect("bootstrap registry");
493
494 let err = MemoryApi::declare_with_key(101, "crate_a", "slot", "app.crate_a.slot.v1")
495 .expect_err("late declaration should fail");
496 assert!(matches!(
497 err,
498 MemoryRegistryError::RegistrationAfterBootstrap {
499 ranges: 0,
500 registrations: 1,
501 }
502 ));
503 }
504
505 #[test]
506 fn register_memory_rejects_before_bootstrap_validation() {
507 reset_for_tests();
508
509 let Err(err) = MemoryApi::register(100, "crate_a", "slot") else {
510 panic!("opening before bootstrap must fail")
511 };
512 assert!(matches!(err, MemoryRegistryError::RegistryNotBootstrapped));
513 }
514
515 #[test]
516 fn register_memory_rejects_new_claim_after_bootstrap_seal() {
517 reset_for_tests();
518 MemoryApi::bootstrap_owner_range("crate_a", 100, 102).expect("bootstrap registry");
519
520 let Err(err) = MemoryApi::register(101, "crate_a", "slot") else {
521 panic!("new registration after bootstrap must fail")
522 };
523 assert!(matches!(
524 err,
525 MemoryRegistryError::RegistrationAfterBootstrap {
526 ranges: 0,
527 registrations: 1,
528 }
529 ));
530 }
531
532 #[test]
533 fn bootstrap_pending_flushes_deferred_state() {
534 reset_for_tests();
535 defer_reserve_range("crate_a", 100, 102).expect("defer range");
536 defer_register(101, "crate_a", "slot").expect("defer register");
537
538 MemoryApi::bootstrap_pending().expect("bootstrap pending");
539
540 assert!(MemoryRegistry::export_ranges().contains(&(
541 "crate_a".to_string(),
542 MemoryRange {
543 start: 100,
544 end: 102
545 }
546 )));
547 let entries = MemoryRegistry::export();
548 assert!(entries.iter().any(|(id, entry)| {
549 *id == 101 && entry.crate_name == "crate_a" && entry.label == "slot"
550 }));
551 }
552
553 #[test]
554 fn inspect_memory_returns_reserved_owner_without_label() {
555 reset_for_tests();
556 MemoryApi::bootstrap_owner_range("crate_a", 100, 102).expect("bootstrap registry");
557
558 let inspection = MemoryApi::inspect(101).expect("reserved slot should inspect");
559 assert_eq!(inspection.owner, "crate_a");
560 assert_eq!(
561 inspection.range,
562 MemoryRange {
563 start: 100,
564 end: 102
565 }
566 );
567 assert_eq!(inspection.label, None);
568 }
569
570 #[test]
571 fn inspect_memory_returns_registered_label() {
572 reset_for_tests();
573 defer_reserve_range("crate_a", 100, 102).expect("defer range");
574 defer_register(101, "crate_a", "slot").expect("defer register");
575 MemoryApi::bootstrap_pending().expect("bootstrap registry");
576
577 let inspection = MemoryApi::inspect(101).expect("registered slot should inspect");
578 assert_eq!(inspection.owner, "crate_a");
579 assert_eq!(
580 inspection.range,
581 MemoryRange {
582 start: 100,
583 end: 102
584 }
585 );
586 assert_eq!(inspection.label.as_deref(), Some("slot"));
587 assert_eq!(
588 inspection.stable_key.as_deref(),
589 Some("legacy.crate_a.slot.v1")
590 );
591 }
592
593 #[test]
594 fn inspect_memory_returns_none_for_unowned_id() {
595 reset_for_tests();
596 assert_eq!(MemoryApi::inspect(99), None);
597 }
598
599 #[test]
600 fn registered_memories_lists_registered_slots_with_owner_context() {
601 reset_for_tests();
602 defer_reserve_range("crate_a", 100, 102).expect("defer range A");
603 defer_reserve_range("crate_b", 110, 112).expect("defer range B");
604 defer_register(101, "crate_a", "slot_a").expect("defer register A");
605 defer_register(111, "crate_b", "slot_b").expect("defer register B");
606 MemoryApi::bootstrap_pending().expect("bootstrap registry");
607
608 let registrations = MemoryApi::registered();
609 assert_eq!(registrations.len(), 3);
610 assert!(registrations.contains(&RegisteredMemory {
611 id: 101,
612 owner: "crate_a".to_string(),
613 range: MemoryRange {
614 start: 100,
615 end: 102
616 },
617 label: "slot_a".to_string(),
618 stable_key: "legacy.crate_a.slot_a.v1".to_string(),
619 schema_version: None,
620 schema_fingerprint: None,
621 }));
622 assert!(registrations.contains(&RegisteredMemory {
623 id: 111,
624 owner: "crate_b".to_string(),
625 range: MemoryRange {
626 start: 110,
627 end: 112
628 },
629 label: "slot_b".to_string(),
630 stable_key: "legacy.crate_b.slot_b.v1".to_string(),
631 schema_version: None,
632 schema_fingerprint: None,
633 }));
634 }
635
636 #[test]
637 fn registered_memories_for_owner_filters_to_owner() {
638 reset_for_tests();
639 defer_reserve_range("crate_a", 100, 102).expect("defer range A");
640 defer_reserve_range("crate_b", 110, 112).expect("defer range B");
641 defer_register(101, "crate_a", "slot_a").expect("defer register A");
642 defer_register(111, "crate_b", "slot_b").expect("defer register B");
643 MemoryApi::bootstrap_pending().expect("bootstrap registry");
644
645 let registrations = MemoryApi::registered_for_owner("crate_a");
646 assert_eq!(
647 registrations,
648 vec![RegisteredMemory {
649 id: 101,
650 owner: "crate_a".to_string(),
651 range: MemoryRange {
652 start: 100,
653 end: 102
654 },
655 label: "slot_a".to_string(),
656 stable_key: "legacy.crate_a.slot_a.v1".to_string(),
657 schema_version: None,
658 schema_fingerprint: None,
659 }]
660 );
661 }
662
663 #[test]
664 fn find_registered_memory_returns_match_for_owner_and_label() {
665 reset_for_tests();
666 defer_reserve_range("crate_a", 100, 102).expect("defer range");
667 defer_register(101, "crate_a", "slot_a").expect("defer register");
668 MemoryApi::bootstrap_pending().expect("bootstrap registry");
669
670 let registration = MemoryApi::find("crate_a", "slot_a").expect("slot should exist");
671 assert_eq!(
672 registration,
673 RegisteredMemory {
674 id: 101,
675 owner: "crate_a".to_string(),
676 range: MemoryRange {
677 start: 100,
678 end: 102
679 },
680 label: "slot_a".to_string(),
681 stable_key: "legacy.crate_a.slot_a.v1".to_string(),
682 schema_version: None,
683 schema_fingerprint: None,
684 }
685 );
686 }
687
688 #[test]
689 fn find_registered_memory_returns_none_when_missing() {
690 reset_for_tests();
691 MemoryApi::bootstrap_owner_range("crate_a", 100, 102).expect("bootstrap registry");
692 assert_eq!(MemoryApi::find("crate_a", "slot_a"), None);
693 }
694
695 #[test]
696 fn ledger_snapshot_reads_historical_records() {
697 reset_for_tests();
698 defer_reserve_range("crate_a", 100, 102).expect("defer range");
699 defer_register(101, "crate_a", "slot").expect("defer register");
700 MemoryApi::bootstrap_pending().expect("bootstrap registry");
701
702 let snapshot = MemoryApi::ledger_snapshot().expect("ledger snapshot");
703 assert_eq!(snapshot.format_id, 1);
704 assert_eq!(snapshot.schema_version, 1);
705 assert_eq!(snapshot.layout_epoch, 1);
706 assert!(snapshot.authorities.iter().any(|authority| {
707 authority.owner == "ic_memory.internal"
708 && authority.range == MemoryRange { start: 0, end: 9 }
709 }));
710 assert!(snapshot.authorities.iter().any(|authority| {
711 authority.owner == "canic.framework"
712 && authority.range == MemoryRange { start: 10, end: 99 }
713 }));
714 assert!(snapshot.authorities.iter().any(|authority| {
715 authority.owner == "applications"
716 && authority.range
717 == MemoryRange {
718 start: 100,
719 end: 254,
720 }
721 }));
722 assert!(snapshot.ranges.iter().any(|(owner, range)| {
723 owner == "crate_a"
724 && *range
725 == MemoryRange {
726 start: 100,
727 end: 102,
728 }
729 }));
730 assert!(snapshot.entries.iter().any(|(id, entry)| {
731 *id == 101
732 && entry.crate_name == "crate_a"
733 && entry.label == "slot"
734 && entry.stable_key == "legacy.crate_a.slot.v1"
735 }));
736 }
737}