From 3c204aa2fb13862deaf7d17935e0c49265d199e5 Mon Sep 17 00:00:00 2001 From: lifegpc Date: Wed, 8 Jun 2022 12:03:59 +0000 Subject: [PATCH] Update --- Cargo.lock | 8 +++++ Cargo.toml | 2 ++ src/downloader/downloader.rs | 13 ++++---- src/downloader/error.rs | 9 ++++++ src/downloader/pd_file/file.rs | 22 +++++++++++++ src/downloader/tasks.rs | 57 ++++++++++++++++++++++++++++++++-- src/ext/atomic.rs | 10 +++--- src/main.rs | 2 ++ 8 files changed, 110 insertions(+), 13 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 12b1d12..73119d6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -665,6 +665,12 @@ dependencies = [ "pin-project-lite", ] +[[package]] +name = "http-content-range" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f0d1a8ef218a86416107794b34cc446958d9203556c312bb41eab4c924c1d2e" + [[package]] name = "httparse" version = "1.7.1" @@ -1152,6 +1158,8 @@ dependencies = [ "getopts", "gettext", "html_parser", + "http", + "http-content-range", "indicatif", "int-enum", "json", diff --git a/Cargo.toml b/Cargo.toml index 86d5e29..1487f71 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -16,6 +16,8 @@ futures-util = "0.3" getopts = "0.2" gettext = "0.4" html_parser = "0.6.3" +http = "0.2" +http-content-range = "0.1" indicatif = "0.17.0-rc.11" int-enum = "0.4" json = "0.12" diff --git a/src/downloader/downloader.rs b/src/downloader/downloader.rs index 348eda8..dddc32a 100644 --- a/src/downloader/downloader.rs +++ b/src/downloader/downloader.rs @@ -26,19 +26,19 @@ use url::Url; /// A file downloader pub struct DownloaderInternal { /// The webclient - client: Arc, + pub client: Arc, /// The download status - pd: Arc, + pub pd: Arc, /// The url of the file - url: Arc, + pub url: Arc, /// The HTTP headers map - headers: Arc>, + pub headers: Arc>, /// The target file file: RwLock>, /// The status of the downloader status: RwLock, /// All tasks - tasks: RwLock>>, + tasks: RwLock>>>, /// Whether to enable mulitple thread mode multi: AtomicBool, } @@ -114,7 +114,7 @@ impl DownloaderInternal { impl DownloaderInternal { /// Add a new task to tasks /// * `task` - Task - pub fn add_task(&self, task: JoinHandle) { + pub fn add_task(&self, task: JoinHandle>) { self.tasks.get_mut().push(task) } @@ -162,6 +162,7 @@ impl Downloader { } } +#[doc(hidden)] macro_rules! define_downloader_fn { {$f:ident, $t:ty, $doc:expr} => { #[inline] diff --git a/src/downloader/error.rs b/src/downloader/error.rs index 8cbba5b..2428250 100644 --- a/src/downloader/error.rs +++ b/src/downloader/error.rs @@ -1,8 +1,17 @@ use crate::downloader::pd_file::PdFileError; +use http::status::StatusCode; #[derive(Debug, derive_more::From)] pub enum DownloaderError { ReqwestError(reqwest::Error), PdFileError(PdFileError), IoError(std::io::Error), + String(String), + ErrorStatusCode(StatusCode), +} + +impl From<&str> for DownloaderError { + fn from(v: &str) -> Self { + Self::String(String::from(v)) + } } diff --git a/src/downloader/pd_file/file.rs b/src/downloader/pd_file/file.rs index df749e4..5b01728 100644 --- a/src/downloader/pd_file/file.rs +++ b/src/downloader/pd_file/file.rs @@ -85,6 +85,22 @@ impl PdFile { } } + /// Set the status to the initailzed status. + pub fn clear(&self) -> Result<(), PdFileError> { + self.status.replace_with2(PdFileStatus::Started); + self.ftype.replace_with2(PdFileType::SingleThread); + self.file_size.qstore(0); + self.downloaded_file_size.qstore(0); + self.part_size.qstore(0); + self.part_datas.get_mut().clear(); + if !self.is_mem_only() { + self.need_saved.qstore(true); + // Rewrite all datas. + self.write()?; + } + Ok(()) + } + /// Close the file. /// This function will return error if write failed. /// If you want to force close the file. Please use [Self::force_close()]. @@ -111,6 +127,12 @@ impl PdFile { self.downloaded_file_size.qload() } + #[inline] + /// The target size of the file. 0 if unknown. + pub fn get_file_size(&self) -> u64 { + self.file_size.qload() + } + /// Return status data of a part /// * `index` - The part index pub fn get_part_data(&self, index: usize) -> Option> { diff --git a/src/downloader/tasks.rs b/src/downloader/tasks.rs index 057b7cd..f869a97 100644 --- a/src/downloader/tasks.rs +++ b/src/downloader/tasks.rs @@ -1,9 +1,62 @@ +use crate::ext::try_err::TryErr; +use crate::gettext; +use super::error::DownloaderError; use super::downloader::DownloaderInternal; +use http_content_range::ContentRange; +use std::ops::Deref; use std::io::Seek; use std::io::Write; use std::sync::Arc; /// Create a download tasks in simple thread mode. -pub async fn create_download_tasks_simple(d: Arc>) -> bool { - false +pub async fn create_download_tasks_simple(d: Arc>) -> Result<(), DownloaderError> { + let start = if d.pd.is_downloading() { + d.pd.get_downloaded_file_size() + } else { + 0 + }; + let file_size = d.pd.get_file_size(); + let mut headers = d.headers.deref().clone(); + if start != 0 { + headers.insert(String::from("Range"), format!("bytes={}-", start)); + } + let mut result = d.client.get(d.url.deref().clone(), headers).try_err(gettext("Failed to get url."))?; + let mut status = result.status(); + if status == 416 { + result = d.client.get(d.url.deref().clone(), d.headers.deref().clone()).try_err(gettext("Failed to get url."))?; + status = result.status(); + } else if status == 206 { + let range = ContentRange::parse_bytes(result.headers()["Content-Range"].as_bytes()); + let need_reget = match range { + ContentRange::Bytes(b) => { + if file_size != 0 && b.complete_length != file_size { + true + } else if start != b.first_byte { + true + } else { + false + } + } + ContentRange::UnboundBytes(b) => { + if start != b.first_byte { + true + } else { + false + } + } + ContentRange::Unknown => { + true + } + _ => { false } + }; + if need_reget { + d.pd.clear()?; + result = d.client.get(d.url.deref().clone(), d.headers.deref().clone()).try_err(gettext("Failed to get url."))?; + status = result.status(); + } + } + if status.as_u16() >= 400 { + return Err(DownloaderError::from(status)); + } + Ok(()) } diff --git a/src/ext/atomic.rs b/src/ext/atomic.rs index ba83b01..04dbc43 100644 --- a/src/ext/atomic.rs +++ b/src/ext/atomic.rs @@ -2,12 +2,12 @@ use std::sync::atomic::Ordering; /// A trait to help to load and store atomic value quickly. pub trait AtomicQuick { - /// Loads a value from the atomic integer + /// Loads a value from the atomic integer in [Ordering::SeqCst] mode. fn qload(&self) -> T; - /// Stores a value into the atomic integer. + /// Stores a value into the atomic integer in [Ordering::SeqCst] mode. fn qstore(&self, value: T); #[inline] - /// Stores a value into the atomic integer. + /// Stores a value into the atomic integer in [Ordering::SeqCst] mode. /// Alias for [Self::qstore] fn qsave(&self, value: T) { self.qstore(value) @@ -19,11 +19,11 @@ macro_rules! impl_atomic_quick_with_atomic { impl AtomicQuick<$type2> for $type1 { #[inline] fn qload(&self) -> $type2 { - self.load(Ordering::Relaxed) + self.load(Ordering::SeqCst) } #[inline] fn qstore(&self, value: $type2) { - self.store(value, Ordering::Relaxed) + self.store(value, Ordering::SeqCst) } } } diff --git a/src/main.rs b/src/main.rs index f417cb4..7b9c69a 100644 --- a/src/main.rs +++ b/src/main.rs @@ -8,6 +8,8 @@ extern crate derive_more; extern crate flagset; extern crate futures_util; extern crate json; +extern crate http; +extern crate http_content_range; extern crate indicatif; extern crate int_enum; #[macro_use]