From 9d4025804d183f506f69d9ae47bdc4bf5fbd6c99 Mon Sep 17 00:00:00 2001
From: Arne Welzel <arne.welzel@corelight.com>
Date: Thu, 12 Sep 2024 22:20:47 +0200
Subject: [PATCH] gtpv1: Replace connection_state_remove() with RemovalHook

---
 scripts/base/packet-protocols/gtpv1/main.zeek | 13 ++++++++++++-
 src/packet_analysis/protocol/gtpv1/GTPv1.cc   |  7 ++++++-
 src/packet_analysis/protocol/gtpv1/events.bif | 10 ++++++++++
 3 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/scripts/base/packet-protocols/gtpv1/main.zeek b/scripts/base/packet-protocols/gtpv1/main.zeek
index 0ad181f0f7..9373a97379 100644
--- a/scripts/base/packet-protocols/gtpv1/main.zeek
+++ b/scripts/base/packet-protocols/gtpv1/main.zeek
@@ -3,11 +3,15 @@ module PacketAnalyzer::GTPV1;
 # This needs to be loaded here so the function is available. Function BIFs normally aren't
 # loaded until after the packet analysis init scripts are run, and then zeek complains it
 # can't find the function.
+@load base/bif/plugins/Zeek_GTPv1.events.bif
 @load base/bif/plugins/Zeek_GTPv1.functions.bif
 
 # Needed for port registration for BPF
 @load base/frameworks/analyzer/main
 
+# Needed to register Conn::RemovalHook
+@load base/protocols/conn/removal-hooks
+
 export {
         ## Default analyzer
         const default_analyzer: PacketAnalyzer::Tag = PacketAnalyzer::ANALYZER_IP &redef;
@@ -21,7 +25,14 @@ event zeek_init() &priority=20
 	PacketAnalyzer::register_for_ports(PacketAnalyzer::ANALYZER_UDP, PacketAnalyzer::ANALYZER_GTPV1, gtpv1_ports);
 	}
 
-event connection_state_remove(c: connection)
+# The analyzer keeps a BinPac interpreter per connection
+# that isn't cleaned due to being stored in a global table.
+hook finalize_gtpv1(c: connection)
 	{
 	remove_gtpv1_connection(c$id);
 	}
+
+event new_gtpv1_state(c: connection)
+	{
+	Conn::register_removal_hook(c, finalize_gtpv1);
+	}
diff --git a/src/packet_analysis/protocol/gtpv1/GTPv1.cc b/src/packet_analysis/protocol/gtpv1/GTPv1.cc
index 9db7b3c9ab..d724725b54 100644
--- a/src/packet_analysis/protocol/gtpv1/GTPv1.cc
+++ b/src/packet_analysis/protocol/gtpv1/GTPv1.cc
@@ -22,9 +22,14 @@ bool GTPv1_Analyzer::AnalyzePacket(size_t len, const uint8_t* data, Packet* pack
     zeek::detail::ConnKey conn_key = conn->Key();
 
     auto cm_it = conn_map.find(conn_key);
-    if ( cm_it == conn_map.end() )
+    if ( cm_it == conn_map.end() ) {
         cm_it = conn_map.insert(cm_it, {conn_key, std::make_unique<binpac::GTPv1::GTPv1_Conn>(this)});
 
+        // Let script land know about the state we created, so it will
+        // register a conn removal hook for cleanup.
+        BifEvent::enqueue_new_gtpv1_state(nullptr, conn);
+    }
+
     try {
         cm_it->second->set_raw_packet(packet);
         cm_it->second->NewData(packet->is_orig, data, data + len);
diff --git a/src/packet_analysis/protocol/gtpv1/events.bif b/src/packet_analysis/protocol/gtpv1/events.bif
index b3bac93a2e..ed2da358bb 100644
--- a/src/packet_analysis/protocol/gtpv1/events.bif
+++ b/src/packet_analysis/protocol/gtpv1/events.bif
@@ -1,3 +1,13 @@
+module GLOBAL;
+
+## Generated when a new GTP analyzer is instantiated for a connection.
+##
+## This event exists to install a connection removal hook to clear
+## internal per-connection GTPv1 state.
+##
+## c: The connection for which the analyzer is instantiated.
+event new_gtpv1_state%(c: connection%);
+
 ## Generated for any GTP message with a GTPv1 header.
 ##
 ## c: The connection over which the message is sent.