netsock/sys/linux/
mod.rs

1mod ext;
2mod netlink_iter;
3mod procfs;
4
5use crate::error::Error;
6use crate::family::AddressFamilyFlags;
7use crate::protocol::ProtocolFlags;
8use crate::socket::SocketInfo;
9use crate::sys::linux::netlink_iter::*;
10use netlink_packet_sock_diag::{AF_INET, AF_INET6, IPPROTO_TCP, IPPROTO_UDP};
11pub use procfs::ProcessCache;
12
13struct ProcessAttached<I> {
14    inner: I,
15    cache: ProcessCache,
16}
17
18impl<I> Iterator for ProcessAttached<I>
19where
20    I: Iterator<Item = Result<SocketInfo, Error>>,
21{
22    type Item = Result<SocketInfo, Error>;
23
24    fn next(&mut self) -> Option<Self::Item> {
25        self.inner.next().map(|result| {
26            result.map(|socket_info| SocketInfo {
27                processes: self.cache.clone_processes(socket_info.inode),
28                ..socket_info
29            })
30        })
31    }
32}
33
34/// Iterates over socket information based on the specified address family and protocol flags.
35///
36/// This function provides an iterator over `SocketInfo` structures, allowing the caller to
37/// iterate through sockets filtered by address family and protocol criteria. It's a higher-level
38/// abstraction over the system's netstat information.
39///
40/// # Parameters
41/// - `af_flags`: An `AddressFamilyFlags` enum specifying the address families to filter by.
42///   This can include flags like `AF_INET` for IPv4 or `AF_INET6` for IPv6.
43/// - `proto_flags`: A `ProtocolFlags` enum specifying the protocols to filter by.
44///   This can include flags like `TCP` or `UDP`.
45///
46/// # Returns
47/// A `Result` containing an iterator over `Result<SocketInfo, Error>`. Each item in the iterator
48/// is a `Result` that either contains a `SocketInfo` struct with details about a socket, or an
49/// `Error` indicating a problem encountered while fetching the socket information.
50///
51/// # Errors
52/// Returns an `Error` if there is a failure in fetching the netstat information, including
53/// failures related to invalid parameters, system call failures, or other OS-level issues.
54///
55/// # Examples
56/// ```
57/// use netsock::family::AddressFamilyFlags;
58/// use netsock::iter_sockets;
59/// use netsock::protocol::ProtocolFlags;
60///
61/// let af_flags = AddressFamilyFlags::IPV4 | AddressFamilyFlags::IPV6;
62/// let proto_flags = ProtocolFlags::TCP | ProtocolFlags::UDP;
63///
64/// if let Ok(socket_iter) = iter_sockets(af_flags, proto_flags) {
65///     for socket_info in socket_iter {
66///         match socket_info {
67///             Ok(info) => println!("Found socket: {:?}", info),
68///             Err(e) => eprintln!("Error fetching socket info: {:?}", e),
69///         }
70///     }
71/// }
72/// ```
73pub fn iter_sockets(
74    af_flags: AddressFamilyFlags,
75    proto_flags: ProtocolFlags,
76) -> Result<impl Iterator<Item = Result<SocketInfo, Error>>, Error> {
77    let sockets = iter_sockets_without_processes(af_flags, proto_flags)?;
78    let cache = ProcessCache::snapshot()?;
79    Ok(attach_processes(sockets, cache))
80}
81
82pub fn iter_sockets_with_cache(
83    af_flags: AddressFamilyFlags,
84    proto_flags: ProtocolFlags,
85    cache: ProcessCache,
86) -> Result<impl Iterator<Item = Result<SocketInfo, Error>>, Error> {
87    let sockets = iter_sockets_without_processes(af_flags, proto_flags)?;
88    Ok(attach_processes(sockets, cache))
89}
90
91/// Iterates over socket information based on the specified address family and protocol flags.
92///
93/// without process info.
94pub fn iter_sockets_without_processes(
95    af_flags: AddressFamilyFlags,
96    proto_flags: ProtocolFlags,
97) -> Result<impl Iterator<Item = Result<SocketInfo, Error>>, Error> {
98    let ipv4 = af_flags.contains(AddressFamilyFlags::IPV4);
99    let ipv6 = af_flags.contains(AddressFamilyFlags::IPV6);
100    let tcp = proto_flags.contains(ProtocolFlags::TCP);
101    let udp = proto_flags.contains(ProtocolFlags::UDP);
102    let mut iterators = Vec::with_capacity(4);
103    if ipv4 {
104        if tcp {
105            iterators.push(NetlinkIterator::new(AF_INET as u8, IPPROTO_TCP as u8)?);
106        }
107        if udp {
108            iterators.push(NetlinkIterator::new(AF_INET as u8, IPPROTO_UDP as u8)?);
109        }
110    }
111    if ipv6 {
112        if tcp {
113            iterators.push(NetlinkIterator::new(AF_INET6 as u8, IPPROTO_TCP as u8)?);
114        }
115        if udp {
116            iterators.push(NetlinkIterator::new(AF_INET6 as u8, IPPROTO_UDP as u8)?);
117        }
118    }
119    Ok(iterators.into_iter().flatten())
120}
121
122fn attach_processes<I>(sockets_info: I, cache: ProcessCache) -> ProcessAttached<I>
123where
124    I: Iterator<Item = Result<SocketInfo, Error>>,
125{
126    ProcessAttached {
127        inner: sockets_info,
128        cache,
129    }
130}