x-forwaded-... headers in proxy module

4 files changed, 47 insertions, 8 deletions

diff --git a/readme.md b/readme.md
index cd5e5d8..f241d66 100644
--- a/readme.md
+++ b/readme.md
@@ -110,6 +110,10 @@ themselves; in that case the request is passed on.
     added to the request. Connection upgrades are handled by direct forwarding
     of network traffic.
   - `backend`: socket address (string) to the backend server
+  - `set_real_ip`: Sets the `X-Real-IP` header. (boolean)
+  - `set_forwarded_for`: Sets the `X-Forwarded-For`, `X-Forwarded-Host`,
+    `X-Forwaded-Proto`, `X-Forwarded-Scheme` and `X-Forwarded-Port` headers.
+    (boolean)
 
 - **module `files`**
   - Provides a simple built-in fileserver. The server handles `accept-ranges`.
diff --git a/src/main.rs b/src/main.rs
index c3e8cc7..fe7856e 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -127,7 +127,7 @@ async fn serve_http(state: Arc<State>) -> Result<()> {
             debug!("connection from {addr}");
             let stream = TokioIo::new(stream);
             let state = state.clone();
-            tokio::spawn(async move { serve_stream(state, stream, addr).await });
+            tokio::spawn(async move { serve_stream(state, stream, addr, false).await });
         }
     }))
     .await;
@@ -163,7 +163,7 @@ async fn serve_https(state: Arc<State>) -> Result<()> {
             tokio::task::spawn(async move {
                 debug!("connection from {addr}");
                 match tls_acceptor.accept(stream).await {
-                    Ok(stream) => serve_stream(state, TokioIo::new(stream), addr).await,
+                    Ok(stream) => serve_stream(state, TokioIo::new(stream), addr, true).await,
                     Err(e) => warn!("error accepting tls: {e}"),
                 };
             });
@@ -178,6 +178,7 @@ pub async fn serve_stream<T: Unpin + Send + 'static + hyper::rt::Read + hyper::r
     state: Arc<State>,
     stream: T,
     addr: SocketAddr,
+    secure: bool,
 ) {
     if let Ok(_semaphore) = state.l_incoming.try_acquire() {
         let builder = hyper_util::server::conn::auto::Builder::new(TokioExecutor::new());
@@ -187,7 +188,7 @@ pub async fn serve_stream<T: Unpin + Send + 'static + hyper::rt::Read + hyper::r
                 let state = state.clone();
                 async move {
                     let config = state.config.read().await.clone();
-                    match service(state, config, req, addr).await {
+                    match service(state, config, req, addr, secure).await {
                         Ok(r) => Ok(r),
                         Err(ServiceError::Hyper(e)) => Err(e),
                         Err(error) => Ok({
@@ -220,6 +221,7 @@ async fn service(
     config: Arc<Config>,
     mut request: Request<Incoming>,
     addr: SocketAddr,
+    secure: bool,
 ) -> Result<hyper::Response<BoxBody<bytes::Bytes, ServiceError>>, ServiceError> {
     // move uri authority used in HTTP/2 to Host header field
     {
@@ -241,7 +243,11 @@ async fn service(
         request.uri()
     );
 
-    let mut context = NodeContext { addr, state };
+    let mut context = NodeContext {
+        addr,
+        state,
+        secure,
+    };
     let mut resp = config
         .handler
         .handle(&mut context, request.map(|body| body.boxed()))
diff --git a/src/modules/mod.rs b/src/modules/mod.rs
index 97dee1a..d3e0dd3 100644
--- a/src/modules/mod.rs
+++ b/src/modules/mod.rs
@@ -55,6 +55,7 @@ pub static MODULES: &[&dyn NodeKind] = &[
 pub struct NodeContext {
     pub state: Arc<State>,
     pub addr: SocketAddr,
+    pub secure: bool,
 }
 
 pub trait NodeKind: Send + Sync + 'static {
diff --git a/src/modules/proxy.rs b/src/modules/proxy.rs
index 763129f..b1af3d9 100644
--- a/src/modules/proxy.rs
+++ b/src/modules/proxy.rs
@@ -15,9 +15,16 @@ pub struct ProxyKind;
 
 #[derive(Debug, Deserialize)]
 struct Proxy {
+    set_forwarded_for: bool,
+    #[serde(default = "ret_true")]
+    set_real_ip: bool,
     backend: SocketAddr,
 }
 
+fn ret_true() -> bool {
+    true
+}
+
 impl NodeKind for ProxyKind {
     fn name(&self) -> &'static str {
         "proxy"
@@ -33,10 +40,31 @@ impl Node for Proxy {
         mut request: NodeRequest,
     ) -> Pin<Box<dyn Future<Output = Result<NodeResponse, ServiceError>> + Send + Sync + 'a>> {
         Box::pin(async move {
-            request.headers_mut().insert(
-                "x-real-ip",
-                HeaderValue::from_str(&format!("{}", context.addr.ip())).unwrap(),
-            );
+            if self.set_real_ip {
+                request.headers_mut().insert(
+                    "x-real-ip",
+                    HeaderValue::from_str(&format!("{}", context.addr.ip())).unwrap(),
+                );
+            }
+            if self.set_forwarded_for {
+                request.headers_mut().insert(
+                    "x-forwarded-for",
+                    HeaderValue::from_str(&format!("{}", context.addr.ip())).unwrap(),
+                );
+                request.headers_mut().insert(
+                    "x-forwarded-port",
+                    HeaderValue::from_str(&context.addr.port().to_string()).unwrap(),
+                );
+                let scheme =
+                    HeaderValue::from_str(if context.secure { "https" } else { "http" }).unwrap();
+                request
+                    .headers_mut()
+                    .insert("x-forwarded-scheme", scheme.clone());
+                request.headers_mut().insert("x-forwarded-proto", scheme);
+                if let Some(host) = request.headers().get("host").cloned() {
+                    request.headers_mut().insert("x-forwarded-host", host);
+                }
+            }
 
             let on_upgrade_downstream = request.extensions_mut().remove::<OnUpgrade>();