add pid

oryx-tui/src/app.rs

@@ -10,9 +10,17 @@ use std::{
     time::Duration,
 };
 
-use crate::{filter::Filter, help::Help, pid};
+use crate::{
+    filter::Filter,
+    help::Help,
+    packet::{
+        network::{IpPacket, IpProto},
+        AppPacket,
+    },
+    pid::{self, ConnectionsInfo, IpMap},
+};
 use crate::{filter::IoChannels, notification::Notification};
-use crate::{packet::AppPacket, section::Section};
+use crate::{packet::NetworkPacket, section::Section};
 
 pub type AppResult<T> = std::result::Result<T, Box<dyn error::Error>>;
 
@@ -31,6 +39,17 @@ pub struct DataEventHandler {
     pub sender: kanal::Sender<[u8; RawPacket::LEN]>,
     pub handler: thread::JoinHandle<()>,
 }
+#[derive(Debug, Clone)]
+pub struct Channels<T> {
+    pub sender: kanal::Sender<T>,
+    pub receiver: kanal::Receiver<T>,
+}
+impl<T> Channels<T> {
+    pub fn new() -> Self {
+        let (sender, receiver) = kanal::unbounded();
+        Self { sender, receiver }
+    }
+}
 
 #[derive(Debug)]
 pub struct App {
@@ -39,6 +58,7 @@ pub struct App {
     pub filter: Filter,
     pub start_sniffing: bool,
     pub packets: Arc<Mutex<Vec<AppPacket>>>,
+    pub connections_info: ConnectionsInfo,
     pub notifications: Vec<Notification>,
     pub section: Section,
     pub data_channel_sender: kanal::Sender<[u8; RawPacket::LEN]>,
@@ -54,35 +74,72 @@ impl Default for App {
 
 impl App {
     pub fn new() -> Self {
-        let pids = pid::ConnectionsInfo::new();
-        let packets = Arc::new(Mutex::new(Vec::with_capacity(AppPacket::LEN * 1024 * 1024)));
+        let app_packets: Arc<Mutex<Vec<AppPacket>>> = Arc::new(Mutex::new(Vec::with_capacity(
+            NetworkPacket::LEN * 1024 * 1024,
+        )));
+        let data_channels = Channels::new();
 
-        let (sender, receiver) = kanal::unbounded();
+        let tcp_map: Arc<Mutex<IpMap>> = Arc::new(Mutex::new(IpMap::new()));
+        let udp_map: Arc<Mutex<IpMap>> = Arc::new(Mutex::new(IpMap::new()));
 
-        let firewall_channels = IoChannels::new();
         thread::spawn({
-            let packets = packets.clone();
+            let app_packets = app_packets.clone();
+            let tcp_map = tcp_map.clone();
+            let udp_map = udp_map.clone();
             move || loop {
-                if let Ok(raw_packet) = receiver.recv() {
-                    let app_packet = AppPacket::from(raw_packet);
-                    let mut packets = packets.lock().unwrap();
-                    if packets.len() == packets.capacity() {
-                        packets.reserve(1024 * 1024);
+                if let Ok(raw_packet) = data_channels.receiver.recv() {
+                    let network_packet = NetworkPacket::from(raw_packet);
+                    let mut app_packets = app_packets.lock().unwrap();
+                    if app_packets.len() == app_packets.capacity() {
+                        app_packets.reserve(1024 * 1024);
                     }
-                    packets.push(app_packet);
+                    let mut app_packet = AppPacket {
+                        packet: network_packet,
+                        pid: None,
+                    };
+                    let pid = match &app_packet.packet {
+                        NetworkPacket::Ip(IpPacket::V4(ipv4packet)) => match ipv4packet.proto {
+                            IpProto::Tcp(_) => {
+                                let ipmap = tcp_map.lock().unwrap().clone();
+                                app_packet.try_get_pid(ipmap)
+                            }
+
+                            IpProto::Udp(_) => {
+                                let ipmap = udp_map.lock().unwrap().clone();
+
+                                app_packet.try_get_pid(ipmap)
+                            }
+
+                            _ => None,
+                        },
+                        _ => None,
+                    };
+                    app_packet.pid = pid;
+                    app_packets.push(app_packet);
                 }
             }
         });
 
+        let firewall_channels = IoChannels::new();
+
+        let udp_map: Arc<Mutex<IpMap>> = Arc::new(Mutex::new(IpMap::new()));
+        let conn_info = pid::ConnectionsInfo::new(tcp_map.clone(), udp_map.clone());
+
         Self {
             running: true,
             help: Help::new(),
             filter: Filter::new(firewall_channels.clone()),
+            connections_info: conn_info,
             start_sniffing: false,
-            packets: packets.clone(),
+            packets: app_packets.clone(),
             notifications: Vec::new(),
-            section: Section::new(packets.clone(), firewall_channels.clone()),
-            data_channel_sender: sender,
+            section: Section::new(
+                app_packets.clone(),
+                tcp_map.clone(),
+                udp_map.clone(),
+                firewall_channels.clone(),
+            ),
+            data_channel_sender: data_channels.sender,
             is_editing: false,
             active_popup: None,
         }

oryx-tui/src/export.rs

@@ -8,11 +8,11 @@ use crate::{
     app::AppResult,
     packet::{
         network::{IpPacket, IpProto},
-        AppPacket,
+        NetworkPacket,
     },
 };
 
-pub fn export(packets: &[AppPacket]) -> AppResult<()> {
+pub fn export(packets: &[NetworkPacket]) -> AppResult<()> {
     let uid = unsafe { libc::geteuid() };
 
     let oryx_export_dir = dirs::home_dir().unwrap().join("oryx");
@@ -40,7 +40,7 @@ pub fn export(packets: &[AppPacket]) -> AppResult<()> {
     )?;
     for packet in packets {
         match packet {
-            AppPacket::Arp(p) => {
+            NetworkPacket::Arp(p) => {
                 writeln!(
                     file,
                     "{:39}  {:^11}  {:39}  {:^11}  ARP",
@@ -50,7 +50,7 @@ pub fn export(packets: &[AppPacket]) -> AppResult<()> {
                     "-"
                 )?;
             }
-            AppPacket::Ip(packet) => match packet {
+            NetworkPacket::Ip(packet) => match packet {
                 IpPacket::V4(ipv4_packet) => match ipv4_packet.proto {
                     IpProto::Tcp(p) => {
                         writeln!(

oryx-tui/src/filter.rs

@@ -26,7 +26,7 @@ use transport::TransportFilter;
 use tui_big_text::{BigText, PixelSize};
 
 use crate::{
-    app::AppResult,
+    app::{AppResult, Channels},
     ebpf::{egress::load_egress, ingress::load_ingress},
     event::Event,
     interface::Interface,
@@ -40,25 +40,12 @@ pub enum FilterChannelSignal {
     Kill,
 }
 
-#[derive(Debug, Clone)]
-pub struct Channels<T> {
-    pub sender: kanal::Sender<T>,
-    pub receiver: kanal::Receiver<T>,
-}
-
 #[derive(Debug, Clone)]
 pub struct IoChannels<T> {
     pub ingress: Channels<T>,
     pub egress: Channels<T>,
 }
 
-impl<T> Channels<T> {
-    pub fn new() -> Self {
-        let (sender, receiver) = kanal::unbounded();
-        Self { sender, receiver }
-    }
-}
-
 impl<T> IoChannels<T> {
     pub fn new() -> Self {
         Self {
@@ -68,12 +55,6 @@ impl<T> IoChannels<T> {
     }
 }
 
-impl<T> Default for Channels<T> {
-    fn default() -> Self {
-        Self::new()
-    }
-}
-
 impl<T> Default for IoChannels<T> {
     fn default() -> Self {
         Self::new()

oryx-tui/src/filter/fuzzy.rs

@@ -98,7 +98,7 @@ impl Fuzzy {
         self.packets = packets
             .iter()
             .copied()
-            .filter(|p| p.to_string().contains(self.filter.value()))
+            .filter(|p| p.packet.to_string().contains(self.filter.value()))
             .collect::<Vec<AppPacket>>();
     }
 
@@ -107,7 +107,7 @@ impl Fuzzy {
             &mut packets
                 .iter()
                 .copied()
-                .filter(|p| p.to_string().contains(self.filter.value()))
+                .filter(|p| p.packet.to_string().contains(self.filter.value()))
                 .collect::<Vec<AppPacket>>(),
         );
     }

oryx-tui/src/lib.rs

@@ -27,4 +27,3 @@ pub mod packet;
 pub mod section;
 
 pub mod pid;
-pub mod test;

oryx-tui/src/packet.rs

@@ -5,22 +5,65 @@ pub mod transport;
 use std::{fmt::Display, mem, net::Ipv4Addr};
 
 use link::{ArpPacket, ArpType, MacAddr};
+use log::info;
 use network::{IcmpPacket, IcmpType, IpPacket, IpProto, Ipv4Packet, Ipv6Packet};
 use network_types::ip::IpHdr;
 use oryx_common::{ProtoHdr, RawPacket};
 use transport::{TcpPacket, UdpPacket};
 
+use crate::pid::IpMap;
+
 #[derive(Debug, Copy, Clone)]
-pub enum AppPacket {
+pub enum NetworkPacket {
     Ip(IpPacket),
     Arp(ArpPacket),
 }
 
-impl AppPacket {
+impl NetworkPacket {
     pub const LEN: usize = mem::size_of::<Self>();
 }
 
-impl Display for AppPacket {
+#[derive(Debug, Copy, Clone)]
+pub struct AppPacket {
+    pub packet: NetworkPacket,
+    pub pid: Option<u32>,
+}
+
+impl AppPacket {
+    fn get_possible_keys(&self) -> Option<[String; 2]> {
+        match self.packet {
+            NetworkPacket::Ip(IpPacket::V4(ipv4packet)) => {
+                let src_ip = ipv4packet.src_ip;
+                let dst_ip = ipv4packet.dst_ip;
+                match ipv4packet.proto {
+                    IpProto::Tcp(tcp) => Some([
+                        format!("{}:{}_{}:{}", src_ip, tcp.src_port, dst_ip, tcp.dst_port),
+                        format!("{}:{}_{}:{}", dst_ip, tcp.dst_port, src_ip, tcp.src_port),
+                    ]),
+                    IpProto::Udp(udp) => Some([
+                        format!("{}:{}_{}:{}", src_ip, udp.src_port, dst_ip, udp.dst_port),
+                        format!("{}:{}_{}:{}", dst_ip, udp.dst_port, src_ip, udp.src_port),
+                    ]),
+                    _ => None,
+                }
+            }
+            _ => None,
+        }
+    }
+    pub fn try_get_pid(&self, ipmap: IpMap) -> Option<u32> {
+        if let Some(keys) = self.get_possible_keys() {
+            for k in keys {
+                if let Some(conn) = ipmap.map.get(&k) {
+                    info!("found pid for {}", k);
+                    return conn.pid;
+                }
+            }
+        }
+        None
+    }
+}
+
+impl Display for NetworkPacket {
     fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
         match self {
             Self::Arp(packet) => write!(f, "{}", packet),
@@ -29,7 +72,7 @@ impl Display for AppPacket {
     }
 }
 
-impl From<[u8; RawPacket::LEN]> for AppPacket {
+impl From<[u8; RawPacket::LEN]> for NetworkPacket {
     fn from(value: [u8; RawPacket::LEN]) -> Self {
         let raw_packet = value.as_ptr() as *const RawPacket;
         match unsafe { &*raw_packet } {
@@ -87,7 +130,7 @@ impl From<[u8; RawPacket::LEN]> for AppPacket {
                         }
                     };
 
-                    AppPacket::Ip(IpPacket::V4(Ipv4Packet {
+                    NetworkPacket::Ip(IpPacket::V4(Ipv4Packet {
                         src_ip,
                         dst_ip,
                         ihl: u8::from_be(ipv4_packet.ihl()),
@@ -153,7 +196,7 @@ impl From<[u8; RawPacket::LEN]> for AppPacket {
                         }
                     };
 
-                    AppPacket::Ip(IpPacket::V6(Ipv6Packet {
+                    NetworkPacket::Ip(IpPacket::V6(Ipv6Packet {
                         traffic_class: ipv6_packet.priority(),
                         flow_label: ipv6_packet.flow_label,
                         payload_length: u16::from_be(ipv6_packet.payload_len),