use std::os::fd::{AsRawFd, FromRawFd, OwnedFd, RawFd};
use std::time::Duration;
use crate::consts::*;
use crate::error::{Error, Result};
pub struct ProcConnector {
fd: OwnedFd,
}
impl ProcConnector {
pub fn new() -> Result<Self> {
let fd = create_socket()?;
let connector = Self { fd };
connector.bind()?;
connector.subscribe()?;
Ok(connector)
}
fn bind(&self) -> Result<()> {
let mut sa: libc::sockaddr_nl = unsafe { std::mem::zeroed() };
sa.nl_family = libc::AF_NETLINK as u16;
sa.nl_pid = 0; sa.nl_groups = CN_IDX_PROC;
let ret = unsafe {
libc::bind(
self.fd.as_raw_fd(),
&sa as *const libc::sockaddr_nl as *const libc::sockaddr,
std::mem::size_of::<libc::sockaddr_nl>() as u32,
)
};
if ret < 0 {
return Err(Error::Os(std::io::Error::last_os_error()));
}
Ok(())
}
pub fn subscribe(&self) -> Result<()> {
self.send_mcast_op(PROC_CN_MCAST_LISTEN)
}
pub fn unsubscribe(&self) -> Result<()> {
self.send_mcast_op(PROC_CN_MCAST_IGNORE)
}
fn send_mcast_op(&self, op: u32) -> Result<()> {
let nlmsg_payload_len = SIZE_CN_MSG + std::mem::size_of::<u32>();
let nlmsg_len = nlmsg_length(nlmsg_payload_len);
let mut buf = vec![0u8; nlmsg_len];
let pid = std::process::id();
let hdr = &mut buf[..SIZE_NLMSGHDR];
hdr[0..4].copy_from_slice(&(nlmsg_len as u32).to_ne_bytes()); hdr[4..6].copy_from_slice(&NLMSG_DONE.to_ne_bytes()); hdr[6..8].copy_from_slice(&0u16.to_ne_bytes()); hdr[8..12].copy_from_slice(&0u32.to_ne_bytes()); hdr[12..16].copy_from_slice(&pid.to_ne_bytes());
let cn_off = nlmsg_hdrlen();
let cn = &mut buf[cn_off..cn_off + SIZE_CN_MSG + 4];
cn[0..4].copy_from_slice(&CN_IDX_PROC.to_ne_bytes());
cn[4..8].copy_from_slice(&CN_VAL_PROC.to_ne_bytes());
cn[8..12].copy_from_slice(&0u32.to_ne_bytes());
cn[12..16].copy_from_slice(&0u32.to_ne_bytes());
cn[16..18]
.copy_from_slice(&(std::mem::size_of::<u32>() as u16).to_ne_bytes());
cn[18..20].copy_from_slice(&0u16.to_ne_bytes());
cn[20..24].copy_from_slice(&op.to_ne_bytes());
let iov = libc::iovec {
iov_base: buf.as_mut_ptr() as *mut libc::c_void,
iov_len: nlmsg_len,
};
let msg_hdr = libc::msghdr {
msg_name: std::ptr::null_mut(),
msg_namelen: 0,
msg_iov: &iov as *const libc::iovec as *mut libc::iovec,
msg_iovlen: 1,
msg_control: std::ptr::null_mut(),
msg_controllen: 0,
msg_flags: 0,
};
let ret = unsafe { libc::sendmsg(self.fd.as_raw_fd(), &msg_hdr, 0) };
if ret < 0 {
return Err(Error::Os(std::io::Error::last_os_error()));
}
Ok(())
}
pub fn recv_raw(&self, buf: &mut [u8]) -> Result<usize> {
let len = unsafe {
libc::recv(
self.fd.as_raw_fd(),
buf.as_mut_ptr() as *mut libc::c_void,
buf.len(),
0,
)
};
if len < 0 {
let err = std::io::Error::last_os_error();
return match err.raw_os_error() {
Some(libc::EINTR) => Err(Error::Interrupted),
_ => Err(Error::Os(err)),
};
}
if len == 0 {
return Err(Error::ConnectionClosed);
}
Ok(len as usize)
}
pub fn recv_raw_timeout(&self, buf: &mut [u8], timeout: Duration) -> Result<Option<usize>> {
let poll_fd = libc::pollfd {
fd: self.fd.as_raw_fd(),
events: libc::POLLIN,
revents: 0,
};
let timeout_ms = timeout
.as_millis()
.try_into()
.unwrap_or(libc::c_int::MAX);
let ret = unsafe { libc::poll(&poll_fd as *const libc::pollfd as *mut _, 1, timeout_ms) };
if ret < 0 {
let err = std::io::Error::last_os_error();
return match err.raw_os_error() {
Some(libc::EINTR) => Err(Error::Interrupted),
_ => Err(Error::Os(err)),
};
}
if ret == 0 {
return Ok(None);
}
self.recv_raw(buf).map(Some)
}
pub fn as_raw_fd(&self) -> RawFd {
self.fd.as_raw_fd()
}
}
impl Drop for ProcConnector {
fn drop(&mut self) {
let _ = self.unsubscribe();
}
}
fn create_socket() -> Result<OwnedFd> {
let fd = unsafe {
let fd = libc::socket(libc::PF_NETLINK, libc::SOCK_DGRAM, NETLINK_CONNECTOR);
if fd < 0 {
return Err(Error::Os(std::io::Error::last_os_error()));
}
OwnedFd::from_raw_fd(fd)
};
let val: libc::c_int = 1;
let ret = unsafe {
libc::setsockopt(
fd.as_raw_fd(),
libc::SOL_NETLINK,
NETLINK_NO_ENOBUFS,
&val as *const libc::c_int as *const libc::c_void,
std::mem::size_of::<libc::c_int>() as u32,
)
};
if ret < 0 {
let _ = std::io::Error::last_os_error();
}
Ok(fd)
}