This commit is contained in:
2022-05-17 15:02:06 +08:00
parent d926defeec
commit 650e533519
12 changed files with 437 additions and 25 deletions

2
Cargo.lock generated
View File

@@ -1142,9 +1142,11 @@ dependencies = [
"json",
"lazy_static",
"link-cplusplus",
"quote",
"regex",
"reqwest",
"spin_on",
"syn",
"tokio",
"urlparse",
"utf16string",

View File

@@ -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 }

View File

@@ -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;

View File

@@ -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.

View File

@@ -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<T: IntEnum> From<IntEnumError<T>> for PdFileError {
fn from(e: IntEnumError<T>) -> Self {
Self::String(format!("{} {}", gettext("Invalid pd file: "), e))
}
}

View File

@@ -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<u8> = 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<P: AsRef<Path> + ?Sized>(path: &P) -> Result<PdFileResult, PdFileError> {
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<P: AsRef<Path> + ?Sized>(path: &P) -> Result<Self, PdFileError> {
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<S: AsRef<str> + ?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();

View File

@@ -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<u8> {
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<R: Read>(reader: &mut R) -> std::io::Result<Self> {
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
///

53
src/ext/io.rs Normal file
View File

@@ -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<Vec<u8>, Self::Error>;
}
impl<T: Read> 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<Vec<u8>, 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)
}
}
}

View File

@@ -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;

82
src/ext/replace.rs Normal file
View File

@@ -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<T> {
/// 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<T> {
/// Replace current value with another value
/// * `another` - another value
///
/// Returns the old value.
fn replace_with2(&self, another: T) -> T;
}
impl<T> ReplaceWith<T> for T {
#[inline]
fn replace_with(&mut self, another: T) -> T {
std::mem::replace(self, another)
}
}
impl<'a, T> ReplaceWith<T> for RwLockWriteGuard<'a, T> {
#[inline]
fn replace_with(&mut self, another: T) -> T {
self.deref_mut().replace_with(another)
}
}
impl<T> ReplaceWith<T> for RwLock<T> {
#[inline]
fn replace_with(&mut self, another: T) -> T {
self.replace_with2(another)
}
}
impl<T> ReplaceWith2<T> for RwLock<T> {
#[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);

View File

@@ -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<U: IntoUrl + Clone>(&self, url: U) -> Option<Response> {
pub async fn adownload_image<U: IntoUrl + Clone>(&self, url: U, pdf: &Option<PdFile>) -> Option<Response> {
self.auto_init();
let r = self.client.aget(url, json::object!{"referer": "https://www.pixiv.net/"}).await;
if r.is_none() {

45
src/proc_macros.rs Normal file
View File

@@ -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()
}