From 88477e0470cd23f11ebb1b278cca9944b569c40a Mon Sep 17 00:00:00 2001
From: "Christian W. Zuckschwerdt" <christian@zuckschwerdt.org>
Date: Sun, 3 Nov 2019 14:13:59 +0100
Subject: [PATCH] Add Handshake::forwarded_addr() to get a trusted forwarded
 address

The X-Forwarded-For header is implicitly untrustworthy,
this adds `Handshake::forwarded_addr()` and `Request::forwarded_addr()`
to get a trusted address at exactly specified proxy depth.
---
 src/handshake.rs | 65 ++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 65 insertions(+)

diff --git a/src/handshake.rs b/src/handshake.rs
index b7520bd..3b6457d 100644
--- a/src/handshake.rs
+++ b/src/handshake.rs
@@ -107,6 +107,32 @@ impl Handshake {
             }
         }))
     }
+
+    /// Get a forwarded IP address of the remote connection.
+    ///
+    /// This is the preferred method of obtaining the client's trusted IP address.
+    /// It will attempt to retrieve the forwarded IP address for `trusted` depth
+    /// of reverse-proxies based on request headers, `trusted=0` returns the direct
+    /// address of the peer.
+    ///
+    /// # Note
+    /// This assumes that the peer is a client. If you are implementing a
+    /// WebSocket client and want to obtain the address of the server, use
+    /// `Handshake::peer_addr` instead.
+    ///
+    /// This method does not ensure that the address is a valid IP address.
+    #[allow(dead_code)]
+    fn forwarded_addr(&self, trusted: usize) -> Result<Option<String>> {
+        if trusted == 0 {
+            if let Some(addr) = self.peer_addr {
+                Ok(Some(addr.ip().to_string()))
+            } else {
+                Ok(None)
+            }
+        } else {
+            Ok(self.request.forwarded_addr(trusted)?.map(String::from))
+        }
+    }
 }
 
 /// The handshake request.
@@ -323,6 +349,45 @@ impl Request {
         Ok(None)
     }
 
+    /// Get the trusted IP address of the client.
+    ///
+    /// This method will attempt to retrieve a specific forwarded IP address of the requester
+    /// in the following manner:
+    ///
+    /// If the `X-Forwarded-For` header exists, this method will return the `trusted` right most
+    /// address in the list.
+    ///
+    /// If the [Forwarded HTTP Header Field](https://tools.ietf.org/html/rfc7239) exits,
+    /// this method will return the `trusted` right most address indicated by the `for` parameter,
+    /// if it exists.
+    ///
+    /// # Note
+    /// This method does not ensure that the address is a valid IP address.
+    #[allow(dead_code)]
+    fn forwarded_addr(&self, trusted: usize) -> Result<Option<&str>> {
+        if trusted == 0 {
+            return Ok(None);
+        }
+        if let Some(x_forward) = self.header("x-forwarded-for") {
+            return Ok(from_utf8(x_forward)?.rsplit(',').nth(trusted - 1));
+        }
+
+        // We only care about the first forwarded header, so header is ok
+        if let Some(forward) = self.header("forwarded") {
+            if let Some(_for) = from_utf8(forward)?
+                .split(';')
+                .find(|f| f.trim().starts_with("for"))
+            {
+                if let Some(_for_eq) = _for.trim().rsplit(',').nth(trusted - 1) {
+                    let mut it = _for_eq.split('=');
+                    it.next();
+                    return Ok(it.next());
+                }
+            }
+        }
+        Ok(None)
+    }
+
     /// Attempt to parse an HTTP request from a buffer. If the buffer does not contain a complete
     /// request, this will return `Ok(None)`.
     pub fn parse(buf: &[u8]) -> Result<Option<Request>> {