From 39890fda5e5167793851b3b6958e92c79e53c224 Mon Sep 17 00:00:00 2001 From: Tim Wojtulewicz Date: Thu, 7 Mar 2024 17:01:23 -0700 Subject: [PATCH] Add a method to register an event handler to a std::function via C++ --- src/EventRegistry.h | 1 + src/Func.cc | 5 +++++ src/Func.h | 1 + src/Stmt.cc | 14 +++++++++++++ src/Stmt.h | 21 +++++++++++++++++++ src/StmtEnums.h | 5 +++-- .../btest/Baseline/plugins.protocol/output | 4 ++++ .../btest/plugins/protocol-plugin/src/Foo.cc | 15 +++++++++++++ .../btest/plugins/protocol-plugin/src/Foo.h | 1 + 9 files changed, 65 insertions(+), 2 deletions(-) diff --git a/src/EventRegistry.h b/src/EventRegistry.h index f164dfe60e..eb9b9b0026 100644 --- a/src/EventRegistry.h +++ b/src/EventRegistry.h @@ -12,6 +12,7 @@ #include #include "zeek/IntrusivePtr.h" +#include "zeek/ZeekArgs.h" namespace zeek { diff --git a/src/Func.cc b/src/Func.cc index fcf0373f86..c08657e3e4 100644 --- a/src/Func.cc +++ b/src/Func.cc @@ -143,6 +143,11 @@ void Func::AddBody(detail::StmtPtr /* new_body */, const std::vector(std::move(body)); + AddBody(stmt, {}, priority); +} + void Func::SetScope(detail::ScopePtr newscope) { scope = std::move(newscope); } FuncPtr Func::DoClone() { diff --git a/src/Func.h b/src/Func.h index fc01ba53d6..a2f3e9c31c 100644 --- a/src/Func.h +++ b/src/Func.h @@ -114,6 +114,7 @@ public: void AddBody(detail::StmtPtr new_body, const std::vector& new_inits, size_t new_frame_size, int priority = 0); void AddBody(detail::StmtPtr new_body, size_t new_frame_size); + void AddBody(detail::StdFunctionStmt::FunctionVariant body, int priority = 0); virtual void SetScope(detail::ScopePtr newscope); virtual detail::ScopePtr GetScope() const { return scope; } diff --git a/src/Stmt.cc b/src/Stmt.cc index 2d36a7632a..e478d02358 100644 --- a/src/Stmt.cc +++ b/src/Stmt.cc @@ -53,6 +53,7 @@ const char* stmt_name(StmtTag t) { "null", "assert", "extern", + "std-function", }; #pragma GCC diagnostic push @@ -2083,4 +2084,17 @@ TraversalCode WhenStmt::Traverse(TraversalCallback* cb) const { HANDLE_TC_STMT_POST(tc); } +ValPtr StdFunctionStmt::Exec(Frame* f, StmtFlowType& flow) { + zeek::Args args = *f->GetFuncArgs(); + + if ( func.index() == 0 ) { + flow = FLOW_NEXT; + std::get<0>(func)(args); + } + else + std::get<1>(func)(args, flow); + + return nullptr; +} + } // namespace zeek::detail diff --git a/src/Stmt.h b/src/Stmt.h index 054bf263eb..13a9d0550c 100644 --- a/src/Stmt.h +++ b/src/Stmt.h @@ -780,4 +780,25 @@ protected: int expected_len; }; +// Statement that calls a std::function. These can be added to a Func body +// to directly all a C++ method. +class StdFunctionStmt : public Stmt { +public: + using FunctionVariant = + std::variant, std::function>; + StdFunctionStmt(FunctionVariant f) : Stmt(STMT_STD_FUNCTION), func(std::move(f)) {} + + ValPtr Exec(Frame* f, StmtFlowType& flow) override; + + StmtPtr Duplicate() override { + reporter->Error("Duplicate() on StdFunctionStmt not implemented"); + return {zeek::NewRef{}, this}; + } + + TraversalCode Traverse(TraversalCallback* cb) const override { return TC_CONTINUE; } + +private: + FunctionVariant func; +}; + } // namespace zeek::detail diff --git a/src/StmtEnums.h b/src/StmtEnums.h index 7cac90deae..db5ab12aa2 100644 --- a/src/StmtEnums.h +++ b/src/StmtEnums.h @@ -31,8 +31,9 @@ enum StmtTag { STMT_ZAM, // a ZAM function body STMT_NULL, STMT_ASSERT, - STMT_EXTERN, // for custom Stmt subclasses provided by plugins -#define NUM_STMTS (int(STMT_EXTERN) + 1) + STMT_EXTERN, // for custom Stmt subclasses provided by plugins + STMT_STD_FUNCTION, // statements that call a std::function from c++ +#define NUM_STMTS (int(STMT_STD_FUNCTION) + 1) }; enum StmtFlowType { diff --git a/testing/btest/Baseline/plugins.protocol/output b/testing/btest/Baseline/plugins.protocol/output index 06f824361c..cc2c789065 100644 --- a/testing/btest/Baseline/plugins.protocol/output +++ b/testing/btest/Baseline/plugins.protocol/output @@ -4,6 +4,10 @@ Demo::Foo - A Foo test analyzer (dynamic, version 1.0.0) [Event] foo_message === +c++ connection_established lambda handler, received 1 arguments +c++ connection_established member handler, received 1 arguments foo_message, [orig_h=::1, orig_p=37927/tcp, resp_h=::1, resp_p=4242/tcp], Hello, Foo!\x0a === +c++ connection_established lambda handler, received 1 arguments +c++ connection_established member handler, received 1 arguments foo_message, [orig_h=::1, orig_p=37927/tcp, resp_h=::1, resp_p=4243/tcp], Hello, Foo!\x0a diff --git a/testing/btest/plugins/protocol-plugin/src/Foo.cc b/testing/btest/plugins/protocol-plugin/src/Foo.cc index 27f0aa9aa4..100ebcb9c4 100644 --- a/testing/btest/plugins/protocol-plugin/src/Foo.cc +++ b/testing/btest/plugins/protocol-plugin/src/Foo.cc @@ -1,19 +1,34 @@ #include "Foo.h" +#include "zeek/EventRegistry.h" #include "zeek/analyzer/protocol/tcp/TCP_Reassembler.h" #include "events.bif.h" #include "foo_pac.h" using namespace btest::plugin::Demo_Foo; +using namespace std::placeholders; Foo::Foo(zeek::Connection* conn) : zeek::analyzer::tcp::TCP_ApplicationAnalyzer("Foo", conn) { interp = new binpac::Foo::Foo_Conn(this); + + auto handler = zeek::event_registry->Lookup("connection_established"); + if ( handler ) { + handler->GetFunc()->AddBody([](const zeek::Args& args) { + printf("c++ connection_established lambda handler, received %zu arguments\n", args.size()); + }); + + handler->GetFunc()->AddBody(std::bind(&Foo::ConnectionEstablishedHandler, this, _1)); + } } Foo::~Foo() { delete interp; } +void Foo::ConnectionEstablishedHandler(const zeek::Args& args) { + printf("c++ connection_established member handler, received %zu arguments\n", args.size()); +} + void Foo::Done() { zeek::analyzer::tcp::TCP_ApplicationAnalyzer::Done(); diff --git a/testing/btest/plugins/protocol-plugin/src/Foo.h b/testing/btest/plugins/protocol-plugin/src/Foo.h index b373f4a747..f1a89bac12 100644 --- a/testing/btest/plugins/protocol-plugin/src/Foo.h +++ b/testing/btest/plugins/protocol-plugin/src/Foo.h @@ -23,6 +23,7 @@ public: virtual void EndpointEOF(bool is_orig); static zeek::analyzer::Analyzer* Instantiate(zeek::Connection* conn) { return new Foo(conn); } + void ConnectionEstablishedHandler(const zeek::Args& args); protected: binpac::Foo::Foo_Conn* interp;