From 82b3a4048feee81169b09a8f4cfc70c429e53354 Mon Sep 17 00:00:00 2001
From: Benjamin Bannier <benjamin.bannier@corelight.com>
Date: Tue, 10 Oct 2023 13:05:35 +0200
Subject: [PATCH] Use LDAP `RemovalHook` instead of implementing
 `connection_state_remove`

---
 scripts/base/protocols/ldap/main.zeek | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/scripts/base/protocols/ldap/main.zeek b/scripts/base/protocols/ldap/main.zeek
index fb3efb5bb3..e4453f58c6 100644
--- a/scripts/base/protocols/ldap/main.zeek
+++ b/scripts/base/protocols/ldap/main.zeek
@@ -30,6 +30,9 @@ export {
   ## Default logging policy hook for LDAP_SEARCH_LOG.
   global log_policy_search: Log::PolicyHook;
 
+  ## LDAP finalization hook.
+  global finalize_ldap: Conn::RemovalHook;
+
   #############################################################################
   # This is the format of ldap.log (ldap operations minus search-related)
   # Each line represents a unique connection+message_id (requests/responses)
@@ -161,6 +164,7 @@ event zeek_init() &priority=5 {
 
 #############################################################################
 function set_session(c: connection, message_id: int, opcode: LDAP::ProtocolOpcode) {
+  Conn::register_removal_hook(c, finalize_ldap);
 
   if (! c?$ldap_messages )
     c$ldap_messages = table();
@@ -363,8 +367,7 @@ event LDAP::bindreq(c: connection,
 }
 
 #############################################################################
-event connection_state_remove(c: connection) {
-
+hook finalize_ldap(c: connection) {
   # log any "pending" unlogged LDAP messages/searches
 
   if ( c?$ldap_messages && (|c$ldap_messages| > 0) ) {