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
34pub 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
91pub 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}