@@ -10,7 +10,7 @@ use crate::{
10
10
11
11
pub ( crate ) fn get_open_sockets ( ) -> OpenSockets {
12
12
let mut open_sockets = HashMap :: new ( ) ;
13
- let mut inode_to_procname = HashMap :: new ( ) ;
13
+ let mut inode_to_proc = HashMap :: new ( ) ;
14
14
15
15
if let Ok ( all_procs) = procfs:: process:: all_processes ( ) {
16
16
for process in all_procs. filter_map ( |res| res. ok ( ) ) {
@@ -20,47 +20,30 @@ pub(crate) fn get_open_sockets() -> OpenSockets {
20
20
let proc_info = ProcessInfo :: new ( & proc_name, stat. pid as u32 ) ;
21
21
for fd in fds. filter_map ( |res| res. ok ( ) ) {
22
22
if let FDTarget :: Socket ( inode) = fd. target {
23
- inode_to_procname . insert ( inode, proc_info. clone ( ) ) ;
23
+ inode_to_proc . insert ( inode, proc_info. clone ( ) ) ;
24
24
}
25
25
}
26
26
}
27
27
}
28
28
29
- if let Ok ( mut tcp) = procfs:: net:: tcp ( ) {
30
- if let Ok ( mut tcp6) = procfs:: net:: tcp6 ( ) {
31
- tcp. append ( & mut tcp6) ;
32
- }
33
- for entry in tcp. into_iter ( ) {
34
- if let Some ( proc_info) = inode_to_procname. get ( & entry. inode ) {
35
- open_sockets. insert (
36
- LocalSocket {
29
+ macro_rules! insert_proto {
30
+ ( $source: expr, $proto: expr) => {
31
+ let entries = $source. into_iter( ) . filter_map( |res| res. ok( ) ) . flatten( ) ;
32
+ for entry in entries {
33
+ if let Some ( proc_info) = inode_to_proc. get( & entry. inode) {
34
+ let socket = LocalSocket {
37
35
ip: entry. local_address. ip( ) ,
38
36
port: entry. local_address. port( ) ,
39
- protocol : Protocol :: Tcp ,
40
- } ,
41
- proc_info. clone ( ) ,
42
- ) ;
43
- } ;
44
- }
37
+ protocol: $proto ,
38
+ } ;
39
+ open_sockets . insert ( socket , proc_info. clone( ) ) ;
40
+ }
41
+ }
42
+ } ;
45
43
}
46
44
47
- if let Ok ( mut udp) = procfs:: net:: udp ( ) {
48
- if let Ok ( mut udp6) = procfs:: net:: udp6 ( ) {
49
- udp. append ( & mut udp6) ;
50
- }
51
- for entry in udp. into_iter ( ) {
52
- if let Some ( proc_info) = inode_to_procname. get ( & entry. inode ) {
53
- open_sockets. insert (
54
- LocalSocket {
55
- ip : entry. local_address . ip ( ) ,
56
- port : entry. local_address . port ( ) ,
57
- protocol : Protocol :: Udp ,
58
- } ,
59
- proc_info. clone ( ) ,
60
- ) ;
61
- } ;
62
- }
63
- }
45
+ insert_proto ! ( [ procfs:: net:: tcp( ) , procfs:: net:: tcp6( ) ] , Protocol :: Tcp ) ;
46
+ insert_proto ! ( [ procfs:: net:: udp( ) , procfs:: net:: udp6( ) ] , Protocol :: Udp ) ;
64
47
65
48
OpenSockets {
66
49
sockets_to_procs : open_sockets,
0 commit comments