iridium/redsunlib
iridium/redsunlib
2
2
Fork 0
Code Issues 2 Pull Requests Packages Projects Releases 4 Wiki Activity

Handle HeaderValue and Uri parsing errors

Browse Source
This commit is contained in:
spikecodes 2021-03-17 16:32:28 -07:00
parent b14b4ff551
commit fb7faf6477
No known key found for this signature in database
GPG Key ID: 004CECFF9B463BCB
3 changed files with 34 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
src/client.rs
View File

@ -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> { async fn stream(url: &str) -> Result<Response<Body>, String> {
// First parameter is target URL (mandatory). // 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. // Prepare the HTTPS connector.
let https = hyper_rustls::HttpsConnector::with_native_roots(); 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 { match client.request(req(url)).await {
Ok(response) => { Ok(response) => {
if response.status().to_string().starts_with('3') { 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 { } else {
Ok(response) Ok(response)
} }

9
src/main.rs
View File

@ -21,7 +21,7 @@ mod utils;
use clap::{App as cli, Arg}; use clap::{App as cli, Arg};
use futures_lite::FutureExt; use futures_lite::FutureExt;
use hyper::{Body, Request, Response}; use hyper::{header::HeaderValue, Body, Request, Response};
mod client; mod client;
use client::proxy; use client::proxy;
@ -73,7 +73,12 @@ async fn resource(body: &str, content_type: &str, cache: bool) -> Result<Respons
.unwrap_or_default(); .unwrap_or_default();
if cache { 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) Ok(res)

20
src/server.rs
View File

@ -1,6 +1,7 @@
use cookie::Cookie; use cookie::Cookie;
use futures_lite::{future::Boxed, Future, FutureExt}; use futures_lite::{future::Boxed, Future, FutureExt};
use hyper::{ use hyper::{
header::HeaderValue,
service::{make_service_fn, service_fn}, service::{make_service_fn, service_fn},
HeaderMap, HeaderMap,
}; };
@ -27,7 +28,10 @@ macro_rules! headers(
{ {
let mut m = hyper::HeaderMap::new(); 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 m
} }
@ -92,14 +96,24 @@ impl ResponseExt for Response<Body> {
} }
fn insert_cookie(&mut self, cookie: Cookie) { 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) { fn remove_cookie(&mut self, name: String) {
let mut cookie = Cookie::named(name); let mut cookie = Cookie::named(name);
cookie.set_path("/"); cookie.set_path("/");
cookie.set_max_age(Duration::second()); 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(_) => (),
}
} }
} }