Fix IPv6 socket detect logic (#383)

* Fix IPv6 socket detect logic

* Write changelog

2 files changed, 17 insertions, 33 deletions

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 5ee803b..93fdb3c 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -11,6 +11,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 * CI: use GitHub API to exempt dependabot from changelog requirement #378 - @cyqsimon
 * Remove unnecessary logging synchronisation #381 - @cyqsimon
 * Apply suggestions from new clippy lint clippy::assigning_clones #382 - @cyqsimon
+* Fix IPv6 socket detect logic #383 - @cyqsimon
 
 ## Added
 * CI: include generated assets in release archive #359 - @cyqsimon
diff --git a/src/os/linux.rs b/src/os/linux.rs
index fdd5ec1..08a4aeb 100644
--- a/src/os/linux.rs
+++ b/src/os/linux.rs
@@ -10,7 +10,7 @@ use crate::{
 
 pub(crate) fn get_open_sockets() -> OpenSockets {
     let mut open_sockets = HashMap::new();
-    let mut inode_to_procname = HashMap::new();
+    let mut inode_to_proc = HashMap::new();
 
     if let Ok(all_procs) = procfs::process::all_processes() {
         for process in all_procs.filter_map(|res| res.ok()) {
@@ -20,47 +20,30 @@ pub(crate) fn get_open_sockets() -> OpenSockets {
             let proc_info = ProcessInfo::new(&proc_name, stat.pid as u32);
             for fd in fds.filter_map(|res| res.ok()) {
                 if let FDTarget::Socket(inode) = fd.target {
-                    inode_to_procname.insert(inode, proc_info.clone());
+                    inode_to_proc.insert(inode, proc_info.clone());
                 }
             }
         }
     }
 
-    if let Ok(mut tcp) = procfs::net::tcp() {
-        if let Ok(mut tcp6) = procfs::net::tcp6() {
-            tcp.append(&mut tcp6);
-        }
-        for entry in tcp.into_iter() {
-            if let Some(proc_info) = inode_to_procname.get(&entry.inode) {
-                open_sockets.insert(
-                    LocalSocket {
+    macro_rules! insert_proto {
+        ($source: expr, $proto: expr) => {
+            let entries = $source.into_iter().filter_map(|res| res.ok()).flatten();
+            for entry in entries {
+                if let Some(proc_info) = inode_to_proc.get(&entry.inode) {
+                    let socket = LocalSocket {
                         ip: entry.local_address.ip(),
                         port: entry.local_address.port(),
-                        protocol: Protocol::Tcp,
-                    },
-                    proc_info.clone(),
-                );
-            };
-        }
+                        protocol: $proto,
+                    };
+                    open_sockets.insert(socket, proc_info.clone());
+                }
+            }
+        };
     }
 
-    if let Ok(mut udp) = procfs::net::udp() {
-        if let Ok(mut udp6) = procfs::net::udp6() {
-            udp.append(&mut udp6);
-        }
-        for entry in udp.into_iter() {
-            if let Some(proc_info) = inode_to_procname.get(&entry.inode) {
-                open_sockets.insert(
-                    LocalSocket {
-                        ip: entry.local_address.ip(),
-                        port: entry.local_address.port(),
-                        protocol: Protocol::Udp,
-                    },
-                    proc_info.clone(),
-                );
-            };
-        }
-    }
+    insert_proto!([procfs::net::tcp(), procfs::net::tcp6()], Protocol::Tcp);
+    insert_proto!([procfs::net::udp(), procfs::net::udp6()], Protocol::Udp);
 
     OpenSockets {
         sockets_to_procs: open_sockets,