From 3f8a33b9d63787097c76adb41600709852082def Mon Sep 17 00:00:00 2001 From: Bnyro Date: Wed, 8 May 2024 18:05:46 +0200 Subject: [PATCH 1/2] refactor: simplify error handling --- src/main.rs | 46 +++++++++++++++------------------------------- src/utils.rs | 4 ++-- 2 files changed, 17 insertions(+), 33 deletions(-) diff --git a/src/main.rs b/src/main.rs index 1b2264e..07e112a 100644 --- a/src/main.rs +++ b/src/main.rs @@ -57,7 +57,7 @@ static CLIENT: Lazy = Lazy::new(|| { .user_agent("Mozilla/5.0 (Windows NT 10.0; rv:102.0) Gecko/20100101 Firefox/102.0"); let proxy = if let Ok(proxy) = env::var("PROXY") { - Some(reqwest::Proxy::all(proxy).unwrap()) + reqwest::Proxy::all(proxy).ok() } else { None }; @@ -75,13 +75,12 @@ static CLIENT: Lazy = Lazy::new(|| { }; if get_env_bool("IPV4_ONLY") { - builder - .local_address(Some("0.0.0.0".parse().unwrap())) - .build() - .unwrap() + builder.local_address("0.0.0.0".parse().ok()) } else { - builder.build().unwrap() + builder } + .build() + .unwrap() }); const ANDROID_USER_AGENT: &str = "com.google.android.youtube/1537338816 (Linux; U; Android 13; en_US; ; Build/TQ2A.230505.002; Cronet/113.0.5672.24)"; @@ -154,13 +153,9 @@ async fn index(req: HttpRequest) -> Result> { let secret = env::var("HASH_SECRET"); if let Ok(secret) = secret { - let qhash = query.get("qhash"); - - if qhash.is_none() { + let Some(qhash) = query.get("qhash") else { return Err("No qhash provided".into()); - } - - let qhash = qhash.unwrap(); + }; if qhash.len() != 8 { return Err("Invalid qhash provided".into()); @@ -218,12 +213,9 @@ async fn index(req: HttpRequest) -> Result> { } } - let res = query.get("host"); - let res = res.map(|s| s.to_string()); - - if res.is_none() { + let Some(host) = query.get("host").map(|s| s.to_string()) else { return Err("No host provided".into()); - } + }; #[cfg(any(feature = "webp", feature = "avif"))] let disallow_image_transcoding = get_env_bool("DISALLOW_IMAGE_TRANSCODING"); @@ -233,14 +225,12 @@ async fn index(req: HttpRequest) -> Result> { #[cfg(feature = "avif")] let avif = query.get("avif") == Some("true"); - let host = res.unwrap(); - let domain = RE_DOMAIN.captures(host.as_str()); - - if domain.is_none() { + let Some(domain) = RE_DOMAIN + .captures(host.as_str()) + .map(|domain| domain.get(1).unwrap().as_str()) + else { return Err("Invalid host provided".into()); - } - - let domain = domain.unwrap().get(1).unwrap().as_str(); + }; if !ALLOWED_DOMAINS.contains(&domain) { return Err("Domain not allowed".into()); @@ -323,13 +313,7 @@ async fn index(req: HttpRequest) -> Result> { request_headers.insert("User-Agent", ANDROID_USER_AGENT.parse().unwrap()); } - let resp = CLIENT.execute(request).await; - - if resp.is_err() { - return Err(resp.err().unwrap().into()); - } - - let resp = resp?; + let resp = CLIENT.execute(request).await?; let mut response = HttpResponse::build(resp.status()); diff --git a/src/utils.rs b/src/utils.rs index 8c415c3..195a4e9 100644 --- a/src/utils.rs +++ b/src/utils.rs @@ -43,9 +43,9 @@ fn finalize_url(path: &str, query: BTreeMap) -> String { } }; - if qhash.is_some() { + if let Some(qhash) = qhash { let mut query = QString::new(query.into_iter().collect::>()); - query.add_pair(("qhash", qhash.unwrap())); + query.add_pair(("qhash", qhash)); return format!("{}?{}", path, query); } } From 4026431403a731ecacec7ecbf340c015d497d670 Mon Sep 17 00:00:00 2001 From: Bnyro Date: Wed, 8 May 2024 21:53:37 +0200 Subject: [PATCH 2/2] refactor: move get_env_bool to utils.rs --- src/main.rs | 13 +++---------- src/utils.rs | 9 ++++++++- 2 files changed, 11 insertions(+), 11 deletions(-) diff --git a/src/main.rs b/src/main.rs index 07e112a..46247be 100644 --- a/src/main.rs +++ b/src/main.rs @@ -34,7 +34,7 @@ async fn main() -> std::io::Result<()> { // get socket/port from env // backwards compat when only UDS is set - if get_env_bool("UDS") { + if utils::get_env_bool("UDS") { let socket_path = env::var("BIND_UNIX").unwrap_or_else(|_| "./socket/actix.sock".to_string()); server.bind_uds(socket_path)? @@ -74,7 +74,7 @@ static CLIENT: Lazy = Lazy::new(|| { builder }; - if get_env_bool("IPV4_ONLY") { + if utils::get_env_bool("IPV4_ONLY") { builder.local_address("0.0.0.0".parse().ok()) } else { builder @@ -126,13 +126,6 @@ fn is_header_allowed(header: &str) -> bool { ) } -fn get_env_bool(key: &str) -> bool { - match env::var(key) { - Ok(val) => val.to_lowercase() == "true" || val == "1", - Err(_) => false, - } -} - async fn index(req: HttpRequest) -> Result> { if req.method() == Method::OPTIONS { let mut response = HttpResponse::Ok(); @@ -218,7 +211,7 @@ async fn index(req: HttpRequest) -> Result> { }; #[cfg(any(feature = "webp", feature = "avif"))] - let disallow_image_transcoding = get_env_bool("DISALLOW_IMAGE_TRANSCODING"); + let disallow_image_transcoding = utils::get_env_bool("DISALLOW_IMAGE_TRANSCODING"); let rewrite = query.get("rewrite") != Some("false"); diff --git a/src/utils.rs b/src/utils.rs index 195a4e9..d5efe91 100644 --- a/src/utils.rs +++ b/src/utils.rs @@ -2,6 +2,7 @@ use qstring::QString; use reqwest::Url; use std::borrow::Cow; use std::collections::BTreeMap; +use std::env; pub fn read_buf(buf: &[u8], pos: &mut usize) -> u8 { let byte = buf[*pos]; @@ -13,7 +14,6 @@ fn finalize_url(path: &str, query: BTreeMap) -> String { #[cfg(feature = "qhash")] { use std::collections::BTreeSet; - use std::env; let qhash = { let secret = env::var("HASH_SECRET"); @@ -94,3 +94,10 @@ pub fn escape_xml(raw: &str) -> Cow<'_, str> { Cow::Owned(escaped) } } + +pub fn get_env_bool(key: &str) -> bool { + match env::var(key) { + Ok(val) => val.to_lowercase() == "true" || val == "1", + Err(_) => false, + } +}