From 03e4323c9500774568b453509422e01c54138fcf Mon Sep 17 00:00:00 2001 From: Avery Harnish Date: Wed, 20 Nov 2019 11:21:51 -0600 Subject: [PATCH 1/3] proxy: Handle response headers --- src/commands/preview/proxy.rs | 37 +++++++++++++++++++++++++++-------- 1 file changed, 29 insertions(+), 8 deletions(-) diff --git a/src/commands/preview/proxy.rs b/src/commands/preview/proxy.rs index dd5d42607..6e297e0c0 100644 --- a/src/commands/preview/proxy.rs +++ b/src/commands/preview/proxy.rs @@ -6,7 +6,7 @@ use chrono::prelude::*; use hyper2::client::{HttpConnector, ResponseFuture}; use hyper2::header::{HeaderMap, HeaderName, HeaderValue}; use hyper2::service::{make_service_fn, service_fn}; -use hyper2::{Body, Client, Request, Server, Uri}; +use hyper2::{Body, Client, Request, Response, Server, Uri}; use hyper_tls2::HttpsConnector; @@ -22,6 +22,7 @@ use crate::settings::target::Target; use crate::commands::preview::upload; const PREVIEW_HOST: &str = "rawhttp.cloudflareworkers.com"; +const HEADER_PREFIX: &str = "cf-ew-raw-"; #[derive(Clone)] struct ProxyConfig { @@ -99,12 +100,32 @@ pub async fn proxy( let proxy_config = proxy_config.clone(); async move { Ok::<_, failure::Error>(service_fn(move |req| { - preview_request( - req, - client.to_owned(), - preview_id.to_owned(), - proxy_config.clone(), - ) + let client = client.to_owned(); + let preview_id = preview_id.to_owned(); + let proxy_config = proxy_config.clone(); + async move { + let resp = preview_request(req, client, preview_id, proxy_config).await?; + + let (mut parts, body) = resp.into_parts(); + + let mut headers = HeaderMap::new(); + + for header in &parts.headers { + let (name, value) = header; + let name = name.as_str(); + if name.starts_with(HEADER_PREFIX) { + let header_name = &name[HEADER_PREFIX.len()..]; + // TODO: remove unwrap + let header_name = + HeaderName::from_bytes(header_name.as_bytes()).unwrap(); + headers.insert(header_name, value.clone()); + } + } + parts.headers = headers; + + let resp = Response::from_parts(parts, body); + Ok::<_, failure::Error>(resp) + } })) } }); @@ -145,7 +166,7 @@ fn preview_request( for header in &parts.headers { let (name, value) = header; - let forward_header = format!("cf-ew-raw-{}", name); + let forward_header = format!("{}{}", HEADER_PREFIX, name); // TODO: remove unwrap let header_name = HeaderName::from_bytes(forward_header.as_bytes()).unwrap(); headers.insert(header_name, value.clone()); From 25dff4384f88d22524c7606476c0a9eb23528934 Mon Sep 17 00:00:00 2001 From: Avery Harnish Date: Thu, 12 Dec 2019 16:29:03 -0600 Subject: [PATCH 2/3] Remove an unwrap --- src/commands/dev/mod.rs | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/commands/dev/mod.rs b/src/commands/dev/mod.rs index 4d69d7ac2..61549a7df 100644 --- a/src/commands/dev/mod.rs +++ b/src/commands/dev/mod.rs @@ -63,9 +63,7 @@ pub async fn dev( let name = name.as_str(); if name.starts_with(HEADER_PREFIX) { let header_name = &name[HEADER_PREFIX.len()..]; - // TODO: remove unwrap - let header_name = - HeaderName::from_bytes(header_name.as_bytes()).unwrap(); + let header_name = HeaderName::from_bytes(header_name.as_bytes())?; headers.insert(header_name, value.clone()); } } From 605972a44c79e81a85c434f3cde61c1e387a6b09 Mon Sep 17 00:00:00 2001 From: Avery Harnish Date: Thu, 12 Dec 2019 16:39:08 -0600 Subject: [PATCH 3/3] Split out response header munging to separate function --- src/commands/dev/mod.rs | 30 ++++++++++++++++++------------ 1 file changed, 18 insertions(+), 12 deletions(-) diff --git a/src/commands/dev/mod.rs b/src/commands/dev/mod.rs index 61549a7df..d309b9bf3 100644 --- a/src/commands/dev/mod.rs +++ b/src/commands/dev/mod.rs @@ -7,6 +7,7 @@ use hyper::client::{HttpConnector, ResponseFuture}; use hyper::header::{HeaderMap, HeaderName, HeaderValue}; use hyper::service::{make_service_fn, service_fn}; use hyper::{Body, Client, Request, Response, Server, Uri}; +use hyper::http::response::Parts; use hyper::http::uri::InvalidUri; use hyper_tls::HttpsConnector; @@ -56,18 +57,7 @@ pub async fn dev( let (mut parts, body) = resp.into_parts(); - let mut headers = HeaderMap::new(); - - for header in &parts.headers { - let (name, value) = header; - let name = name.as_str(); - if name.starts_with(HEADER_PREFIX) { - let header_name = &name[HEADER_PREFIX.len()..]; - let header_name = HeaderName::from_bytes(header_name.as_bytes())?; - headers.insert(header_name, value.clone()); - } - } - parts.headers = headers; + munge_response_headers(&mut parts)?; let resp = Response::from_parts(parts, body); Ok::<_, failure::Error>(resp) @@ -84,6 +74,22 @@ pub async fn dev( Ok(()) } +fn munge_response_headers(parts: &mut Parts) -> Result<(), failure::Error> { + let mut headers = HeaderMap::new(); + + for header in &parts.headers { + let (name, value) = header; + let name = name.as_str(); + if name.starts_with(HEADER_PREFIX) { + let header_name = &name[HEADER_PREFIX.len()..]; + let header_name = HeaderName::from_bytes(header_name.as_bytes())?; + headers.insert(header_name, value.clone()); + } + } + parts.headers = headers; + Ok(()) +} + fn get_preview_url(path_string: &str) -> Result { format!("https://{}{}", PREVIEW_HOST, path_string).parse() }