Add Methods to SocketDigest for Retrieving SO_ORIGINAL_DST Information

pingora-core/src/protocols/digest.rs

@@ -19,7 +19,9 @@ 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_v4, get_original_dest_v6, 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 +69,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 +79,7 @@ impl SocketDigest {
             raw_fd,
             peer_addr: OnceCell::new(),
             local_addr: OnceCell::new(),
+            original_dst: OnceCell::new(),
         }
     }
 
@@ -109,6 +114,22 @@ impl SocketDigest {
             None
         }
     }
+
+    pub fn original_dst(&self) -> Option<&SocketAddr> {
+        self.original_dst
+            .get_or_init(|| {
+                if let Some(SocketAddr::Inet(addr)) = self.local_addr() {
+                    let dst = match addr {
+                        std::net::SocketAddr::V4(_) => get_original_dest_v4(self.raw_fd),
+                        std::net::SocketAddr::V6(_) => get_original_dest_v6(self.raw_fd),
+                    };
+                    dst.ok().flatten().map(SocketAddr::Inet)
+                } else {
+                    None
+                }
+            })
+            .as_ref()
+    }
 }
 
 /// The interface to return timing information

pingora-core/src/protocols/l4/ext.rs

@@ -23,6 +23,8 @@ use pingora_error::{Error, ErrorType::*, OrErr, Result};
 use std::io::{self, ErrorKind};
 use std::mem;
 use std::net::SocketAddr;
+#[cfg(target_os = "linux")]
+use std::net::{SocketAddrV4, SocketAddrV6};
 use std::os::unix::io::{AsRawFd, RawFd};
 use std::time::Duration;
 use tokio::net::{TcpSocket, TcpStream, UnixStream};
@@ -335,6 +337,42 @@ 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_v4(fd: RawFd) -> io::Result<Option<SocketAddr>> {
+    let addr = get_opt_sized::<libc::sockaddr_in>(fd, libc::SOL_IP, libc::SO_ORIGINAL_DST)?;
+    Ok(Some(
+        SocketAddrV4::new(
+            u32::from_be(addr.sin_addr.s_addr).into(),
+            u16::from_be(addr.sin_port),
+        )
+        .into(),
+    ))
+}
+
+#[cfg(not(target_os = "linux"))]
+pub fn get_original_dest_v4(_fd: RawFd) -> io::Result<Option<SocketAddr>> {
+    Ok(None)
+}
+
+#[cfg(target_os = "linux")]
+pub fn get_original_dest_v6(fd: RawFd) -> io::Result<Option<SocketAddr>> {
+    let addr = get_opt_sized::<libc::sockaddr_in6>(fd, libc::SOL_IPV6, libc::IP6T_SO_ORIGINAL_DST)?;
+    Ok(Some(
+        SocketAddrV6::new(
+            addr.sin6_addr.s6_addr.into(),
+            u16::from_be(addr.sin6_port),
+            addr.sin6_flowinfo,
+            addr.sin6_scope_id,
+        )
+        .into(),
+    ))
+}
+
+#[cfg(not(target_os = "linux"))]
+pub fn get_original_dest_v6(_fd: RawFd) -> io::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.