refactor(stm): applied more suggestions and reverted the use of HOME variable

Author: damrobi

mithril-stm/Cargo.toml

@@ -30,6 +30,7 @@ future_snark = [
     "dep:rand",
     "dep:sha2",
     "dep:rand_chacha",
+    "dep:reqwest",
 ]
 num-integer-backend = ["dep:num-bigint", "dep:num-integer", "dep:num-rational", "dep:num-traits"]
 rug-backend = ["rug/default"]
@@ -55,11 +56,10 @@ rand = { version = "0.10.1", optional = true }
 rand_chacha = { workspace = true, optional = true }
 rand_core = { workspace = true, features = ["std"] }
 rayon = { workspace = true }
-reqwest = { workspace = true, default-features = true, features = ["blocking"] }
+reqwest = { workspace = true, default-features = true, features = ["blocking"], optional = true }
 serde = { workspace = true }
 sha2 = { version = "0.10.9", optional = true }
 subtle = "2.6.1"
-tempfile = "3.27.0"
 thiserror = { workspace = true }
 
 [target.'cfg(any(target_family = "wasm", target_env = "musl", windows))'.dependencies]

mithril-stm/src/circuits/mod.rs

@@ -15,3 +15,6 @@ pub(crate) mod test_utils;
 pub(crate) use halo2::witness::{
     CircuitInstance, CircuitMerkleTreeLeaf, CircuitWitness, MerklePath,
 };
+
+/// Constant holding the current path of the cached values related to the circuits
+const MITHRIL_CIRCUIT_CACHE_FOLDER: &str = "mithril-circuit";

mithril-stm/src/circuits/trusted_setup.rs

@@ -1,16 +1,16 @@
 use std::{
     fs::File,
-    io::{BufReader, BufWriter, Read, Write},
+    io::{BufReader, BufWriter, Write},
     path::PathBuf,
+    time::Duration,
 };
 
 use anyhow::{Context, anyhow};
 use midnight_curves::Bls12;
 use midnight_proofs::{poly::kzg::params::ParamsKZG, utils::SerdeFormat};
 use sha2::{Digest, Sha256};
-use tempfile::NamedTempFile;
 
-use crate::StmResult;
+use crate::{StmResult, circuits::MITHRIL_CIRCUIT_CACHE_FOLDER};
 
 /// TODO: remove allow(dead_code) when the constants are used or remove the constatnts
 #[allow(dead_code)]
@@ -22,35 +22,36 @@ use crate::StmResult;
 /// the proper value available here: https://github.com/midnightntwrk/midnight-trusted-setup/blob/main/MIDNIGHT_SRS_CATALOG.md
 const SRS_HASH_K22: &str = "e8ad5eed936d657a0fb59d2a55ba19f81a3083bb3554ef88f464f5377e9b2c2f";
 #[allow(dead_code)]
-/// Constant storing local path of the SRS of degree 22 used to create proof in production
-const SRS_PATH_K22: &str = ".local/local_trusted_setup/";
-#[allow(dead_code)]
 /// Constant storing URL to download the SRS of degree 22 used to create proof in production
 const SRS_URL_K22: &str = "https://srs.midnight.network/midnight-srs-2p22";
 
-/// An handle to manage the trusted setup srs. It stores the local path of the srs file
+/// A structure to manage the trusted setup SRS. It stores the local path of the SRS file
 /// and information to download the file and verify integrity if it is missing.
-pub struct TrustedSetupHandle {
-    /// Path of the local srs file
+pub struct TrustedSetupProvider {
+    /// Path of the local SRS file
     local_srs_path: PathBuf,
-    /// Expected hash of the downloaded srs file
+    /// Expected hash of the downloaded SRS file
     srs_expected_hash: String,
-    /// URL where to download the srs file if it is not present locally
+    /// URL where to download the SRS file if it is not present locally
     url_to_download_srs: String,
+    /// The timeout limit when trying to download the SRS file
+    download_timeout_limit: Duration,
 }
 
 /// TODO: remove allow(dead_code) when used
 #[allow(dead_code)]
-impl TrustedSetupHandle {
-    fn new<P, S>(local_srs_path: P, srs_expected_hash: S, url_to_download_srs: S) -> Self
-    where
-        P: Into<PathBuf>,
-        S: Into<String>,
-    {
+impl TrustedSetupProvider {
+    fn new<P: Into<PathBuf>, S: Into<String>>(
+        local_srs_path: P,
+        srs_expected_hash: S,
+        url_to_download_srs: S,
+        download_timeout_limit: Duration,
+    ) -> Self {
         Self {
             local_srs_path: local_srs_path.into(),
             srs_expected_hash: srs_expected_hash.into(),
             url_to_download_srs: url_to_download_srs.into(),
+            download_timeout_limit,
         }
     }
 
@@ -62,18 +63,6 @@ impl TrustedSetupHandle {
         hex::encode(hasher.finalize())
     }
 
-    /// Reads the local SRS file at the stored path and checks its SHA256 hash against the expected value.
-    // This function can be removed if it is not used.
-    fn verify_file_sha256_hash(&self) -> StmResult<bool> {
-        let mut file = File::open(&self.local_srs_path)
-            .with_context(|| "Loading of the SRS file should have succeeded!")?;
-        let mut srs_buffer = vec![];
-        file.read_to_end(&mut srs_buffer)
-            .with_context(|| "Reading the SRS file should have succeeded!")?;
-
-        Ok(self.verify_bytes_sha256_hash(&srs_buffer))
-    }
-
     /// Checks SHA256 hash of the given bytes against the stored expected value.
     fn verify_bytes_sha256_hash(&self, srs_bytes: &[u8]) -> bool {
         let recomputed_hash = Self::compute_hash(srs_bytes);
@@ -85,7 +74,7 @@ impl TrustedSetupHandle {
     fn download_srs_file(&self) -> StmResult<Vec<u8>> {
         let response = reqwest::blocking::Client::builder()
             // TODO: For now a timeout but this should be updated depending on the behavior we want
-            .timeout(std::time::Duration::from_secs(600))
+            .timeout(self.download_timeout_limit)
             .build()?
             .get(&self.url_to_download_srs)
             .header("User-Agent", "mithril-stm")
@@ -96,8 +85,9 @@ impl TrustedSetupHandle {
         Ok(bytes.to_vec())
     }
 
-    /// Saves the given bytes in a file at the stored path while creating
-    /// the directories of the path if needed.
+    /// Saves the given bytes in a temporary file then atomically moves it to the stored path
+    /// while creating the directories of the path if needed.
+    /// If the writing is interu
     fn store_srs_bytes_to_file(&self, srs_bytes: &[u8]) -> StmResult<()> {
         let parent = self
             .local_srs_path
@@ -106,12 +96,14 @@ impl TrustedSetupHandle {
         std::fs::create_dir_all(parent)
             .with_context(|| "Subdirectory creation should have succeeded!")?;
 
-        let mut temporary_file = NamedTempFile::new_in(parent)
-            .with_context(|| "Failed to generate temporary file to store the SRS")?;
+        let temp_path = self.local_srs_path.with_extension(".temp");
+        // Create/clean the file
+        let mut temporary_file = File::create(&temp_path)?;
         BufWriter::new(&mut temporary_file).write_all(srs_bytes)?;
-        temporary_file
-            .persist(&self.local_srs_path)
-            .with_context(|| "Failed to atomically rename SRS temp file.")?;
+
+        // atomic renaming
+        std::fs::rename(temp_path, &self.local_srs_path)?;
+
         Ok(())
     }
 
@@ -143,19 +135,18 @@ impl TrustedSetupHandle {
             .with_context(|| format!("Failed to open SRS file at {:?}", self.local_srs_path))?;
         let mut reader = BufReader::new(file);
 
-        ParamsKZG::read_custom(&mut reader, SerdeFormat::RawBytes)
+        ParamsKZG::read_custom(&mut reader, SerdeFormat::RawBytesUnchecked)
             .with_context(|| "Failed to deserialize SRS from file")
     }
 }
 
-impl Default for TrustedSetupHandle {
+impl Default for TrustedSetupProvider {
     fn default() -> Self {
         Self::new(
-            PathBuf::from(std::env::var("HOME").unwrap_or_else(|_| ".".to_string()))
-                .join(SRS_PATH_K22)
-                .join("midnight-srs-2p22"),
+            std::env::temp_dir().join(MITHRIL_CIRCUIT_CACHE_FOLDER).join("srs"),
             SRS_HASH_K22,
             SRS_URL_K22,
+            Duration::from_secs(600),
         )
     }
 }
@@ -215,7 +206,8 @@ mod tests {
     fn both_bytes_encoding_work_to_load_srs_from_file() {
         let mut tmp_file = NamedTempFile::new().unwrap();
         std::fs::write(&mut tmp_file, SRS_K1.as_slice()).unwrap();
-        let srs_manager = TrustedSetupHandle::new(tmp_file.path(), "", "");
+        let srs_manager =
+            TrustedSetupProvider::new(tmp_file.path(), "", "", Duration::from_secs(600));
         let loaded_srs = srs_manager.check_available_file_then_load_srs().unwrap();
         let srs_rawbytes: ParamsKZG<Bls12> =
             ParamsKZG::read_custom(&mut SRS_K1.as_slice(), SerdeFormat::RawBytes).unwrap();
@@ -243,57 +235,44 @@ mod tests {
         assert_eq!(raw_bytes_unchecked_buffer, raw_bytes_buffer);
     }
 
-    #[test]
-    fn hash_of_correct_srs_file_verifies() {
-        let srs_file = NamedTempFile::new_in("/tmp").unwrap();
-        std::fs::write(&srs_file, SRS_K1).unwrap();
-
-        let result = TrustedSetupHandle::new(srs_file.path(), SRS_HASH_K1, "")
-            .verify_file_sha256_hash()
-            .unwrap();
-
-        assert!(result);
-    }
-
     #[test]
     fn verification_of_hash_of_invalid_srs_file_fails() {
         let mut tampered_bytes = SRS_K1.to_vec();
         tampered_bytes[0] = tampered_bytes[0].wrapping_add(1);
 
-        let result =
-            TrustedSetupHandle::new("", SRS_HASH_K1, "").verify_bytes_sha256_hash(&tampered_bytes);
+        let result = TrustedSetupProvider::new("", SRS_HASH_K1, "", Duration::from_secs(600))
+            .verify_bytes_sha256_hash(&tampered_bytes);
 
         assert!(!result);
     }
 
     #[test]
     fn hash_of_correct_bytes_verifies() {
-        let result = TrustedSetupHandle::new("", SRS_HASH_K1, "").verify_bytes_sha256_hash(SRS_K1);
+        let result = TrustedSetupProvider::new("", SRS_HASH_K1, "", Duration::from_secs(600))
+            .verify_bytes_sha256_hash(SRS_K1);
 
         assert!(result);
     }
 
     #[test]
-    fn verification_of_hash_of_invalid_bytes_fails() {
+    fn existing_file_on_disk_skips_download_and_verification() {
         let srs_file = NamedTempFile::new_in("/tmp").unwrap();
-        let mut tampered = SRS_K1.to_vec();
-        tampered[0] = tampered[0].wrapping_add(1);
-        std::fs::write(&srs_file, &tampered).unwrap();
+        std::fs::write(&srs_file, [0, 1, 2, 3, 4]).unwrap();
 
-        let result = TrustedSetupHandle::new(srs_file.path(), SRS_HASH_K1, "")
-            .verify_file_sha256_hash()
-            .unwrap();
+        let result = TrustedSetupProvider::new(srs_file.path(), "", "", Duration::from_secs(600))
+            .ensure_srs_file_is_available();
 
-        assert!(!result);
+        assert!(result.is_ok());
     }
 
     #[test]
-    fn existing_file_on_disk_skips_download_and_verification() {
+    fn interrupted_writing_of_srs_resumes_properly_at_next_try() {
         let srs_file = NamedTempFile::new_in("/tmp").unwrap();
-        std::fs::write(&srs_file, [0, 1, 2, 3, 4]).unwrap();
+        let temp_srs_file = srs_file.path().with_extension("temp");
+        std::fs::write(&temp_srs_file, [0, 1, 2, 3, 4]).unwrap();
 
-        let result =
-            TrustedSetupHandle::new(srs_file.path(), "", "").ensure_srs_file_is_available();
+        let result = TrustedSetupProvider::new(srs_file.path(), "", "", Duration::from_secs(600))
+            .ensure_srs_file_is_available();
 
         assert!(result.is_ok());
     }
@@ -302,7 +281,7 @@ mod tests {
         use super::*;
 
         #[test]
-        fn missing_srs_file_triggers_download_verifiation_and_storage() {
+        fn missing_srs_file_triggers_download_verification_and_storage() {
             let server = MockServer::start();
             let mock = server.mock(|when, then| {
                 when.method(httpmock::Method::GET).path("/srs");
@@ -312,9 +291,14 @@ mod tests {
             let temp_dir = tempfile::tempdir().unwrap();
             let srs_path = temp_dir.path().join("missing_file_trigger_dl");
 
-            TrustedSetupHandle::new(&srs_path, SRS_HASH_K1, &server.url("/srs"))
-                .ensure_srs_file_is_available()
-                .unwrap();
+            TrustedSetupProvider::new(
+                &srs_path,
+                SRS_HASH_K1,
+                &server.url("/srs"),
+                Duration::from_secs(600),
+            )
+            .ensure_srs_file_is_available()
+            .unwrap();
 
             mock.assert();
             assert!(srs_path.exists());
@@ -331,8 +315,13 @@ mod tests {
             let temp_dir = tempfile::tempdir().unwrap();
             let srs_path = temp_dir.path().join("dl_wrong_hash");
 
-            let result = TrustedSetupHandle::new(&srs_path, SRS_HASH_K1, &server.url("/srs"))
-                .ensure_srs_file_is_available();
+            let result = TrustedSetupProvider::new(
+                &srs_path,
+                SRS_HASH_K1,
+                &server.url("/srs"),
+                Duration::from_secs(600),
+            )
+            .ensure_srs_file_is_available();
 
             assert!(result.is_err());
             assert!(!srs_path.exists());
@@ -349,8 +338,13 @@ mod tests {
             let temp_dir = tempfile::tempdir().unwrap();
             let srs_path = temp_dir.path().join("server_error_fails");
 
-            let result = TrustedSetupHandle::new(&srs_path, SRS_HASH_K1, &server.url("/srs"))
-                .ensure_srs_file_is_available();
+            let result = TrustedSetupProvider::new(
+                &srs_path,
+                SRS_HASH_K1,
+                &server.url("/srs"),
+                Duration::from_secs(600),
+            )
+            .ensure_srs_file_is_available();
 
             assert!(result.is_err());
         }