From acd03fc29b447f20ba2944355f37de148839eadc Mon Sep 17 00:00:00 2001
From: Brooks Prumo <brooks@solana.com>
Date: Mon, 20 Dec 2021 14:22:34 -0600
Subject: [PATCH] Check file size of snapshot_version when unarchiving snapshot
 (backport #21925) (#21987)

---
 runtime/src/snapshot_utils.rs | 52 ++++++++++++++++++++++++++++++++---
 1 file changed, 48 insertions(+), 4 deletions(-)

diff --git a/runtime/src/snapshot_utils.rs b/runtime/src/snapshot_utils.rs
index 48852dee1d..3661e900a0 100644
--- a/runtime/src/snapshot_utils.rs
+++ b/runtime/src/snapshot_utils.rs
@@ -41,6 +41,7 @@ pub const SNAPSHOT_STATUS_CACHE_FILE_NAME: &str = "status_cache";
 
 pub const MAX_SNAPSHOTS: usize = 8; // Save some snapshots but not too many
 const MAX_SNAPSHOT_DATA_FILE_SIZE: u64 = 32 * 1024 * 1024 * 1024; // 32 GiB
+const MAX_SNAPSHOT_VERSION_FILE_SIZE: u64 = 8; // byte
 const VERSION_STRING_V1_2_0: &str = "1.2.0";
 const DEFAULT_SNAPSHOT_VERSION: SnapshotVersion = SnapshotVersion::V1_2_0;
 const TMP_SNAPSHOT_PREFIX: &str = "tmp-snapshot-";
@@ -633,11 +634,10 @@ pub fn bank_from_archive<P: AsRef<Path> + std::marker::Sync>(
     let unpacked_snapshots_dir = unpack_dir.as_ref().join("snapshots");
     let unpacked_version_file = unpack_dir.as_ref().join("version");
 
-    let mut snapshot_version = String::new();
-    File::open(unpacked_version_file).and_then(|mut f| f.read_to_string(&mut snapshot_version))?;
+    let snapshot_version = snapshot_version_from_file(unpacked_version_file)?;
 
     let bank = rebuild_bank_from_snapshots(
-        snapshot_version.trim(),
+        &snapshot_version,
         frozen_account_pubkeys,
         &unpacked_snapshots_dir,
         account_paths,
@@ -669,6 +669,28 @@ pub fn bank_from_archive<P: AsRef<Path> + std::marker::Sync>(
     Ok((bank, timings))
 }
 
+/// Reads the `snapshot_version` from a file. Before opening the file, its size
+/// is compared to `MAX_SNAPSHOT_VERSION_FILE_SIZE`. If the size exceeds this
+/// threshold, it is not opened and an error is returned.
+fn snapshot_version_from_file(path: impl AsRef<Path>) -> Result<String> {
+    // Check file size.
+    let file_size = fs::metadata(&path)?.len();
+    if file_size > MAX_SNAPSHOT_VERSION_FILE_SIZE {
+        let error_message = format!(
+            "snapshot version file too large: {} has {} bytes (max size is {} bytes)",
+            path.as_ref().display(),
+            file_size,
+            MAX_SNAPSHOT_VERSION_FILE_SIZE,
+        );
+        return Err(get_io_error(&error_message));
+    }
+
+    // Read snapshot_version from file.
+    let mut snapshot_version = String::new();
+    File::open(path).and_then(|mut f| f.read_to_string(&mut snapshot_version))?;
+    Ok(snapshot_version.trim().to_string())
+}
+
 pub fn get_snapshot_archive_path(
     snapshot_output_dir: PathBuf,
     snapshot_hash: &(Slot, Hash),
@@ -1116,7 +1138,8 @@ mod tests {
         super::*,
         assert_matches::assert_matches,
         bincode::{deserialize_from, serialize_into},
-        std::mem::size_of,
+        std::{convert::TryFrom, mem::size_of},
+        tempfile::NamedTempFile,
     };
 
     #[test]
@@ -1224,6 +1247,27 @@ mod tests {
         assert_matches!(result, Err(SnapshotError::Io(ref message)) if message.to_string().starts_with("invalid snapshot data file"));
     }
 
+    #[test]
+    fn test_snapshot_version_from_file_under_limit() {
+        let file_content = format!("v{}", DEFAULT_SNAPSHOT_VERSION);
+        let mut file = NamedTempFile::new().unwrap();
+        file.write_all(file_content.as_bytes()).unwrap();
+        let version_from_file = snapshot_version_from_file(file.path()).unwrap();
+        assert_eq!(version_from_file, file_content);
+    }
+
+    #[test]
+    fn test_snapshot_version_from_file_over_limit() {
+        let over_limit_size = usize::try_from(MAX_SNAPSHOT_VERSION_FILE_SIZE + 1).unwrap();
+        let file_content = vec![7u8; over_limit_size];
+        let mut file = NamedTempFile::new().unwrap();
+        file.write_all(&file_content).unwrap();
+        assert_matches!(
+            snapshot_version_from_file(file.path()),
+            Err(SnapshotError::Io(ref message)) if message.to_string().starts_with("snapshot version file too large")
+        );
+    }
+
     #[test]
     fn test_parse_snapshot_archive_filename() {
         assert_eq!(