x-forwaded-... headers in proxy module

1 files changed, 32 insertions, 4 deletions

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>();