From eb826bd337705ab03e06bc0d45a3ef8ae66aad6b Mon Sep 17 00:00:00 2001
From: Jon Siwek <jsiwek@corelight.com>
Date: Tue, 14 Jul 2020 22:11:04 -0700
Subject: [PATCH] GH-734: Improve handling of lambdas that escape enclosing
 frame

---
 src/Frame.cc                                  | 14 +++++++++
 .../Baseline/language.lambda-escaping/out     |  4 +++
 .../Baseline/language.lambda-zeek-init/out    |  1 +
 testing/btest/language/lambda-escaping.zeek   | 31 +++++++++++++++++++
 testing/btest/language/lambda-zeek-init.zeek  | 12 +++++++
 5 files changed, 62 insertions(+)
 create mode 100644 testing/btest/Baseline/language.lambda-escaping/out
 create mode 100644 testing/btest/Baseline/language.lambda-zeek-init/out
 create mode 100644 testing/btest/language/lambda-escaping.zeek
 create mode 100644 testing/btest/language/lambda-zeek-init.zeek

diff --git a/src/Frame.cc b/src/Frame.cc
index 4328cf49a3..ffb1e30968 100644
--- a/src/Frame.cc
+++ b/src/Frame.cc
@@ -128,6 +128,20 @@ const zeek::ValPtr& Frame::GetElementByID(const zeek::detail::ID* id) const
 
 void Frame::Reset(int startIdx)
 	{
+	if ( functions_with_closure_frame_reference )
+		{
+		for ( auto& func : *functions_with_closure_frame_reference )
+			{
+			// A lambda could be escaping its enclosing Frame at this point so
+			// it needs to claim some ownership (or copy) of the Frame in
+			// order to be of any further use.
+			func->StrengthenClosureReference(this);
+			Unref(func);
+			}
+
+		functions_with_closure_frame_reference.reset();
+		}
+
 	for ( int i = startIdx; i < size; ++i )
 		ClearElement(i);
 	}
diff --git a/testing/btest/Baseline/language.lambda-escaping/out b/testing/btest/Baseline/language.lambda-escaping/out
new file mode 100644
index 0000000000..20fe27d57e
--- /dev/null
+++ b/testing/btest/Baseline/language.lambda-escaping/out
@@ -0,0 +1,4 @@
+101
+101
+101
+101
diff --git a/testing/btest/Baseline/language.lambda-zeek-init/out b/testing/btest/Baseline/language.lambda-zeek-init/out
new file mode 100644
index 0000000000..a9c8fe8292
--- /dev/null
+++ b/testing/btest/Baseline/language.lambda-zeek-init/out
@@ -0,0 +1 @@
+103
diff --git a/testing/btest/language/lambda-escaping.zeek b/testing/btest/language/lambda-escaping.zeek
new file mode 100644
index 0000000000..7552d39cb2
--- /dev/null
+++ b/testing/btest/language/lambda-escaping.zeek
@@ -0,0 +1,31 @@
+# @TEST-EXEC: zeek -b %INPUT >out
+# @TEST-EXEC: btest-diff out
+
+global my_func: function();
+
+event other_event()
+	{
+	local baz = 42;
+	my_func();
+	}
+
+event zeek_init() &priority=+10
+	{
+	local outer = 101;
+
+	local lambda = function()
+		{ print outer; };
+
+	lambda();
+
+	my_func = lambda;
+	my_func();
+
+	event other_event();
+	}
+
+event zeek_init() &priority=-10
+	{
+	local qux = 13;
+	my_func();
+	}
diff --git a/testing/btest/language/lambda-zeek-init.zeek b/testing/btest/language/lambda-zeek-init.zeek
new file mode 100644
index 0000000000..2576831c22
--- /dev/null
+++ b/testing/btest/language/lambda-zeek-init.zeek
@@ -0,0 +1,12 @@
+# @TEST-EXEC: zeek -b %INPUT >out
+# @TEST-EXEC: btest-diff out
+
+event zeek_init() &priority=+10
+	{
+	local outer = 101;
+
+	local lambda = function()
+		{ print outer + 2; };
+
+	lambda();
+	}