diff --git a/Cargo.lock b/Cargo.lock index a585d4e..a2b66b9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -389,13 +389,12 @@ dependencies = [ [[package]] name = "nhentairs" -version = "0.5.4" +version = "0.5.5" dependencies = [ "reqwest", "serde", "serde_json", "tokio", - "tokio-stream", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 50f4e97..8aa6571 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "nhentairs" -version = "0.5.4" +version = "0.5.5" authors = ["blank X "] edition = "2018" @@ -12,6 +12,5 @@ lto = true [dependencies] serde = { version = "1.0", features = ["derive"] } serde_json = "1.0" -reqwest = { version = "0.11", features = ["stream"] } -tokio = { version = "1.0", features = ["rt-multi-thread", "sync"] } -tokio-stream = "0.1" +reqwest = "0.11" +tokio = { version = "1.0", features = ["rt-multi-thread", "sync", "time"] } diff --git a/src/commands/download.rs b/src/commands/download.rs index 7a7fc54..4578699 100644 --- a/src/commands/download.rs +++ b/src/commands/download.rs @@ -7,11 +7,13 @@ use std::path::Path; use std::process::exit; use tokio::sync::Mutex; use tokio::task::JoinHandle; +use tokio::time::{sleep, Duration}; use std::fs::{rename, create_dir, write}; extern crate tokio; extern crate reqwest; const DOWNLOAD_WORKERS: usize = 5; +const FAIL_DOWNLOAD_WAIT_TIME: u64 = 5000; pub async fn run(args: env::Args) { let sauces = utils::get_arg_sauces(args).unwrap(); @@ -90,7 +92,17 @@ pub async fn run(args: env::Args) { eprintln!("[DW{}] Downloading {} to {}", worker_id, url, file_path); let mut tmp_file_path = file_path.clone(); tmp_file_path.push_str(".tmp"); - utils::download_file(cloned_client, &url, &tmp_file_path).await.unwrap(); + loop { + match utils::download_file(cloned_client.clone(), &url, &tmp_file_path).await { + Ok(success) => { + if success { + break; + } + }, + Err(err) => eprintln!("[DW{}] Failed to download {} due to {}, sleeping for {}ms", worker_id, file_path, err, FAIL_DOWNLOAD_WAIT_TIME) + }; + sleep(Duration::from_millis(FAIL_DOWNLOAD_WAIT_TIME)).await; + } rename(&tmp_file_path, &file_path).unwrap(); eprintln!("[DW{}] {} downloaded", worker_id, file_path); } diff --git a/src/utils.rs b/src/utils.rs index f0c1b6f..d5c40f8 100644 --- a/src/utils.rs +++ b/src/utils.rs @@ -3,7 +3,6 @@ use crate::structs; use std::env; use std::fs::File; use std::io::Write; -use tokio_stream::StreamExt; extern crate serde_json; extern crate reqwest; @@ -25,16 +24,19 @@ pub async fn get_search_info(client: reqwest::Client, search_query: &str) -> Res Ok(serde_json::from_str(&resp.text().await?).unwrap()) } -pub async fn download_file(client: reqwest::Client, url: &str, file_name: &str) -> Result<(), reqwest::Error> { +pub async fn download_file(client: reqwest::Client, url: &str, file_name: &str) -> Result { let resp = client.get(url) .send() .await?; - let mut file = File::create(&file_name).unwrap(); - let mut stream = resp.bytes_stream(); - while let Some(item) = stream.next().await { - file.write(&item?).unwrap(); - } - Ok(()) + Ok(match resp.headers().get("Content-Type") { + Some(header) if header.to_str().unwrap_or_default().starts_with("image/") => { + let bytes = resp.bytes().await?; + let mut file = File::create(&file_name).unwrap(); + file.write_all(&bytes).unwrap(); + true + }, + _ => false + }) } pub fn get_arg_sauces(args: env::Args) -> Result, String> {