From b75a42197698e1d5bb1c2d30b069b4bd35c2310f Mon Sep 17 00:00:00 2001 From: lifegpc Date: Sun, 23 Oct 2022 03:38:55 +0000 Subject: [PATCH] Fix #280 Update README --- README.md | 7 +- src/cookies.rs | 193 +++++++++++++++++++++++++++++++++++++++++++++-- src/pixiv_web.rs | 17 ----- src/webclient.rs | 26 +++---- 4 files changed, 196 insertions(+), 47 deletions(-) diff --git a/README.md b/README.md index 76525e0..a3d207e 100644 --- a/README.md +++ b/README.md @@ -4,12 +4,7 @@ A pixiv downloader written in Rust. * Write exif metatata to picture. * Merge ugoira(GIF) pictures to video files. ### TODO -- [ ] Add support for multiple threads download. -- [ ] Batch download. -- [ ] Fanbox support. -### TODO - CVE -- [ ] Remove parse_duration dependency which used in unittest. -## Rust features flags +See [issues](https://github.com/lifegpc/pixiv_downloader/issues) or [projects](https://github.com/lifegpc/pixiv_downloader/projects) ### all Enable all unconflicted features, this will enable [`exif`](#exif) and [`ugoira`](#ugoira). ### exif diff --git a/src/cookies.rs b/src/cookies.rs index de7c038..2efef3c 100644 --- a/src/cookies.rs +++ b/src/cookies.rs @@ -1,14 +1,20 @@ +use crate::ext::replace::ReplaceWith; +use crate::ext::rw_lock::GetRwLock; use crate::gettext; use chrono::DateTime; use chrono::TimeZone; use chrono::Utc; use reqwest::IntoUrl; +use std::collections::HashMap; +#[cfg(test)] +use std::fs::create_dir; use std::fs::{remove_file, File}; use std::io::BufRead; use std::io::BufReader; use std::io::Write; use std::iter::Iterator; -use std::path::Path; +use std::path::{Path, PathBuf}; +use std::sync::{Arc, RwLock}; trait ToNetscapeStr { fn to_netscape_str(&self) -> &'static str; @@ -311,16 +317,27 @@ impl CookieJar { self.cookies.clear(); } - pub fn read(&mut self, file_name: &str) -> bool { + #[allow(dead_code)] + pub fn get + ?Sized>(&self, name: &S) -> Option<&Cookie> { + let name = name.as_ref(); + for i in self.cookies.iter() { + if i._name == name { + return Some(i); + } + } + None + } + + pub fn read + ?Sized>(&mut self, file_name: &P) -> bool { self.cookies.clear(); - let p = Path::new(file_name); + let p = file_name.as_ref(); if !p.exists() { - println!("{} {}", gettext("Can not find file:"), file_name); + println!("{} {}", gettext("Can not find file:"), p.display()); return false; } let re = File::open(p); if re.is_err() { - println!("{} {}", gettext("Can not open file:"), file_name); + println!("{} {}", gettext("Can not open file:"), p.display()); return false; } let f = re.unwrap(); @@ -361,8 +378,8 @@ impl CookieJar { true } - pub fn save(&mut self, file_name: &str) -> bool { - let p = Path::new(file_name); + pub fn save + ?Sized>(&mut self, file_name: &P) -> bool { + let p = file_name.as_ref(); self.check_expired(); if p.exists() { let re = remove_file(p); @@ -391,3 +408,165 @@ impl CookieJar { self.cookies.iter() } } + +struct CookieJarManager { + jars: RwLock>, usize)>>, +} + +impl CookieJarManager { + pub fn new() -> Self { + Self { + jars: RwLock::new(HashMap::new()), + } + } + + pub fn get_cookie_jar + ?Sized>( + &self, + path: &P, + ) -> Result>, ()> { + let path = path.as_ref().to_owned(); + let mut jars = self.jars.get_mut(); + match jars.get_mut(&path) { + Some((jar, count)) => { + *count += 1; + Ok(jar.clone()) + } + None => { + let mut jar = CookieJar::new(); + if jar.read(&path) { + let jar = Arc::new(RwLock::new(jar)); + jars.insert(path, (jar.clone(), 1)); + Ok(jar) + } else { + Err(()) + } + } + } + } + + pub fn drop_cookie_jar + ?Sized>(&self, path: &P) { + let path = path.as_ref().to_owned(); + let mut jars = self.jars.get_mut(); + match jars.get_mut(&path) { + Some((jar, count)) => { + *count -= 1; + if *count == 0 { + if !jar.get_mut().save(&path) { + println!( + "{} {}", + gettext("Warning: Failed to save cookies file:"), + path.display() + ); + } + jars.remove(&path); + } + } + None => {} + } + } +} + +lazy_static! { + static ref MANAGER: CookieJarManager = CookieJarManager::new(); +} + +#[derive(Clone, Debug)] +/// A cookie jar that make sure there are one cookie interface for a cookies file +pub struct ManagedCookieJar { + pub jar: Arc>, + path: Option, +} + +impl ManagedCookieJar { + pub fn new() -> Self { + Self { + jar: Arc::new(RwLock::new(CookieJar::new())), + path: None, + } + } + + pub fn read + ?Sized>(&mut self, path: &P) -> bool { + let path = path.as_ref(); + let path = match path.canonicalize() { + Ok(p) => p, + Err(_) => path.to_owned(), + }; + match self.path.as_ref() { + Some(p) => { + if p == &path { + return true; + } + MANAGER.drop_cookie_jar(p); + let jar = match MANAGER.get_cookie_jar(&path) { + Ok(jar) => jar, + Err(()) => return false, + }; + self.jar.replace_with(jar); + self.path.replace(path); + true + } + None => { + let jar = match MANAGER.get_cookie_jar(&path) { + Ok(jar) => jar, + Err(()) => return false, + }; + self.jar.replace_with(jar); + self.path.replace(path); + true + } + } + } +} + +impl Drop for ManagedCookieJar { + fn drop(&mut self) { + if let Some(path) = self.path.as_ref() { + MANAGER.drop_cookie_jar(path); + } + } +} + +#[test] +fn test_managed_cookie_jar() { + let p = Path::new("./test"); + if !p.exists() { + let re = create_dir("./test"); + assert!(re.is_ok() || p.exists()); + } + { + let mut jar = CookieJar::new(); + jar.add(Cookie::new( + "test", + "de", + "example.com", + true, + "/", + false, + None, + )); + jar.save("./test/cookies.txt"); + } + { + let mut jar = ManagedCookieJar::new(); + assert!(jar.read("./test/cookies.txt")); + assert_eq!(jar.jar.get_mut().get("test").unwrap()._value, "de"); + let mut jar2 = ManagedCookieJar::new(); + assert!(jar2.read("./test/cookies.txt")); + jar2.jar.get_mut().add(Cookie::new( + "test2", + "de", + "example.com", + true, + "/", + false, + None, + )); + assert_eq!(jar.jar.get_mut().get("test2").unwrap()._value, "de"); + } + { + let mut jar = CookieJar::new(); + assert!(jar.read("./test/cookies.txt")); + assert_eq!(jar.get("test").unwrap()._value, "de"); + assert_eq!(jar.get("test2").unwrap()._value, "de"); + } +} diff --git a/src/pixiv_web.rs b/src/pixiv_web.rs index 33f8c90..196f6d1 100644 --- a/src/pixiv_web.rs +++ b/src/pixiv_web.rs @@ -297,20 +297,3 @@ impl PixivWebClient { false } } - -impl Drop for PixivWebClient { - fn drop(&mut self) { - if self.is_inited() { - let c = get_helper().cookies(); - if c.is_some() { - if !self.client.save_cookies(c.as_ref().unwrap()) { - println!( - "{} {}", - gettext("Warning: Failed to save cookies file:"), - c.as_ref().unwrap() - ); - } - } - } - } -} diff --git a/src/webclient.rs b/src/webclient.rs index 636e79d..3ce7564 100644 --- a/src/webclient.rs +++ b/src/webclient.rs @@ -1,5 +1,5 @@ use crate::cookies::Cookie; -use crate::cookies::CookieJar; +use crate::cookies::ManagedCookieJar; use crate::ext::atomic::AtomicQuick; use crate::ext::json::ToJson; use crate::ext::rw_lock::GetRwLock; @@ -59,10 +59,10 @@ impl ToHeaders for JsonValue { /// * `c` - Cookies /// * `url` - URL pub fn gen_cookie_header(c: &WebClient, url: U) -> String { - c.get_cookies_as_mut().check_expired(); + c.get_cookies_as_mut().jar.get_mut().check_expired(); let mut s = String::from(""); let u = url.as_str(); - for a in c.get_cookies().iter() { + for a in c.get_cookies().jar.get_ref().iter() { if a.matched(u) { if s.len() > 0 { s += " "; @@ -81,7 +81,7 @@ pub struct WebClient { /// HTTP Headers headers: RwLock>, /// Cookies - cookies: RwLock, + cookies: RwLock, /// Verbose logging verbose: Arc, /// Retry times, 0 means disable, < 0 means always retry @@ -98,18 +98,18 @@ impl WebClient { Self { client, headers: RwLock::new(HashMap::new()), - cookies: RwLock::new(CookieJar::new()), + cookies: RwLock::new(ManagedCookieJar::new()), verbose: Arc::new(AtomicBool::new(false)), retry: Arc::new(AtomicI64::new(3)), retry_interval: RwLock::new(None), } } - pub fn get_cookies_as_mut<'a>(&'a self) -> RwLockWriteGuard<'a, CookieJar> { + pub fn get_cookies_as_mut<'a>(&'a self) -> RwLockWriteGuard<'a, ManagedCookieJar> { self.cookies.get_mut() } - pub fn get_cookies<'a>(&'a self) -> RwLockReadGuard<'a, CookieJar> { + pub fn get_cookies<'a>(&'a self) -> RwLockReadGuard<'a, ManagedCookieJar> { self.cookies.get_ref() } @@ -153,7 +153,7 @@ impl WebClient { let c = Cookie::from_set_cookie(u.as_str(), val); match c { Some(c) => { - self.get_cookies_as_mut().add(c); + self.get_cookies_as_mut().jar.get_mut().add(c); } None => { println!("{}", gettext("Failed to parse Set-Cookie header.")); @@ -177,19 +177,11 @@ impl WebClient { let mut c = self.get_cookies_as_mut(); let r = c.read(file_name); if !r { - c.clear(); + c.jar.get_mut().clear(); } r } - /// Save cookies to file - /// * `file_name`: File name - /// - /// returns true if saved successfully. - pub fn save_cookies(&self, file_name: &str) -> bool { - self.get_cookies_as_mut().save(file_name) - } - /// Set new HTTP header /// * `key` - The key of the new HTTP header /// * `value` - The value of the new HTTP value