From c49b2075a30460a319ce53da300e4943d45d5add Mon Sep 17 00:00:00 2001
From: Alex Auvolat <alex@adnab.me>
Date: Wed, 8 Dec 2021 13:28:07 +0100
Subject: [PATCH] Try to fix things

---
 src/cert_store.rs    |  6 +++++-
 src/https.rs         |  2 ++
 src/main.rs          | 11 ++++++-----
 src/proxy_config.rs  |  7 +++++++
 src/reverse_proxy.rs | 27 ++++++++++++++++++++++++++-
 5 files changed, 46 insertions(+), 7 deletions(-)

diff --git a/src/cert_store.rs b/src/cert_store.rs
index 2834795..6cc3ea9 100644
--- a/src/cert_store.rs
+++ b/src/cert_store.rs
@@ -25,7 +25,11 @@ pub struct CertStore {
 }
 
 impl CertStore {
-	pub fn new(consul: Consul, rx_proxy_config: watch::Receiver<Arc<ProxyConfig>>, letsencrypt_email: String) -> Arc<Self> {
+	pub fn new(
+		consul: Consul,
+		rx_proxy_config: watch::Receiver<Arc<ProxyConfig>>,
+		letsencrypt_email: String,
+	) -> Arc<Self> {
 		Arc::new(Self {
 			consul,
 			certs: RwLock::new(HashMap::new()),
diff --git a/src/https.rs b/src/https.rs
index ae737d9..33b8030 100644
--- a/src/https.rs
+++ b/src/https.rs
@@ -117,6 +117,7 @@ async fn handle(
 
 		let to_addr = format!("http://{}", proxy_to.target_addr);
 		info!("Proxying {} {} -> {}", host, path, to_addr);
+		trace!("Request: {:?}", req);
 
 		let mut response = reverse_proxy::call(remote_addr.ip(), &to_addr, req).await?;
 
@@ -126,6 +127,7 @@ async fn handle(
 				HeaderValue::from_str(value)?,
 			);
 		}
+		trace!("Response: {:?}", response);
 
 		Ok(response)
 	} else {
diff --git a/src/main.rs b/src/main.rs
index bea09bd..d495fb2 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -54,10 +54,7 @@ struct Opt {
 	pub https_bind_addr: SocketAddr,
 
 	/// E-mail address for Let's Encrypt certificate requests
-	#[structopt(
-		long = "letsencrypt-email",
-		env = "TRICOT_LETSENCRYPT_EMAIL",
-	)]
+	#[structopt(long = "letsencrypt-email", env = "TRICOT_LETSENCRYPT_EMAIL")]
 	pub letsencrypt_email: String,
 }
 
@@ -75,7 +72,11 @@ async fn main() {
 	let consul = consul::Consul::new(&opt.consul_addr, &opt.consul_kv_prefix, &opt.node_name);
 	let mut rx_proxy_config = proxy_config::spawn_proxy_config_task(consul.clone());
 
-	let cert_store = cert_store::CertStore::new(consul.clone(), rx_proxy_config.clone(), opt.letsencrypt_email.clone());
+	let cert_store = cert_store::CertStore::new(
+		consul.clone(),
+		rx_proxy_config.clone(),
+		opt.letsencrypt_email.clone(),
+	);
 	tokio::spawn(cert_store.clone().watch_proxy_config());
 
 	tokio::spawn(http::serve_http(opt.http_bind_addr, consul.clone()));
diff --git a/src/proxy_config.rs b/src/proxy_config.rs
index dd7a357..a15c143 100644
--- a/src/proxy_config.rs
+++ b/src/proxy_config.rs
@@ -136,6 +136,8 @@ fn parse_tricot_add_header_tag(tag: &str) -> Option<(String, String)> {
 }
 
 fn parse_consul_catalog(catalog: &ConsulNodeCatalog) -> Vec<ProxyEntry> {
+	trace!("Parsing node catalog: {:#?}", catalog);
+
 	let mut entries = vec![];
 
 	for (_, svc) in catalog.services.iter() {
@@ -168,6 +170,11 @@ fn parse_consul_catalog(catalog: &ConsulNodeCatalog) -> Vec<ProxyEntry> {
 		}
 	}
 
+	trace!("Result of parsing catalog:");
+	for ent in entries.iter() {
+		trace!("    {}", ent);
+	}
+
 	entries
 }
 
diff --git a/src/reverse_proxy.rs b/src/reverse_proxy.rs
index 046808f..3d2a1fc 100644
--- a/src/reverse_proxy.rs
+++ b/src/reverse_proxy.rs
@@ -2,7 +2,10 @@
 //! See there for original Copyright notice
 
 use anyhow::Result;
+use log::*;
+use std::convert::TryInto;
 
+use http::header::HeaderName;
 use hyper::header::{HeaderMap, HeaderValue};
 use hyper::{Body, Client, Request, Response, Uri};
 use lazy_static::lazy_static;
@@ -63,7 +66,9 @@ fn create_proxied_request<B>(
 	forward_url: &str,
 	request: Request<B>,
 ) -> Result<Request<B>> {
-	let mut builder = Request::builder().uri(forward_uri(forward_url, &request)?);
+	let mut builder = Request::builder()
+		.method(request.method())
+		.uri(forward_uri(forward_url, &request)?);
 
 	*builder.headers_mut().unwrap() = remove_hop_headers(request.headers());
 
@@ -95,6 +100,21 @@ fn create_proxied_request<B>(
 		}
 	}
 
+	if let Some(conn) = request.headers().get("connection") {
+		if conn.to_str()?.to_lowercase() == "upgrade" {
+			if let Some(upgrade) = request.headers().get("upgrade") {
+				builder.headers_mut().unwrap().insert(
+					HeaderName::from_bytes(b"connection")?,
+					"Upgrade".try_into()?,
+				);
+				builder
+					.headers_mut()
+					.unwrap()
+					.insert(HeaderName::from_bytes(b"upgrade")?, upgrade.clone());
+			}
+		}
+	}
+
 	Ok(builder.body(request.into_body())?)
 }
 
@@ -105,8 +125,13 @@ pub async fn call(
 ) -> Result<Response<Body>> {
 	let proxied_request = create_proxied_request(client_ip, &forward_uri, request)?;
 
+	trace!("Proxied request: {:?}", proxied_request);
+
 	let client = Client::new();
 	let response = client.request(proxied_request).await?;
+
+	trace!("Inner response: {:?}", response);
+
 	let proxied_response = create_proxied_response(response);
 	Ok(proxied_response)
 }