From 88c86cc7d4f093e633604bee0dd096a29f066937 Mon Sep 17 00:00:00 2001
From: Jan Grashoefer <jan.grashoefer@corelight.com>
Date: Fri, 21 Apr 2023 17:38:19 +0200
Subject: [PATCH] Add hook into cluster connection setup.

---
 scripts/base/frameworks/cluster/main.zeek        | 11 ++++++-----
 .../frameworks/cluster/setup-connections.zeek    | 16 ++++++++++++++++
 2 files changed, 22 insertions(+), 5 deletions(-)

diff --git a/scripts/base/frameworks/cluster/main.zeek b/scripts/base/frameworks/cluster/main.zeek
index 3f3b222ce2..ee2cf71681 100644
--- a/scripts/base/frameworks/cluster/main.zeek
+++ b/scripts/base/frameworks/cluster/main.zeek
@@ -191,6 +191,12 @@ export {
 		id: string                &optional;
 	};
 
+	## Record to represent a cluster node including its name.
+	type NamedNode: record {
+		name: string;
+		node: Node;
+	};
+
 	## This function can be called at any time to determine if the cluster
 	## framework is being enabled for this run.
 	##
@@ -281,11 +287,6 @@ export {
 # Track active nodes per type.
 global active_node_ids: table[NodeType] of set[string];
 
-type NamedNode: record {
-	name: string;
-	node: Node;
-};
-
 function nodes_with_type(node_type: NodeType): vector of NamedNode
 	{
 	local rval: vector of NamedNode = vector();
diff --git a/scripts/base/frameworks/cluster/setup-connections.zeek b/scripts/base/frameworks/cluster/setup-connections.zeek
index 55ab8a460a..848555419c 100644
--- a/scripts/base/frameworks/cluster/setup-connections.zeek
+++ b/scripts/base/frameworks/cluster/setup-connections.zeek
@@ -7,6 +7,15 @@
 
 module Cluster;
 
+export {
+	## This hook is called when the local node connects to other nodes based on
+	## the given cluster layout. Breaking from the hook will prevent connection
+	## establishment.
+	##
+	## connectee: The node to connect to.
+	global connect_node_hook: hook(connectee: NamedNode);
+}
+
 function connect_peer(node_type: NodeType, node_name: string)
 	{
 	local nn = nodes_with_type(node_type);
@@ -17,12 +26,15 @@ function connect_peer(node_type: NodeType, node_name: string)
 
 		if ( n$name != node_name )
 			next;
+		if ( ! hook connect_node_hook(n) )
+			return;
 
 		local status = Broker::peer(cat(n$node$ip), n$node$p,
 		                            Cluster::retry_interval);
 		Cluster::log(fmt("initiate peering with %s:%s, retry=%s, status=%s",
 		                 n$node$ip, n$node$p, Cluster::retry_interval,
 		                 status));
+		return;
 		}
 	}
 
@@ -33,6 +45,10 @@ function connect_peers_with_type(node_type: NodeType)
 	for ( i in nn )
 		{
 		local n = nn[i];
+
+		if ( ! hook connect_node_hook(n) )
+			next;
+
 		local status = Broker::peer(cat(n$node$ip), n$node$p,
 		                            Cluster::retry_interval);
 		Cluster::log(fmt("initiate peering with %s:%s, retry=%s, status=%s",