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();