diff --git a/Cargo.lock b/Cargo.lock index a17982f..795c9be 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1142,9 +1142,11 @@ dependencies = [ "json", "lazy_static", "link-cplusplus", + "quote", "regex", "reqwest", "spin_on", + "syn", "tokio", "urlparse", "utf16string", diff --git a/Cargo.toml b/Cargo.toml index dfc07c3..2a54a6f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -5,6 +5,11 @@ edition = "2018" # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html +[lib] +name = "proc_macros" +path = "src/proc_macros.rs" +proc-macro = true + [dependencies] atty = "0.2" c_fixed_string = { version = "0.2", optional = true } @@ -20,10 +25,12 @@ indicatif = "0.17.0-rc.10" int-enum = "0.4" json = "0.12" lazy_static = "1.4" +quote = "1" regex = "1" reqwest = { version = "0.11", features = ["brotli", "deflate", "gzip", "rustls-tls", "socks", "stream"] } RustyXML = "0.3" spin_on = "0.1.1" +syn = "1" tokio = { version = "1.18", features = ["rt", "macros", "rt-multi-thread", "time"] } urlparse = "0.7" utf16string = { version= "0.2", optional = true } diff --git a/src/download.rs b/src/download.rs index dd7d6a2..74b9583 100644 --- a/src/download.rs +++ b/src/download.rs @@ -6,6 +6,8 @@ use crate::data::exif::add_exifdata_to_image; use crate::data::json::JSONDataFile; #[cfg(feature = "ugoira")] use crate::data::video::get_video_metadata; +use crate::downloader::pd_file::enums::PdFileResult; +use crate::downloader::pd_file::file::PdFile; use crate::gettext; use crate::opthelper::OptHelper; use crate::pixiv_link::PixivID; @@ -65,35 +67,85 @@ impl Main { } let file_name = file_name.unwrap(); let file_name = base.join(file_name); - if file_name.exists() { - match helper.overwrite() { - Some(overwrite) => { - if !overwrite { - #[cfg(feature = "exif")] - { - if helper.update_exif() { - if add_exifdata_to_image(&file_name, &datas, np).is_err() { - println!( - "{} {}", - gettext("Failed to add exif data to image:"), - file_name.to_str().unwrap_or("(null)") - ); + let pdf = match PdFile::open(&file_name) { + Ok(f) => { + match f { + PdFileResult::TargetExisted => { + match helper.overwrite() { + Some(overwrite) => { + if !overwrite { + #[cfg(feature = "exif")] + { + if helper.update_exif() { + if add_exifdata_to_image(&file_name, &datas, np).is_err() { + println!( + "{} {}", + gettext("Failed to add exif data to image:"), + file_name.to_str().unwrap_or("(null)") + ); + } + } + } + return 0; + } + } + None => { + if !ask_need_overwrite(file_name.to_str().unwrap()) { + return 0; } } } - return 0; - } - } - None => { - if !ask_need_overwrite(file_name.to_str().unwrap()) { - return 0; + match PdFile::open(&file_name) { + Ok(v) => { + match v { + PdFileResult::Ok(e) => { Some(e) } + _ => { None } + } + } + Err(e) => { + println!("{}", e); + None + } + } } + PdFileResult::Ok(e) => { Some(e) } + PdFileResult::ExistedOk(e) => { Some(e) } } } - } + Err(e) => { + println!("{}", e); + if file_name.exists() { + match helper.overwrite() { + Some(overwrite) => { + if !overwrite { + #[cfg(feature = "exif")] + { + if helper.update_exif() { + if add_exifdata_to_image(&file_name, &datas, np).is_err() { + println!( + "{} {}", + gettext("Failed to add exif data to image:"), + file_name.to_str().unwrap_or("(null)") + ); + } + } + } + return 0; + } + } + None => { + if !ask_need_overwrite(file_name.to_str().unwrap()) { + return 0; + } + } + } + } + None + } + }; let r; { - r = pw.adownload_image(link).await; + r = pw.adownload_image(link, &pdf).await; if r.is_none() { println!("{} {}", gettext("Failed to download image:"), link); return 1; diff --git a/src/downloader/pd_file/enums.rs b/src/downloader/pd_file/enums.rs index c4ecdc3..c5da1a9 100644 --- a/src/downloader/pd_file/enums.rs +++ b/src/downloader/pd_file/enums.rs @@ -1,3 +1,4 @@ +use crate::downloader::pd_file::file::PdFile; use int_enum::IntEnum; /// The status of the downloaded file. @@ -18,6 +19,12 @@ impl PdFileStatus { pub fn is_completed(&self) -> bool { *self == PdFileStatus::Downloaded } + + #[inline] + /// Returns true if the download is in progress. + pub fn is_downloading(&self) -> bool { + *self == PdFileStatus::Downloading + } } /// The type of the downloader. @@ -30,6 +37,20 @@ pub enum PdFileType { MultiThread = 1, } +#[derive(Debug)] +/// The result when try opening pd file. +pub enum PdFileResult { + /// The pd file is not existed, and the new pd file is created. + /// In this case, need download whole file. + Ok(PdFile), + /// The pd file is not existed but the target file is existed. + /// In most case, this means the download already completed. + TargetExisted, + /// The pd file is existed. + /// In this case, can continue to download. + ExistedOk(PdFile), +} + impl PdFileType { #[inline] /// Returns true if is multiple thread mode. diff --git a/src/downloader/pd_file/error.rs b/src/downloader/pd_file/error.rs index f1bfa43..15c2ada 100644 --- a/src/downloader/pd_file/error.rs +++ b/src/downloader/pd_file/error.rs @@ -1,11 +1,17 @@ use crate::gettext; +use int_enum::IntEnum; +use int_enum::IntEnumError; use std::convert::From; use std::fmt::Display; +use std::string::FromUtf8Error; #[derive(Debug, derive_more::From)] pub enum PdFileError { IoError(std::io::Error), String(String), + InvalidPdFile, + Unsupported, + Utf8Error(FromUtf8Error), } impl Display for PdFileError { @@ -18,6 +24,16 @@ impl Display for PdFileError { Self::String(e) => { f.write_str(e)?; } + Self::InvalidPdFile => { + f.write_str(gettext("Invalid pd file."))?; + } + Self::Unsupported => { + f.write_str(gettext("The pd file is newer version, please update the program."))?; + } + Self::Utf8Error(e) => { + f.write_str(gettext("Failed to decode UTF-8: "))?; + e.fmt(f)?; + } } Ok(()) } @@ -28,3 +44,9 @@ impl From<&str> for PdFileError { PdFileError::String(String::from(value)) } } + +impl From> for PdFileError { + fn from(e: IntEnumError) -> Self { + Self::String(format!("{} {}", gettext("Invalid pd file: "), e)) + } +} diff --git a/src/downloader/pd_file/file.rs b/src/downloader/pd_file/file.rs index 3b63d0c..535c572 100644 --- a/src/downloader/pd_file/file.rs +++ b/src/downloader/pd_file/file.rs @@ -1,7 +1,10 @@ use crate::downloader::pd_file::error::PdFileError; +use crate::downloader::pd_file::enums::PdFileResult; use crate::downloader::pd_file::enums::PdFileStatus; use crate::downloader::pd_file::enums::PdFileType; use crate::downloader::pd_file::version::PdFileVersion; +use crate::ext::io::StructRead; +use crate::ext::replace::ReplaceWith2; use crate::ext::rw_lock::GetRwLock; use crate::ext::try_err::TryErr; use crate::ext::try_err::TryErr2; @@ -14,6 +17,7 @@ use std::fs::create_dir; #[cfg(test)] use std::fs::metadata; use std::fs::remove_file; +use std::io::Read; use std::io::Seek; use std::io::SeekFrom; use std::io::Write; @@ -31,6 +35,7 @@ lazy_static! { static ref MAGIC_WORDS: Vec = vec![0x50, 0x44, 0xff, 0xff]; } +#[derive(Debug)] /// The pd file pub struct PdFile { /// The version of the current file. @@ -53,6 +58,8 @@ pub struct PdFile { downloaded_file_size: AtomicU64, /// The size of the each part. Ignored in single thread mode. part_size: AtomicU32, + /// Only stored in memory. + mem_only: AtomicBool, } impl PdFile { @@ -69,6 +76,7 @@ impl PdFile { file_size: AtomicU64::new(0), downloaded_file_size: AtomicU64::new(0), part_size: AtomicU32::new(0), + mem_only: AtomicBool::new(true), } } @@ -92,18 +100,107 @@ impl PdFile { } } + #[inline] + /// Returns the size of the downloaded data + pub fn get_downloaded_file_size(&self) -> u64 { + self.downloaded_file_size.load(Ordering::Relaxed) + } + #[inline] /// Returns true if the download is completed. pub fn is_completed(&self) -> bool { self.status.get_ref().is_completed() } + #[inline] + /// Returns true if the download is in progress. + pub fn is_downloading(&self) -> bool { + self.status.get_ref().is_downloading() + } + + #[inline] + /// Returns true if stored in memory only. + fn is_mem_only(&self) -> bool { + self.mem_only.load(Ordering::Relaxed) + } + + #[inline] + /// Returns true if is multiple thread mode. + pub fn is_multi_threads(&self) -> bool { + self.ftype.get_ref().is_multi() + } + #[inline] /// Returns true if needed to save to file. fn is_need_saved(&self) -> bool { self.need_saved.load(Ordering::Relaxed) } + /// Open a new [PdFile] if download is needed. + /// * `path` - The path of the file which want to download. + pub fn open + ?Sized>(path: &P) -> Result { + let p = path.as_ref(); + let mut pb = PathBuf::from(p); + let mut file_name = pb.file_name().try_err(gettext("Path need have a file name."))?.to_owned(); + file_name.push(".pd"); + pb.set_file_name(&file_name); + if p.exists() { + if pb.exists() { + let f = Self::read_from_file(p)?; + if f.is_completed() { + return Ok(PdFileResult::TargetExisted); + } + Ok(PdFileResult::ExistedOk(f)) + } else { + Ok(PdFileResult::TargetExisted) + } + } else { + let f = PdFile::new(); + f.open_with_create_file(&pb)?; + f.set_file_name(p.file_name().try_err(gettext("Path need have a file name."))?.to_str().unwrap_or("(null)"))?; + Ok(PdFileResult::Ok(f)) + } + } + + /// Create a new [PdFile] instance from the pd file. + /// * `path` - The path to the pd file. + /// + /// Returns errors or a new instance. + pub fn read_from_file + ?Sized>(path: &P) -> Result { + let p = path.as_ref(); + let mut f = File::open(p)?; + f.seek(SeekFrom::Start(0))?; + let mut buf = [0u8, 0, 0, 0]; + f.read_exact(&mut buf)?; + if MAGIC_WORDS.as_ref() == buf { + return Err(PdFileError::InvalidPdFile); + } + let version = PdFileVersion::read_from(&mut f)?; + if !version.is_supported() { + return Err(PdFileError::Unsupported); + } + let file_name_len = f.read_le_u32()?; + let status: PdFileStatus = PdFileStatus::from_int(f.read_le_u8()?)?; + let ftype: PdFileType = PdFileType::from_int(f.read_le_u8()?)?; + let file_size = f.read_le_u64()?; + let downloaded_file_size = f.read_le_u64()?; + let part_size = f.read_le_u32()?; + let file_name = String::from_utf8(f.read_bytes(file_name_len as usize)?)?; + Ok(Self { + version, + need_saved: AtomicBool::new(false), + file: RwLock::new(Some(f)), + file_path: RwLock::new(Some(p.to_path_buf())), + file_name: RwLock::new(Some(file_name)), + status: RwLock::new(status), + ftype: RwLock::new(ftype), + file_size: AtomicU64::new(file_size), + downloaded_file_size: AtomicU64::new(downloaded_file_size), + part_size: AtomicU32::new(part_size), + mem_only: AtomicBool::new(false), + }) + } + /// Create a new file and prepare to write data to it. /// If file alreay exists, will remove it first. /// * `path` - The path to the pd file. @@ -115,6 +212,7 @@ impl PdFile { let f = File::create(p)?; self.file.get_mut().replace(f); self.file_path.get_mut().replace(PathBuf::from(p)); + self.mem_only.store(false, Ordering::Relaxed); self.need_saved.store(true, Ordering::Relaxed); Ok(()) } @@ -142,6 +240,12 @@ impl PdFile { } } + #[inline] + /// Set status to alreay downloaded. + fn set_completed(&self) { + self.status.replace_with2(PdFileStatus::Downloaded); + } + /// Set the file name /// * `file_name` - The file name. Should not be empty. pub fn set_file_name + ?Sized>(&self, file_name: &S) -> Result<(), PdFileError> { @@ -150,9 +254,11 @@ impl PdFile { Err(gettext("File name should not be empty."))? } else { self.file_name.get_mut().replace(String::from(fname)); - self.need_saved.store(true, Ordering::Relaxed); - // Rewrite all datas. - self.write()?; + if !self.is_mem_only() { + self.need_saved.store(true, Ordering::Relaxed); + // Rewrite all datas. + self.write()?; + } Ok(()) } } @@ -186,6 +292,9 @@ impl PdFile { impl Drop for PdFile { fn drop(&mut self) { + if self.is_mem_only() { + return; + } if self.is_completed() { self.force_close(); self.remove_pd_file_with_err_msg(); diff --git a/src/downloader/pd_file/version.rs b/src/downloader/pd_file/version.rs index 0b7a82d..11781fa 100644 --- a/src/downloader/pd_file/version.rs +++ b/src/downloader/pd_file/version.rs @@ -4,6 +4,7 @@ use std::cmp::PartialEq; use std::cmp::PartialOrd; use std::convert::AsRef; use std::convert::TryFrom; +use std::io::Read; use std::io::Write; #[derive(Clone, Copy, Debug, PartialEq, Eq)] @@ -37,6 +38,11 @@ impl PdFileVersion { } } + /// Check the version is supported or not. + pub fn is_supported(&self) -> bool { + *self <= [1, 0] + } + /// Get version bytes pub fn to_bytes(&self) -> Vec { let mut data = Vec::new(); @@ -45,6 +51,16 @@ impl PdFileVersion { data } + /// Create a new instance of the [PdFileVersion] from reader + /// * `reader` - The reader which implement the [Read] trait + /// + /// Returns io Error or [PdFileVersion] instance. + pub fn read_from(reader: &mut R) -> std::io::Result { + let mut buf = [0u8; 2]; + reader.read_exact(&mut buf)?; + Ok(Self::from_bytes(&buf, 0).unwrap()) + } + /// Write version bytes to writer. /// * `writer` - The writer which implement the [Write] trait /// diff --git a/src/ext/io.rs b/src/ext/io.rs new file mode 100644 index 0000000..d925a34 --- /dev/null +++ b/src/ext/io.rs @@ -0,0 +1,53 @@ +use proc_macros::define_struct_reader_fn; +use proc_macros::impl_struct_reader_read; +use std::io::Read; + +/// Read number. +pub trait StructRead { + /// The error type + type Error; + define_struct_reader_fn!(u8); + define_struct_reader_fn!(i8); + define_struct_reader_fn!(u16); + define_struct_reader_fn!(i16); + define_struct_reader_fn!(u32); + define_struct_reader_fn!(i32); + define_struct_reader_fn!(u64); + define_struct_reader_fn!(i64); + define_struct_reader_fn!(usize); + define_struct_reader_fn!(isize); + define_struct_reader_fn!(u128); + define_struct_reader_fn!(i128); + /// Read exact number of bytes. + /// * `size` - The number of bytes + /// + /// Returns io error or the bytes. + fn read_bytes(&mut self, size: usize) -> Result, Self::Error>; +} + +impl StructRead for T { + type Error = std::io::Error; + impl_struct_reader_read!(u8); + impl_struct_reader_read!(i8); + impl_struct_reader_read!(u16); + impl_struct_reader_read!(i16); + impl_struct_reader_read!(u32); + impl_struct_reader_read!(i32); + impl_struct_reader_read!(u64); + impl_struct_reader_read!(i64); + impl_struct_reader_read!(usize); + impl_struct_reader_read!(isize); + impl_struct_reader_read!(u128); + impl_struct_reader_read!(i128); + + fn read_bytes(&mut self, size: usize) -> Result, Self::Error> { + let mut h = self.take(size as u64); + let mut r = Vec::new(); + let s = h.read_to_end(&mut r)?; + if s != size { + Err(std::io::Error::from(std::io::ErrorKind::UnexpectedEof)) + } else { + Ok(r) + } + } +} diff --git a/src/ext/mod.rs b/src/ext/mod.rs index a88cd4f..8c87a1b 100644 --- a/src/ext/mod.rs +++ b/src/ext/mod.rs @@ -1,9 +1,11 @@ pub mod cstr; #[cfg(feature = "flagset")] pub mod flagset; +pub mod io; pub mod json; #[cfg(any(feature = "exif", feature = "avdict", feature = "ugoira"))] pub mod rawhandle; +pub mod replace; pub mod rw_lock; pub mod try_err; pub mod use_or_not; diff --git a/src/ext/replace.rs b/src/ext/replace.rs new file mode 100644 index 0000000..ec8a67d --- /dev/null +++ b/src/ext/replace.rs @@ -0,0 +1,82 @@ +use crate::ext::rw_lock::GetRwLock; +use std::ops::DerefMut; +use std::sync::RwLock; +use std::sync::RwLockWriteGuard; + +/// Replace current value with another value +pub trait ReplaceWith { + /// Replace current value with another value + /// * `another` - another value + /// + /// Returns the old value. + fn replace_with(&mut self, another: T) -> T; +} + +/// Replace current value with another value +/// +/// If you want to mutably borrows, please use [ReplaceWith] instead. +pub trait ReplaceWith2 { + /// Replace current value with another value + /// * `another` - another value + /// + /// Returns the old value. + fn replace_with2(&self, another: T) -> T; +} + +impl ReplaceWith for T { + #[inline] + fn replace_with(&mut self, another: T) -> T { + std::mem::replace(self, another) + } +} + +impl<'a, T> ReplaceWith for RwLockWriteGuard<'a, T> { + #[inline] + fn replace_with(&mut self, another: T) -> T { + self.deref_mut().replace_with(another) + } +} + +impl ReplaceWith for RwLock { + #[inline] + fn replace_with(&mut self, another: T) -> T { + self.replace_with2(another) + } +} + +impl ReplaceWith2 for RwLock { + #[inline] + fn replace_with2(&self, another: T) -> T { + self.get_mut().replace_with(another) + } +} + +macro_rules! impl_replace_with_atomic { + ($type1:ty, $type2:ty) => { + impl ReplaceWith<$type2> for $type1 { + #[inline] + fn replace_with(&mut self, another: $type2) -> $type2 { + self.replace_with2(another) + } + } + impl ReplaceWith2<$type2> for $type1 { + fn replace_with2(&self, another: $type2) -> $type2 { + let ori = self.load(std::sync::atomic::Ordering::Relaxed); + self.store(another, std::sync::atomic::Ordering::Relaxed); + ori + } + } + }; +} + +impl_replace_with_atomic!(std::sync::atomic::AtomicBool, bool); +impl_replace_with_atomic!(std::sync::atomic::AtomicI8, i8); +impl_replace_with_atomic!(std::sync::atomic::AtomicU8, u8); +impl_replace_with_atomic!(std::sync::atomic::AtomicI16, i16); +impl_replace_with_atomic!(std::sync::atomic::AtomicU16, u16); +impl_replace_with_atomic!(std::sync::atomic::AtomicI32, i32); +impl_replace_with_atomic!(std::sync::atomic::AtomicU32, u32); +impl_replace_with_atomic!(std::sync::atomic::AtomicI64, i64); +impl_replace_with_atomic!(std::sync::atomic::AtomicU64, u64); +impl_replace_with_atomic!(std::sync::atomic::AtomicIsize, isize); +impl_replace_with_atomic!(std::sync::atomic::AtomicUsize, usize); diff --git a/src/pixiv_web.rs b/src/pixiv_web.rs index 70ee279..8f48061 100644 --- a/src/pixiv_web.rs +++ b/src/pixiv_web.rs @@ -1,3 +1,4 @@ +use crate::downloader::pd_file::file::PdFile; use crate::ext::rw_lock::GetRwLock; use crate::gettext; use crate::opthelper::OptHelper; @@ -204,7 +205,7 @@ impl PixivWebClient { Some(r) } - pub async fn adownload_image(&self, url: U) -> Option { + pub async fn adownload_image(&self, url: U, pdf: &Option) -> Option { self.auto_init(); let r = self.client.aget(url, json::object!{"referer": "https://www.pixiv.net/"}).await; if r.is_none() { diff --git a/src/proc_macros.rs b/src/proc_macros.rs new file mode 100644 index 0000000..7e3931a --- /dev/null +++ b/src/proc_macros.rs @@ -0,0 +1,45 @@ +extern crate quote; +extern crate syn; + +use proc_macro::TokenStream; +use quote::quote; +use syn::Ident; +use syn::parse_macro_input; + +#[proc_macro] +pub fn define_struct_reader_fn(item: TokenStream) -> TokenStream { + let i = parse_macro_input!(item as Ident); + let lefname = format!("read_le_{}", i); + let lefname = Ident::new(&lefname, i.span()); + let befname = format!("read_be_{}", i); + let befname = Ident::new(&befname, i.span()); + let stream = quote! { + #[doc = concat!("Read [", stringify!(#i), "] in little endian.")] + fn #lefname(&mut self) -> Result<#i, Self::Error>; + #[doc = concat!("Read [", stringify!(#i), "] in big endian.")] + fn #befname(&mut self) -> Result<#i, Self::Error>; + }; + stream.into() +} + +#[proc_macro] +pub fn impl_struct_reader_read(item: TokenStream) -> TokenStream { + let i = parse_macro_input!(item as Ident); + let lefname = format!("read_le_{}", i); + let lefname = Ident::new(&lefname, i.span()); + let befname = format!("read_be_{}", i); + let befname = Ident::new(&befname, i.span()); + let stream = quote! { + fn #lefname(&mut self) -> Result<#i, Self::Error> { + let mut buf = [0u8; std::mem::size_of::<#i>()]; + self.read_exact(&mut buf)?; + Ok(<#i>::from_le_bytes(buf)) + } + fn #befname(&mut self) -> Result<#i, Self::Error> { + let mut buf = [0u8; std::mem::size_of::<#i>()]; + self.read_exact(&mut buf)?; + Ok(<#i>::from_be_bytes(buf)) + } + }; + stream.into() +}