[framework-snaps] Update snapshot loading (MystenLabs#21193)

Since deployed protocol versions are not guaranteed to be sequential,
snapshots are not guaranteed to be sequential (e.g., 45 may exist, 46
may not, but 47 may).

Previously when trying to load a snapshot for a protocol version we
would only load a snapshot if there was one for that exact version,
otherwise we'd default to the newest version of the framework. This is
incorrect, and @dariorussi has run into problems with this.

This updates snapshot loading so if a snapshot for a particular version
`v` is not found, the least version greater than the given version is
returned.

## Test plan 

CI
 
---

## Release notes

Check each box that your changes affect. If none of the boxes relate to
your changes, release notes aren't required.

For each box you select, include information after the relevant heading
that describes the impact of your changes that a user might notice and
any actions they must take to implement updates.

- [ ] Protocol: 
- [ ] Nodes (Validators and Full nodes): 
- [ ] gRPC:
- [ ] JSON-RPC: 
- [ ] GraphQL: 
- [ ] CLI: 
- [ ] Rust SDK:

---------

Co-authored-by: ronny-mysten <[email protected]>

Author: tzakian

crates/sui-framework-snapshot/src/lib.rs

@@ -2,7 +2,7 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use serde::{Deserialize, Serialize};
-use std::collections::BTreeMap;
+use std::collections::{BTreeMap, BTreeSet};
 use std::{fs, io::Read, path::PathBuf};
 use sui_framework::{SystemPackage, SystemPackageMetadata};
 use sui_types::base_types::ObjectID;
@@ -100,8 +100,7 @@ pub fn update_bytecode_snapshot_manifest(
 }
 
 pub fn load_bytecode_snapshot(protocol_version: u64) -> anyhow::Result<Vec<SystemPackage>> {
-    let mut snapshot_path = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
-    snapshot_path.extend(["bytecode_snapshot", protocol_version.to_string().as_str()]);
+    let snapshot_path = snapshot_path_for_version(protocol_version)?;
     let mut snapshots: BTreeMap<ObjectID, SystemPackage> = fs::read_dir(&snapshot_path)?
         .flatten()
         .map(|entry| {
@@ -128,3 +127,33 @@ pub fn load_bytecode_snapshot(protocol_version: u64) -> anyhow::Result<Vec<Syste
 pub fn manifest_path() -> PathBuf {
     PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("manifest.json")
 }
+
+/// Given a protocol version:
+/// * The path to the snapshot directory for that version is returned, if it exists.
+/// * If the version is greater than the latest snapshot version, then `Ok(None)` is returned.
+/// * If the version does not exist, but there are snapshots present with versions greater than
+///   `version`, then the smallest snapshot number greater than `version` is returned.
+fn snapshot_path_for_version(version: u64) -> anyhow::Result<PathBuf> {
+    let snapshot_dir = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("bytecode_snapshot");
+    let mut snapshots = BTreeSet::new();
+
+    for entry in fs::read_dir(&snapshot_dir)? {
+        let entry = entry?;
+        let path = entry.path();
+        if path.is_dir() {
+            if let Some(snapshot_number) = path
+                .file_name()
+                .and_then(|n| n.to_str())
+                .and_then(|n| n.parse::<u64>().ok())
+            {
+                snapshots.insert(snapshot_number);
+            }
+        }
+    }
+
+    snapshots
+        .range(version..)
+        .next()
+        .map(|v| snapshot_dir.join(v.to_string()))
+        .ok_or_else(|| anyhow::anyhow!("No snapshot found for version {}", version))
+}