tc full packet test

backend/src/utilities/packet_tracker.py

@@ -14,6 +14,7 @@ from src.Models.etherType import EtherTypeEnum, ethertype_from_int
 from src.Models.ip_protocol import protocol_from_number
 from src.config import settings
 from src.utilities.packet_identity import build_packet_uid
+from src.utilities.packet_mark import packet_id_from_mark, verdict_from_mark
 
 logger = logging.getLogger("packet_tracker")
 
@@ -42,6 +43,8 @@ class PacketTracker:
             "persisted_merged": 0,
             "persisted_with_raw": 0,
             "persisted_without_raw": 0,
+            "persisted_kernel_mark": 0,
+            "persisted_legacy_hash": 0,
         }
         self._lock = threading.Lock()
         self._stop_event = threading.Event()
@@ -55,46 +58,46 @@ class PacketTracker:
     def observe_packet(self, pkt_info: Dict[str, Any]) -> str:
         """Merge parsed packet information into a pending packet entry."""
         now_ts = time.time()
-        packet_uid = pkt_info.get("packet_uid") or build_packet_uid(pkt_info)
-        pkt_info["packet_uid"] = packet_uid
+        correlation_key = self._ensure_correlation(pkt_info)
         pkt_info["raw_present"] = pkt_info.get("raw") is not None
         pkt_info["capture_sources"] = ["af_packet"]
 
         with self._lock:
-            entry = self._entries.get(packet_uid)
+            entry = self._entries.get(correlation_key)
             if entry is None:
-                entry = self._new_entry(packet_uid, now_ts)
-                self._entries[packet_uid] = entry
+                entry = self._new_entry(correlation_key, now_ts)
+                self._entries[correlation_key] = entry
 
             self._merge_packet_info(entry, pkt_info, now_ts)
             self._maybe_promote_reject_from_reply(pkt_info, now_ts)
             self._maybe_mark_complete(entry)
-        return packet_uid
+        return correlation_key
 
     def observe_telemetry(self, event: Dict[str, Any]) -> Optional[str]:
         """Merge ingress/egress/verdict telemetry into a pending packet entry."""
-        packet_uid = event.get("packet_uid")
-        if not packet_uid:
-            try:
-                packet_uid = build_packet_uid(event)
-            except Exception:
-                logger.debug("Telemetry event missing packet identity: %s", event)
-                return None
+        correlation_key = self._ensure_correlation(event)
+        if not correlation_key:
+            logger.debug("Telemetry event missing packet identity: %s", event)
+            return None
 
         now_ts = time.time()
         with self._lock:
-            entry = self._entries.get(packet_uid)
+            entry = self._entries.get(correlation_key)
             if entry is None:
-                entry = self._new_entry(packet_uid, now_ts)
-                self._entries[packet_uid] = entry
+                entry = self._new_entry(correlation_key, now_ts)
+                self._entries[correlation_key] = entry
 
             payload = entry["payload"]
-            payload["packet_uid"] = packet_uid
+            payload["correlation_key"] = correlation_key
+            payload["packet_id"] = event.get("packet_id") or payload.get("packet_id")
+            payload["packet_uid"] = event.get("packet_uid") or payload.get("packet_uid")
+            payload["correlation_source"] = event.get("correlation_source") or payload.get("correlation_source")
+            payload["skb_mark"] = event.get("skb_mark") or payload.get("skb_mark")
             payload["telemetry_metadata"] = event
             payload["last_observed_at"] = now_ts
             self._add_capture_source(payload, "telemetry")
             for key, value in event.items():
-                if value is None or key in {"event_type", "reason", "reason_code", "iface", "packet_uid"}:
+                if value is None or key in {"event_type", "reason", "reason_code", "iface", "packet_uid", "correlation_key"}:
                     continue
                 if payload.get(key) is None:
                     payload[key] = value
@@ -112,6 +115,7 @@ class PacketTracker:
                     payload["protocol"] = int(payload["protocol_raw"])
             event_type = event.get("event_type")
             iface = event.get("iface")
+            verdict_hint = event.get("verdict_hint")
 
             if event_type == "ingress":
                 payload["ingress_if"] = iface
@@ -124,11 +128,11 @@ class PacketTracker:
                 payload["verdict_confidence"] = "high"
                 payload["verdict_seen_at"] = _utcnow()
             elif event_type == "drop":
-                payload["verdict"] = "drop"
-                payload["verdict_reason"] = event.get("reason") or "kfree_skb"
+                payload["verdict"] = verdict_hint or "drop"
+                payload["verdict_reason"] = event.get("reason") or ("mark-verdict" if verdict_hint else "kfree_skb")
                 payload["verdict_confidence"] = "high"
                 payload["verdict_seen_at"] = _utcnow()
-            elif event_type == "reject":
+            elif event_type == "reject" or verdict_hint == "reject":
                 payload["verdict"] = "reject"
                 payload["verdict_reason"] = event.get("reason") or "netfilter-reject"
                 payload["verdict_confidence"] = event.get("verdict_confidence") or "medium"
@@ -137,18 +141,23 @@ class PacketTracker:
             entry["last_observed_at"] = now_ts
             entry["dirty"] = True
             self._maybe_mark_complete(entry)
-        return packet_uid
+        return correlation_key
 
-    def _new_entry(self, packet_uid: str, now_ts: float) -> Dict[str, Any]:
+    def _new_entry(self, correlation_key: str, now_ts: float) -> Dict[str, Any]:
         return {
-            "packet_uid": packet_uid,
+            "correlation_key": correlation_key,
             "payload": {
-                "packet_uid": packet_uid,
+                "correlation_key": correlation_key,
+                "correlation_source": None,
+                "packet_id": None,
+                "packet_uid": None,
+                "skb_mark": None,
                 "verdict": "pending",
                 "verdict_reason": None,
                 "verdict_confidence": None,
                 "raw_present": False,
                 "capture_sources": [],
+                "capture_metadata": None,
                 "telemetry_metadata": None,
             },
             "persisted": False,
@@ -160,6 +169,34 @@ class PacketTracker:
             "last_persisted_at": 0.0,
         }
 
+    def _ensure_correlation(self, payload: Dict[str, Any]) -> Optional[str]:
+        skb_mark = payload.get("skb_mark")
+        if payload.get("packet_id") is None and skb_mark is not None:
+            payload["packet_id"] = packet_id_from_mark(skb_mark)
+        if payload.get("verdict_hint") is None and skb_mark is not None:
+            payload["verdict_hint"] = verdict_from_mark(skb_mark)
+
+        packet_id = payload.get("packet_id")
+        if packet_id not in (None, ""):
+            packet_id = str(packet_id)
+            payload["packet_id"] = packet_id
+            payload["correlation_key"] = f"pid:{packet_id}"
+            if not payload.get("correlation_source"):
+                payload["correlation_source"] = "kernel_mark"
+            return payload["correlation_key"]
+
+        packet_uid = payload.get("packet_uid")
+        if packet_uid in (None, ""):
+            try:
+                packet_uid = build_packet_uid(payload)
+            except Exception:
+                return None
+        payload["packet_uid"] = packet_uid
+        payload["correlation_key"] = f"uid:{packet_uid}"
+        if not payload.get("correlation_source"):
+            payload["correlation_source"] = "legacy_hash"
+        return payload["correlation_key"]
+
     def _add_capture_source(self, payload: Dict[str, Any], source: str) -> None:
         capture_sources = payload.setdefault("capture_sources", [])
         if source not in capture_sources:
@@ -182,7 +219,14 @@ class PacketTracker:
             changed = True
 
         iface = pkt_info.get("iface")
-        if iface and not payload.get("ingress_if"):
+        if iface and pkt_info.get("capture_iface") and not payload.get("capture_iface"):
+            payload["capture_iface"] = pkt_info.get("capture_iface")
+            changed = True
+        elif iface and pkt_info.get("capture_metadata") and not payload.get("capture_iface"):
+            payload["capture_iface"] = iface
+            changed = True
+
+        if iface and not pkt_info.get("capture_metadata") and not payload.get("ingress_if"):
             payload["ingress_if"] = iface
             payload["ingress_seen_at"] = _utcnow()
             changed = True
@@ -267,11 +311,11 @@ class PacketTracker:
         while not self._stop_event.is_set():
             time.sleep(0.05)
             due_entries: List[Dict[str, Any]] = []
-            expired_uids: List[str] = []
+            expired_keys: List[str] = []
             now_ts = time.time()
 
             with self._lock:
-                for packet_uid, entry in list(self._entries.items()):
+                for correlation_key, entry in list(self._entries.items()):
                     age = now_ts - entry["last_observed_at"]
                     if not entry["finalized"] and age >= self._finalize_delay_seconds:
                         entry["finalized"] = True
@@ -290,7 +334,7 @@ class PacketTracker:
                     if should_flush:
                         due_entries.append(
                             {
-                                "packet_uid": entry["packet_uid"],
+                                "correlation_key": entry["correlation_key"],
                                 "payload": dict(entry["payload"]),
                             }
                         )
@@ -298,10 +342,10 @@ class PacketTracker:
                         if entry["finalized"] and not entry["stats_recorded"]:
                             self._record_stats(entry["payload"])
                             entry["stats_recorded"] = True
-                        expired_uids.append(packet_uid)
+                        expired_keys.append(correlation_key)
 
-                for packet_uid in expired_uids:
-                    self._entries.pop(packet_uid, None)
+                for correlation_key in expired_keys:
+                    self._entries.pop(correlation_key, None)
 
             for entry in due_entries:
                 self._persist(entry)
@@ -318,13 +362,13 @@ class PacketTracker:
             fut = asyncio.run_coroutine_threadsafe(web_db.upsert_packet(payload), web_loop)
             fut.result(timeout=settings.packet_tracker_persist_timeout_seconds)
             with self._lock:
-                current = self._entries.get(entry["packet_uid"])
+                current = self._entries.get(entry["correlation_key"])
                 if current is not None:
                     current["persisted"] = True
                     current["dirty"] = False
                     current["last_persisted_at"] = time.time()
         except Exception:
-            logger.exception("Failed to persist packet %s", entry["packet_uid"])
+            logger.exception("Failed to persist packet %s", entry["correlation_key"])
 
     def _record_stats(self, payload: Dict[str, Any]) -> None:
         capture_sources = set(payload.get("capture_sources") or [])
@@ -333,6 +377,10 @@ class PacketTracker:
             self._stats["persisted_with_raw"] += 1
         else:
             self._stats["persisted_without_raw"] += 1
+        if payload.get("correlation_source") == "kernel_mark":
+            self._stats["persisted_kernel_mark"] = self._stats.get("persisted_kernel_mark", 0) + 1
+        else:
+            self._stats["persisted_legacy_hash"] = self._stats.get("persisted_legacy_hash", 0) + 1
 
         if capture_sources == {"af_packet"}:
             self._stats["persisted_af_packet_only"] += 1
@@ -352,6 +400,8 @@ class PacketTracker:
         active_af_packet_only = 0
         active_telemetry_only = 0
         active_merged = 0
+        active_kernel_mark = 0
+        active_legacy_hash = 0
         for entry in active_entries:
             capture_sources = set(entry["payload"].get("capture_sources") or [])
             if capture_sources == {"af_packet"}:
@@ -360,6 +410,10 @@ class PacketTracker:
                 active_telemetry_only += 1
             else:
                 active_merged += 1
+            if entry["payload"].get("correlation_source") == "kernel_mark":
+                active_kernel_mark += 1
+            else:
+                active_legacy_hash += 1
 
         return {
             "active_total": active_total,
@@ -368,6 +422,8 @@ class PacketTracker:
             "active_af_packet_only": active_af_packet_only,
             "active_telemetry_only": active_telemetry_only,
             "active_merged": active_merged,
+            "active_kernel_mark": active_kernel_mark,
+            "active_legacy_hash": active_legacy_hash,
             "cumulative": stats,
         }