From 76f6691f8a8c6ae457e31565c04a15d89f92ef03 Mon Sep 17 00:00:00 2001 From: lifegpc Date: Sun, 15 May 2022 12:26:58 +0800 Subject: [PATCH] fix support for language switch update document --- src/ext/json.rs | 24 +++++++++++++++++ src/ext/mod.rs | 1 + src/ext/rw_lock.rs | 35 ++++++++++++++++++++++++ src/pixiv_web.rs | 17 ++++++++---- src/webclient.rs | 67 +++++++++++++++++++++++++++++++++++++--------- 5 files changed, 127 insertions(+), 17 deletions(-) create mode 100644 src/ext/rw_lock.rs diff --git a/src/ext/json.rs b/src/ext/json.rs index ba3a054..43715ca 100644 --- a/src/ext/json.rs +++ b/src/ext/json.rs @@ -1,4 +1,7 @@ use json::JsonValue; +use std::ops::Deref; +use std::sync::RwLockReadGuard; +use std::sync::RwLockWriteGuard; pub trait ToJson { fn to_json(&self) -> Option; @@ -28,6 +31,27 @@ impl ToJson for &T { } } +impl ToJson for Option { + fn to_json(&self) -> Option { + match self { + Some(d) => { d.to_json() } + None => { None } + } + } +} + +impl ToJson for RwLockReadGuard<'_, T> { + fn to_json(&self) -> Option { + self.deref().to_json() + } +} + +impl ToJson for RwLockWriteGuard<'_, T> { + fn to_json(&self) -> Option { + self.deref().to_json() + } +} + pub trait FromJson where Self: Sized { type Err; fn from_json(v: T) -> Result; diff --git a/src/ext/mod.rs b/src/ext/mod.rs index 3a5957a..a88cd4f 100644 --- a/src/ext/mod.rs +++ b/src/ext/mod.rs @@ -4,5 +4,6 @@ pub mod flagset; pub mod json; #[cfg(any(feature = "exif", feature = "avdict", feature = "ugoira"))] pub mod rawhandle; +pub mod rw_lock; pub mod try_err; pub mod use_or_not; diff --git a/src/ext/rw_lock.rs b/src/ext/rw_lock.rs new file mode 100644 index 0000000..7989963 --- /dev/null +++ b/src/ext/rw_lock.rs @@ -0,0 +1,35 @@ +use spin_on::spin_on; +use std::sync::RwLock; +use std::sync::RwLockReadGuard; +use std::sync::RwLockWriteGuard; +use std::time::Duration; + +pub trait GetRwLock { + type Target; + fn get_ref(&self) -> RwLockReadGuard; + fn get_mut(&self) -> RwLockWriteGuard; +} + +impl GetRwLock for RwLock { + type Target = T; + fn get_ref<'a>(&'a self) -> RwLockReadGuard<'a, Self::Target> { + loop { + match self.try_read() { + Ok(f) => { return f; } + Err(_) => { + spin_on(tokio::time::sleep(Duration::new(0, 1_000_000))); + } + } + } + } + fn get_mut<'a>(&'a self) -> RwLockWriteGuard<'a, Self::Target> { + loop { + match self.try_write() { + Ok(f) => { return f; } + Err(_) => { + spin_on(tokio::time::sleep(Duration::new(0, 1_000_000))); + } + } + } + } +} diff --git a/src/pixiv_web.rs b/src/pixiv_web.rs index caa62cb..70ee279 100644 --- a/src/pixiv_web.rs +++ b/src/pixiv_web.rs @@ -1,3 +1,4 @@ +use crate::ext::rw_lock::GetRwLock; use crate::gettext; use crate::opthelper::OptHelper; use crate::parser::metadata::MetaDataParser; @@ -21,7 +22,10 @@ pub struct PixivWebClient { pub helper: OptHelper, /// true if in is initialized inited: Arc, + /// pixiv global data data: RwLock>, + /// Get basic params + params: RwLock>, } impl PixivWebClient { @@ -31,6 +35,7 @@ impl PixivWebClient { helper: OptHelper::new(m.cmd.as_ref().unwrap().clone(), m.settings.as_ref().unwrap().clone()), inited: Arc::new(AtomicBool::new(false)), data: RwLock::new(None), + params: RwLock::new(None), } } @@ -79,8 +84,10 @@ impl PixivWebClient { let l = self.helper.language(); if l.is_some() { self.client.set_header("Accept-Language", l.as_ref().unwrap()); + self.params.get_mut().replace(json::object! { "lang": l.as_ref().unwrap().replace("-", "_").as_str() }); } else { self.client.set_header("Accept-Language", "ja"); + self.params.get_mut().replace(json::object! { "lang": "ja" }); } self.client.set_verbose(self.helper.verbose()); let retry = self.helper.retry(); @@ -103,7 +110,7 @@ impl PixivWebClient { pub fn check_login(&self) -> bool { self.auto_init(); - let r = self.client.get("https://www.pixiv.net/", None); + let r = self.client.get_with_param("https://www.pixiv.net/", self.params.get_ref(), None); if r.is_none() { return false; } @@ -215,7 +222,7 @@ impl PixivWebClient { pub fn get_artwork_ajax(&self, id: u64) -> Option { self.auto_init(); - let r = self.client.get(format!("https://www.pixiv.net/ajax/illust/{}", id), None); + let r = self.client.get_with_param(format!("https://www.pixiv.net/ajax/illust/{}", id), self.params.get_ref(), None); if r.is_none() { return None; } @@ -229,7 +236,7 @@ impl PixivWebClient { pub fn get_artwork(&self, id: u64) -> Option { self.auto_init(); - let r = self.client.get(format!("https://www.pixiv.net/artworks/{}", id), None); + let r = self.client.get_with_param(format!("https://www.pixiv.net/artworks/{}", id), self.params.get_ref(), None); if r.is_none() { return None; } @@ -259,7 +266,7 @@ impl PixivWebClient { pub fn get_illust_pages(&self, id: u64) -> Option { self.auto_init(); - let r = self.client.get(format!("https://www.pixiv.net/ajax/illust/{}/pages", id), None); + let r = self.client.get_with_param(format!("https://www.pixiv.net/ajax/illust/{}/pages", id), self.params.get_ref(), None); if r.is_none() { return None; } @@ -273,7 +280,7 @@ impl PixivWebClient { pub fn get_ugoira(&self, id: u64) -> Option { self.auto_init(); - let r = self.client.get(format!("https://www.pixiv.net/ajax/illust/{}/ugoira_meta", id), None); + let r = self.client.get_with_param(format!("https://www.pixiv.net/ajax/illust/{}/ugoira_meta", id), self.params.get_ref(), None); if r.is_none() { return None; } diff --git a/src/webclient.rs b/src/webclient.rs index 7636635..771588f 100644 --- a/src/webclient.rs +++ b/src/webclient.rs @@ -2,6 +2,7 @@ extern crate spin_on; use crate::cookies::Cookie; use crate::cookies::CookieJar; +use crate::ext::json::ToJson; use crate::gettext; use crate::list::NonTailList; use crate::opthelper::OptHelper; @@ -28,7 +29,9 @@ use std::sync::atomic::AtomicU64; use std::sync::atomic::Ordering; use std::time::Duration; +/// Convert data to HTTP headers map pub trait ToHeaders { + /// return HTTP headers map fn to_headers(&self) -> Option>; } @@ -82,10 +85,13 @@ pub fn gen_cookie_header(c: &WebClient, url: U) -> String { /// A Web Client pub struct WebClient { + /// Basic Web Client client: Client, /// HTTP Headers headers: RwLock>, + /// Cookies cookies: RwLock, + /// Verbose logging verbose: Arc, /// Retry times, 0 means disable retry: Arc, @@ -94,6 +100,7 @@ pub struct WebClient { } impl WebClient { + /// Create a new instance of client pub fn new() -> Self { Self { client: Client::new(), @@ -204,6 +211,8 @@ impl WebClient { self.verbose.load(Ordering::Relaxed) } + /// Used to handle Set-Cookie header in an [Response] + /// * `r` - reference to an [Response] pub fn handle_set_cookie(&self, r: &Response) { let u = r.url(); let h = r.headers(); @@ -229,6 +238,12 @@ impl WebClient { } } + /// Read cookies from file. + /// * `file_name`: File name + /// + /// returns true if readed successfully. + /// # Note + /// If read failed, will clean all entries in the current [CookieJar] pub fn read_cookies(&self, file_name: &str) -> bool { let mut c = self.get_cookies_as_mut(); let r = c.read(file_name); @@ -238,10 +253,19 @@ impl WebClient { 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 + /// + /// Returns the old HTTP header value if presented. pub fn set_header(&self, key: &str, value: &str) -> Option { self.get_headers_as_mut().insert(String::from(key), String::from(value)) } @@ -254,36 +278,42 @@ impl WebClient { pub fn set_verbose(&self, verbose: bool) { self.verbose.store(verbose, Ordering::Relaxed) } + /// Send GET requests with parameters /// * `param` - GET parameters. Should be a JSON object/array. If value in map is not a string, will dump it /// # Examples /// ``` /// let client = WebClient::new(); - /// client.verbose = true; - /// client.get_with_param("https://test.com/a", json::object!{"data": "param1"}); - /// client.get_with_param("https://test.com/a", json::object!{"daa": {"ad": "test"}}); - /// client.get_with_param("https://test.com/a", json::array![["daa", "param1"]]); + /// client.set_verbose(true); + /// client.get_with_param("https://test.com/a", json::object!{"data": "param1"}, None); + /// client.get_with_param("https://test.com/a", json::object!{"daa": {"ad": "test"}}, None); + /// client.get_with_param("https://test.com/a", json::array![["daa", "param1"]], None); /// ``` /// It will GET `https://test.com/a?data=param1`, `https://test.com/a?daa=%7B%22ad%22%3A%22test%22%7D`, `https://test.com/a?daa=param1` - pub fn get_with_param(&self, url: U, param: JsonValue) -> Option { + pub fn get_with_param(&self, url: U, param: J, headers: H) -> Option { let u = url.into_url(); if u.is_err() { println!("{} \"{}\"", gettext("Can not parse URL:"), u.unwrap_err()); return None; } let mut u = u.unwrap(); - if !param.is_object() && !param.is_array() { + let obj = param.to_json(); + if obj.is_none() { + return self.get(u, headers); + } + let obj = obj.unwrap(); + if !obj.is_object() && !obj.is_array() { println!( "{} \"{}\"", gettext("Parameters should be object or array:"), - param + obj ); return None; } { let mut query = u.query_pairs_mut(); - if param.is_object() { - for (k, v) in param.entries() { + if obj.is_object() { + for (k, v) in obj.entries() { let s: String; if v.is_string() { s = String::from(v.as_str().unwrap()); @@ -293,7 +323,7 @@ impl WebClient { query.append_pair(k, s.as_str()); } } else { - for v in param.members() { + for v in obj.members() { if !v.is_object() { println!("{} \"{}\"", gettext("Parameters should be array:"), v); return None; @@ -323,9 +353,10 @@ impl WebClient { } } } - self.get(u.as_str(), None) + self.get(u.as_str(), headers) } + /// Send Get Requests pub fn get(&self, url: U, headers: H) -> Option { let mut count = 0u64; let retry = self.get_retry(); @@ -350,6 +381,7 @@ impl WebClient { None } + /// Send Get Requests pub async fn aget(&self, url: U, headers: H) -> Option { let mut count = 0u64; let retry = self.get_retry(); @@ -371,7 +403,7 @@ impl WebClient { None } - /// Send GET requests + /// Send GET requests without retry pub fn _get(&self, url: U, headers: H) -> Option { let r = self._aget(url, headers); let r = r.send(); @@ -391,6 +423,7 @@ impl WebClient { Some(r) } + /// Send GET requests without retry pub async fn _aget2(&self, url: U, headers: H) -> Option { let r = self._aget(url, headers); let r = r.send().await; @@ -409,6 +442,7 @@ impl WebClient { Some(r) } + /// Generate a requests pub fn _aget(&self, url: U, headers: H) -> RequestBuilder { let s = url.as_str(); if self.get_verbose() { @@ -432,6 +466,13 @@ impl WebClient { r } + /// Download a stream + /// * `file_name` - File name + /// * `r` - Response + /// * `opt` - Options + /// * `progess_bars` - Multiple progress bars + /// + /// Note: If file already exists, will remove existing file first. pub async fn adownload_stream + ?Sized>(file_name: &S, r: Response, opt: &OptHelper, progress_bars: Option>) -> Result<(), ()> { let content_length = r.content_length(); let use_progress_bar = match &content_length { @@ -513,6 +554,8 @@ impl WebClient { /// Download a stream /// * `file_name` - File name /// * `r` - Response + /// * `opt` - Options + /// /// Note: If file already exists, will remove existing file first. pub fn download_stream + ?Sized>(file_name: &S, r: Response, opt: &OptHelper) -> Result<(), ()> { let content_length = r.content_length();