diff --git a/src/supervisor/supervisor.bif b/src/supervisor/supervisor.bif index b97c562deb..6bc6db17d6 100644 --- a/src/supervisor/supervisor.bif +++ b/src/supervisor/supervisor.bif @@ -98,3 +98,15 @@ function Supervisor::__is_supervisor%(%): bool %{ return val_mgr->GetBool(zeek::supervisor_mgr != nullptr); %} + +function Supervisor::__stem_pid%(%): int + %{ + if ( zeek::supervisor_mgr ) + return val_mgr->GetInt(zeek::supervisor_mgr->StemPID()); + + if ( zeek::Supervisor::ThisNode() ) + return val_mgr->GetInt(zeek::Supervisor::ThisNode()->parent_pid); + + builtin_error("supervisor mode not enabled and not a supervised node"); + return val_mgr->GetInt(-1); + %} diff --git a/testing/btest/Baseline/supervisor.revive-leaf/zeek.node.out b/testing/btest/Baseline/supervisor.revive-leaf/zeek.node.out new file mode 100644 index 0000000000..101b306cf0 --- /dev/null +++ b/testing/btest/Baseline/supervisor.revive-leaf/zeek.node.out @@ -0,0 +1,2 @@ +supervised node zeek_init() +supervised node zeek_done() diff --git a/testing/btest/Baseline/supervisor.revive-leaf/zeek.supervisor.out b/testing/btest/Baseline/supervisor.revive-leaf/zeek.supervisor.out new file mode 100644 index 0000000000..7e3ffd50b4 --- /dev/null +++ b/testing/btest/Baseline/supervisor.revive-leaf/zeek.supervisor.out @@ -0,0 +1,7 @@ +supervisor zeek_init() +supervisor connected to peer +supervisor lost peer +supervisor connected to peer +supervisor lost peer +supervisor connected to peer +supervisor zeek_done() diff --git a/testing/btest/Baseline/supervisor.revive-stem/zeek.node.out b/testing/btest/Baseline/supervisor.revive-stem/zeek.node.out new file mode 100644 index 0000000000..101b306cf0 --- /dev/null +++ b/testing/btest/Baseline/supervisor.revive-stem/zeek.node.out @@ -0,0 +1,2 @@ +supervised node zeek_init() +supervised node zeek_done() diff --git a/testing/btest/Baseline/supervisor.revive-stem/zeek.supervisor.out b/testing/btest/Baseline/supervisor.revive-stem/zeek.supervisor.out new file mode 100644 index 0000000000..7e3ffd50b4 --- /dev/null +++ b/testing/btest/Baseline/supervisor.revive-stem/zeek.supervisor.out @@ -0,0 +1,7 @@ +supervisor zeek_init() +supervisor connected to peer +supervisor lost peer +supervisor connected to peer +supervisor lost peer +supervisor connected to peer +supervisor zeek_done() diff --git a/testing/btest/supervisor/revive-leaf.zeek b/testing/btest/supervisor/revive-leaf.zeek new file mode 100644 index 0000000000..448970283b --- /dev/null +++ b/testing/btest/supervisor/revive-leaf.zeek @@ -0,0 +1,70 @@ +# @TEST-PORT: BROKER_PORT +# @TEST-EXEC: btest-bg-run zeek zeek -j -b %INPUT +# @TEST-EXEC: btest-bg-wait 20 +# @TEST-EXEC: btest-diff zeek/supervisor.out +# @TEST-EXEC: btest-diff zeek/node.out + +# So the supervised node doesn't terminate right away. +redef exit_only_after_terminate=T; + +global supervisor_output_file: file; +global node_output_file: file; +global topic = "test-topic"; +global peers_added = 0; + +event kill_self() + { + system(fmt("kill %s", getpid())); + } + +event zeek_init() + { + if ( Supervisor::is_supervisor() ) + { + Broker::subscribe(topic); + Broker::listen("127.0.0.1", to_port(getenv("BROKER_PORT"))); + supervisor_output_file = open("supervisor.out"); + print supervisor_output_file, "supervisor zeek_init()"; + local sn = Supervisor::NodeConfig($name="grault"); + local res = Supervisor::create(sn); + + if ( res != "" ) + print supervisor_output_file, res; + } + else + { + Broker::subscribe(topic); + Broker::peer("127.0.0.1", to_port(getenv("BROKER_PORT"))); + node_output_file = open("node.out"); + print node_output_file, "supervised node zeek_init()"; + } + } + +event Broker::peer_added(endpoint: Broker::EndpointInfo, msg: string) + { + ++peers_added; + + if ( Supervisor::is_supervisor() ) + { + print supervisor_output_file, "supervisor connected to peer"; + + if ( peers_added == 3 ) + terminate(); + else + Broker::publish(topic, kill_self); + } + } + +event Broker::peer_lost(endpoint: Broker::EndpointInfo, msg: string) + { + if ( Supervisor::is_supervisor() ) + print supervisor_output_file, "supervisor lost peer"; + } + +event zeek_done() + { + if ( Supervisor::is_supervisor() ) + print supervisor_output_file, "supervisor zeek_done()"; + else + print node_output_file, "supervised node zeek_done()"; + } diff --git a/testing/btest/supervisor/revive-stem.zeek b/testing/btest/supervisor/revive-stem.zeek new file mode 100644 index 0000000000..37dc3e5053 --- /dev/null +++ b/testing/btest/supervisor/revive-stem.zeek @@ -0,0 +1,65 @@ +# @TEST-PORT: BROKER_PORT +# @TEST-EXEC: btest-bg-run zeek zeek -j -b %INPUT +# @TEST-EXEC: btest-bg-wait 20 +# @TEST-EXEC: btest-diff zeek/supervisor.out +# @TEST-EXEC: btest-diff zeek/node.out + +# So the supervised node doesn't terminate right away. +redef exit_only_after_terminate=T; + +global supervisor_output_file: file; +global node_output_file: file; +global topic = "test-topic"; +global peers_added = 0; + +event zeek_init() + { + if ( Supervisor::is_supervisor() ) + { + Broker::subscribe(topic); + Broker::listen("127.0.0.1", to_port(getenv("BROKER_PORT"))); + supervisor_output_file = open("supervisor.out"); + print supervisor_output_file, "supervisor zeek_init()"; + local sn = Supervisor::NodeConfig($name="grault"); + local res = Supervisor::create(sn); + + if ( res != "" ) + print supervisor_output_file, res; + } + else + { + Broker::subscribe(topic); + Broker::peer("127.0.0.1", to_port(getenv("BROKER_PORT"))); + node_output_file = open("node.out"); + print node_output_file, "supervised node zeek_init()"; + } + } + +event Broker::peer_added(endpoint: Broker::EndpointInfo, msg: string) + { + ++peers_added; + + if ( Supervisor::is_supervisor() ) + { + print supervisor_output_file, "supervisor connected to peer"; + + if ( peers_added == 3 ) + terminate(); + else + system(fmt("kill %s", Supervisor::__stem_pid())); + } + } + +event Broker::peer_lost(endpoint: Broker::EndpointInfo, msg: string) + { + if ( Supervisor::is_supervisor() ) + print supervisor_output_file, "supervisor lost peer"; + } + +event zeek_done() + { + if ( Supervisor::is_supervisor() ) + print supervisor_output_file, "supervisor zeek_done()"; + else + print node_output_file, "supervised node zeek_done()"; + }