feat: add CloudFront header to Insecure and Secure

imbolc · Apr 8, 2024 · ed4052a · ed4052a
1 parent dcbe03e
commit ed4052a
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 2 deletions.
diff --git a/src/insecure.rs b/src/insecure.rs
@@ -1,3 +1,5 @@
+#[cfg(feature = "aws-cloudfront")]
+use crate::rudimental::CloudFrontViewerAddress;
 use crate::rudimental::{
     CfConnectingIp, FlyClientIp, Forwarded, MultiIpHeader, SingleIpHeader, TrueClientIp,
     XForwardedFor, XRealIp,
@@ -43,12 +45,17 @@ impl InsecureClientIp {
         headers: &HeaderMap<HeaderValue>,
         extensions: &Extensions,
     ) -> Result<Self, Rejection> {
-        XForwardedFor::maybe_leftmost_ip(headers)
+        let maybe_ip = XForwardedFor::maybe_leftmost_ip(headers)
             .or_else(|| Forwarded::maybe_leftmost_ip(headers))
             .or_else(|| XRealIp::maybe_ip_from_headers(headers))
             .or_else(|| FlyClientIp::maybe_ip_from_headers(headers))
             .or_else(|| TrueClientIp::maybe_ip_from_headers(headers))
-            .or_else(|| CfConnectingIp::maybe_ip_from_headers(headers))
+            .or_else(|| CfConnectingIp::maybe_ip_from_headers(headers));
+
+        #[cfg(feature = "aws-cloudfront")]
+        let maybe_ip = maybe_ip.or_else(|| CloudFrontViewerAddress::maybe_ip_from_headers(headers));
+
+        maybe_ip
             .or_else(|| maybe_connect_info(extensions))
             .map(Self)
             .ok_or((

diff --git a/src/secure.rs b/src/secure.rs
@@ -1,3 +1,5 @@
+#[cfg(feature = "aws-cloudfront")]
+use crate::rudimental::CloudFrontViewerAddress;
 use crate::rudimental::{
     CfConnectingIp, FlyClientIp, Forwarded, MultiIpHeader, SingleIpHeader, StringRejection,
     TrueClientIp, XForwardedFor, XRealIp,
@@ -44,6 +46,9 @@ pub enum SecureClientIpSource {
     CfConnectingIp,
     /// IP from the [`axum::extract::ConnectInfo`]
     ConnectInfo,
+    /// IP from the `CloudFront-Viewer-Address` header
+    #[cfg(feature = "aws-cloudfront")]
+    CloudFrontViewerAddress,
 }
 
 impl SecureClientIpSource {
@@ -77,6 +82,8 @@ impl FromStr for SecureClientIpSource {
             "TrueClientIp" => Self::TrueClientIp,
             "CfConnectingIp" => Self::CfConnectingIp,
             "ConnectInfo" => Self::ConnectInfo,
+            #[cfg(feature = "aws-cloudfront")]
+            "CloudFrontViewerAddress" => Self::CloudFrontViewerAddress,
             _ => return Err(ParseSecureClientIpSourceError(s.to_string())),
         })
     }
@@ -100,6 +107,10 @@ impl SecureClientIp {
             SecureClientIpSource::FlyClientIp => FlyClientIp::ip_from_headers(headers),
             SecureClientIpSource::TrueClientIp => TrueClientIp::ip_from_headers(headers),
             SecureClientIpSource::CfConnectingIp => CfConnectingIp::ip_from_headers(headers),
+            #[cfg(feature = "aws-cloudfront")]
+            SecureClientIpSource::CloudFrontViewerAddress => {
+                CloudFrontViewerAddress::ip_from_headers(headers)
+            }
             SecureClientIpSource::ConnectInfo => extensions
                 .get::<ConnectInfo<SocketAddr>>()
                 .map(|ConnectInfo(addr)| addr.ip())