From 2dc98acd1ffef1ed8714577fa09711b882460158 Mon Sep 17 00:00:00 2001
From: Arne Welzel <arne.welzel@corelight.com>
Date: Fri, 14 Feb 2025 17:25:39 -0800
Subject: [PATCH] packet_analysis: Track data spans of packet analyzers

Do not just track the analyzer instance in the stack, but also the
data span it is given. This allows to extract more information on-demand
during event processing.

TrackAnalyzer() is technically a public API, but no one should use it
outside of the Analyzer's Forward methods itself.
---
 src/packet_analysis/Analyzer.cc |  4 ++--
 src/packet_analysis/Manager.cc  | 12 +++++++++++-
 src/packet_analysis/Manager.h   | 22 ++++++++++++++++++++--
 3 files changed, 33 insertions(+), 5 deletions(-)

diff --git a/src/packet_analysis/Analyzer.cc b/src/packet_analysis/Analyzer.cc
index e61ef83238..17f6987b8d 100644
--- a/src/packet_analysis/Analyzer.cc
+++ b/src/packet_analysis/Analyzer.cc
@@ -114,7 +114,7 @@ bool Analyzer::ForwardPacket(size_t len, const uint8_t* data, Packet* packet, ui
     DBG_LOG(DBG_PACKET_ANALYSIS, "Analysis in %s succeeded, next layer identifier is %#x.", GetAnalyzerName(),
             identifier);
 
-    packet_mgr->TrackAnalyzer(inner_analyzer.get());
+    packet_mgr->TrackAnalyzer(inner_analyzer.get(), len, data);
     return inner_analyzer->AnalyzePacket(len, data, packet);
 }
 
@@ -130,7 +130,7 @@ bool Analyzer::ForwardPacket(size_t len, const uint8_t* data, Packet* packet) co
         return false;
     }
 
-    packet_mgr->TrackAnalyzer(inner_analyzer.get());
+    packet_mgr->TrackAnalyzer(inner_analyzer.get(), len, data);
     return inner_analyzer->AnalyzePacket(len, data, packet);
 }
 
diff --git a/src/packet_analysis/Manager.cc b/src/packet_analysis/Manager.cc
index c04cb7b17e..3c4c3437b1 100644
--- a/src/packet_analysis/Manager.cc
+++ b/src/packet_analysis/Manager.cc
@@ -232,7 +232,7 @@ zeek::VectorValPtr Manager::BuildAnalyzerHistory() const {
     auto history = zeek::make_intrusive<zeek::VectorVal>(zeek::id::string_vec);
 
     for ( unsigned int i = 0; i < analyzer_stack.size(); i++ ) {
-        auto analyzer_name = analyzer_stack[i]->GetAnalyzerName();
+        auto analyzer_name = analyzer_stack[i].analyzer->GetAnalyzerName();
         history->Assign(i, make_intrusive<StringVal>(analyzer_name));
     }
 
@@ -249,3 +249,13 @@ void Manager::ReportUnknownProtocol(const std::string& analyzer, uint32_t protoc
         }
     }
 }
+
+std::vector<zeek::Span<const uint8_t>> Manager::GetAnalyzerData(const AnalyzerPtr& analyzer) {
+    std::vector<zeek::Span<const uint8_t>> result;
+    for ( const auto [sa, span] : analyzer_stack ) {
+        if ( sa == analyzer.get() )
+            result.push_back(span);
+    }
+
+    return result;
+}
diff --git a/src/packet_analysis/Manager.h b/src/packet_analysis/Manager.h
index 5497382714..00f36c2b8e 100644
--- a/src/packet_analysis/Manager.h
+++ b/src/packet_analysis/Manager.h
@@ -182,8 +182,21 @@ public:
      * call to ProcessPacket() but maintained throughout calls to ProcessInnerPacket().
      *
      * @param analyzer The analyzer to track.
+     * @param len The remaining length of the data in the packet being processed.
+     * @param data A pointer to the data
      */
-    void TrackAnalyzer(Analyzer* analyzer) { analyzer_stack.push_back(analyzer); }
+    void TrackAnalyzer(const Analyzer* analyzer, size_t len, const uint8_t* data) {
+        analyzer_stack.push_back({analyzer, {data, len}});
+    }
+
+    /**
+     * Get all tracked data spans for a given analyzer instance.
+     *
+     * @analyzer The analyzer instance.
+     *
+     * @returns An array of data spans.
+     */
+    std::vector<zeek::Span<const uint8_t>> GetAnalyzerData(const AnalyzerPtr& analyzer);
 
 private:
     /**
@@ -234,7 +247,12 @@ private:
     uint64_t total_not_processed = 0;
     iosource::PktDumper* unprocessed_dumper = nullptr;
 
-    std::vector<Analyzer*> analyzer_stack;
+    struct StackEntry {
+        const Analyzer* analyzer;
+        zeek::Span<const uint8_t> data; // Start of this layer, limited by span's size.
+    };
+
+    std::vector<StackEntry> analyzer_stack;
 };
 
 } // namespace packet_analysis