refresh map when new socket is created

pythops · Nov 17, 2024 · 1c9b6e6 · 1c9b6e6
1 parent b90cad5
commit 1c9b6e6
Show file tree

Hide file tree

Showing 4 changed files with 42 additions and 7 deletions.
diff --git a/oryx-tui/src/ebpf/pid.rs b/oryx-tui/src/ebpf/pid.rs
@@ -16,7 +16,7 @@ use log::error;
 use crate::{
     event::Event,
     notification::{Notification, NotificationLevel},
-    pid::ConnectionMap,
+    pid::{tcp::TcpConnectionMap, udp::UdpConnectionMap, ConnectionMap},
 };
 use mio::{unix::SourceFd, Events, Interest, Poll, Token};
 
@@ -115,9 +115,44 @@ pub fn load_pid(
                             let pid = u32::from_ne_bytes(pid);
 
                             let fd_dir = format!("/proc/{}/fd", pid);
-                            if let Ok(_fds) = fs::read_dir(&fd_dir) {
+                            if let Ok(fds) = fs::read_dir(&fd_dir) {
                                 let mut map = pid_map.lock().unwrap();
-                                *map = ConnectionMap::new();
+                                let tcp_inode_map = TcpConnectionMap::inode_map();
+                                let udp_inode_map = UdpConnectionMap::inode_map();
+
+                                for fd in fds.flatten() {
+                                    let link_path = fd.path();
+
+                                    if let Ok(link_target) = fs::read_link(&link_path) {
+                                        if let Some(inode_str) = link_target.to_str() {
+                                            if inode_str.starts_with("socket:[")
+                                                && inode_str.ends_with(']')
+                                            {
+                                                if let Ok(inode) = inode_str[8..inode_str.len() - 1]
+                                                    .parse::<usize>()
+                                                {
+                                                    if let Some(connection_hash) =
+                                                        tcp_inode_map.get(&inode)
+                                                    {
+                                                        map.tcp.map.insert(
+                                                            *connection_hash,
+                                                            pid.try_into().unwrap(),
+                                                        );
+                                                    }
+
+                                                    if let Some(connection_hash) =
+                                                        udp_inode_map.get(&inode)
+                                                    {
+                                                        map.udp.map.insert(
+                                                            *connection_hash,
+                                                            pid.try_into().unwrap(),
+                                                        );
+                                                    }
+                                                }
+                                            }
+                                        }
+                                    }
+                                }
                             }
                         }
                     }

diff --git a/oryx-tui/src/pid.rs b/oryx-tui/src/pid.rs
@@ -65,8 +65,8 @@ fn decode_hex_port(hex_str: &str) -> Result<u16, ParseIntError> {
 
 #[derive(Clone, Debug)]
 pub struct ConnectionMap {
-    tcp: TcpConnectionMap,
-    udp: UdpConnectionMap,
+    pub tcp: TcpConnectionMap,
+    pub udp: UdpConnectionMap,
 }
 
 impl ConnectionMap {

diff --git a/oryx-tui/src/pid/tcp.rs b/oryx-tui/src/pid/tcp.rs
@@ -13,7 +13,7 @@ pub struct TcpConnectionMap {
 }
 
 impl TcpConnectionMap {
-    fn inode_map() -> HashMap<usize, u64> {
+    pub fn inode_map() -> HashMap<usize, u64> {
         let mut map = HashMap::new();
         let mut file = File::open("/proc/net/tcp").unwrap();
         let mut buffer = String::new();

diff --git a/oryx-tui/src/pid/udp.rs b/oryx-tui/src/pid/udp.rs
@@ -13,7 +13,7 @@ pub struct UdpConnectionMap {
 }
 
 impl UdpConnectionMap {
-    fn inode_map() -> HashMap<usize, u64> {
+    pub fn inode_map() -> HashMap<usize, u64> {
         let mut map = HashMap::new();
         let mut file = File::open("/proc/net/udp").unwrap();
         let mut buffer = String::new();