Handle HeaderValue and Uri parsing errors

3 years ago · fb7faf6477
parent b14b4ff551
commit fb7faf6477
3 changed files with 34 additions and 7 deletions
--- a/src/client.rs
+++ b/src/client.rs
@ -19,7 +19,7 @@ pub async fn proxy(req: Request<Body>, format: &str) -> Result<Response<Body>, S

 async fn stream(url: &str) -> Result<Response<Body>, 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<Result<Response<Body>, 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)
 				}
--- 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<Respons
 		.unwrap_or_default();

 	if cache {
-		res.headers_mut().insert("Cache-Control", "public, max-age=1209600, s-maxage=86400".parse().unwrap());
+		match HeaderValue::from_str("public, max-age=1209600, s-maxage=86400") {
+			Ok(val) => {
+				res.headers_mut().insert("Cache-Control", val);
+			}
+			Err(_) => (),
+		}
 	}

 	Ok(res)
--- 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<Body> {
 	}

 	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(_) => (),
+		}
 	}
 }