From c2c57289cff695bfb68fdb9dbbcba195860b8a35 Mon Sep 17 00:00:00 2001
From: John Doty <john@d0ty.me>
Date: Sun, 4 Aug 2024 08:14:19 -0700
Subject: [PATCH] Refactor refresh

Moving different mechanisms into different conditionally-compiled
modules. This way it can be extended, e.g. with docker lookups, MacOS
support, etc.
---
 src/server/refresh.rs        | 65 +++++++++++-------------------------
 src/server/refresh/procfs.rs | 54 ++++++++++++++++++++++++++++++
 2 files changed, 73 insertions(+), 46 deletions(-)
 create mode 100644 src/server/refresh/procfs.rs

diff --git a/src/server/refresh.rs b/src/server/refresh.rs
index a62dfa7..b1b95b8 100644
--- a/src/server/refresh.rs
+++ b/src/server/refresh.rs
@@ -1,61 +1,34 @@
-use crate::message::PortDesc;
 use anyhow::Result;
+use log::{error, warn};
+use std::collections::HashMap;
 
-#[cfg(not(target_os = "linux"))]
-pub fn get_entries() -> Result<Vec<PortDesc>> {
-    use anyhow::bail;
-    bail!("Not supported on this operating system");
-}
+use crate::message::PortDesc;
 
 #[cfg(target_os = "linux")]
+mod procfs;
+
 pub fn get_entries() -> Result<Vec<PortDesc>> {
-    use procfs::process::FDTarget;
-    use std::collections::HashMap;
+    #[allow(unused)]
+    let mut attempts = 0;
 
-    let all_procs = procfs::process::all_processes()?;
+    let mut result: HashMap<u16, PortDesc> = HashMap::new();
 
-    // build up a map between socket inodes and process stat info. Ignore any
-    // error we encounter as it probably means we have no access to that
-    // process or something.
-    let mut map: HashMap<u64, String> = HashMap::new();
-    for process in all_procs.flatten() {
-        if !process.is_alive() {
-            continue; // Ignore zombies.
-        }
-
-        if let (Ok(fds), Ok(cmd)) = (process.fd(), process.cmdline()) {
-            for fd in fds.flatten() {
-                if let FDTarget::Socket(inode) = fd.target {
-                    map.insert(inode, cmd.join(" "));
+    #[cfg(target_os = "linux")]
+    {
+        attempts += 1;
+        match procfs::get_entries() {
+            Ok(m) => {
+                for (p, d) in m {
+                    result.entry(p).or_insert(d);
                 }
             }
+            Err(e) => error!("Error reading procfs: {e:?}"),
         }
     }
 
-    let mut h: HashMap<u16, PortDesc> = HashMap::new();
-
-    // Go through all the listening IPv4 and IPv6 sockets and take the first
-    // instance of listening on each port *if* the address is loopback or
-    // unspecified. (TODO: Do we want this restriction really?)
-    let tcp = procfs::net::tcp()?;
-    let tcp6 = procfs::net::tcp6()?;
-    for tcp_entry in tcp.into_iter().chain(tcp6) {
-        if tcp_entry.state == procfs::net::TcpState::Listen
-            && (tcp_entry.local_address.ip().is_loopback()
-                || tcp_entry.local_address.ip().is_unspecified())
-            && !h.contains_key(&tcp_entry.local_address.port())
-        {
-            if let Some(cmd) = map.get(&tcp_entry.inode) {
-                h.insert(
-                    tcp_entry.local_address.port(),
-                    PortDesc {
-                        port: tcp_entry.local_address.port(),
-                        desc: cmd.clone(),
-                    },
-                );
-            }
-        }
+    if attempts == 0 {
+        warn!("Port scanning is not supported for this server");
     }
 
-    Ok(h.into_values().collect())
+    Ok(result.into_values().collect())
 }
diff --git a/src/server/refresh/procfs.rs b/src/server/refresh/procfs.rs
new file mode 100644
index 0000000..aa1a79f
--- /dev/null
+++ b/src/server/refresh/procfs.rs
@@ -0,0 +1,54 @@
+use anyhow::Result;
+use procfs::process::FDTarget;
+use std::collections::HashMap;
+
+use crate::message::PortDesc;
+
+pub fn get_entries() -> Result<HashMap<u16, PortDesc>> {
+    let all_procs = procfs::process::all_processes()?;
+
+    // build up a map between socket inodes and process stat info. Ignore any
+    // error we encounter as it probably means we have no access to that
+    // process or something.
+    let mut map: HashMap<u64, String> = HashMap::new();
+    for process in all_procs.flatten() {
+        if !process.is_alive() {
+            continue; // Ignore zombies.
+        }
+
+        if let (Ok(fds), Ok(cmd)) = (process.fd(), process.cmdline()) {
+            for fd in fds.flatten() {
+                if let FDTarget::Socket(inode) = fd.target {
+                    map.insert(inode, cmd.join(" "));
+                }
+            }
+        }
+    }
+
+    let mut h: HashMap<u16, PortDesc> = HashMap::new();
+
+    // Go through all the listening IPv4 and IPv6 sockets and take the first
+    // instance of listening on each port *if* the address is loopback or
+    // unspecified. (TODO: Do we want this restriction really?)
+    let tcp = procfs::net::tcp()?;
+    let tcp6 = procfs::net::tcp6()?;
+    for tcp_entry in tcp.into_iter().chain(tcp6) {
+        if tcp_entry.state == procfs::net::TcpState::Listen
+            && (tcp_entry.local_address.ip().is_loopback()
+                || tcp_entry.local_address.ip().is_unspecified())
+            && !h.contains_key(&tcp_entry.local_address.port())
+        {
+            if let Some(cmd) = map.get(&tcp_entry.inode) {
+                h.insert(
+                    tcp_entry.local_address.port(),
+                    PortDesc {
+                        port: tcp_entry.local_address.port(),
+                        desc: cmd.clone(),
+                    },
+                );
+            }
+        }
+    }
+
+    Ok(h)
+}