zeek/scripts/base/protocols/conn/removal-hooks.zeek
Jon Siwek 05cf511f18 GH-1119: add base/protcols/conn/removal-hooks.zeek
This adds two new functions: `Conn::register_removal_hook()` and
`Conn::unregister_removal_hook()` for registering a hook function to be
called back during `connection_state_remove`.  The benefit of using hook
callback approach is better scalability: the overhead of unrelated
protocols having to dispatch no-op `connection_state_remove` handlers is
avoided.
2020-09-11 12:12:10 -07:00

79 lines
2.5 KiB
Text

##! Adds a framework for registering "connection removal hooks".
##! All registered hooks for a given connection get run within the
##! :zeek:see:`connection_state_remove` event for that connection.
##! This functionality is useful from a performance/scaling concern:
##! if every new protocol-analysis script uses
##! :zeek:see:`connection_state_remove` to implement its finalization/cleanup
##! logic, then all connections take the performance hit of dispatching that
##! event, even if they aren't related to that specific protocol.
module Conn;
export {
## A hook function for use with either :zeek:see:`Conn::register_removal_hook`
## or :zeek:see:`Conn::unregister_removal_hook`. The :zeek:see:`connection`
## argument refers to the connection currently being removed within a
## :zeek:see:`connection_state_remove` event.
type RemovalHook: hook(c: connection);
## Register a hook that will later be called during a connection's
## :zeek:see:`connection_state_remove` event.
##
## c: The associated connection whose :zeek:see:`connection_state_remove`
## event should trigger a callback to *hk*.
##
## hk: The hook function to use as a callback.
##
## Returns: false if the provided hook was previously registered, else true.
global register_removal_hook: function(c: connection, hk: RemovalHook): bool;
## Unregister a hook that would have been called during a connection's
## :zeek:see:`connection_state_remove` event such that it will no longer
## be called.
##
## c: The associated connection whose :zeek:see:`connection_state_remove`
## event could have triggered a callback to *hk*.
##
## hk: The hook function that would have been used as a callback.
##
## Returns: true if the provided hook was previously registered, else false.
global unregister_removal_hook: function(c: connection, hk: RemovalHook): bool;
}
redef record connection += {
removal_hooks: set[RemovalHook] &optional;
};
function register_removal_hook(c: connection, hk: RemovalHook): bool
{
if ( c?$removal_hooks )
{
if ( hk in c$removal_hooks )
return F;
add c$removal_hooks[hk];
return T;
}
c$removal_hooks = set(hk);
return T;
}
function unregister_removal_hook(c: connection, hk: RemovalHook): bool
{
if ( ! c?$removal_hooks )
return F;
if ( hk !in c$removal_hooks )
return F;
delete c$removal_hooks[hk];
return T;
}
event connection_state_remove(c: connection) &priority=-3
{
if ( c?$removal_hooks )
for ( removal_hook in c$removal_hooks )
hook removal_hook(c);
}