From 3d16ed1f600dfa8b74c21f81ca3c15fedd578e20 Mon Sep 17 00:00:00 2001
From: Thea Rossman <thearossman@gmail.com>
Date: Thu, 5 Sep 2024 00:17:40 +0000
Subject: [PATCH] Add additional state for UDP connections aging out

- UDP performance is bad for connection-level subscriptions, due
  to UDP connections being re-inserted in connection table.
- This is one proposed solution: keep UDP connections around in
  the connection table in a "dropped" state until the timerwheel
  ages them out.

Testing done:
- Ran log_quic and log_dns on both live traffic (anon ips) and pcaps
- Output on live traffic was reasonable; output on pcaps was correct
- Repeated with log_* with connection-level subscriptions (Connection)
---
 core/Cargo.toml                      |  2 +-
 core/src/conntrack/conn/conn_info.rs | 23 +++++++++++++++++------
 core/src/conntrack/conn/mod.rs       |  2 +-
 core/src/conntrack/mod.rs            |  8 +++++++-
 4 files changed, 26 insertions(+), 9 deletions(-)

diff --git a/core/Cargo.toml b/core/Cargo.toml
index c1e32141..814270d0 100644
--- a/core/Cargo.toml
+++ b/core/Cargo.toml
@@ -5,7 +5,7 @@ edition = "2021"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [build-dependencies]
-bindgen = "0.69.4" 
+bindgen = "0.69.4"
 cc = "1.0.79"
 
 [dependencies]
diff --git a/core/src/conntrack/conn/conn_info.rs b/core/src/conntrack/conn/conn_info.rs
index 2547a7b5..2b5a52d4 100644
--- a/core/src/conntrack/conn/conn_info.rs
+++ b/core/src/conntrack/conn/conn_info.rs
@@ -1,6 +1,7 @@
 use crate::conntrack::conn_id::FiveTuple;
 use crate::conntrack::pdu::L4Pdu;
 use crate::filter::FilterResult;
+use crate::protocols::packet::udp::UDP_PROTOCOL;
 use crate::protocols::stream::{
     ConnData, ParseResult, ParserRegistry, ProbeRegistryResult, Session,
 };
@@ -47,7 +48,7 @@ where
             ConnState::Tracking => {
                 self.on_track(pdu, subscription);
             }
-            ConnState::Remove => {
+            ConnState::Remove | ConnState::Dropped => {
                 drop(pdu);
             }
         }
@@ -69,7 +70,7 @@ where
                         self.on_parse(pdu, subscription);
                     }
                     FilterResult::NoMatch => {
-                        self.state = ConnState::Remove;
+                        self.state = self.get_drop_state();
                     }
                 }
             }
@@ -82,10 +83,10 @@ where
                         self.state = self.get_match_state(0);
                     }
                     FilterResult::MatchNonTerminal(_idx) => {
-                        self.state = ConnState::Remove;
+                        self.state = self.get_drop_state();
                     }
                     FilterResult::NoMatch => {
-                        self.state = ConnState::Remove;
+                        self.state = self.get_drop_state();
                     }
                 }
             }
@@ -108,7 +109,7 @@ where
                     }
                 } else {
                     log::error!("Done parse but no mru");
-                    self.state = ConnState::Remove;
+                    self.state = self.get_drop_state();
                 }
             }
             ParseResult::Continue(id) => {
@@ -134,11 +135,18 @@ where
 
     fn get_nomatch_state(&self, session_id: usize) -> ConnState {
         if session_id == 0 && T::Subscribed::level() == Level::Connection {
-            ConnState::Remove
+            self.get_drop_state()
         } else {
             self.cdata.conn_parser.session_nomatch_state()
         }
     }
+
+    fn get_drop_state(&self) -> ConnState {
+        if self.cdata.five_tuple.proto == UDP_PROTOCOL {
+            return ConnState::Dropped;
+        }
+        ConnState::Remove
+    }
 }
 
 #[derive(Debug, PartialEq, Eq, Clone, Copy)]
@@ -151,4 +159,7 @@ pub enum ConnState {
     Tracking,
     /// Connection will be removed
     Remove,
+    /// Unmatched UDP connection; waiting to be aged out by timerwheel.
+    /// Prevents dropped UDP conns from being re-inserted in table
+    Dropped,
 }
diff --git a/core/src/conntrack/conn/mod.rs b/core/src/conntrack/conn/mod.rs
index 0ab98cb5..218a6431 100644
--- a/core/src/conntrack/conn/mod.rs
+++ b/core/src/conntrack/conn/mod.rs
@@ -160,7 +160,7 @@ where
             ConnState::Tracking => {
                 self.info.sdata.on_terminate(subscription);
             }
-            ConnState::Remove => {
+            ConnState::Remove | ConnState::Dropped => {
                 // do nothing
             }
         }
diff --git a/core/src/conntrack/mod.rs b/core/src/conntrack/mod.rs
index 4febca9c..ea9980ad 100644
--- a/core/src/conntrack/mod.rs
+++ b/core/src/conntrack/mod.rs
@@ -83,8 +83,14 @@ where
         match self.table.raw_entry_mut().from_key(&conn_id) {
             RawEntryMut::Occupied(mut occupied) => {
                 let conn = occupied.get_mut();
-                let dir = conn.packet_dir(&ctxt);
                 conn.last_seen_ts = Instant::now();
+                if conn.state() == ConnState::Dropped {
+                    // Allow connection to age out.
+                    // last_seen_ts is updated to avoid aging out long-lived UDP
+                    // connections prematurely
+                    return;
+                }
+                let dir = conn.packet_dir(&ctxt);
                 conn.inactivity_window = match &conn.l4conn {
                     L4Conn::Tcp(_) => self.config.tcp_inactivity_timeout,
                     L4Conn::Udp(_) => self.config.udp_inactivity_timeout,