use std::collections::HashMap;
use std::io;
use std::net::TcpListener;
use std::os::fd::{AsRawFd, OwnedFd, RawFd};
use std::path::{Path, PathBuf};
use std::sync::Mutex;
pub fn bind_listener(bind_addr: &str) -> io::Result<OwnedFd> {
let listener = TcpListener::bind(bind_addr)?;
Ok(OwnedFd::from(listener))
}
#[cfg(target_os = "linux")]
pub fn bind_listener_in_netns(bind_addr: &str, netns_path: &Path) -> io::Result<OwnedFd> {
use nix::sched::{setns, CloneFlags};
use std::os::fd::AsFd;
let bind_addr = bind_addr.to_string();
let netns_path = netns_path.to_path_buf();
std::thread::scope(|scope| {
scope
.spawn(move || -> io::Result<OwnedFd> {
let ns = std::fs::File::open(&netns_path).map_err(|e| {
io::Error::new(
e.kind(),
format!("opening netns {}: {e}", netns_path.display()),
)
})?;
setns(ns.as_fd(), CloneFlags::CLONE_NEWNET)
.map_err(|e| io::Error::other(format!("setns(CLONE_NEWNET): {e}")))?;
let listener = TcpListener::bind(&bind_addr)?;
Ok(OwnedFd::from(listener))
})
.join()
.map_err(|_| io::Error::other("netns bind thread panicked"))?
})
}
#[cfg(unix)]
pub fn send_listener_fds(upgrade_sock: &Path, binds: &[(&str, RawFd)]) -> io::Result<()> {
use nix::sys::socket::{
connect, sendmsg, socket, AddressFamily, ControlMessage, MsgFlags, SockFlag, SockType,
UnixAddr,
};
use std::time::Duration;
if binds.is_empty() {
return Ok(());
}
let payload = binds
.iter()
.map(|(key, _)| *key)
.collect::<Vec<_>>()
.join(" ");
let raw_fds: Vec<RawFd> = binds.iter().map(|(_, fd)| *fd).collect();
let addr = UnixAddr::new(upgrade_sock).map_err(|e| {
io::Error::other(format!("upgrade sock addr {}: {e}", upgrade_sock.display()))
})?;
let sock = socket(
AddressFamily::Unix,
SockType::Stream,
SockFlag::empty(),
None,
)
.map_err(|e| io::Error::other(format!("socket(AF_UNIX): {e}")))?;
const MAX_RETRY: usize = 10;
const RETRY_INTERVAL: Duration = Duration::from_millis(200);
let mut attempt = 0;
loop {
match connect(sock.as_raw_fd(), &addr) {
Ok(()) => break,
Err(nix::errno::Errno::ENOENT)
| Err(nix::errno::Errno::ECONNREFUSED)
| Err(nix::errno::Errno::EACCES)
if attempt < MAX_RETRY =>
{
attempt += 1;
std::thread::sleep(RETRY_INTERVAL);
}
Err(e) => {
return Err(io::Error::other(format!(
"connect({}): {e}",
upgrade_sock.display()
)))
}
}
}
let iov = [io::IoSlice::new(payload.as_bytes())];
let cmsg = [ControlMessage::ScmRights(&raw_fds)];
sendmsg::<UnixAddr>(sock.as_raw_fd(), &iov, &cmsg, MsgFlags::empty(), None)
.map_err(|e| io::Error::other(format!("sendmsg(SCM_RIGHTS): {e}")))?;
Ok(())
}
struct HeldSockets {
binds: Vec<(String, OwnedFd)>,
#[cfg(target_os = "linux")]
_netns: Option<OwnedFd>,
netns: Option<PathBuf>,
}
#[derive(Default)]
pub struct SocketCustodian {
held: Mutex<HashMap<String, HeldSockets>>,
}
impl std::fmt::Debug for SocketCustodian {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let held = self.held.lock().map(|h| h.len()).unwrap_or(0);
f.debug_struct("SocketCustodian")
.field("workloads_held", &held)
.finish()
}
}
impl SocketCustodian {
pub fn new() -> Self {
Self::default()
}
pub fn bind_and_hold(
&self,
ident: &str,
bind_addr: &str,
netns_path: Option<&Path>,
) -> io::Result<()> {
let fd = match netns_path {
#[cfg(target_os = "linux")]
Some(path) => bind_listener_in_netns(bind_addr, path)?,
#[cfg(not(target_os = "linux"))]
Some(_) => {
return Err(io::Error::other(
"netns-scoped socket custody is only supported on Linux",
))
}
None => bind_listener(bind_addr)?,
};
#[cfg(target_os = "linux")]
let netns_fd = match netns_path {
Some(path) => Some(OwnedFd::from(std::fs::File::open(path)?)),
None => None,
};
let mut held = self.held.lock().unwrap();
let entry = held
.entry(ident.to_string())
.or_insert_with(|| HeldSockets {
binds: Vec::new(),
#[cfg(target_os = "linux")]
_netns: None,
netns: None,
});
if entry.binds.iter().any(|(b, _)| b == bind_addr) {
return Err(io::Error::other(format!(
"socket custodian already holds {bind_addr} for workload {ident}"
)));
}
entry.binds.push((bind_addr.to_string(), fd));
#[cfg(target_os = "linux")]
if entry._netns.is_none() {
entry._netns = netns_fd;
}
if entry.netns.is_none() {
entry.netns = netns_path.map(Path::to_path_buf);
}
Ok(())
}
pub fn holds(&self, ident: &str) -> bool {
self.held.lock().unwrap().contains_key(ident)
}
pub fn held_netns(&self, ident: &str) -> Option<PathBuf> {
self.held.lock().unwrap().get(ident)?.netns.clone()
}
#[cfg(unix)]
pub fn hand_off(&self, ident: &str, upgrade_sock: &Path) -> io::Result<()> {
let held = self.held.lock().unwrap();
let entry = held.get(ident).ok_or_else(|| {
io::Error::other(format!(
"socket custodian holds nothing for workload {ident}"
))
})?;
let binds: Vec<(&str, RawFd)> = entry
.binds
.iter()
.map(|(key, fd)| (key.as_str(), fd.as_raw_fd()))
.collect();
send_listener_fds(upgrade_sock, &binds)
}
pub fn release(&self, ident: &str) {
self.held.lock().unwrap().remove(ident);
}
pub fn held_binds(&self, ident: &str) -> Option<Vec<String>> {
self.held
.lock()
.unwrap()
.get(ident)
.map(|h| h.binds.iter().map(|(b, _)| b.clone()).collect())
}
pub fn held_raw_fds(&self, ident: &str) -> Option<Vec<RawFd>> {
self.held
.lock()
.unwrap()
.get(ident)
.map(|h| h.binds.iter().map(|(_, fd)| fd.as_raw_fd()).collect())
}
}
#[cfg(all(test, unix))]
mod tests {
use super::*;
use std::os::fd::FromRawFd;
use std::os::unix::net::UnixListener;
fn receive_fds(listener: &UnixListener) -> (String, Vec<RawFd>) {
use nix::sys::socket::{recvmsg, ControlMessageOwned, MsgFlags, UnixAddr};
let (stream, _) = listener.accept().unwrap();
let mut payload = [0u8; 2048];
let mut iov = [io::IoSliceMut::new(&mut payload)];
let mut cmsg_buf = nix::cmsg_space!([RawFd; 32]);
let msg = recvmsg::<UnixAddr>(
stream.as_raw_fd(),
&mut iov,
Some(&mut cmsg_buf),
MsgFlags::empty(),
)
.unwrap();
let mut fds = Vec::new();
for cmsg in msg.cmsgs().unwrap() {
if let ControlMessageOwned::ScmRights(v) = cmsg {
fds.extend(v);
}
}
let bytes = msg.bytes;
let keys = String::from_utf8(payload[..bytes].to_vec()).unwrap();
(keys, fds)
}
#[test]
fn send_listener_fd_round_trips_pingora_wire_format() {
let listener_fd = bind_listener("127.0.0.1:0").unwrap();
let local = {
let borrowed = unsafe { std::net::TcpListener::from_raw_fd(listener_fd.as_raw_fd()) };
let addr = borrowed.local_addr().unwrap();
std::mem::forget(borrowed); addr
};
let key = local.to_string();
let dir = tempfile::tempdir().unwrap();
let sock_path = dir.path().join("upgrade.sock");
let ulistener = UnixListener::bind(&sock_path).unwrap();
let sock_path_recv = sock_path.clone();
let key_expect = key.clone();
let recv = std::thread::spawn(move || {
let _ = &sock_path_recv;
receive_fds(&ulistener)
});
send_listener_fds(&sock_path, &[(key.as_str(), listener_fd.as_raw_fd())]).unwrap();
let (got_keys, got_fds) = recv.join().unwrap();
assert_eq!(
got_keys, key_expect,
"bind-address key must survive verbatim"
);
assert_eq!(got_fds.len(), 1, "exactly one fd transferred");
let received = unsafe { std::net::TcpListener::from_raw_fd(got_fds[0]) };
assert_eq!(
received.local_addr().unwrap().port(),
local.port(),
"receiver adopted kamaji's listening socket"
);
drop(received);
}
#[test]
fn custodian_holds_and_releases() {
let cust = SocketCustodian::new();
assert!(!cust.holds("ingress"));
cust.bind_and_hold("ingress", "127.0.0.1:0", None).unwrap();
assert!(cust.holds("ingress"));
assert_eq!(cust.held_binds("ingress").unwrap().len(), 1);
let held_key = cust.held_binds("ingress").unwrap()[0].clone();
let dup = cust.bind_and_hold("ingress", &held_key, None);
assert!(dup.is_err(), "re-binding the same key is rejected");
cust.release("ingress");
assert!(!cust.holds("ingress"));
cust.release("ingress");
}
#[test]
fn a_host_bound_listener_reports_no_netns() {
let cust = SocketCustodian::new();
cust.bind_and_hold("ingress", "127.0.0.1:0", None).unwrap();
assert!(cust.holds("ingress"));
assert_eq!(cust.held_netns("ingress"), None);
assert_eq!(cust.held_netns("never-deployed"), None);
cust.release("ingress");
}
#[test]
fn hand_off_without_holding_errs() {
let cust = SocketCustodian::new();
let dir = tempfile::tempdir().unwrap();
let err = cust
.hand_off("nobody", &dir.path().join("x.sock"))
.unwrap_err();
assert!(err.to_string().contains("holds nothing"));
}
}