cache connection added

Cargo.toml

@@ -17,6 +17,7 @@ regex = "1.11.1"
 failure = "0.1.8"
 log = "0.4.22"
 lazy_static = "1.5.0"
+ttl_cache = "0.5.1"
 
 [lib]
 name = "passivetcp"

src/lib.rs

@@ -13,38 +13,44 @@ use crate::db::Database;
 use crate::p0f_output::{MTUOutput, P0fOutput, SynAckTCPOutput, UptimeOutput};
 use crate::packet::SignatureDetails;
 use crate::signature_matcher::SignatureMatcher;
+use std::net::IpAddr;
+use ttl_cache::TtlCache;
+
+#[derive(Debug, Hash, Eq, PartialEq, Clone)]
+struct Connection {
+    src_ip: IpAddr,
+    src_port: u16,
+    dst_ip: IpAddr,
+    dst_port: u16,
+}
 
 struct SynData {
     ts1: u32,
     recv_ms: u64,
 }
-struct UptimeData {
-    client: Option<SynData>,
-    server: Option<SynData>,
-    sendsyn: bool,
-}
-
 pub struct P0f<'a> {
     pub matcher: SignatureMatcher<'a>,
-    uptime_data: UptimeData,
+    cache: TtlCache<Connection, SynData>,
 }
 
 impl<'a> P0f<'a> {
     pub fn new(database: &'a Database) -> Self {
         let matcher = SignatureMatcher::new(database);
-        let uptime_data = UptimeData {
-            client: None,
-            server: None,
-            sendsyn: false,
-        };
+        let cache: TtlCache<Connection, SynData> = TtlCache::new(100); // TODO: capacity by param...
+                                                                       /*let uptime_data = UptimeData {
+                                                                           client: None,
+                                                                           server: None,
+                                                                           sendsyn: false,
+                                                                       };*/
         Self {
             matcher,
-            uptime_data,
+            cache,
+            //uptime_data,
         }
     }
 
     pub fn analyze_tcp(&mut self, packet: &[u8]) -> P0fOutput {
-        if let Ok(signature_details) = SignatureDetails::extract(packet, &mut self.uptime_data) {
+        if let Ok(signature_details) = SignatureDetails::extract(packet, &mut self.cache) {
             if signature_details.is_client {
                 let mtu: Option<MTUOutput> = if let Some(mtu) = signature_details.mtu {
                     if let Some((link, _matched_mtu)) = self.matcher.matching_by_mtu(&mtu) {

src/packet.rs

@@ -1,6 +1,6 @@
 use crate::tcp::{IpVersion, PayloadSize, Quirk, Signature, TcpOption, Ttl, WindowSize};
 use crate::uptime::{check_ts_tcp, Uptime};
-use crate::{mtu, UptimeData};
+use crate::{mtu, Connection, SynData};
 use failure::{bail, err_msg, Error};
 use pnet::packet::{
     ethernet::{EtherType, EtherTypes, EthernetPacket},
@@ -13,6 +13,7 @@ use pnet::packet::{
 };
 use std::convert::TryInto;
 use std::net::IpAddr;
+use ttl_cache::TtlCache;
 
 #[derive(Clone)]
 pub struct IpPort {
@@ -29,39 +30,43 @@ pub struct SignatureDetails {
     pub is_client: bool,
 }
 impl SignatureDetails {
-    pub fn extract(packet: &[u8], uptime_data: &mut UptimeData) -> Result<Self, Error> {
+    pub fn extract(
+        packet: &[u8],
+        cache: &mut TtlCache<Connection, SynData>,
+    ) -> Result<Self, Error> {
         EthernetPacket::new(packet)
             .ok_or_else(|| err_msg("ethernet packet too short"))
-            .and_then(|packet| {
-                visit_ethernet(packet.get_ethertype(), uptime_data, packet.payload())
-            })
+            .and_then(|packet| visit_ethernet(packet.get_ethertype(), cache, packet.payload()))
     }
 }
 
 fn visit_ethernet(
     ethertype: EtherType,
-    uptime_data: &mut UptimeData,
+    cache: &mut TtlCache<Connection, SynData>,
     payload: &[u8],
 ) -> Result<SignatureDetails, Error> {
     match ethertype {
         EtherTypes::Vlan => VlanPacket::new(payload)
             .ok_or_else(|| err_msg("vlan packet too short"))
-            .and_then(|packet| visit_vlan(uptime_data, packet)),
+            .and_then(|packet| visit_vlan(cache, packet)),
 
         EtherTypes::Ipv4 => Ipv4Packet::new(payload)
             .ok_or_else(|| err_msg("ipv4 packet too short"))
-            .and_then(|packet| visit_ipv4(uptime_data, packet)),
+            .and_then(|packet| visit_ipv4(cache, packet)),
 
         EtherTypes::Ipv6 => Ipv6Packet::new(payload)
             .ok_or_else(|| err_msg("ipv6 packet too short"))
-            .and_then(|packet| visit_ipv6(uptime_data, packet)),
+            .and_then(|packet| visit_ipv6(cache, packet)),
 
         ty => bail!("unsupport ethernet type: {}", ty),
     }
 }
 
-fn visit_vlan(uptime_data: &mut UptimeData, packet: VlanPacket) -> Result<SignatureDetails, Error> {
-    visit_ethernet(packet.get_ethertype(), uptime_data, packet.payload())
+fn visit_vlan(
+    cache: &mut TtlCache<Connection, SynData>,
+    packet: VlanPacket,
+) -> Result<SignatureDetails, Error> {
+    visit_ethernet(packet.get_ethertype(), cache, packet.payload())
 }
 
 /// Congestion encountered
@@ -75,7 +80,10 @@ fn is_client(tcp_flags: u8) -> bool {
     tcp_flags & TcpFlags::SYN != 0 && tcp_flags & TcpFlags::ACK == 0
 }
 
-fn visit_ipv4(uptime_data: &mut UptimeData, packet: Ipv4Packet) -> Result<SignatureDetails, Error> {
+fn visit_ipv4(
+    cache: &mut TtlCache<Connection, SynData>,
+    packet: Ipv4Packet,
+) -> Result<SignatureDetails, Error> {
     if packet.get_next_level_protocol() != IpNextHeaderProtocols::Tcp {
         bail!(
             "unsuppport IPv4 packet with non-TCP payload: {}",
@@ -124,7 +132,7 @@ fn visit_ipv4(uptime_data: &mut UptimeData, packet: Ipv4Packet) -> Result<Signat
         .ok_or_else(|| err_msg("TCP packet too short"))
         .and_then(|tcp_packet| {
             visit_tcp(
-                uptime_data,
+                cache,
                 &tcp_packet,
                 version,
                 ttl,
@@ -137,7 +145,10 @@ fn visit_ipv4(uptime_data: &mut UptimeData, packet: Ipv4Packet) -> Result<Signat
         })
 }
 
-fn visit_ipv6(uptime_data: &mut UptimeData, packet: Ipv6Packet) -> Result<SignatureDetails, Error> {
+fn visit_ipv6(
+    cache: &mut TtlCache<Connection, SynData>,
+    packet: Ipv6Packet,
+) -> Result<SignatureDetails, Error> {
     if packet.get_next_header() != IpNextHeaderProtocols::Tcp {
         bail!(
             "unsuppport IPv6 packet with non-TCP payload: {}",
@@ -167,7 +178,7 @@ fn visit_ipv6(uptime_data: &mut UptimeData, packet: Ipv6Packet) -> Result<Signat
         .ok_or_else(|| err_msg("TCP packet too short"))
         .and_then(|tcp_packet| {
             visit_tcp(
-                uptime_data,
+                cache,
                 &tcp_packet,
                 version,
                 ttl,
@@ -195,15 +206,15 @@ fn guess_dist(ttl: u8) -> u8 {
 //TODO: WIP: observable tcp params
 #[allow(clippy::too_many_arguments)]
 fn visit_tcp(
-    uptime_data: &mut UptimeData,
+    cache: &mut TtlCache<Connection, SynData>,
     tcp: &TcpPacket,
     version: IpVersion,
     ittl: Ttl,
     ip_package_header_length: u8,
     olen: u8,
     mut quirks: Vec<Quirk>,
-    client_ip: IpAddr,
-    server_ip: IpAddr,
+    source_ip: IpAddr,      // TODO: rename ALL to source
+    destination_ip: IpAddr, // TODO: rename ALL to destination
 ) -> Result<SignatureDetails, Error> {
     use TcpFlags::*;
 
@@ -322,7 +333,13 @@ fn visit_tcp(
 
                 if data.len() >= 8 {
                     let ts_val: u32 = u32::from_ne_bytes(data[..4].try_into()?);
-                    uptime = check_ts_tcp(uptime_data, is_client, ts_val, tcp_type);
+                    let connection: Connection = Connection {
+                        src_ip: source_ip,
+                        src_port: tcp.get_source(),
+                        dst_ip: destination_ip,
+                        dst_port: tcp.get_destination(),
+                    };
+                    uptime = check_ts_tcp(cache, &connection, is_client, ts_val);
                 }
 
                 /*if data.len() != 10 {
@@ -378,11 +395,11 @@ fn visit_tcp(
         mtu,
         uptime,
         client: IpPort {
-            ip: client_ip,
+            ip: source_ip,
             port: client_port,
         },
         server: IpPort {
-            ip: server_ip,
+            ip: destination_ip,
             port: server_port,
         },
         is_client,

src/uptime.rs

@@ -1,6 +1,6 @@
-use crate::{SynData, UptimeData};
-use pnet::packet::tcp::TcpFlags::{ACK, SYN};
-use std::time::{SystemTime, UNIX_EPOCH};
+use crate::{Connection, SynData};
+use std::time::{Duration, SystemTime, UNIX_EPOCH};
+use ttl_cache::TtlCache;
 
 const MIN_TWAIT: u64 = 25;
 const MAX_TWAIT: u64 = 10 * 60 * 1000;
@@ -24,25 +24,33 @@ fn get_unix_time_ms() -> u64 {
 }
 
 pub fn check_ts_tcp(
-    uptime_data: &mut UptimeData,
-    to_server: bool,
+    cache: &mut TtlCache<Connection, SynData>,
+    connection: &Connection,
+    is_client: bool,
     ts_val: u32,
-    tcp_type: u8,
 ) -> Option<Uptime> {
-    // If no timestamp or if the connection is still in SYN state, return early
-    if ts_val == 0 || uptime_data.sendsyn {
-        return None;
-    }
-
-    // Select the correct side (client or server) based on the `to_server` flag
-    let last_syn_data: Option<&SynData> = if to_server {
-        uptime_data.client.as_ref()
+    let syn_data: Option<&SynData> = if !is_client {
+        let client_connection = Connection {
+            src_ip: connection.dst_ip.clone(),
+            src_port: connection.dst_port,
+            dst_ip: connection.src_ip.clone(),
+            dst_port: connection.src_port,
+        };
+        cache.get(&client_connection)
     } else {
-        uptime_data.server.as_ref()
+        cache.insert(
+            connection.clone(),
+            SynData {
+                ts1: ts_val,
+                recv_ms: get_unix_time_ms(),
+            },
+            Duration::new(60, 0),
+        );
+        None
     };
 
     // If there's no valid SYN data yet, return early
-    let last_syn_data = last_syn_data?;
+    let last_syn_data = syn_data?;
     let ms_diff = get_unix_time_ms().saturating_sub(last_syn_data.recv_ms);
     let ts_diff = ts_val.saturating_sub(last_syn_data.ts1);
 
@@ -67,14 +75,6 @@ pub fn check_ts_tcp(
 
     // Check if the frequency is within acceptable bounds
     if ffreq < MIN_TSCALE || ffreq > MAX_TSCALE {
-        // Allow invalid readings on SYN, might be caused by IP sharing or OS changes
-        if tcp_type != SYN {
-            if to_server {
-                uptime_data.client = None; // Mark client as invalid
-            } else {
-                uptime_data.server = None; // Mark server as invalid
-            }
-        }
         return None;
     }
 
@@ -92,20 +92,6 @@ pub fn check_ts_tcp(
     let up_min = ts_val / freq / 60;
     let up_mod_days = 0xFFFFFFFF / (freq * 60 * 60 * 24);
 
-    // Store the last timestamp for client or server, depending on direction
-    if tcp_type == SYN {
-        uptime_data.sendsyn = true;
-        uptime_data.client = Some(SynData {
-            ts1: ts_val,
-            recv_ms: get_unix_time_ms(),
-        });
-    } else if tcp_type == ACK {
-        uptime_data.server = Some(SynData {
-            ts1: ts_val,
-            recv_ms: get_unix_time_ms(),
-        });
-    }
-
     Some(Uptime {
         days: up_min / 60 / 24,
         hours: (up_min / 60) % 24,