add request header override option

src/proxy/backend.rs

@@ -0,0 +1,206 @@
+use super::UpstreamOption;
+use crate::log::*;
+use parking_lot::Mutex;
+use rand::Rng;
+use rustc_hash::{FxHashMap as HashMap, FxHashSet as HashSet};
+use std::{
+  fs::File,
+  io::{self, BufReader, Cursor, Read},
+  path::PathBuf,
+  sync::{
+    atomic::{AtomicUsize, Ordering},
+    Arc,
+  },
+};
+use tokio_rustls::rustls::{Certificate, PrivateKey, ServerConfig};
+
+pub struct Backends {
+  pub apps: HashMap<String, Backend>, // TODO: hyper::uriで抜いたhostで引っ掛ける。Stringでいいのか？
+  pub default_app: Option<String>,    // for plaintext http
+}
+
+pub struct Backend {
+  pub app_name: String,
+  pub server_name: String,
+  pub reverse_proxy: ReverseProxy,
+
+  // tls settings
+  pub tls_cert_path: Option<PathBuf>,
+  pub tls_cert_key_path: Option<PathBuf>,
+  pub https_redirection: Option<bool>,
+  pub server_config: Mutex<Option<ServerConfig>>,
+}
+
+#[derive(Debug, Clone)]
+pub struct ReverseProxy {
+  pub default_upstream: Option<Upstream>,
+  pub upstream: HashMap<String, Upstream>,
+}
+
+#[allow(dead_code)]
+#[derive(Debug, Clone)]
+pub enum LoadBalance {
+  RoundRobin,
+  Random,
+}
+impl Default for LoadBalance {
+  fn default() -> Self {
+    Self::RoundRobin
+  }
+}
+
+#[derive(Debug, Clone)]
+pub struct Upstream {
+  pub uri: Vec<hyper::Uri>,
+  pub lb: LoadBalance,
+  pub cnt: UpstreamCount, // counter for load balancing
+  pub opts: HashSet<UpstreamOption>,
+}
+
+#[derive(Debug, Clone, Default)]
+pub struct UpstreamCount(Arc<AtomicUsize>);
+
+impl Upstream {
+  pub fn get(&self) -> Option<&hyper::Uri> {
+    match self.lb {
+      LoadBalance::RoundRobin => {
+        let idx = self.increment_cnt();
+        self.uri.get(idx)
+      }
+      LoadBalance::Random => {
+        let mut rng = rand::thread_rng();
+        let max = self.uri.len() - 1;
+        self.uri.get(rng.gen_range(0..max))
+      }
+    }
+  }
+
+  fn current_cnt(&self) -> usize {
+    self.cnt.0.load(Ordering::Relaxed)
+  }
+
+  fn increment_cnt(&self) -> usize {
+    if self.current_cnt() < self.uri.len() - 1 {
+      self.cnt.0.fetch_add(1, Ordering::Relaxed)
+    } else {
+      self.cnt.0.fetch_and(0, Ordering::Relaxed)
+    }
+  }
+}
+
+impl Backend {
+  pub fn get_tls_server_config(&self) -> Option<ServerConfig> {
+    let lock = self.server_config.lock();
+    let opt_clone = lock.clone();
+    if let Some(sc) = opt_clone {
+      return Some(sc);
+    }
+    None
+  }
+  pub async fn update_server_config(&self) -> io::Result<()> {
+    debug!("Update TLS server config");
+    let certs_path = self.tls_cert_path.as_ref().unwrap();
+    let certs_keys_path = self.tls_cert_key_path.as_ref().unwrap();
+    let certs: Vec<_> = {
+      let certs_path_str = certs_path.display().to_string();
+      let mut reader = BufReader::new(File::open(certs_path).map_err(|e| {
+        io::Error::new(
+          e.kind(),
+          format!(
+            "Unable to load the certificates [{}]: {}",
+            certs_path_str, e
+          ),
+        )
+      })?);
+      rustls_pemfile::certs(&mut reader).map_err(|_| {
+        io::Error::new(
+          io::ErrorKind::InvalidInput,
+          "Unable to parse the certificates",
+        )
+      })?
+    }
+    .drain(..)
+    .map(Certificate)
+    .collect();
+    let certs_keys: Vec<_> = {
+      let certs_keys_path_str = certs_keys_path.display().to_string();
+      let encoded_keys = {
+        let mut encoded_keys = vec![];
+        File::open(certs_keys_path)
+          .map_err(|e| {
+            io::Error::new(
+              e.kind(),
+              format!(
+                "Unable to load the certificate keys [{}]: {}",
+                certs_keys_path_str, e
+              ),
+            )
+          })?
+          .read_to_end(&mut encoded_keys)?;
+        encoded_keys
+      };
+      let mut reader = Cursor::new(encoded_keys);
+      let pkcs8_keys = rustls_pemfile::pkcs8_private_keys(&mut reader).map_err(|_| {
+        io::Error::new(
+          io::ErrorKind::InvalidInput,
+          "Unable to parse the certificates private keys (PKCS8)",
+        )
+      })?;
+      reader.set_position(0);
+      let mut rsa_keys = rustls_pemfile::rsa_private_keys(&mut reader).map_err(|_| {
+        io::Error::new(
+          io::ErrorKind::InvalidInput,
+          "Unable to parse the certificates private keys (RSA)",
+        )
+      })?;
+      let mut keys = pkcs8_keys;
+      keys.append(&mut rsa_keys);
+      if keys.is_empty() {
+        return Err(io::Error::new(
+          io::ErrorKind::InvalidInput,
+          "No private keys found - Make sure that they are in PKCS#8/PEM format",
+        ));
+      }
+      keys.drain(..).map(PrivateKey).collect()
+    };
+
+    let mut server_config = certs_keys
+      .into_iter()
+      .find_map(|certs_key| {
+        let server_config_builder = ServerConfig::builder()
+          .with_safe_defaults()
+          .with_no_client_auth();
+        if let Ok(found_config) = server_config_builder.with_single_cert(certs.clone(), certs_key) {
+          Some(found_config)
+        } else {
+          None
+        }
+      })
+      .ok_or_else(|| {
+        io::Error::new(
+          io::ErrorKind::InvalidInput,
+          "Unable to find a valid certificate and key",
+        )
+      })?;
+
+    #[cfg(feature = "h3")]
+    {
+      server_config.alpn_protocols = vec![
+        b"h3".to_vec(),
+        b"hq-29".to_vec(), // quinn draft example TODO: remove later
+        b"h2".to_vec(),
+        b"http/1.1".to_vec(),
+      ];
+    }
+    #[cfg(not(feature = "h3"))]
+    {
+      server_config.alpn_protocols = vec![b"h2".to_vec(), b"http/1.1".to_vec()];
+    }
+
+    let mut config_store = self.server_config.lock();
+    *config_store = Some(server_config);
+
+    // server_config;
+    Ok(())
+  }
+}

src/proxy/backend_opt.rs

@@ -0,0 +1,16 @@
+use crate::error::*;
+
+#[derive(Debug, Clone, Hash, Eq, PartialEq)]
+pub enum UpstreamOption {
+  OverrideHost,
+  // TODO: Adds more options for heder override
+}
+impl TryFrom<&str> for UpstreamOption {
+  type Error = anyhow::Error;
+  fn try_from(val: &str) -> Result<Self> {
+    match val {
+      "override_host" => Ok(Self::OverrideHost),
+      _ => Err(anyhow!("Unsupported header option")),
+    }
+  }
+}

src/proxy/mod.rs

@@ -1,7 +1,11 @@
+mod backend;
+mod backend_opt;
 #[cfg(feature = "h3")]
 mod proxy_h3;
 mod proxy_handler;
 mod proxy_main;
 mod proxy_tls;
 
+pub use backend::*;
+pub use backend_opt::UpstreamOption;
 pub use proxy_main::Proxy;

src/proxy/proxy_handler.rs

@@ -1,5 +1,5 @@
 // Highly motivated by https://github.com/felipenoris/hyper-reverse-proxy
-use super::Proxy;
+use super::{Proxy, Upstream, UpstreamOption};
 use crate::{constants::*, error::*, log::*};
 use hyper::{
   client::connect::Connect,
@@ -68,14 +68,14 @@ where
     // Find reverse proxy for given path and choose one of upstream host
     // TODO: More flexible path matcher
     let path = req.uri().path();
-    let upstream_uri = if let Some(upstream) = backend.reverse_proxy.upstream.get(path) {
-      upstream.get()
+    let upstream = if let Some(upstream) = backend.reverse_proxy.upstream.get(path) {
+      upstream
     } else if let Some(default_upstream) = &backend.reverse_proxy.default_upstream {
-      default_upstream.get()
+      default_upstream
     } else {
       return http_error(StatusCode::NOT_FOUND);
     };
-    let upstream_scheme_host = if let Some(u) = upstream_uri {
+    let upstream_scheme_host = if let Some(u) = upstream.get() {
       u
     } else {
       return http_error(StatusCode::INTERNAL_SERVER_ERROR);
@@ -92,6 +92,7 @@ where
       upstream_scheme_host,
       path_and_query,
       &upgrade_in_request,
+      upstream,
     ) {
       req
     } else {
@@ -184,6 +185,7 @@ fn generate_request_forwarded<B: core::fmt::Debug>(
   upstream_scheme_host: &Uri,
   path_and_query: String,
   upgrade: &Option<String>,
+  upstream: &Upstream,
 ) -> Result<Request<B>> {
   debug!("Generate request to be forwarded");
 
@@ -206,11 +208,27 @@ fn generate_request_forwarded<B: core::fmt::Debug>(
   remove_hop_header(headers);
   // X-Forwarded-For
   add_forwarding_header(headers, client_addr)?;
+
   // Add te: trailer if te_trailer
   if te_trailers {
     headers.insert("te", "trailer".parse().unwrap());
   }
 
+  // add "host" header of original server_name if not exist (default)
+  if req.headers().get(hyper::header::HOST).is_none() {
+    let org_host = req.uri().host().unwrap_or("none").to_owned();
+    req.headers_mut().insert(
+      hyper::header::HOST,
+      HeaderValue::from_str(org_host.as_str()).unwrap(),
+    );
+  };
+
+  // apply upstream-specific headers given in upstream_option
+  let headers = req.headers_mut();
+  println!("before {:?}", headers);
+  apply_upstream_options_to_header(headers, client_addr, upstream_scheme_host, upstream)?;
+  println!("after {:?}", req);
+
   // update uri in request
   *req.uri_mut() = Uri::builder()
     .scheme(upstream_scheme_host.scheme().unwrap().as_str())
@@ -223,7 +241,7 @@ fn generate_request_forwarded<B: core::fmt::Debug>(
     req.headers_mut().insert("upgrade", v.parse().unwrap());
     req
       .headers_mut()
-      .insert("connection", HeaderValue::from_str("upgrade")?);
+      .insert(hyper::header::CONNECTION, HeaderValue::from_str("upgrade")?);
   }
 
   // Change version to http/1.1 when destination scheme is http
@@ -237,8 +255,29 @@ fn generate_request_forwarded<B: core::fmt::Debug>(
   Ok(req)
 }
 
+fn apply_upstream_options_to_header(
+  headers: &mut HeaderMap,
+  _client_addr: SocketAddr,
+  upstream_scheme_host: &Uri,
+  upstream: &Upstream,
+) -> Result<()> {
+  upstream.opts.iter().for_each(|opt| match opt {
+    UpstreamOption::OverrideHost => {
+      let upstream_host = upstream_scheme_host.host().unwrap();
+      headers
+        .insert(
+          hyper::header::HOST,
+          HeaderValue::from_str(upstream_host).unwrap(),
+        )
+        .unwrap();
+    }
+  });
+  Ok(())
+}
+
 fn add_forwarding_header(headers: &mut HeaderMap, client_addr: SocketAddr) -> Result<()> {
-  // TODO: Option対応？
+  // default process
+  // optional process defined by upstream_option is applied in fn apply_upstream_options
   let client_ip = client_addr.ip();
   match headers.entry("x-forwarded-for") {
     hyper::header::Entry::Vacant(entry) => {

src/proxy/proxy_main.rs

@@ -1,5 +1,6 @@
 // use super::proxy_handler::handle_request;
-use crate::{backend::Backends, error::*, globals::Globals, log::*};
+use super::Backends;
+use crate::{error::*, globals::Globals, log::*};
 use hyper::{
   client::connect::Connect, server::conn::Http, service::service_fn, Body, Client, Request,
 };