diff --git a/proc_macros/proc_macros.rs b/proc_macros/proc_macros.rs index 17c7bc1..c928243 100644 --- a/proc_macros/proc_macros.rs +++ b/proc_macros/proc_macros.rs @@ -215,6 +215,7 @@ struct FilterHttpMethods { pub req: Ident, pub typ: Expr, pub handle_options: LitBool, + pub ctx: Option, pub methods: Vec, } @@ -231,6 +232,12 @@ impl Parse for FilterHttpMethods { return Err(syn::Error::new(input.span(), "Failed to parse boolean.")); } }; + let ctx = if handle_options.value() { + token::Comma::parse(input)?; + Some(Expr::parse(input)?) + } else { + None + }; loop { if input.cursor().eof() { break; @@ -243,6 +250,7 @@ impl Parse for FilterHttpMethods { req, typ, handle_options, + ctx, methods, }) } @@ -257,6 +265,7 @@ pub fn filter_http_methods(item: TokenStream) -> TokenStream { req, typ, handle_options, + ctx, methods, } = parse_macro_input!(item as FilterHttpMethods); let mut header_value = Vec::new(); @@ -274,9 +283,68 @@ pub fn filter_http_methods(item: TokenStream) -> TokenStream { let allow_header = LitStr::new(allow_header.as_str(), req.span()); if enable_options { streams.push(quote!(&hyper::Method::OPTIONS => { - return Ok(hyper::Response::builder().status(200).header("Allow", #allow_header).body(#typ).unwrap()); + let builder = hyper::Response::builder(); + let headers = #req.headers(); + let origin = match headers.get(hyper::header::ORIGIN) { + Some(origin) => match origin.to_str() { + Ok(origin) => Some(origin.to_owned()), + Err(_) => None, + }, + None => None, + }; + match origin { + Some(origin) => { + match #ctx.cors.matches(origin.as_str()) { + crate::server::cors::CorsResult::Allowed => { + let builder = builder.header("Access-Control-Allow-Origin", origin.as_str()); + return Ok(builder.status(200).header("Allow", #allow_header).body(#typ).unwrap()); + } + crate::server::cors::CorsResult::AllowedAll => { + let builder = builder.header("Access-Control-Allow-Origin", "*"); + return Ok(builder.status(200).header("Allow", #allow_header).body(#typ).unwrap()); + } + _ => { + return Ok(builder.status(400).header("Allow", #allow_header).body(#typ).unwrap()); + } + } + } + None => { + return Ok(builder.status(200).header("Allow", #allow_header).body(#typ).unwrap()); + } + } })); } + let post_stream = if enable_options { + quote!( + let mut builder = hyper::Response::builder(); + let headers = #req.headers(); + let origin = match headers.get(hyper::header::ORIGIN) { + Some(origin) => match origin.to_str() { + Ok(origin) => Some(origin.to_owned()), + Err(_) => None, + }, + None => None, + }; + match origin { + Some(origin) => { + match #ctx.cors.matches(origin.as_str()) { + crate::server::cors::CorsResult::Allowed => { + builder.headers_mut().unwrap().insert(hyper::header::ACCESS_CONTROL_ALLOW_ORIGIN, origin.parse().unwrap()); + } + crate::server::cors::CorsResult::AllowedAll => { + builder.headers_mut().unwrap().insert(hyper::header::ACCESS_CONTROL_ALLOW_ORIGIN, "*".parse().unwrap()); + } + _ => { + return Ok(builder.status(403).body(#typ).unwrap()); + } + } + } + None => {} + } + ) + } else { + quote!() + }; let stream = quote! { match #req.method() { #(#streams)* @@ -284,6 +352,7 @@ pub fn filter_http_methods(item: TokenStream) -> TokenStream { return Ok(hyper::Response::builder().status(405).header("Allow", #allow_header).body(#typ).unwrap()) } } + #post_stream }; stream.into() } diff --git a/src/opthelper.rs b/src/opthelper.rs index 13fbb39..f18793d 100644 --- a/src/opthelper.rs +++ b/src/opthelper.rs @@ -10,6 +10,10 @@ use crate::opt::size::parse_u32_size; use crate::opt::use_progress_bar::UseProgressBar; use crate::opts::CommandOpts; use crate::retry_interval::parse_retry_interval_from_json; +#[cfg(feature = "server")] +use crate::server::cors::parse_cors_entries; +#[cfg(feature = "server")] +use crate::server::cors::CorsEntry; use crate::settings::SettingStore; use std::convert::TryFrom; #[cfg(feature = "server")] @@ -38,6 +42,8 @@ pub struct OptHelper { _use_progress_bar: RwLock>, /// Proxy settings _proxy_chain: RwLock, + #[cfg(feature = "server")] + _cors_entries: RwLock>, } impl OptHelper { @@ -59,6 +65,12 @@ impl OptHelper { } } + #[cfg(feature = "server")] + #[inline] + pub fn cors_entries(&self) -> Vec { + self._cors_entries.get_ref().clone() + } + /// Return max retry count of each part when downloading in multiple thread mode. pub fn download_part_retry(&self) -> Option { if self.opt.get_ref().download_part_retry.is_some() { @@ -188,6 +200,11 @@ impl OptHelper { self._proxy_chain .replace_with2(ProxyChain::try_from(settings.get("proxy").unwrap()).unwrap()); } + #[cfg(feature = "server")] + if settings.have("cors-entries") { + self._cors_entries + .replace_with2(parse_cors_entries(&settings.get("cors-entries").unwrap()).unwrap()); + } self.opt.replace_with2(opt); self.settings.replace_with2(settings); } @@ -323,6 +340,8 @@ impl Default for OptHelper { _author_name_filters: RwLock::new(Vec::new()), _use_progress_bar: RwLock::new(None), _proxy_chain: RwLock::new(ProxyChain::default()), + #[cfg(feature = "server")] + _cors_entries: RwLock::new(Vec::new()), } } } diff --git a/src/server/context.rs b/src/server/context.rs new file mode 100644 index 0000000..cd3c390 --- /dev/null +++ b/src/server/context.rs @@ -0,0 +1,14 @@ +use super::cors::CorsContext; +use std::default::Default; + +pub struct ServerContext { + pub cors: CorsContext, +} + +impl Default for ServerContext { + fn default() -> Self { + Self { + cors: CorsContext::default(), + } + } +} diff --git a/src/server/cors.rs b/src/server/cors.rs index f5a5572..80939c4 100644 --- a/src/server/cors.rs +++ b/src/server/cors.rs @@ -1,10 +1,13 @@ use crate::gettext; +use crate::opthelper::get_helper; use http::uri::InvalidUri; use http::uri::Scheme; use http::Uri; +use json::JsonValue; use std::cmp::PartialEq; use std::convert::From; use std::convert::TryFrom; +use std::default::Default; use std::fmt::Display; use std::net::SocketAddr; use std::str::FromStr; @@ -165,6 +168,28 @@ impl TryFrom for CorsEntry { } } +impl TryFrom<&JsonValue> for CorsEntry { + type Error = CorsError; + fn try_from(value: &JsonValue) -> Result { + match value.as_str() { + Some(s) => Self::try_from(s), + None => Err(CorsError::from(gettext("Object is not a string."))), + } + } +} + +pub fn parse_cors_entries(v: &JsonValue) -> Result, CorsError> { + if v.is_array() { + let mut list = Vec::new(); + for i in v.members() { + list.push(CorsEntry::try_from(i)?); + } + Ok(list) + } else { + Err(CorsError::from(gettext("Object is not an array."))) + } +} + #[derive(Clone, Debug, PartialEq)] /// Current host pub struct CorsHost { @@ -322,6 +347,75 @@ impl From<&str> for CorsError { } } +#[derive(Clone, Debug, derive_more::Display, PartialEq, Eq, PartialOrd, Ord)] +pub enum CorsResult { + Allowed, + AllowedAll, + NotAllowed, +} + +pub struct CorsContext { + allow_all: bool, + entries: Vec, + hosts: Vec, +} + +impl CorsContext { + pub fn matches(&self, host: T) -> CorsResult + where + CorsEntry: PartialEq, + CorsHost: PartialEq, + { + if self.is_in_hosts(host.clone()) || self.is_in_entries(host.clone()) { + CorsResult::Allowed + } else if self.is_all_allowed() { + CorsResult::AllowedAll + } else { + CorsResult::NotAllowed + } + } + + /// Returns true if every domain allowed + pub fn is_all_allowed(&self) -> bool { + self.allow_all + } + + pub fn is_in_entries(&self, host: T) -> bool + where + CorsEntry: PartialEq, + { + for ent in self.entries.iter() { + if *ent == host { + return true; + } + } + false + } + + pub fn is_in_hosts(&self, host: T) -> bool + where + CorsHost: PartialEq, + { + for h in self.hosts.iter() { + if *h == host { + return true; + } + } + false + } +} + +impl Default for CorsContext { + fn default() -> Self { + let helper = get_helper(); + Self { + allow_all: false, + entries: helper.cors_entries(), + hosts: Vec::new(), + } + } +} + #[test] fn test_cors_entry() { let ent = CorsEntry::try_from("https://.test.com").unwrap(); diff --git a/src/server/mod.rs b/src/server/mod.rs index 4d276f9..e10833b 100644 --- a/src/server/mod.rs +++ b/src/server/mod.rs @@ -1,3 +1,4 @@ +pub mod context; /// CORS Handle pub mod cors; /// Routes diff --git a/src/server/route.rs b/src/server/route.rs index 05dd1fc..a5f03f6 100644 --- a/src/server/route.rs +++ b/src/server/route.rs @@ -1,8 +1,10 @@ +use super::context::ServerContext; use super::traits::MatchRoute; use super::traits::ResponseFor; use super::version::VersionRoute; use hyper::Body; use hyper::Request; +use std::sync::Arc; pub type RouteType = dyn MatchRoute + Send + Sync; pub type ResponseForType = dyn ResponseFor + Send + Sync; @@ -18,10 +20,14 @@ impl ServerRoutes { Self { routes } } - pub fn match_route(&self, req: &Request) -> Option> { + pub fn match_route( + &self, + req: &Request, + ctx: &Arc, + ) -> Option> { for i in self.routes.iter() { if i.match_route(req) { - return Some(i.get_route()); + return Some(i.get_route(Arc::clone(ctx))); } } None diff --git a/src/server/service.rs b/src/server/service.rs index 7cf42c8..1a69901 100644 --- a/src/server/service.rs +++ b/src/server/service.rs @@ -1,3 +1,4 @@ +use super::context::ServerContext; use super::route::ServerRoutes; use hyper::server::conn::AddrIncoming; use hyper::server::Server; @@ -13,12 +14,13 @@ use std::task::Context; use std::task::Poll; pub struct PixivDownloaderSvc { + context: Arc, routes: Arc, } impl PixivDownloaderSvc { - pub fn new(routes: Arc) -> Self { - Self { routes } + pub fn new(routes: Arc, context: Arc) -> Self { + Self { routes, context } } } @@ -33,7 +35,7 @@ impl Service> for PixivDownloaderSvc { fn call(&mut self, req: Request) -> Self::Future { println!("{} {}", req.method(), req.uri()); - match self.routes.match_route(&req) { + match self.routes.match_route(&req, &self.context) { Some(route) => Box::pin(async move { match route.response(req).await { Ok(data) => Ok(data), @@ -54,12 +56,14 @@ impl Service> for PixivDownloaderSvc { } pub struct PixivDownloaderMakeSvc { + context: Arc, routes: Arc, } impl PixivDownloaderMakeSvc { pub fn new() -> Self { Self { + context: Arc::new(ServerContext::default()), routes: Arc::new(ServerRoutes::new()), } } @@ -76,7 +80,8 @@ impl Service for PixivDownloaderMakeSvc { fn call(&mut self, _: T) -> Self::Future { let routes = Arc::clone(&self.routes); - let fut = async move { Ok(PixivDownloaderSvc::new(routes)) }; + let context = Arc::clone(&self.context); + let fut = async move { Ok(PixivDownloaderSvc::new(routes, context)) }; Box::pin(fut) } } diff --git a/src/server/traits.rs b/src/server/traits.rs index 5963523..3158f85 100644 --- a/src/server/traits.rs +++ b/src/server/traits.rs @@ -1,11 +1,13 @@ +use super::context::ServerContext; use crate::error::PixivDownloaderError; use hyper::Body; use hyper::Request; use hyper::Response; use json::JsonValue; +use std::sync::Arc; pub trait MatchRoute { - fn get_route(&self) -> Box + Send + Sync>; + fn get_route(&self, ctx: Arc) -> Box + Send + Sync>; fn match_route(&self, req: &Request) -> bool; } diff --git a/src/server/version.rs b/src/server/version.rs index 2ee18fc..8c40c9e 100644 --- a/src/server/version.rs +++ b/src/server/version.rs @@ -1,3 +1,4 @@ +use super::context::ServerContext; use super::route::ResponseForType; use super::traits::MatchRoute; use super::traits::ResponseJsonFor; @@ -8,14 +9,15 @@ use hyper::Response; use json::JsonValue; use proc_macros::filter_http_methods; use regex::Regex; +use std::sync::Arc; pub struct VersionContext { - _unused: [u8; 0], + ctx: Arc, } impl VersionContext { - pub fn new() -> Self { - Self { _unused: [] } + pub fn new(ctx: Arc) -> Self { + Self { ctx } } } @@ -25,8 +27,10 @@ impl ResponseJsonFor for VersionContext { &self, req: Request, ) -> Result, PixivDownloaderError> { - filter_http_methods!(req, json::object! {}, true, GET, OPTIONS, POST); - Ok(Response::new(json::object! {"version": [0, 0, 1, 0]})) + filter_http_methods!(req, json::object! {}, true, self.ctx, GET, OPTIONS, POST); + Ok(builder + .body(json::object! {"version": [0, 0, 1, 0]}) + .unwrap()) } } @@ -43,8 +47,8 @@ impl VersionRoute { } impl MatchRoute for VersionRoute { - fn get_route(&self) -> Box { - Box::new(VersionContext::new()) + fn get_route(&self, ctx: Arc) -> Box { + Box::new(VersionContext::new(ctx)) } fn match_route(&self, req: &http::Request) -> bool { diff --git a/src/settings_list.rs b/src/settings_list.rs index f3bd24d..ecf6922 100644 --- a/src/settings_list.rs +++ b/src/settings_list.rs @@ -7,6 +7,8 @@ use crate::settings::JsonValueType; use crate::opt::author_name_filter::check_author_name_filters; use crate::opt::proxy::check_proxy; use crate::opt::size::parse_u32_size; +#[cfg(feature = "server")] +use crate::server::cors::parse_cors_entries; use json::JsonValue; #[cfg(feature = "server")] use std::net::SocketAddr; @@ -36,6 +38,8 @@ pub fn get_settings_list() -> Vec { SettingDes::new("proxy", gettext("Proxy settings."), JsonValueType::Array, Some(check_proxy)).unwrap(), #[cfg(feature = "server")] SettingDes::new("server", gettext("Server address."), JsonValueType::Str, Some(check_socket_addr)).unwrap(), + #[cfg(feature = "server")] + SettingDes::new("cors-entries", gettext("The domains allowed to send CORS requests."), JsonValueType::Array, Some(check_cors_entries)).unwrap(), ] } @@ -44,6 +48,17 @@ fn check_i64(obj: &JsonValue) -> bool { r.is_some() } +#[cfg(feature = "server")] +fn check_cors_entries(obj: &JsonValue) -> bool { + match parse_cors_entries(obj) { + Ok(_) => true, + Err(e) => { + println!("{}", e); + false + } + } +} + #[cfg(feature = "server")] fn check_socket_addr(obj: &JsonValue) -> bool { match obj.as_str() {