From 9de1dd16d5d55c84c27e562f7b4ded91db82a985 Mon Sep 17 00:00:00 2001
From: Tim Wojtulewicz <tim@corelight.com>
Date: Tue, 19 Mar 2024 09:07:45 -0700
Subject: [PATCH] Remove variant from StdFunctionStmt

The variant ended up conflicting with std::bind, which resulted in failures
on the btest invoking it. Change back to a single function that takes a
flow, and default it to a value in Exec.
---
 src/Func.cc                                      | 4 ++--
 src/Func.h                                       | 2 +-
 src/Stmt.cc                                      | 9 +++------
 src/Stmt.h                                       | 7 +++----
 testing/btest/plugins/protocol-plugin/src/Foo.cc | 6 +++---
 testing/btest/plugins/protocol-plugin/src/Foo.h  | 3 ++-
 6 files changed, 14 insertions(+), 17 deletions(-)

diff --git a/src/Func.cc b/src/Func.cc
index c08657e3e4..87783a30a9 100644
--- a/src/Func.cc
+++ b/src/Func.cc
@@ -129,7 +129,7 @@ void Func::AddBody(const detail::FunctionIngredients& ingr, detail::StmtPtr new_
 void Func::AddBody(detail::StmtPtr new_body, const std::vector<detail::IDPtr>& new_inits, size_t new_frame_size,
                    int priority) {
     std::set<EventGroupPtr> groups;
-    AddBody(new_body, new_inits, new_frame_size, priority, groups);
+    AddBody(std::move(new_body), new_inits, new_frame_size, priority, groups);
 }
 
 void Func::AddBody(detail::StmtPtr new_body, size_t new_frame_size) {
@@ -143,7 +143,7 @@ void Func::AddBody(detail::StmtPtr /* new_body */, const std::vector<detail::IDP
     Internal("Func::AddBody called");
 }
 
-void Func::AddBody(detail::StdFunctionStmt::FunctionVariant body, int priority) {
+void Func::AddBody(std::function<void(const zeek::Args&, detail::StmtFlowType&)> body, int priority) {
     auto stmt = zeek::make_intrusive<detail::StdFunctionStmt>(std::move(body));
     AddBody(stmt, {}, priority);
 }
diff --git a/src/Func.h b/src/Func.h
index a2f3e9c31c..8b31976e69 100644
--- a/src/Func.h
+++ b/src/Func.h
@@ -114,7 +114,7 @@ public:
     void AddBody(detail::StmtPtr new_body, const std::vector<detail::IDPtr>& 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);
+    void AddBody(std::function<void(const zeek::Args&, detail::StmtFlowType&)> 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 e478d02358..d2def610c3 100644
--- a/src/Stmt.cc
+++ b/src/Stmt.cc
@@ -2087,12 +2087,9 @@ TraversalCode WhenStmt::Traverse(TraversalCallback* cb) const {
 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);
+    // Set this to NEXT by default. The function can override that if it wants.
+    flow = FLOW_NEXT;
+    func(args, flow);
 
     return nullptr;
 }
diff --git a/src/Stmt.h b/src/Stmt.h
index 13a9d0550c..03653b6fbc 100644
--- a/src/Stmt.h
+++ b/src/Stmt.h
@@ -784,9 +784,8 @@ protected:
 // to directly all a C++ method.
 class StdFunctionStmt : public Stmt {
 public:
-    using FunctionVariant =
-        std::variant<std::function<void(const zeek::Args&)>, std::function<void(const zeek::Args&, StmtFlowType&)>>;
-    StdFunctionStmt(FunctionVariant f) : Stmt(STMT_STD_FUNCTION), func(std::move(f)) {}
+    StdFunctionStmt(std::function<void(const zeek::Args&, StmtFlowType&)> f)
+        : Stmt(STMT_STD_FUNCTION), func(std::move(f)) {}
 
     ValPtr Exec(Frame* f, StmtFlowType& flow) override;
 
@@ -798,7 +797,7 @@ public:
     TraversalCode Traverse(TraversalCallback* cb) const override { return TC_CONTINUE; }
 
 private:
-    FunctionVariant func;
+    std::function<void(const zeek::Args&, StmtFlowType&)> func;
 };
 
 } // namespace zeek::detail
diff --git a/testing/btest/plugins/protocol-plugin/src/Foo.cc b/testing/btest/plugins/protocol-plugin/src/Foo.cc
index 100ebcb9c4..5c27f3a3a9 100644
--- a/testing/btest/plugins/protocol-plugin/src/Foo.cc
+++ b/testing/btest/plugins/protocol-plugin/src/Foo.cc
@@ -15,17 +15,17 @@ Foo::Foo(zeek::Connection* conn) : zeek::analyzer::tcp::TCP_ApplicationAnalyzer(
 
     auto handler = zeek::event_registry->Lookup("connection_established");
     if ( handler ) {
-        handler->GetFunc()->AddBody([](const zeek::Args& args) {
+        handler->GetFunc()->AddBody([](const zeek::Args& args, zeek::detail::StmtFlowType& flow) {
             printf("c++ connection_established lambda handler, received %zu arguments\n", args.size());
         });
 
-        handler->GetFunc()->AddBody(std::bind(&Foo::ConnectionEstablishedHandler, this, _1));
+        handler->GetFunc()->AddBody(std::bind(&Foo::ConnectionEstablishedHandler, this, _1, _2));
     }
 }
 
 Foo::~Foo() { delete interp; }
 
-void Foo::ConnectionEstablishedHandler(const zeek::Args& args) {
+void Foo::ConnectionEstablishedHandler(const zeek::Args& args, zeek::detail::StmtFlowType& flow) {
     printf("c++ connection_established member handler, received %zu arguments\n", args.size());
 }
 
diff --git a/testing/btest/plugins/protocol-plugin/src/Foo.h b/testing/btest/plugins/protocol-plugin/src/Foo.h
index f1a89bac12..7e59f2e705 100644
--- a/testing/btest/plugins/protocol-plugin/src/Foo.h
+++ b/testing/btest/plugins/protocol-plugin/src/Foo.h
@@ -1,6 +1,7 @@
 
 #pragma once
 
+#include "Stmt.h"
 #include "analyzer/protocol/pia/PIA.h"
 #include "analyzer/protocol/tcp/TCP.h"
 
@@ -23,7 +24,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);
+    void ConnectionEstablishedHandler(const zeek::Args& args, zeek::detail::StmtFlowType& flow);
 
 protected:
     binpac::Foo::Foo_Conn* interp;