Skip to content

Commit

Permalink
Add Methods to SocketDigest for Retrieving SO_ORIGINAL_DST Information
Browse files Browse the repository at this point in the history
  • Loading branch information
ermakov-oleg committed Sep 20, 2024
1 parent 6c8e7aa commit a3f2aff
Show file tree
Hide file tree
Showing 2 changed files with 55 additions and 1 deletion.
16 changes: 15 additions & 1 deletion pingora-core/src/protocols/digest.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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 {
Expand All @@ -75,6 +77,7 @@ impl SocketDigest {
raw_fd,
peer_addr: OnceCell::new(),
local_addr: OnceCell::new(),
original_dst: OnceCell::new(),
}
}

Expand Down Expand Up @@ -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
Expand Down
40 changes: 40 additions & 0 deletions pingora-core/src/protocols/l4/ext.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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")
.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.
Expand Down

0 comments on commit a3f2aff

Please sign in to comment.