iota_framework_snapshot/
lib.rs

1// Copyright (c) Mysten Labs, Inc.
2// Modifications Copyright (c) 2024 IOTA Stiftung
3// SPDX-License-Identifier: Apache-2.0
4
5use std::{
6    collections::{BTreeMap, BTreeSet},
7    fs,
8    io::Read,
9    path::{Path, PathBuf},
10};
11
12use iota_framework::SystemPackage;
13pub use iota_framework_snapshot_manifest::*;
14use iota_protocol_config::ProtocolVersion;
15use iota_sdk_types::ObjectId;
16
17const SYSTEM_PACKAGE_PUBLISH_ORDER: &[ObjectId] = &[
18    ObjectId::STD,
19    ObjectId::FRAMEWORK,
20    ObjectId::SYSTEM,
21    ObjectId::STARDUST,
22];
23
24/// Returns the list of system packages in the order they should be published.
25/// If the protocol version is < 9 then include also the bridge package.
26pub fn get_system_package_publish_order(protocol_version: u64) -> Vec<ObjectId> {
27    let mut publish_order = SYSTEM_PACKAGE_PUBLISH_ORDER.to_vec();
28    if protocol_version < 9 {
29        publish_order.insert(3, ObjectId::GENESIS_BRIDGE);
30    }
31    publish_order
32}
33
34pub fn load_bytecode_snapshot(protocol_version: u64) -> anyhow::Result<Vec<SystemPackage>> {
35    let snapshot_path = snapshot_path_for_version(protocol_version)?;
36    let mut snapshots: BTreeMap<ObjectId, SystemPackage> = fs::read_dir(&snapshot_path)?
37        .flatten()
38        .map(|entry| {
39            let file_name = entry.file_name().to_str().unwrap().to_string();
40            let mut file = fs::File::open(snapshot_path.clone().join(file_name))?;
41            let mut buffer = Vec::new();
42            file.read_to_end(&mut buffer)?;
43            let package: SystemPackage = bcs::from_bytes(&buffer)?;
44            Ok((package.id, package))
45        })
46        .collect::<anyhow::Result<_>>()?;
47
48    // system packages need to be restored in a specific order
49    let snapshots_publish_order = get_system_package_publish_order(protocol_version);
50    assert!(snapshots.len() <= snapshots_publish_order.len());
51    let mut snapshot_objects = Vec::new();
52    for package_id in &snapshots_publish_order {
53        if let Some(object) = snapshots.remove(package_id) {
54            snapshot_objects.push(object);
55        }
56    }
57    Ok(snapshot_objects)
58}
59
60/// Returns the path of the snapshot directory holding the framework in effect
61/// at `version`, which is the newest snapshot taken at or before it. A
62/// protocol version that changed nothing in the framework has no snapshot of
63/// its own and resolves to the last one taken before it.
64///
65/// Returns an error if `version` is the max protocol version supported by this
66/// build and no snapshot has been taken for it yet, so that callers fall back
67/// to the framework compiled into this build rather than to an older snapshot.
68/// An error is also returned if `version` predates every snapshot.
69fn snapshot_path_for_version(version: u64) -> anyhow::Result<PathBuf> {
70    let snapshot_dir = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("bytecode_snapshot");
71    let snapshots = read_snapshot_versions(&snapshot_dir)?;
72    let selected = select_snapshot_version(version, ProtocolVersion::MAX.as_u64(), &snapshots)?;
73
74    Ok(snapshot_dir.join(selected.to_string()))
75}
76
77/// Returns the protocol versions that `snapshot_dir` holds a snapshot for.
78fn read_snapshot_versions(snapshot_dir: &Path) -> anyhow::Result<BTreeSet<u64>> {
79    let mut snapshots = BTreeSet::new();
80
81    for entry in fs::read_dir(snapshot_dir)? {
82        let entry = entry?;
83        let path = entry.path();
84        if path.is_dir() {
85            if let Some(snapshot_number) = path
86                .file_name()
87                .and_then(|n| n.to_str())
88                .and_then(|n| n.parse::<u64>().ok())
89            {
90                snapshots.insert(snapshot_number);
91            }
92        }
93    }
94
95    Ok(snapshots)
96}
97
98/// Picks the newest snapshot taken at or before `version` out of `snapshots`.
99///
100/// `max_protocol_version` never resolves to an earlier snapshot: while it has
101/// no snapshot of its own, the framework in effect at that version is the one
102/// compiled into this build, which is not on disk.
103fn select_snapshot_version(
104    version: u64,
105    max_protocol_version: u64,
106    snapshots: &BTreeSet<u64>,
107) -> anyhow::Result<u64> {
108    if version == max_protocol_version && !snapshots.contains(&version) {
109        anyhow::bail!("No snapshot found for version {version}");
110    }
111
112    snapshots
113        .range(..=version)
114        .next_back()
115        .copied()
116        .ok_or_else(|| anyhow::anyhow!("No snapshot found for version {version}"))
117}
118
119#[cfg(test)]
120mod tests {
121    use super::*;
122
123    #[test]
124    fn select_snapshot_of_the_requested_version() {
125        let snapshots = BTreeSet::from([1, 2, 3]);
126        assert_eq!(select_snapshot_version(2, 3, &snapshots).unwrap(), 2);
127    }
128
129    #[test]
130    fn select_snapshot_preceding_a_version_without_one() {
131        let snapshots = BTreeSet::from([33, 35]);
132        assert_eq!(select_snapshot_version(34, 36, &snapshots).unwrap(), 33);
133    }
134
135    #[test]
136    fn select_snapshot_of_the_requested_version_past_a_gap() {
137        let snapshots = BTreeSet::from([33, 35]);
138        assert_eq!(select_snapshot_version(35, 36, &snapshots).unwrap(), 35);
139    }
140
141    #[test]
142    fn select_no_snapshot_for_the_max_version_without_one() {
143        let snapshots = BTreeSet::from([33, 35]);
144        assert!(select_snapshot_version(36, 36, &snapshots).is_err());
145    }
146
147    #[test]
148    fn select_snapshot_of_the_max_version_when_it_has_one() {
149        let snapshots = BTreeSet::from([33, 34, 35]);
150        assert_eq!(select_snapshot_version(35, 35, &snapshots).unwrap(), 35);
151    }
152
153    #[test]
154    fn select_no_snapshot_for_a_version_preceding_all_of_them() {
155        let snapshots = BTreeSet::from([10, 11]);
156        assert!(select_snapshot_version(5, 12, &snapshots).is_err());
157    }
158}
iota_framework_snapshot/lib.rs

iota_framework_snapshot/
lib.rs
