use std::io::BufWriter;
use std::ops::Deref;
use std::path::PathBuf;
use std::sync::{Arc, RwLock};
use std::time::Duration;
use bon::bon;
use bytes::Bytes;
use cookie_store::{CookieStore, RawCookie, RawCookieParseError};
use reqwest::header::{ACCEPT, ACCEPT_LANGUAGE, CONNECTION, HeaderMap, HeaderValue};
use reqwest::{Certificate, Client, Proxy, StatusCode, redirect};
use tokio::fs;
use url::Url;
use crate::Error;
pub(crate) fn check_status<T>(code: StatusCode, msg: T) -> Result<(), Error>
where
T: AsRef<str>,
{
if code != StatusCode::OK {
return Err(Error::Http {
code,
msg: msg.as_ref().trim().to_string(),
});
}
Ok(())
}
#[must_use]
pub(crate) struct HTTPClient {
app_name: &'static str,
cookie_provider: Option<Arc<Jar>>,
client: Client,
}
#[bon]
impl HTTPClient {
const COOKIE_FILE_NAME: &'static str = "cookie.json";
const COOKIE_FILE_PASSWORD: &'static str = "gafqad-4Ratne-dirqom";
const COOKIE_FILE_NONCE: &'static str = "novel-rs-cookie";
#[builder]
pub(crate) async fn new(
app_name: &'static str,
cookie: Option<bool>,
user_agent: Option<String>,
accept: Option<HeaderValue>,
accept_language: Option<HeaderValue>,
allow_compress: Option<bool>,
proxy: Option<Url>,
no_proxy: Option<bool>,
cert_path: Option<PathBuf>,
headers: Option<HeaderMap>,
retry_url: Option<Url>,
) -> Result<Self, Error> {
let mut cookie_provider = None;
if cookie.is_some_and(|x| x) {
cookie_provider = Some(Arc::new(
HTTPClient::create_cookie_provider(app_name).await?,
));
}
let mut headers = headers.unwrap_or_default();
if let Some(accept) = accept {
headers.insert(ACCEPT, accept);
}
if let Some(accept_language) = accept_language {
headers.insert(ACCEPT_LANGUAGE, accept_language);
}
headers.insert(CONNECTION, HeaderValue::from_static("keep-alive"));
let mut client_builder = Client::builder()
.default_headers(headers)
.redirect(redirect::Policy::default())
.http2_keep_alive_interval(Duration::from_secs(5));
if !is_ci::cached() {
client_builder = client_builder
.connect_timeout(Duration::from_secs(30))
.timeout(Duration::from_secs(60));
}
if let Some(user_agent) = user_agent {
client_builder = client_builder.user_agent(user_agent);
} else {
client_builder = client_builder.user_agent("Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/26.2 Safari/605.1.15");
}
if let Some(jar) = &cookie_provider {
client_builder = client_builder.cookie_provider(Arc::clone(jar));
}
if allow_compress.is_some_and(|x| !x) {
client_builder = client_builder.no_gzip();
client_builder = client_builder.no_brotli();
client_builder = client_builder.no_deflate();
client_builder = client_builder.no_zstd();
}
if let Some(proxy) = proxy {
client_builder = client_builder.proxy(Proxy::all(proxy)?);
}
if no_proxy.is_some_and(|x| x) {
client_builder = client_builder.no_proxy();
}
if let Some(cert_path) = cert_path {
let cert = Certificate::from_pem(&fs::read(cert_path).await?)?;
client_builder = client_builder.tls_certs_merge([cert]);
}
if let Some(retry_url) = retry_url
&& let Some(retry_host) = retry_url.host_str()
{
let retries = reqwest::retry::for_host(retry_host.to_string()).classify_fn(|req_rep| {
if matches!(
req_rep.status(),
Some(StatusCode::REQUEST_TIMEOUT)
| Some(StatusCode::TOO_MANY_REQUESTS)
| Some(StatusCode::SERVICE_UNAVAILABLE)
| Some(StatusCode::GATEWAY_TIMEOUT)
) {
req_rep.retryable()
} else {
req_rep.success()
}
});
client_builder = client_builder.retry(retries);
}
Ok(Self {
app_name,
cookie_provider,
client: client_builder.build()?,
})
}
pub(crate) fn add_cookie(&self, cookie_str: &str, url: &Url) -> Result<(), Error> {
if let Some(cookie_provider) = &self.cookie_provider {
cookie_provider.0.write().unwrap().parse(cookie_str, url)?;
} else {
return Err(Error::NovelApi("Cookie provider is not set".to_string()));
}
Ok(())
}
pub(crate) fn save_cookies(&self) -> Result<(), Error> {
if let Some(cookie_provider) = &self.cookie_provider {
let mut writer = BufWriter::new(Vec::new());
let cookies = cookie_provider.0.read().unwrap();
cookie_store::serde::json::save(&cookies, &mut writer)?;
let result = simdutf8::basic::from_utf8(writer.buffer())?.to_string();
if !result.is_empty() {
let cookie_path = HTTPClient::cookie_path(self.app_name)?;
tracing::info!("Save the cookie file at: `{}`", cookie_path.display());
super::encrypt_and_save_to_file(
result,
cookie_path,
HTTPClient::COOKIE_FILE_PASSWORD,
HTTPClient::COOKIE_FILE_NONCE,
)?;
}
}
Ok(())
}
async fn create_cookie_provider(app_name: &str) -> Result<Jar, Error> {
let cookie_path = HTTPClient::cookie_path(app_name)?;
let cookie_store = if fs::try_exists(&cookie_path).await? {
tracing::info!("The cookie file is located at: `{}`", cookie_path.display());
if let Ok(json) = super::decrypt_from_file(
&cookie_path,
HTTPClient::COOKIE_FILE_PASSWORD,
HTTPClient::COOKIE_FILE_NONCE,
) {
if let Ok(cookie_store) = cookie_store::serde::json::load(json.as_bytes()) {
cookie_store
} else {
tracing::error!("Fail to load the cookie file, a new one will be created");
CookieStore::default()
}
} else {
tracing::error!("Fail to decrypt the cookie file, a new one will be created");
CookieStore::default()
}
} else {
tracing::info!(
"The cookie file will be created at: `{}`",
cookie_path.display()
);
fs::create_dir_all(cookie_path.parent().unwrap()).await?;
CookieStore::default()
};
Ok(Jar::new(cookie_store))
}
fn cookie_path(app_name: &str) -> Result<PathBuf, Error> {
let mut config_path = crate::config_dir_path(app_name)?;
config_path.push(HTTPClient::COOKIE_FILE_NAME);
Ok(config_path)
}
}
impl Deref for HTTPClient {
type Target = Client;
fn deref(&self) -> &Self::Target {
&self.client
}
}
impl Drop for HTTPClient {
fn drop(&mut self) {
if let Err(err) = self.save_cookies() {
tracing::error!("Fail to save cookie: {err}");
}
}
}
struct Jar(RwLock<CookieStore>);
impl Jar {
fn new(cookie_store: CookieStore) -> Jar {
Jar(RwLock::new(cookie_store))
}
}
impl reqwest::cookie::CookieStore for Jar {
fn set_cookies(&self, cookie_headers: &mut dyn Iterator<Item = &HeaderValue>, url: &url::Url) {
let mut write = self.0.write().unwrap();
set_cookies(&mut write, cookie_headers, url);
}
fn cookies(&self, url: &url::Url) -> Option<HeaderValue> {
let read = self.0.read().unwrap();
cookies(&read, url)
}
}
fn set_cookies(
cookie_store: &mut CookieStore,
cookie_headers: &mut dyn Iterator<Item = &HeaderValue>,
url: &url::Url,
) {
let cookies = cookie_headers.filter_map(|val| {
std::str::from_utf8(val.as_bytes())
.map_err(RawCookieParseError::from)
.and_then(RawCookie::parse)
.map(|c| c.into_owned())
.ok()
});
cookie_store.store_response_cookies(cookies, url);
}
fn cookies(cookie_store: &CookieStore, url: &url::Url) -> Option<HeaderValue> {
let s = cookie_store
.get_request_values(url)
.map(|(name, value)| format!("{name}={value}"))
.collect::<Vec<_>>()
.join("; ");
if s.is_empty() {
return None;
}
HeaderValue::from_maybe_shared(Bytes::from(s)).ok()
}