feat: support zstd archives

Cargo.toml

@@ -16,3 +16,4 @@ reqwest = { version = "0.11.23", features = ["json", "stream", "blocking"] }
 rusqlite = { version = "0.30.0", features = ["bundled"] }
 url = "2.5.0"
 zip = "0.6.6"
+zstd = "0.13.0"

src/checksum.rs

@@ -14,6 +14,9 @@ fn replace_sql_zip_with_md5(url: &Url) -> Result<Url> {
   if url_str.ends_with(".sql.zip") {
     let new_url_str = url_str.replace(".sql.zip", ".sql.md5");
     Ok(Url::parse(&new_url_str)?)
+  } else if url_str.ends_with(".sql.zstd") {
+    let new_url_str = url_str.replace(".sql.zstd", ".sql.md5");
+    Ok(Url::parse(&new_url_str)?)
   } else {
     anyhow::bail!("URL does not end with .sql.zip")
   }

src/main.rs

@@ -8,18 +8,19 @@ mod checksum;
 mod download;
 mod go_spacemesh;
 mod parsers;
+mod reader_with_bytes;
 mod reader_with_progress;
 mod sql;
+mod unpack;
 mod utils;
-mod zip;
 
 use checksum::*;
 use download::download_with_retries;
 use go_spacemesh::get_version;
 use parsers::*;
 use sql::get_last_layer_from_db;
+use unpack::{unpack_zip, unpack_zstd};
 use utils::*;
-use zip::unpack;
 
 #[derive(Parser, Debug)]
 #[clap(author, version, about, long_about = None)]
@@ -159,13 +160,14 @@ fn main() -> anyhow::Result<()> {
       let dir_path = node_data;
       let temp_file_path = dir_path.join("state.download");
       let redirect_file_path = dir_path.join("state.url");
-      let archive_file_path = dir_path.join("state.zip");
+      let archive_zip_file_path = dir_path.join("state.zip");
+      let archive_zstd_file_path = dir_path.join("state.zstd");
       let unpacked_file_path = dir_path.join("state_downloaded.sql");
       let final_file_path = dir_path.join("state.sql");
       let wal_file_path = dir_path.join("state.sql-wal");
 
       // Download archive if needed
-      if !archive_file_path.exists() {
+      let archive_file_path = if !archive_zip_file_path.exists() && !archive_zstd_file_path.exists() {
         println!("Downloading the latest database...");
         let url = if redirect_file_path.exists() {
           std::fs::read_to_string(&redirect_file_path)?
@@ -174,7 +176,7 @@ fn main() -> anyhow::Result<()> {
           let go_path_str = go_path
             .to_str()
             .expect("Cannot resolve path to go-spacemesh");
-          let path = format!("{}/state.zip", &get_version(go_path_str)?);
+          let path = format!("{}/state.zstd", &get_version(go_path_str)?);
           let url = build_url(&download_url, &path);
           url.to_string()
         };
@@ -189,10 +191,28 @@ fn main() -> anyhow::Result<()> {
           process::exit(1);
         }
 
-        // Rename `state.download` -> `state.zip`
+        let archive_file_path = if url.ends_with(".zip") {
+          archive_zip_file_path
+        } else {
+          archive_zstd_file_path
+        };
+
+        // Rename `state.download` -> `state.zstd`
         std::fs::rename(&temp_file_path, &archive_file_path)?;
         println!("Archive downloaded!");
-      }
+        archive_file_path
+      } else if archive_zip_file_path.exists() {
+        archive_zip_file_path
+      } else {
+        archive_zstd_file_path
+      };
+
+      let archive_url = std::fs::read_to_string(&redirect_file_path)?;
+      let unpack = if archive_url.ends_with(".zip") {
+        unpack_zip
+      } else {
+        unpack_zstd
+      };
 
       // Unzip
       match unpack(&archive_file_path, &unpacked_file_path) {

src/reader_with_bytes.rs

@@ -0,0 +1,33 @@
+use std::io::{self, Read};
+
+const MB: u64 = 1024 * 1024;
+
+pub struct ReaderWithBytes<R: Read> {
+  reader: R,
+  bytes_read: u64,
+  last_reported: u64,
+}
+
+impl<R: Read> ReaderWithBytes<R> {
+  pub fn new(reader: R) -> Self {
+    ReaderWithBytes {
+      reader,
+      bytes_read: 0,
+      last_reported: 0,
+    }
+  }
+}
+
+impl<R: Read> Read for ReaderWithBytes<R> {
+  fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
+    let bytes_read = self.reader.read(buf)?;
+    self.bytes_read += bytes_read as u64;
+
+    if self.bytes_read / MB > self.last_reported / MB {
+      println!("Unpacking... {} MB extracted", self.bytes_read / MB);
+      self.last_reported = self.bytes_read;
+    }
+
+    Ok(bytes_read)
+  }
+}

src/zip.rs → src/unpack.rs

@@ -1,10 +1,12 @@
 use anyhow::Result;
 use std::fs::File;
-use std::io::{BufReader, Error};
+use std::io::{BufReader, BufWriter, Error};
 use std::path::Path;
 use zip::read::ZipFile;
 use zip::ZipArchive;
+use zstd::stream::read::Decoder;
 
+use crate::reader_with_bytes::ReaderWithBytes;
 use crate::reader_with_progress::ReaderWithProgress;
 
 fn find_file_in_archive<'a>(
@@ -29,7 +31,27 @@ fn find_file_in_archive<'a>(
   ))
 }
 
-pub fn unpack(archive_path: &Path, output_path: &Path) -> Result<()> {
+pub fn unpack_zstd(archive_path: &Path, output_path: &Path) -> Result<()> {
+  let file = File::open(archive_path)?;
+  let reader = BufReader::new(file);
+  let decoder = Decoder::new(reader)?;
+
+  let outpath = Path::new(output_path);
+  if let Some(p) = outpath.parent() {
+    std::fs::create_dir_all(p)?;
+  }
+  let outfile = File::create(outpath)?;
+  let mut writer = BufWriter::new(outfile);
+
+  let mut reader = ReaderWithBytes::new(decoder);
+
+  std::io::copy(&mut reader, &mut writer)?;
+  println!("Unpacking complete!");
+
+  Ok(())
+}
+
+pub fn unpack_zip(archive_path: &Path, output_path: &Path) -> Result<()> {
   let file = File::open(archive_path)?;
   let mut zip = ZipArchive::new(file)?;
 

src/utils.rs

@@ -56,7 +56,7 @@ pub fn backup_file(original_path: &PathBuf) -> Result<PathBuf> {
 }
 
 fn extract_number_from_url(url: &Url) -> Result<u64> {
-  let re = Regex::new(r"/(\d+)\.sql\.zip$")?;
+  let re = Regex::new(r"/(\d+)\.sql\.(zip|zstd)$")?;
   let path = url.path();
   let caps = re
     .captures(path)
@@ -95,11 +95,17 @@ mod tests {
   use url::Url;
 
   #[test]
-  fn test_extract_number_valid() {
+  fn test_extract_number_zip_valid() {
     let url = Url::parse("https://quicksync-downloads.spacemesh.network/10/61579.sql.zip").unwrap();
     assert_eq!(extract_number_from_url(&url).unwrap(), 61579);
   }
 
+  #[test]
+  fn test_extract_number_zstd_valid() {
+    let url = Url::parse("https://quicksync-downloads.spacemesh.network/10/61579.sql.zstd").unwrap();
+    assert_eq!(extract_number_from_url(&url).unwrap(), 61579);
+  }
+
   #[test]
   fn test_extract_number_invalid() {
     let url = Url::parse("https://quicksync.spacemesh.network/state.zip").unwrap();