From 2a62acd60b2c24c82ca1119f62a27c08508bf730 Mon Sep 17 00:00:00 2001
From: ermakov-oleg <ermakovolegs@gmail.com>
Date: Fri, 20 Sep 2024 21:35:05 +0000
Subject: Add Methods to SocketDigest for Retrieving SO_ORIGINAL_DST
 Information --- Merge branch 'cloudflare:main' into original-dst

Co-authored-by: ermakov-oleg <ermakov-oleg@users.noreply.github.com>
Includes-commit: a3f2affceb1fcc41742c72088f993e05600abd2f
Includes-commit: a73893b243b68108033ecda26b0d36939ad36891
Replicated-from: https://github.com/cloudflare/pingora/pull/387
---
 .bleep                               |  2 +-
 pingora-core/src/protocols/digest.rs | 16 ++++++++++++++-
 pingora-core/src/protocols/l4/ext.rs | 40 ++++++++++++++++++++++++++++++++++++
 3 files changed, 56 insertions(+), 2 deletions(-)

diff --git a/.bleep b/.bleep
index e659893..1ac3b4f 100644
--- a/.bleep
+++ b/.bleep
@@ -1 +1 @@
-08724db8653366f6eee60a3236ed9c2987402fc0
\ No newline at end of file
+73cbdf49b1e1279ebba51525cad01952d4dc04d2
\ No newline at end of file
diff --git a/pingora-core/src/protocols/digest.rs b/pingora-core/src/protocols/digest.rs
index 88720e5..9e23741 100644
--- a/pingora-core/src/protocols/digest.rs
+++ b/pingora-core/src/protocols/digest.rs
@@ -19,7 +19,7 @@ use std::time::{Duration, SystemTime};
 
 use once_cell::sync::OnceCell;
 
-use super::l4::ext::{get_recv_buf, get_tcp_info, TCP_INFO};
+use super::l4::ext::{get_original_dest, get_recv_buf, get_tcp_info, TCP_INFO};
 use super::l4::socket::SocketAddr;
 use super::raw_connect::ProxyDigest;
 use super::tls::digest::SslDigest;
@@ -70,6 +70,8 @@ pub struct SocketDigest {
     pub peer_addr: OnceCell<Option<SocketAddr>>,
     /// Local socket address
     pub local_addr: OnceCell<Option<SocketAddr>>,
+    /// Original destination address
+    pub original_dst: OnceCell<Option<SocketAddr>>,
 }
 
 impl SocketDigest {
@@ -79,6 +81,7 @@ impl SocketDigest {
             raw_fd,
             peer_addr: OnceCell::new(),
             local_addr: OnceCell::new(),
+            original_dst: OnceCell::new(),
         }
     }
 
@@ -158,6 +161,17 @@ impl SocketDigest {
             None
         }
     }
+
+    pub fn original_dst(&self) -> Option<&SocketAddr> {
+        self.original_dst
+            .get_or_init(|| {
+                get_original_dest(self.raw_fd)
+                    .ok()
+                    .flatten()
+                    .map(SocketAddr::Inet)
+            })
+            .as_ref()
+    }
 }
 
 /// The interface to return timing information
diff --git a/pingora-core/src/protocols/l4/ext.rs b/pingora-core/src/protocols/l4/ext.rs
index 5437198..fae93f7 100644
--- a/pingora-core/src/protocols/l4/ext.rs
+++ b/pingora-core/src/protocols/l4/ext.rs
@@ -378,6 +378,46 @@ pub fn get_socket_cookie(_fd: RawFd) -> io::Result<u64> {
     Ok(0) // SO_COOKIE is a Linux concept
 }
 
+#[cfg(target_os = "linux")]
+pub fn get_original_dest(fd: RawFd) -> Result<Option<SocketAddr>> {
+    use super::socket;
+    use pingora_error::OkOrErr;
+    use std::net::{SocketAddrV4, SocketAddrV6};
+
+    let sock = socket::SocketAddr::from_raw_fd(fd, false);
+    let addr = sock
+        .as_ref()
+        .and_then(|s| s.as_inet())
+        .or_err(SocketError, "failed get original dest, invalid IP socket")?;
+
+    let dest = if addr.is_ipv4() {
+        get_opt_sized::<libc::sockaddr_in>(fd, libc::SOL_IP, libc::SO_ORIGINAL_DST).map(|addr| {
+            SocketAddr::V4(SocketAddrV4::new(
+                u32::from_be(addr.sin_addr.s_addr).into(),
+                u16::from_be(addr.sin_port),
+            ))
+        })
+    } else {
+        get_opt_sized::<libc::sockaddr_in6>(fd, libc::SOL_IPV6, libc::IP6T_SO_ORIGINAL_DST).map(
+            |addr| {
+                SocketAddr::V6(SocketAddrV6::new(
+                    addr.sin6_addr.s6_addr.into(),
+                    u16::from_be(addr.sin6_port),
+                    addr.sin6_flowinfo,
+                    addr.sin6_scope_id,
+                ))
+            },
+        )
+    };
+    dest.or_err(SocketError, "failed to get original dest")
+        .map(Some)
+}
+
+#[cfg(not(target_os = "linux"))]
+pub fn get_original_dest(_fd: RawFd) -> Result<Option<SocketAddr>> {
+    Ok(None)
+}
+
 /// connect() to the given address while optionally binding to the specific source address and port range.
 ///
 /// The `set_socket` callback can be used to tune the socket before `connect()` is called.
-- 
cgit v1.2.3