mod ext;
mod netlink_iter;
mod procfs;
use crate::error::Error;
use crate::family::AddressFamilyFlags;
use crate::protocol::ProtocolFlags;
use crate::socket::SocketInfo;
use crate::sys::linux::netlink_iter::*;
use netlink_packet_sock_diag::{AF_INET, AF_INET6, IPPROTO_TCP, IPPROTO_UDP};
pub use procfs::ProcessCache;
struct ProcessAttached<I> {
inner: I,
cache: ProcessCache,
}
impl<I> Iterator for ProcessAttached<I>
where
I: Iterator<Item = Result<SocketInfo, Error>>,
{
type Item = Result<SocketInfo, Error>;
fn next(&mut self) -> Option<Self::Item> {
self.inner.next().map(|result| {
result.map(|socket_info| SocketInfo {
processes: self.cache.clone_processes(socket_info.inode),
..socket_info
})
})
}
}
pub fn iter_sockets(
af_flags: AddressFamilyFlags,
proto_flags: ProtocolFlags,
) -> Result<impl Iterator<Item = Result<SocketInfo, Error>>, Error> {
let sockets = iter_sockets_without_processes(af_flags, proto_flags)?;
let cache = ProcessCache::snapshot()?;
Ok(attach_processes(sockets, cache))
}
pub fn iter_sockets_with_cache(
af_flags: AddressFamilyFlags,
proto_flags: ProtocolFlags,
cache: ProcessCache,
) -> Result<impl Iterator<Item = Result<SocketInfo, Error>>, Error> {
let sockets = iter_sockets_without_processes(af_flags, proto_flags)?;
Ok(attach_processes(sockets, cache))
}
pub fn iter_sockets_without_processes(
af_flags: AddressFamilyFlags,
proto_flags: ProtocolFlags,
) -> Result<impl Iterator<Item = Result<SocketInfo, Error>>, Error> {
let ipv4 = af_flags.contains(AddressFamilyFlags::IPV4);
let ipv6 = af_flags.contains(AddressFamilyFlags::IPV6);
let tcp = proto_flags.contains(ProtocolFlags::TCP);
let udp = proto_flags.contains(ProtocolFlags::UDP);
let mut iterators = Vec::with_capacity(4);
if ipv4 {
if tcp {
iterators.push(NetlinkIterator::new(AF_INET, IPPROTO_TCP)?);
}
if udp {
iterators.push(NetlinkIterator::new(AF_INET, IPPROTO_UDP)?);
}
}
if ipv6 {
if tcp {
iterators.push(NetlinkIterator::new(AF_INET6, IPPROTO_TCP)?);
}
if udp {
iterators.push(NetlinkIterator::new(AF_INET6, IPPROTO_UDP)?);
}
}
Ok(iterators.into_iter().flatten())
}
fn attach_processes<I>(sockets_info: I, cache: ProcessCache) -> ProcessAttached<I>
where
I: Iterator<Item = Result<SocketInfo, Error>>,
{
ProcessAttached {
inner: sockets_info,
cache,
}
}