Add support for custom headers

1 files changed, 11 insertions, 5 deletions

diff --git a/src/https.rs b/src/https.rs
index a62ebea..43a93e2 100644
--- a/src/https.rs
+++ b/src/https.rs
@@ -8,6 +8,7 @@ use futures::FutureExt;
 use hyper::server::conn::Http;
 use hyper::service::service_fn;
 use hyper::{Body, Request, Response, StatusCode};
+use http::header::{HeaderName, HeaderValue};
 use tokio::net::TcpListener;
 use tokio::sync::watch;
 use tokio_rustls::TlsAcceptor;
@@ -17,11 +18,10 @@ use crate::proxy_config::ProxyConfig;
 use crate::reverse_proxy;
 
 pub async fn serve_https(
+	bind_addr: SocketAddr,
 	cert_store: Arc<CertStore>,
 	proxy_config: watch::Receiver<Arc<ProxyConfig>>,
 ) -> Result<()> {
-	let addr = format!("0.0.0.0:1443");
-
 	let mut cfg = rustls::ServerConfig::builder()
 		.with_safe_defaults()
 		.with_no_client_auth()
@@ -31,9 +31,9 @@ pub async fn serve_https(
 	let tls_cfg = Arc::new(cfg);
 	let tls_acceptor = Arc::new(TlsAcceptor::from(tls_cfg));
 
-	println!("Starting to serve on https://{}.", addr);
+	info!("Starting to serve on https://{}.", bind_addr);
 
-	let tcp = TcpListener::bind(&addr).await?;
+	let tcp = TcpListener::bind(bind_addr).await?;
 	loop {
 		let (socket, remote_addr) = tcp.accept().await?;
 
@@ -118,7 +118,13 @@ async fn handle(
 		let to_addr = format!("http://{}", proxy_to.target_addr);
 		info!("Proxying {} {} -> {}", host, path, to_addr);
 
-		reverse_proxy::call(remote_addr.ip(), &to_addr, req).await
+		let mut response = reverse_proxy::call(remote_addr.ip(), &to_addr, req).await?;
+
+		for (header, value) in proxy_to.add_headers.iter() {
+			response.headers_mut().insert(HeaderName::from_bytes(header.as_bytes())?, HeaderValue::from_str(value)?);
+		}
+
+		Ok(response)
 	} else {
 		info!("Proxying {} {} -> NOT FOUND", host, path);