From 13e5dd2558d17e20b73867f4c99c76ee776f56d6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Bartosz=20R=C3=B3=C5=BCa=C5=84ski?= Date: Tue, 20 Aug 2024 12:25:46 +0200 Subject: [PATCH 1/4] add basic tests for downloading --- Cargo.lock | 208 ++++++++++++++++++++++++++++++++++++++++++++---- Cargo.toml | 4 + src/download.rs | 95 +++++++++++++++++++--- 3 files changed, 282 insertions(+), 25 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 0a0ec08..4cd21d8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -134,6 +134,16 @@ dependencies = [ "derive_arbitrary", ] +[[package]] +name = "assert-json-diff" +version = "2.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47e4f2b81832e72834d7518d8487a0396a28cc408186a2e8854c0f98011faf12" +dependencies = [ + "serde", + "serde_json", +] + [[package]] name = "atomic-waker" version = "1.1.2" @@ -249,7 +259,7 @@ dependencies = [ "js-sys", "num-traits", "wasm-bindgen", - "windows-targets", + "windows-targets 0.52.6", ] [[package]] @@ -308,6 +318,16 @@ version = "1.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d3fd119d74b830634cea2a0f58bbd0d54540518a14397557951e79340abc28c0" +[[package]] +name = "colored" +version = "2.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cbf2150cce219b664a8a70df7a1f933836724b503f8a413af9365b4dcc4d90b8" +dependencies = [ + "lazy_static", + "windows-sys 0.48.0", +] + [[package]] name = "constant_time_eq" version = "0.3.0" @@ -703,6 +723,12 @@ version = "1.9.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0fcc0b4a115bf80b728eb8ea024ad5bd707b615bfed49e0665b6e0f86fd082d9" +[[package]] +name = "httpdate" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df3b46402a9d5adb4c86a0cf463f42e19994e3ee891101b1841f30a545cb49a9" + [[package]] name = "hyper" version = "1.4.1" @@ -716,6 +742,7 @@ dependencies = [ "http", "http-body", "httparse", + "httpdate", "itoa", "pin-project-lite", "smallvec", @@ -864,6 +891,12 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "lazy_static" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" + [[package]] name = "libc" version = "0.2.158" @@ -887,6 +920,16 @@ version = "0.4.14" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "78b3ae25bc7c8c38cec158d1f2757ee79e9b3740fbc7ccf0e59e4b08d793fa89" +[[package]] +name = "lock_api" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "07af8b9cdd281b7915f413fa73f29ebd5d55d0d3f0155584dade1ff18cea1b17" +dependencies = [ + "autocfg", + "scopeguard", +] + [[package]] name = "lockfree-object-pool" version = "0.1.6" @@ -957,6 +1000,30 @@ dependencies = [ "windows-sys 0.52.0", ] +[[package]] +name = "mockito" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09b34bd91b9e5c5b06338d392463e1318d683cf82ec3d3af4014609be6e2108d" +dependencies = [ + "assert-json-diff", + "bytes", + "colored", + "futures-util", + "http", + "http-body", + "http-body-util", + "hyper", + "hyper-util", + "log", + "rand", + "regex", + "serde_json", + "serde_urlencoded", + "similar", + "tokio", +] + [[package]] name = "native-tls" version = "0.2.12" @@ -1048,6 +1115,29 @@ dependencies = [ "vcpkg", ] +[[package]] +name = "parking_lot" +version = "0.12.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f1bf18183cf54e8d6059647fc3063646a1801cf30896933ec2311622cc4b9a27" +dependencies = [ + "lock_api", + "parking_lot_core", +] + +[[package]] +name = "parking_lot_core" +version = "0.9.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e401f977ab385c9e4e3ab30627d6f26d00e2c73eef317493c4ec6d468726cf8" +dependencies = [ + "cfg-if", + "libc", + "redox_syscall", + "smallvec", + "windows-targets 0.52.6", +] + [[package]] name = "pbkdf2" version = "0.12.2" @@ -1135,11 +1225,13 @@ dependencies = [ "clap", "duration-string", "md5", + "mockito", "regex", "reqwest", "rusqlite", "serde", "serde_json", + "tempfile", "url", "zip", "zstd", @@ -1184,6 +1276,15 @@ dependencies = [ "getrandom", ] +[[package]] +name = "redox_syscall" +version = "0.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2a908a6e00f1fdd0dfd9c0eb08ce85126f6d8bbda50017e74bc4a4b7d4a926a4" +dependencies = [ + "bitflags", +] + [[package]] name = "regex" version = "1.10.6" @@ -1362,6 +1463,12 @@ dependencies = [ "windows-sys 0.52.0", ] +[[package]] +name = "scopeguard" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" + [[package]] name = "security-framework" version = "2.11.1" @@ -1452,6 +1559,12 @@ version = "0.3.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d66dc143e6b11c1eddc06d5c423cfc97062865baf299914ab64caa38182078fe" +[[package]] +name = "similar" +version = "2.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1de1d4f81173b03af4c0cbed3c898f6bff5b870e4a7f5d6f4057d62a7a4b686e" + [[package]] name = "slab" version = "0.4.9" @@ -1613,6 +1726,7 @@ dependencies = [ "bytes", "libc", "mio", + "parking_lot", "pin-project-lite", "socket2", "windows-sys 0.52.0", @@ -1877,7 +1991,7 @@ version = "0.52.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "33ab640c8d7e35bf8ba19b884ba838ceb4fba93a4e8c65a9059d08afcfc683d9" dependencies = [ - "windows-targets", + "windows-targets 0.52.6", ] [[package]] @@ -1888,7 +2002,7 @@ checksum = "e400001bb720a623c1c69032f8e3e4cf09984deec740f007dd2b03ec864804b0" dependencies = [ "windows-result", "windows-strings", - "windows-targets", + "windows-targets 0.52.6", ] [[package]] @@ -1897,7 +2011,7 @@ version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1d1043d8214f791817bab27572aaa8af63732e11bf84aa21a45a78d6c317ae0e" dependencies = [ - "windows-targets", + "windows-targets 0.52.6", ] [[package]] @@ -1907,7 +2021,16 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4cd9b125c486025df0eabcb585e62173c6c9eddcec5d117d3b6e8c30e2ee4d10" dependencies = [ "windows-result", - "windows-targets", + "windows-targets 0.52.6", +] + +[[package]] +name = "windows-sys" +version = "0.48.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "677d2418bec65e3338edb076e806bc1ec15693c5d0104683f2efe857f61056a9" +dependencies = [ + "windows-targets 0.48.5", ] [[package]] @@ -1916,7 +2039,7 @@ version = "0.52.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "282be5f36a8ce781fad8c8ae18fa3f9beff57ec1b52cb3de0789201425d9a33d" dependencies = [ - "windows-targets", + "windows-targets 0.52.6", ] [[package]] @@ -1925,7 +2048,22 @@ version = "0.59.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b" dependencies = [ - "windows-targets", + "windows-targets 0.52.6", +] + +[[package]] +name = "windows-targets" +version = "0.48.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a2fa6e2155d7247be68c096456083145c183cbbbc2764150dda45a87197940c" +dependencies = [ + "windows_aarch64_gnullvm 0.48.5", + "windows_aarch64_msvc 0.48.5", + "windows_i686_gnu 0.48.5", + "windows_i686_msvc 0.48.5", + "windows_x86_64_gnu 0.48.5", + "windows_x86_64_gnullvm 0.48.5", + "windows_x86_64_msvc 0.48.5", ] [[package]] @@ -1934,28 +2072,46 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" dependencies = [ - "windows_aarch64_gnullvm", - "windows_aarch64_msvc", - "windows_i686_gnu", + "windows_aarch64_gnullvm 0.52.6", + "windows_aarch64_msvc 0.52.6", + "windows_i686_gnu 0.52.6", "windows_i686_gnullvm", - "windows_i686_msvc", - "windows_x86_64_gnu", - "windows_x86_64_gnullvm", - "windows_x86_64_msvc", + "windows_i686_msvc 0.52.6", + "windows_x86_64_gnu 0.52.6", + "windows_x86_64_gnullvm 0.52.6", + "windows_x86_64_msvc 0.52.6", ] +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.48.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b38e32f0abccf9987a4e3079dfb67dcd799fb61361e53e2882c3cbaf0d905d8" + [[package]] name = "windows_aarch64_gnullvm" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" +[[package]] +name = "windows_aarch64_msvc" +version = "0.48.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc35310971f3b2dbbf3f0690a219f40e2d9afcf64f9ab7cc1be722937c26b4bc" + [[package]] name = "windows_aarch64_msvc" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" +[[package]] +name = "windows_i686_gnu" +version = "0.48.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a75915e7def60c94dcef72200b9a8e58e5091744960da64ec734a6c6e9b3743e" + [[package]] name = "windows_i686_gnu" version = "0.52.6" @@ -1968,24 +2124,48 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" +[[package]] +name = "windows_i686_msvc" +version = "0.48.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f55c233f70c4b27f66c523580f78f1004e8b5a8b659e05a4eb49d4166cca406" + [[package]] name = "windows_i686_msvc" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" +[[package]] +name = "windows_x86_64_gnu" +version = "0.48.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53d40abd2583d23e4718fddf1ebec84dbff8381c07cae67ff7768bbf19c6718e" + [[package]] name = "windows_x86_64_gnu" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.48.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b7b52767868a23d5bab768e390dc5f5c55825b6d30b86c844ff2dc7414044cc" + [[package]] name = "windows_x86_64_gnullvm" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" +[[package]] +name = "windows_x86_64_msvc" +version = "0.48.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed94fce61571a4006852b7389a063ab983c02eb1bb37b47f8272ce92d06d9538" + [[package]] name = "windows_x86_64_msvc" version = "0.52.6" diff --git a/Cargo.toml b/Cargo.toml index 91803ec..8d0fae3 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -19,3 +19,7 @@ serde_json = "1.0.114" url = "2.5.0" zip = "2.2.0" zstd = "0.13.0" + +[dev-dependencies] +mockito = "1.5.0" +tempfile = "3.12.0" diff --git a/src/download.rs b/src/download.rs index b1c7791..f0b2ac0 100644 --- a/src/download.rs +++ b/src/download.rs @@ -1,5 +1,6 @@ use anyhow::{anyhow, Result}; use reqwest::blocking::Client; +use reqwest::StatusCode; use std::collections::VecDeque; use std::fs::{self, OpenOptions}; use std::io::{Read, Seek, SeekFrom, Write}; @@ -10,7 +11,7 @@ use crate::eta::Eta; use crate::read_error_response::read_error_response; use crate::user_agent::APP_USER_AGENT; -pub fn download_file(url: &str, file_path: &Path, redirect_path: &Path) -> Result<()> { +fn download_file(url: &str, file_path: &Path, redirect_path: &Path) -> Result<()> { if let Some(dir) = file_path.parent() { fs::create_dir_all(dir)?; } @@ -32,19 +33,21 @@ pub fn download_file(url: &str, file_path: &Path, redirect_path: &Path) -> Resul .header("Range", format!("bytes={}-", file_size)) .send()?; + let code = response.status(); + match code { + StatusCode::PARTIAL_CONTENT => {} + _ if code.is_success() => { + anyhow::bail!("expected {}, but got {}", StatusCode::PARTIAL_CONTENT, code); + } + _ => { + let err = read_error_response(response.text()?); + anyhow::bail!("failed to download from {url}: {code} {err}"); + } + } let final_url = response.url().clone(); fs::write(redirect_path, final_url.as_str())?; - let status = response.status(); - if !status.is_success() { - let err = read_error_response(response.text()?); - anyhow::bail!(format!( - "Failed to download from {}: {} {}", - final_url, status, err - )); - } - let content_len = response .headers() .get(reqwest::header::CONTENT_LENGTH) @@ -53,7 +56,6 @@ pub fn download_file(url: &str, file_path: &Path, redirect_path: &Path) -> Resul .unwrap_or(0); let total_size = content_len + file_size; - file.seek(SeekFrom::End(0))?; const MEASUREMENT_SIZE: usize = 500; @@ -141,3 +143,74 @@ pub fn download_with_retries( } } } + +#[cfg(test)] +mod tests { + use std::fs; + + #[test] + fn rejects_not_206() { + let mut server = mockito::Server::new(); + + let mock = server.mock("GET", "/").with_status(200).create(); + + let tmpdir = tempfile::tempdir().unwrap(); + let file_path = tmpdir.path().join("file.bin"); + let redirect_path = tmpdir.path().join("redirect.txt"); + + let result = super::download_file(&server.url(), &file_path, &redirect_path); + let err = result.unwrap_err(); + assert_eq!( + err.to_string(), + "expected 206 Partial Content, but got 200 OK" + ); + + mock.assert(); + } + + #[test] + fn fails_when_server_fails() { + let mut server = mockito::Server::new(); + + let mock = server.mock("GET", "/").with_status(500).create(); + + let tmpdir = tempfile::tempdir().unwrap(); + let file_path = tmpdir.path().join("file.bin"); + let redirect_path = tmpdir.path().join("redirect.txt"); + + let result = super::download_file(&server.url(), &file_path, &redirect_path); + let err = result.unwrap_err(); + assert!(err.to_string().contains("failed to download from")); + + mock.assert(); + } + + #[test] + fn downloads_file() { + let mut server = mockito::Server::new(); + + let binary = b"1234567890"; + + let mock = server + .mock("GET", "/file") + .with_status(206) + .with_header("Content-Length", &format!("{}", binary.len())) + .with_body(binary) + .create(); + + let tmpdir = tempfile::tempdir().unwrap(); + let file_path = tmpdir.path().join("file.bin"); + let redirect_path = tmpdir.path().join("redirect.txt"); + + let url = server.url() + "/file"; + super::download_file(&url, &file_path, &redirect_path).unwrap(); + + let content = fs::read(file_path).unwrap(); + assert_eq!(content, binary); + + let redirect_url = fs::read_to_string(redirect_path).unwrap(); + assert_eq!(redirect_url, url); + + mock.assert(); + } +} From 30a4064c0d4da865c5d7315ec526cde5098977b4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Bartosz=20R=C3=B3=C5=BCa=C5=84ski?= Date: Tue, 20 Aug 2024 13:55:54 +0200 Subject: [PATCH 2/4] Use &Path in `get_version()` --- src/go_spacemesh.rs | 14 +++++--------- src/main.rs | 10 ++-------- src/utils.rs | 4 ---- 3 files changed, 7 insertions(+), 21 deletions(-) diff --git a/src/go_spacemesh.rs b/src/go_spacemesh.rs index 4aa11e1..56df5a2 100644 --- a/src/go_spacemesh.rs +++ b/src/go_spacemesh.rs @@ -1,21 +1,17 @@ use anyhow::Result; -use std::{io::ErrorKind, process::Command}; +use std::{io::ErrorKind, path::Path, process::Command}; -use crate::utils::trim_version; - -pub fn get_version(path: &str) -> Result { +pub fn get_version(path: &Path) -> Result { let output = Command::new(path) .arg("version") .output() .map_err(|error| match error.kind() { ErrorKind::NotFound => { - anyhow::anyhow!("Executable not found at path: {}", path) + anyhow::anyhow!("executable not found at path: {}", path.display()) } - other_error => panic!("Unexpected error: {:?}", other_error), + other_error => anyhow::anyhow!("unexpected error: {other_error}"), })?; let version = String::from_utf8(output.stdout)?; - let trimmed = trim_version(version.trim()).to_string(); - - Ok(trimmed) + Ok(version.split('+').next().unwrap().to_string()) } diff --git a/src/main.rs b/src/main.rs index 1f922ff..a83c624 100644 --- a/src/main.rs +++ b/src/main.rs @@ -153,10 +153,7 @@ fn main() -> anyhow::Result<()> { println!("Current network layer: {}", time_layer); let go_path = resolve_path(&go_spacemesh_path).unwrap(); - let go_path_str = go_path - .to_str() - .expect("Cannot resolve path to go-spacemesh"); - let go_version = get_version(go_path_str)?; + let go_version = get_version(&go_path)?; let quicksync_layer = fetch_latest_available_layer(&download_url, &go_version)?; println!("Latest layer in cloud: {}", quicksync_layer); Ok(()) @@ -190,10 +187,7 @@ fn main() -> anyhow::Result<()> { std::fs::read_to_string(&redirect_file_path)? } else { let go_path = resolve_path(&go_spacemesh_path).unwrap(); - let go_path_str = go_path - .to_str() - .expect("Cannot resolve path to go-spacemesh"); - let path = format!("{}/state.zst", &get_version(go_path_str)?); + let path = format!("{}/state.zst", &get_version(&go_path)?); let url = build_url(&download_url, &path); url.to_string() }; diff --git a/src/utils.rs b/src/utils.rs index 52e6de6..87fa3a2 100644 --- a/src/utils.rs +++ b/src/utils.rs @@ -25,10 +25,6 @@ pub fn resolve_path(relative_path: &PathBuf) -> Result { Ok(resolved_path) } -pub fn trim_version(version: &str) -> &str { - version.split('+').next().unwrap_or(version) -} - pub fn build_url(base: &Url, path: &str) -> Url { let mut url = base.clone(); url From 5b37caca82968b415cca46ff5d0edbd8023302a7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Bartosz=20R=C3=B3=C5=BCa=C5=84ski?= Date: Tue, 20 Aug 2024 17:41:53 +0200 Subject: [PATCH 3/4] add more downloading tests and fix redirect bug after retry --- Cargo.lock | 1 + Cargo.toml | 1 + src/download.rs | 197 +++++++++++++++++++++++++++++++++++++----------- src/main.rs | 62 ++++++++------- src/utils.rs | 26 ++----- 5 files changed, 199 insertions(+), 88 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 4cd21d8..21c728c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1226,6 +1226,7 @@ dependencies = [ "duration-string", "md5", "mockito", + "rand", "regex", "reqwest", "rusqlite", diff --git a/Cargo.toml b/Cargo.toml index 8d0fae3..83cbb8e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -22,4 +22,5 @@ zstd = "0.13.0" [dev-dependencies] mockito = "1.5.0" +rand = "0.8.5" tempfile = "3.12.0" diff --git a/src/download.rs b/src/download.rs index f0b2ac0..0567f24 100644 --- a/src/download.rs +++ b/src/download.rs @@ -2,7 +2,6 @@ use anyhow::{anyhow, Result}; use reqwest::blocking::Client; use reqwest::StatusCode; use std::collections::VecDeque; -use std::fs::{self, OpenOptions}; use std::io::{Read, Seek, SeekFrom, Write}; use std::path::Path; use std::time::Instant; @@ -11,26 +10,22 @@ use crate::eta::Eta; use crate::read_error_response::read_error_response; use crate::user_agent::APP_USER_AGENT; -fn download_file(url: &str, file_path: &Path, redirect_path: &Path) -> Result<()> { - if let Some(dir) = file_path.parent() { - fs::create_dir_all(dir)?; - } - - let mut file = OpenOptions::new() - .create(true) - .read(true) - .append(true) - .open(file_path)?; +fn download_file(url: &str, file: &mut W, redirect_path: &Path) -> Result<()> { + let offset = file.seek(SeekFrom::End(0))?; - let file_size = file.metadata()?.len(); + let url = if redirect_path.try_exists().unwrap_or(false) { + std::fs::read_to_string(redirect_path)? + } else { + url.to_string() + }; let client = Client::builder() .user_agent(APP_USER_AGENT) .timeout(std::time::Duration::from_secs(30)) .build()?; let mut response = client - .get(url) - .header("Range", format!("bytes={}-", file_size)) + .get(&url) + .header("Range", format!("bytes={offset}-")) .send()?; let code = response.status(); @@ -46,7 +41,7 @@ fn download_file(url: &str, file_path: &Path, redirect_path: &Path) -> Result<() } let final_url = response.url().clone(); - fs::write(redirect_path, final_url.as_str())?; + std::fs::write(redirect_path, final_url.as_str())?; let content_len = response .headers() @@ -55,15 +50,14 @@ fn download_file(url: &str, file_path: &Path, redirect_path: &Path) -> Result<() .and_then(|ct_len| ct_len.parse::().ok()) .unwrap_or(0); - let total_size = content_len + file_size; - file.seek(SeekFrom::End(0))?; + let total_size = content_len + offset; const MEASUREMENT_SIZE: usize = 500; let mut last_reported_progress: Option = None; let start = Instant::now(); let mut measurements = VecDeque::with_capacity(MEASUREMENT_SIZE); - let mut just_downloaded: u64 = 0; + let mut just_downloaded = 0; let mut buffer = [0; 16 * 1024]; loop { @@ -74,7 +68,7 @@ fn download_file(url: &str, file_path: &Path, redirect_path: &Path) -> Result<() Ok(bytes_read) => { file.write_all(&buffer[..bytes_read])?; just_downloaded += bytes_read as u64; - let downloaded = file_size + just_downloaded; + let downloaded = offset + just_downloaded; let elapsed = start.elapsed().as_secs_f64(); let speed = if elapsed > 0.0 { @@ -118,25 +112,20 @@ fn download_file(url: &str, file_path: &Path, redirect_path: &Path) -> Result<() Ok(()) } -pub fn download_with_retries( +pub(crate) fn download_with_retries( url: &str, - file_path: &Path, + file: &mut W, redirect_path: &Path, max_retries: u32, ) -> Result<()> { let mut attempts = 0; loop { - match download_file(url, file_path, redirect_path) { + attempts += 1; + match download_file(url, file, redirect_path) { Ok(()) => return Ok(()), - Err(e) if attempts < max_retries => { - println!( - "Download error: {}. Attempt {} / {}", - e, - attempts + 1, - max_retries - ); - attempts += 1; + Err(e) if attempts <= max_retries => { + println!("Download error: {e}. Attempt {attempts} / {max_retries}",); std::thread::sleep(std::time::Duration::from_secs(5)); } Err(e) => return Err(anyhow!(e)), @@ -146,19 +135,25 @@ pub fn download_with_retries( #[cfg(test)] mod tests { - use std::fs; + use std::{ + cmp::min, + fs, + io::{Error, ErrorKind, Read, Seek}, + iter, + }; + + use rand::{Rng, SeedableRng}; #[test] fn rejects_not_206() { let mut server = mockito::Server::new(); - let mock = server.mock("GET", "/").with_status(200).create(); let tmpdir = tempfile::tempdir().unwrap(); - let file_path = tmpdir.path().join("file.bin"); let redirect_path = tmpdir.path().join("redirect.txt"); + let mut file = tempfile::tempfile().unwrap(); - let result = super::download_file(&server.url(), &file_path, &redirect_path); + let result = super::download_file(&server.url(), &mut file, &redirect_path); let err = result.unwrap_err(); assert_eq!( err.to_string(), @@ -171,14 +166,13 @@ mod tests { #[test] fn fails_when_server_fails() { let mut server = mockito::Server::new(); - let mock = server.mock("GET", "/").with_status(500).create(); let tmpdir = tempfile::tempdir().unwrap(); - let file_path = tmpdir.path().join("file.bin"); let redirect_path = tmpdir.path().join("redirect.txt"); + let mut file = tempfile::tempfile().unwrap(); - let result = super::download_file(&server.url(), &file_path, &redirect_path); + let result = super::download_file(&server.url(), &mut file, &redirect_path); let err = result.unwrap_err(); assert!(err.to_string().contains("failed to download from")); @@ -187,25 +181,24 @@ mod tests { #[test] fn downloads_file() { - let mut server = mockito::Server::new(); - let binary = b"1234567890"; + let mut server = mockito::Server::new(); let mock = server .mock("GET", "/file") .with_status(206) - .with_header("Content-Length", &format!("{}", binary.len())) .with_body(binary) .create(); let tmpdir = tempfile::tempdir().unwrap(); - let file_path = tmpdir.path().join("file.bin"); + let mut file = tempfile::tempfile().unwrap(); let redirect_path = tmpdir.path().join("redirect.txt"); let url = server.url() + "/file"; - super::download_file(&url, &file_path, &redirect_path).unwrap(); - let content = fs::read(file_path).unwrap(); + super::download_file(&url, &mut file, &redirect_path).unwrap(); + file.seek(std::io::SeekFrom::Start(0)).unwrap(); + let content = file.bytes().collect::, _>>().unwrap(); assert_eq!(content, binary); let redirect_url = fs::read_to_string(redirect_path).unwrap(); @@ -213,4 +206,122 @@ mod tests { mock.assert(); } + + #[test] + fn follows_redirect_and_persists_it() { + let binary = b"1234567890"; + + let mut server = mockito::Server::new(); + let redirected_url = server.url() + "/redirected"; + let mock_redirect = server + .mock("GET", "/file") + .with_status(301) + .with_header("location", &redirected_url) + .create(); + + let mock = server + .mock("GET", "/redirected") + .with_status(206) + .with_body(binary) + .create(); + + let tmpdir = tempfile::tempdir().unwrap(); + let mut file = tempfile::tempfile().unwrap(); + let redirect_path = tmpdir.path().join("redirect.txt"); + + let url = server.url() + "/file"; + + super::download_file(&url, &mut file, &redirect_path).unwrap(); + file.seek(std::io::SeekFrom::Start(0)).unwrap(); + let content = file.bytes().collect::, _>>().unwrap(); + assert_eq!(content, binary); + + let redirect_url = fs::read_to_string(redirect_path).unwrap(); + assert_eq!(redirect_url, redirected_url); + + mock_redirect.assert(); + mock.assert(); + } + + #[test] + fn retries_after_failure() { + let mut server = mockito::Server::new(); + + let mut rng = rand::rngs::StdRng::seed_from_u64(11); + let binary: Vec = iter::repeat_with(|| rng.gen()).take(2_000).collect(); + let binary = std::sync::Arc::new(binary); + let binary_clone = binary.clone(); + + let mock_redirect = server + .mock("GET", "/file") + .with_status(301) + .with_header("location", &(server.url() + "/redirected")) + .create(); + + let mock = server + .mock("GET", "/redirected") + .with_status(206) + .with_body_from_request(move |req| { + let range_hdr = req.header("Range").first().unwrap().to_str().unwrap(); + let range = range_hdr.split('=').nth(1).unwrap(); + let start = range.strip_suffix('-').unwrap().parse::().unwrap(); + + binary_clone[start..].to_vec() + }) + .create() + .expect(2); + + let tmpdir = tempfile::tempdir().unwrap(); + let redirect_path = tmpdir.path().join("redirect.txt"); + + // a mock file that fails once after writing the first bytes on the first attempt + struct FileMock { + bytes: Vec, + failed: bool, + } + + impl std::io::Write for FileMock { + fn write(&mut self, buf: &[u8]) -> std::io::Result { + match (self.bytes.len(), self.failed) { + (0, _) => { + // write only the first 50 bytes + let size = min(50, buf.len()); + self.bytes.extend_from_slice(&buf[..size]); + Ok(size) + } + (1.., false) => { + self.failed = true; + Err(Error::new(ErrorKind::Other, "failed writing")) + } + _ => { + self.bytes.extend_from_slice(buf); + Ok(buf.len()) + } + } + } + + fn flush(&mut self) -> std::io::Result<()> { + Ok(()) + } + } + + impl std::io::Seek for FileMock { + fn seek(&mut self, _: std::io::SeekFrom) -> std::io::Result { + Ok(self.bytes.len() as u64) + } + } + + let mut file = FileMock { + bytes: Vec::new(), + failed: false, + }; + + let url = server.url() + "/file"; + super::download_with_retries(&url, &mut file, &redirect_path, 1).unwrap(); + + mock_redirect.assert(); + mock.assert(); + + assert_eq!(file.bytes, *binary); + } } diff --git a/src/main.rs b/src/main.rs index a83c624..7d27255 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,7 +1,9 @@ use chrono::Duration; use clap::{Parser, Subcommand}; -use std::path::PathBuf; +use std::fs::OpenOptions; +use std::path::Path; use std::process; +use std::{env, path::PathBuf}; use url::Url; mod checksum; @@ -17,6 +19,7 @@ mod unpack; mod user_agent; mod utils; +use anyhow::Context; use checksum::*; use download::download_with_retries; use go_spacemesh::get_version; @@ -121,6 +124,11 @@ fn backup_or_fail(file_path: PathBuf) { } } +fn resolve_path(relative_path: &Path) -> anyhow::Result { + let current_dir = env::current_dir()?; + Ok(current_dir.join(relative_path)) +} + fn main() -> anyhow::Result<()> { let cli = Cli::parse(); @@ -166,7 +174,7 @@ fn main() -> anyhow::Result<()> { Commands::Download { node_data, go_spacemesh_path, - download_url, + mut download_url, max_retries, } => { let dir_path = node_data; @@ -179,43 +187,45 @@ fn main() -> anyhow::Result<()> { let wal_file_path = dir_path.join("state.sql-wal"); // Download archive if needed - let archive_file_path = if !archive_zip_file_path.try_exists().unwrap_or(false) - && !archive_zstd_file_path.try_exists().unwrap_or(false) - { + let archive_file_path = if archive_zip_file_path.try_exists().unwrap_or(false) { + archive_zip_file_path + } else if archive_zstd_file_path.try_exists().unwrap_or(false) { + archive_zstd_file_path + } else { + let archive_file_path = archive_zstd_file_path; println!("Downloading the latest database..."); let url = if redirect_file_path.try_exists().unwrap_or(false) { std::fs::read_to_string(&redirect_file_path)? } else { - let go_path = resolve_path(&go_spacemesh_path).unwrap(); - let path = format!("{}/state.zst", &get_version(&go_path)?); - let url = build_url(&download_url, &path); - url.to_string() + let go_path = resolve_path(&go_spacemesh_path).context("checking node version")?; + let version = get_version(&go_path)?; + download_url + .path_segments_mut() + .map_err(|e| anyhow::anyhow!("parsing download url: {e:?}"))? + .extend(&[&version, "state.zst"]); + download_url.to_string() }; - if let Err(e) = - download_with_retries(&url, &temp_file_path, &redirect_file_path, max_retries) - { - eprintln!( - "Failed to download a file after {} attempts: {}", - max_retries, e - ); - process::exit(1); + if let Some(dir) = temp_file_path.parent() { + std::fs::create_dir_all(dir)?; } - let archive_file_path = if url.ends_with(".zip") { - archive_zip_file_path - } else { - archive_zstd_file_path - }; + let mut file = OpenOptions::new() + .create(true) + .read(true) + .append(true) + .open(&temp_file_path)?; + + if let Err(e) = download_with_retries(&url, &mut file, &redirect_file_path, max_retries) { + eprintln!("Failed to download a file after {max_retries} attempts: {e}",); + process::exit(1); + } + drop(file); // Rename `state.download` -> `state.zst` std::fs::rename(&temp_file_path, &archive_file_path)?; println!("Archive downloaded!"); archive_file_path - } else if archive_zip_file_path.try_exists().unwrap_or(false) { - archive_zip_file_path - } else { - archive_zstd_file_path }; if redirect_file_path.try_exists().unwrap_or(false) { diff --git a/src/utils.rs b/src/utils.rs index 87fa3a2..bd812ca 100644 --- a/src/utils.rs +++ b/src/utils.rs @@ -2,7 +2,7 @@ use anyhow::{anyhow, Result}; use chrono::{DateTime, Duration, Utc}; use regex::Regex; use reqwest::{blocking::Client, redirect}; -use std::{env, path::PathBuf}; +use std::path::{Path, PathBuf}; use url::Url; use crate::user_agent::APP_USER_AGENT; @@ -19,22 +19,7 @@ pub fn calculate_latest_layer( Ok(delta.num_milliseconds() / layer_duration.num_milliseconds()) } -pub fn resolve_path(relative_path: &PathBuf) -> Result { - let current_dir = env::current_dir()?; - let resolved_path = current_dir.join(relative_path); - Ok(resolved_path) -} - -pub fn build_url(base: &Url, path: &str) -> Url { - let mut url = base.clone(); - url - .path_segments_mut() - .expect("cannot be base") - .extend(path.split('/')); - url -} - -pub fn backup_file(original_path: &PathBuf) -> Result { +pub fn backup_file(original_path: &Path) -> Result { if !original_path.exists() { anyhow::bail!("No file to make a backup"); } @@ -76,8 +61,11 @@ pub fn fetch_latest_available_layer(download_url: &Url, go_version: &str) -> Res .timeout(std::time::Duration::from_secs(30)) .build()?; - let path = format!("{}/state.zst", go_version); - let url = build_url(download_url, &path); + let mut url = download_url.clone(); + url + .path_segments_mut() + .unwrap() + .extend(&[go_version, "state.zst"]); let response = client.head(url).send()?; From f5723de7b406ed67b5b14c451057180cdb51f296 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Bartosz=20R=C3=B3=C5=BCa=C5=84ski?= Date: Tue, 20 Aug 2024 17:47:00 +0200 Subject: [PATCH 4/4] shorten retry delay in tests --- src/download.rs | 13 +++++++++++-- src/main.rs | 10 +++++++++- 2 files changed, 20 insertions(+), 3 deletions(-) diff --git a/src/download.rs b/src/download.rs index 0567f24..e8c19fa 100644 --- a/src/download.rs +++ b/src/download.rs @@ -117,6 +117,7 @@ pub(crate) fn download_with_retries( file: &mut W, redirect_path: &Path, max_retries: u32, + retry_delay: std::time::Duration, ) -> Result<()> { let mut attempts = 0; @@ -126,7 +127,7 @@ pub(crate) fn download_with_retries( Ok(()) => return Ok(()), Err(e) if attempts <= max_retries => { println!("Download error: {e}. Attempt {attempts} / {max_retries}",); - std::thread::sleep(std::time::Duration::from_secs(5)); + std::thread::sleep(retry_delay); } Err(e) => return Err(anyhow!(e)), } @@ -135,6 +136,7 @@ pub(crate) fn download_with_retries( #[cfg(test)] mod tests { + use core::time; use std::{ cmp::min, fs, @@ -317,7 +319,14 @@ mod tests { }; let url = server.url() + "/file"; - super::download_with_retries(&url, &mut file, &redirect_path, 1).unwrap(); + super::download_with_retries( + &url, + &mut file, + &redirect_path, + 1, + time::Duration::from_millis(1), + ) + .unwrap(); mock_redirect.assert(); mock.assert(); diff --git a/src/main.rs b/src/main.rs index e90f2de..b44ec93 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,6 +1,7 @@ use chrono::Duration; use clap::{Parser, Subcommand}; use std::fs::OpenOptions; +use std::io::Write; use std::path::Path; use std::process; use std::{env, path::PathBuf}; @@ -215,8 +216,15 @@ fn main() -> anyhow::Result<()> { .append(true) .open(&temp_file_path)?; - if let Err(e) = download_with_retries(&url, &mut file, &redirect_file_path, max_retries) { + if let Err(e) = download_with_retries( + &url, + &mut file, + &redirect_file_path, + max_retries, + std::time::Duration::from_secs(5), + ) { eprintln!("Failed to download a file after {max_retries} attempts: {e}",); + file.flush()?; process::exit(1); } drop(file);