From fb7faf6477b3a220b213e45b76f9ad891a99b725 Mon Sep 17 00:00:00 2001 From: spikecodes <19519553+spikecodes@users.noreply.github.com> Date: Wed, 17 Mar 2021 16:32:28 -0700 Subject: [PATCH] Handle HeaderValue and Uri parsing errors --- src/client.rs | 12 ++++++++++-- src/main.rs | 9 +++++++-- src/server.rs | 20 +++++++++++++++++--- 3 files changed, 34 insertions(+), 7 deletions(-) diff --git a/src/client.rs b/src/client.rs index cd5fe15..d60f71d 100644 --- a/src/client.rs +++ b/src/client.rs @@ -19,7 +19,7 @@ pub async fn proxy(req: Request, format: &str) -> Result, S async fn stream(url: &str) -> Result, String> { // First parameter is target URL (mandatory). - let url = Uri::from_str(url).unwrap(); + let url = Uri::from_str(url).map_err(|_| "Couldn't parse URL".to_string())?; // Prepare the HTTPS connector. let https = hyper_rustls::HttpsConnector::with_native_roots(); @@ -71,7 +71,15 @@ fn request(url: String) -> Boxed, String>> { match client.request(req(url)).await { Ok(response) => { if response.status().to_string().starts_with('3') { - request(response.headers().get("Location").unwrap().to_str().unwrap_or_default().to_string()).await + request( + response + .headers() + .get("Location") + .map(|val| val.to_str().unwrap_or_default()) + .unwrap_or_default() + .to_string(), + ) + .await } else { Ok(response) } diff --git a/src/main.rs b/src/main.rs index b516bb4..daa3320 100644 --- a/src/main.rs +++ b/src/main.rs @@ -21,7 +21,7 @@ mod utils; use clap::{App as cli, Arg}; use futures_lite::FutureExt; -use hyper::{Body, Request, Response}; +use hyper::{header::HeaderValue, Body, Request, Response}; mod client; use client::proxy; @@ -73,7 +73,12 @@ async fn resource(body: &str, content_type: &str, cache: bool) -> Result { + res.headers_mut().insert("Cache-Control", val); + } + Err(_) => (), + } } Ok(res) diff --git a/src/server.rs b/src/server.rs index 6d7ce8e..ad542af 100644 --- a/src/server.rs +++ b/src/server.rs @@ -1,6 +1,7 @@ use cookie::Cookie; use futures_lite::{future::Boxed, Future, FutureExt}; use hyper::{ + header::HeaderValue, service::{make_service_fn, service_fn}, HeaderMap, }; @@ -27,7 +28,10 @@ macro_rules! headers( { let mut m = hyper::HeaderMap::new(); $( - m.insert($key, hyper::header::HeaderValue::from_str($value).unwrap()); + match hyper::header::HeaderValue::from_str($value) { + Ok(val) => { m.insert($key, val); } + Err(_) => () + } )+ m } @@ -92,14 +96,24 @@ impl ResponseExt for Response { } fn insert_cookie(&mut self, cookie: Cookie) { - self.headers_mut().append("Set-Cookie", cookie.to_string().parse().unwrap()); + match HeaderValue::from_str(&cookie.to_string()) { + Ok(val) => { + self.headers_mut().append("Set-Cookie", val); + } + Err(_) => (), + } } fn remove_cookie(&mut self, name: String) { let mut cookie = Cookie::named(name); cookie.set_path("/"); cookie.set_max_age(Duration::second()); - self.headers_mut().append("Set-Cookie", cookie.to_string().parse().unwrap()); + match HeaderValue::from_str(&cookie.to_string()) { + Ok(val) => { + self.headers_mut().append("Set-Cookie", val); + } + Err(_) => (), + } } }