set forwarding headers

1 files changed, 25 insertions, 6 deletions

diff --git a/src/main.rs b/src/main.rs
index 7f6ec88..0e6e4e9 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -6,14 +6,17 @@ use http_body_util::{combinators::BoxBody, BodyExt};
 use hyper::{
     body::Incoming,
     header::{HOST, UPGRADE},
-    http::{uri::PathAndQuery, HeaderValue},
+    http::{
+        uri::{PathAndQuery, Scheme},
+        HeaderValue,
+    },
     server::conn::http1,
     service::service_fn,
     upgrade::OnUpgrade,
     Request, Response, StatusCode, Uri,
 };
 use log::{debug, error, info, warn};
-use std::{fs::File, io::BufReader, path::Path, sync::Arc};
+use std::{fs::File, io::BufReader, net::SocketAddr, path::Path, sync::Arc};
 use tokio::{
     io::{AsyncRead, AsyncWrite},
     net::{TcpListener, TcpStream},
@@ -56,7 +59,7 @@ async fn serve_http(config: Arc<Config>) -> Result<()> {
     loop {
         let (stream, addr) = listener.accept().await.context("accepting connection")?;
         debug!("connection from {addr}");
-        serve_stream(config.clone(), stream)
+        serve_stream(config.clone(), stream, addr)
     }
 }
 async fn serve_https(config: Arc<Config>) -> Result<()> {
@@ -82,7 +85,7 @@ async fn serve_https(config: Arc<Config>) -> Result<()> {
         let (stream, addr) = listener.accept().await.context("accepting connection")?;
         debug!("connection from {addr}");
         match tls_acceptor.accept(stream).await {
-            Ok(stream) => serve_stream(config.clone(), stream),
+            Ok(stream) => serve_stream(config.clone(), stream, addr),
             Err(e) => warn!("error accepting tls: {e}"),
         };
     }
@@ -91,6 +94,7 @@ async fn serve_https(config: Arc<Config>) -> Result<()> {
 pub fn serve_stream<T: AsyncRead + AsyncWrite + Unpin + Send + 'static>(
     config: Arc<Config>,
     stream: T,
+    addr: SocketAddr,
 ) {
     tokio::task::spawn(async move {
         let conn = http1::Builder::new()
@@ -99,7 +103,7 @@ pub fn serve_stream<T: AsyncRead + AsyncWrite + Unpin + Send + 'static>(
                 service_fn(move |req| {
                     let config = config.clone();
                     async move {
-                        match service(config, req).await {
+                        match service(config, req, addr).await {
                             Ok(r) => Ok(r),
                             Err(ServiceError::Hyper(e)) => Err(e),
                             Err(error) => Ok({
@@ -139,8 +143,10 @@ fn load_private_key(path: &Path) -> anyhow::Result<rustls::PrivateKey> {
 async fn service(
     config: Arc<Config>,
     mut req: Request<Incoming>,
+    addr: SocketAddr,
 ) -> Result<hyper::Response<BoxBody<bytes::Bytes, hyper::Error>>, ServiceError> {
-    debug!("<- {:?} {}", req.headers().get(HOST), req.uri());
+    let scheme_secure = req.uri().scheme() == Some(&Scheme::HTTPS);
+    debug!("{addr} ~> {:?} {}", req.headers().get(HOST), req.uri());
     *req.uri_mut() = Uri::builder()
         .scheme("http")
         .authority("backend")
@@ -154,6 +160,19 @@ async fn service(
         .build()
         .unwrap();
 
+    req.headers_mut().insert(
+        "x-forwarded-for",
+        HeaderValue::from_str(&format!("{addr}")).unwrap(),
+    );
+    req.headers_mut().insert(
+        "x-forwarded-proto",
+        if scheme_secure {
+            HeaderValue::from_static("https")
+        } else {
+            HeaderValue::from_static("http")
+        },
+    );
+
     let route = config
         .hosts
         .get(remove_port(