From a9829f5325b2b5137816880a6c82c39e131ca5a3 Mon Sep 17 00:00:00 2001
From: ermakov-oleg <ermakovolegs@gmail.com>
Date: Mon, 16 Sep 2024 17:03:28 +0200
Subject: [PATCH] Add Methods to SocketDigest for Retrieving SO_ORIGINAL_DST
 Information

---
 pingora-core/src/protocols/digest.rs | 16 ++++++++++-
 pingora-core/src/protocols/l4/ext.rs | 40 ++++++++++++++++++++++++++++
 2 files changed, 55 insertions(+), 1 deletion(-)

diff --git a/pingora-core/src/protocols/digest.rs b/pingora-core/src/protocols/digest.rs
index 3150306fd..f23d5820d 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;
@@ -67,6 +67,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 {
@@ -75,6 +77,7 @@ impl SocketDigest {
             raw_fd,
             peer_addr: OnceCell::new(),
             local_addr: OnceCell::new(),
+            original_dst: OnceCell::new(),
         }
     }
 
@@ -109,6 +112,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 4b65f84f3..f0455fbdc 100644
--- a/pingora-core/src/protocols/l4/ext.rs
+++ b/pingora-core/src/protocols/l4/ext.rs
@@ -335,6 +335,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")
+        .into()
+}
+
+#[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.