Fix reference counting issues related to lambdas/closures

For example, circular references between a lambda function the frame
it's stored within and/or its closure could cause memory leaks.

This also fixes other various reference-count ownership issues that
could lead to memory errors.

There may still be some potential/undiscovered issues because the "outer
ID" finding logic doesn't look quite right as the AST traversal descends
within nested lambdas and considers their locals as "outer", but
possibly the other logic for locating values in closures or cloning
closures just works around that behavior.

src/Brofiler.cc

@@ -15,6 +15,17 @@ Brofiler::Brofiler()
 
 Brofiler::~Brofiler()
 	{
+	for ( auto& s : stmts )
+		Unref(s);
+	}
+
+void Brofiler::AddStmt(Stmt* s)
+	{
+	if ( ignoring != 0 )
+		return;
+
+	::Ref(s);
+	stmts.push_back(s);
 	}
 
 bool Brofiler::ReadStats()
@@ -109,7 +120,7 @@ bool Brofiler::WriteStats()
 		return false;
 		}
 
-	for ( list<const Stmt*>::const_iterator it = stmts.begin();
+	for ( list<Stmt*>::const_iterator it = stmts.begin();
 	      it != stmts.end(); ++it )
 		{
 		ODesc location_info;

src/Brofiler.h

@@ -38,13 +38,13 @@ public:
 	void IncIgnoreDepth() { ignoring++; }
 	void DecIgnoreDepth() { ignoring--; }
 
-	void AddStmt(const Stmt* s) { if ( ignoring == 0 ) stmts.push_back(s); }
+	void AddStmt(Stmt* s);
 
 private:
 	/**
 	 * The current, global Brofiler instance creates this list at parse-time.
 	 */
-	list<const Stmt*> stmts;
+	list<Stmt*> stmts;
 
 	/**
 	 * Indicates whether new statments will not be considered as part of

src/Expr.cc

@@ -4339,6 +4339,7 @@ LambdaExpr::LambdaExpr(std::unique_ptr<function_ingredients> arg_ing,
 
 	// Install a dummy version of the function globally for use only
 	// when broker provides a closure.
+	::Ref(ingredients->body);
 	BroFunc* dummy_func = new BroFunc(
 		ingredients->id,
 		ingredients->body,
@@ -4378,6 +4379,7 @@ LambdaExpr::LambdaExpr(std::unique_ptr<function_ingredients> arg_ing,
 	dummy_func->SetName(my_name.c_str());
 
 	Val* v = new Val(dummy_func);
+	Unref(dummy_func);
 	id->SetVal(v); // id will unref v when its done.
 	id->SetType(ingredients->id->Type()->Ref());
 	id->SetConst();
@@ -4385,6 +4387,7 @@ LambdaExpr::LambdaExpr(std::unique_ptr<function_ingredients> arg_ing,
 
 Val* LambdaExpr::Eval(Frame* f) const
 	{
+	::Ref(ingredients->body);
 	BroFunc* lamb = new BroFunc(
 		ingredients->id,
 		ingredients->body,
@@ -4398,7 +4401,9 @@ Val* LambdaExpr::Eval(Frame* f) const
 	// Allows for lookups by the receiver.
 	lamb->SetName(my_name.c_str());
 
-	return new Val(lamb);
+	auto rval = new Val(lamb);
+	Unref(lamb);
+	return rval;
 	}
 
 void LambdaExpr::ExprDescribe(ODesc* d) const

src/Frame.cc

@@ -31,20 +31,54 @@ Frame::Frame(int arg_size, const BroFunc* func, const val_list* fn_args)
 
 Frame::~Frame()
 	{
+	for ( auto& func : functions_with_closure_frame_reference )
+		{
+		func->StrengthenClosureReference(this);
+		Unref(func);
+		}
+
 	// Deleting a Frame that is a view is a no-op.
 	Unref(trigger);
+
+	if ( ! weak_closure_ref )
 		Unref(closure);
 
 	for ( auto& i : outer_ids )
 		Unref(i);
 
 	Release();
+
+	delete [] weak_refs;
 	}
 
-void Frame::SetElement(int n, Val* v)
+void Frame::AddFunctionWithClosureRef(BroFunc* func)
 	{
-	Unref(frame[n]);
+	::Ref(func);
+	functions_with_closure_frame_reference.emplace_back(func);
+	}
+
+void Frame::SetElement(int n, Val* v, bool weak_ref)
+	{
+	UnrefElement(n);
 	frame[n] = v;
+
+	if ( weak_ref )
+		{
+		if ( ! weak_refs )
+			{
+			weak_refs = new bool[size];
+
+			for ( auto i = 0; i < size; ++i )
+				weak_refs[i] = false;
+			}
+
+		weak_refs[n] = true;
+		}
+	else
+		{
+		if ( weak_refs )
+			weak_refs[n] = false;
+		}
 	}
 
 void Frame::SetElement(const ID* id, Val* v)
@@ -62,8 +96,15 @@ void Frame::SetElement(const ID* id, Val* v)
 	if ( offset_map.size() )
 		{
 		auto where = offset_map.find(std::string(id->Name()));
+
 		if ( where != offset_map.end() )
-			SetElement(where->second, v);
+			{
+			// Need to add a Ref to 'v' since the SetElement() for
+			// id->Offset() below is otherwise responsible for keeping track
+			// of the implied reference count of the passed-in 'v' argument.
+			// i.e. if we end up storing it twice, we need an addition Ref.
+			SetElement(where->second, v->Ref());
+			}
 		}
 
 	SetElement(id->Offset(), v);
@@ -92,7 +133,7 @@ void Frame::Reset(int startIdx)
 	{
 	for ( int i = startIdx; i < size; ++i )
 		{
-		Unref(frame[i]);
+		UnrefElement(i);
 		frame[i] = 0;
 		}
 	}
@@ -100,7 +141,7 @@ void Frame::Reset(int startIdx)
 void Frame::Release()
 	{
 	for ( int i = 0; i < size; ++i )
-		Unref(frame[i]);
+		UnrefElement(i);
 
 	delete [] frame;
 	}
@@ -145,7 +186,34 @@ Frame* Frame::Clone() const
 	return other;
 	}
 
-Frame* Frame::SelectiveClone(const id_list& selection) const
+static bool val_is_func(Val* v, BroFunc* func)
+	{
+	if ( v->Type()->Tag() != TYPE_FUNC )
+		return false;
+
+	return v->AsFunc() == func;
+	}
+
+static Val* clone_if_not_func(Val** frame, int offset, BroFunc* func,
+                              Frame* other)
+	{
+	auto v = frame[offset];
+
+	if ( ! v )
+		return nullptr;
+
+	if ( val_is_func(v, func) )
+		{
+		other->SetElement(offset, v, true);
+		return v;
+		}
+
+	auto rval = v->Clone();
+	other->SetElement(offset, rval);
+	return rval;
+	}
+
+Frame* Frame::SelectiveClone(const id_list& selection, BroFunc* func) const
 	{
 	if ( selection.length() == 0 )
 		return nullptr;
@@ -171,7 +239,7 @@ Frame* Frame::SelectiveClone(const id_list& selection) const
 			auto where = offset_map.find(std::string(id->Name()));
 			if ( where != offset_map.end() )
 				{
-				other->frame[where->second] = frame[where->second]->Clone();
+				clone_if_not_func(frame, where->second, func, other);
 				continue;
 				}
 			}
@@ -179,7 +247,7 @@ Frame* Frame::SelectiveClone(const id_list& selection) const
 		if ( ! frame[id->Offset()] )
 			reporter->InternalError("Attempted to clone an id ('%s') with no associated value.", id->Name());
 
-		other->frame[id->Offset()] = frame[id->Offset()]->Clone();
+		clone_if_not_func(frame, id->Offset(), func, other);
 		}
 
 	/**
@@ -379,6 +447,7 @@ std::pair<bool, Frame*> Frame::Unserialize(const broker::vector& data)
 	// Frame takes ownership of unref'ing elements in outer_ids
 	rf->outer_ids = std::move(outer_ids);
 	rf->closure = closure;
+	rf->weak_closure_ref = false;
 
 	for ( int i = 0; i < frame_size; ++i )
 		{
@@ -437,7 +506,7 @@ void Frame::CaptureClosure(Frame* c, id_list arg_outer_ids)
 
 	closure = c;
 	if ( closure )
-		Ref(closure);
+		weak_closure_ref = true;
 
 	/**
 	 * Want to capture closures by copy?

src/Frame.h

@@ -44,8 +44,11 @@ public:
 	 *
 	 * @param n the index to set
 	 * @param v the value to set it to
+	 * @param weak_ref whether the frame owns the value and should unref
+	 * it upon destruction.  Used to break circular references between
+	 * lambda functions and closure frames.
 	 */
-	void SetElement(int n, Val* v);
+	void SetElement(int n, Val* v, bool weak_ref = false);
 
 	/**
 	 * Associates *id* and *v* in the frame. Future lookups of
@@ -149,7 +152,7 @@ public:
 	 * *selection* have been cloned. All other values are made to be
 	 * null.
 	 */
-	Frame* SelectiveClone(const id_list& selection) const;
+	Frame* SelectiveClone(const id_list& selection, BroFunc* func) const;
 
 	/**
 	 * Serializes the Frame into a Broker representation.
@@ -215,8 +218,28 @@ public:
 	void SetDelayed()	{ delayed = true; }
 	bool HasDelayed() const	{ return delayed; }
 
+	/**
+	 * Track a new function that refers to this frame for use as a closure.
+	 * This frame's destructor will then upgrade that functions reference
+	 * from weak to strong (by making a copy).  The initial use of
+	 * weak references prevents unbreakable circular references that
+	 * otherwise cause memory leaks.
+	 */
+	void AddFunctionWithClosureRef(BroFunc* func);
 
 private:
+
+	/**
+	 * Unrefs the value at offset 'n' frame unless it's a weak reference.
+	 */
+	void UnrefElement(int n)
+		{
+		if ( weak_refs && weak_refs[n] )
+			return;
+
+		Unref(frame[n]);
+		}
+
 	/** Have we captured this id? */
 	bool IsOuterID(const ID* in) const;
 
@@ -242,8 +265,13 @@ private:
 	/** Associates ID's offsets with values. */
 	Val** frame;
 
+	/** Values that are weakly referenced by the frame.  Used to
+	 * prevent circular reference memory leaks in lambda/closures */
+	bool* weak_refs = nullptr;
+
 	/** The enclosing frame of this frame. */
 	Frame* closure;
+	bool weak_closure_ref = false;
 
 	/** ID's used in this frame from the enclosing frame. */
 	id_list outer_ids;
@@ -268,6 +296,8 @@ private:
 	Trigger* trigger;
 	const CallExpr* call;
 	bool delayed;
+
+	std::vector<BroFunc*> functions_with_closure_frame_reference;
 };
 
 /**

src/Func.cc

@@ -132,6 +132,7 @@ Func* Func::DoClone()
 	{
 	// By default, ok just to return a reference. Func does not have any state
 	// that is different across instances.
+	::Ref(this);
 	return this;
 	}
 
@@ -286,6 +287,8 @@ BroFunc::~BroFunc()
 	{
 	std::for_each(bodies.begin(), bodies.end(),
 		[](Body& b) { Unref(b.stmts); });
+
+	if ( ! weak_closure_ref )
 		Unref(closure);
 	}
 
@@ -490,14 +493,35 @@ void BroFunc::AddClosure(id_list ids, Frame* f)
 	SetClosureFrame(f);
 	}
 
+bool BroFunc::StrengthenClosureReference(Frame* f)
+	{
+	if ( closure != f )
+		return false;
+
+	if ( ! weak_closure_ref )
+		return false;
+
+	closure = closure->SelectiveClone(outer_ids, this);
+	weak_closure_ref = false;
+	return true;
+	}
+
 void BroFunc::SetClosureFrame(Frame* f)
 	{
 	if ( closure )
 		reporter->InternalError("Tried to override closure for BroFunc %s.",
 					Name());
 
+	// Have to use weak references initially because otherwise Ref'ing the
+	// original frame creates a circular reference: the function holds a
+	// reference to the frame and the frame contains a reference to this
+	// function value.  And we can't just do a shallow clone of the frame
+	// up front because the closure semantics in Zeek allow mutating
+	// the outer frame.
+
 	closure = f;
-	Ref(closure);
+	weak_closure_ref = true;
+	f->AddFunctionWithClosureRef(this);
 	}
 
 bool BroFunc::UpdateClosure(const broker::vector& data)
@@ -510,9 +534,10 @@ bool BroFunc::UpdateClosure(const broker::vector& data)
 	if ( new_closure )
 		new_closure->SetFunction(this);
 
-	if ( closure )
+	if ( ! weak_closure_ref )
 		Unref(closure);
 
+	weak_closure_ref = false;
 	closure = new_closure;
 
 	return true;
@@ -528,7 +553,8 @@ Func* BroFunc::DoClone()
 	CopyStateInto(other);
 
 	other->frame_size = frame_size;
-	other->closure = closure ? closure->SelectiveClone(outer_ids) : nullptr;
+	other->closure = closure ? closure->SelectiveClone(outer_ids, this) : nullptr;
+	other->weak_closure_ref = false;
 	other->outer_ids = outer_ids;
 
 	return other;
@@ -811,3 +837,68 @@ bool check_built_in_call(BuiltinFunc* f, CallExpr* call)
 
 	return true;
 	}
+
+// Gets a function's priority from its Scope's attributes. Errors if it sees any
+// problems.
+static int get_func_priority(const attr_list& attrs)
+	{
+	int priority = 0;
+
+	for ( const auto& a : attrs )
+		{
+		if ( a->Tag() == ATTR_DEPRECATED )
+			continue;
+
+		if ( a->Tag() != ATTR_PRIORITY )
+			{
+			a->Error("illegal attribute for function body");
+			continue;
+			}
+
+		Val* v = a->AttrExpr()->Eval(0);
+		if ( ! v )
+			{
+			a->Error("cannot evaluate attribute expression");
+			continue;
+			}
+
+		if ( ! IsIntegral(v->Type()->Tag()) )
+			{
+			a->Error("expression is not of integral type");
+			continue;
+			}
+
+		priority = v->InternalInt();
+		}
+
+	return priority;
+	}
+
+function_ingredients::function_ingredients(Scope* scope, Stmt* body)
+	{
+	frame_size = scope->Length();
+	inits = scope->GetInits();
+
+	this->scope = scope;
+	::Ref(this->scope);
+	id = scope->ScopeID();
+	::Ref(id);
+
+	auto attrs = scope->Attrs();
+
+	priority = (attrs ? get_func_priority(*attrs) : 0);
+	this->body = body;
+	::Ref(this->body);
+	}
+
+function_ingredients::~function_ingredients()
+	{
+	Unref(id);
+	Unref(body);
+	Unref(scope);
+
+	for ( const auto& i : *inits )
+		Unref(i);
+
+	delete inits;
+	}

src/Func.h

@@ -114,6 +114,12 @@ public:
 	 */
 	bool UpdateClosure(const broker::vector& data);
 
+	/**
+	 * If the function's closure is a weak reference to the given frame,
+	 * upgrade to a strong reference of a shallow clone of that frame.
+	 */
+	bool StrengthenClosureReference(Frame* f);
+
 	/**
 	 * Serializes this function's closure.
 	 *
@@ -154,6 +160,7 @@ private:
 	id_list outer_ids;
 	// The frame the BroFunc was initialized in.
 	Frame* closure = nullptr;
+	bool weak_closure_ref = false;
 };
 
 typedef Val* (*built_in_func)(Frame* frame, val_list* args);
@@ -192,13 +199,20 @@ struct CallInfo {
 // Struct that collects all the specifics defining a Func. Used for BroFuncs
 // with closures.
 struct function_ingredients {
+
+	// Gathers all of the information from a scope and a function body needed
+	// to build a function.
+	function_ingredients(Scope* scope, Stmt* body);
+
+	~function_ingredients();
+
 	ID* id;
 	Stmt* body;
 	id_list* inits;
 	int frame_size;
 	int priority;
 	Scope* scope;
-	};
+};
 
 extern vector<CallInfo> call_stack;
 

src/Val.cc

@@ -120,7 +120,12 @@ Val* Val::DoClone(CloneState* state)
 		// Derived classes are responsible for this. Exception:
 		// Functions and files. There aren't any derived classes.
 		if ( type->Tag() == TYPE_FUNC )
-			return new Val(AsFunc()->DoClone());
+			{
+			auto c = AsFunc()->DoClone();
+			auto rval = new Val(c);
+			Unref(c);
+			return rval;
+			}
 
 		if ( type->Tag() == TYPE_FILE )
 			{

src/Var.cc

@@ -416,13 +416,30 @@ public:
 		: scope(s) { }
 
 	virtual TraversalCode PreExpr(const Expr*);
+	virtual TraversalCode PostExpr(const Expr*);
 
 	Scope* scope;
 	vector<const NameExpr*> outer_id_references;
+	int lambda_depth = 0;
+	// Note: think we really ought to toggle this to false to prevent
+	// considering locals within inner-lambdas as "outer", but other logic
+	// for "selective cloning" and locating IDs in the closure chain may
+	// depend on current behavior and also needs to be changed.
+	bool search_inner_lambdas = true;
 };
 
 TraversalCode OuterIDBindingFinder::PreExpr(const Expr* expr)
 	{
+	if ( expr->Tag() == EXPR_LAMBDA )
+		++lambda_depth;
+
+	if ( lambda_depth > 0 && ! search_inner_lambdas )
+		// Don't inspect the bodies of inner lambdas as they will have their
+		// own traversal to find outer IDs and we don't want to detect
+		// references to local IDs inside and accidentally treat them as
+		// "outer" since they can't be found in current scope.
+		return TC_CONTINUE;
+
 	if ( expr->Tag() != EXPR_NAME )
 		return TC_CONTINUE;
 
@@ -438,45 +455,20 @@ TraversalCode OuterIDBindingFinder::PreExpr(const Expr* expr)
 	return TC_CONTINUE;
 	}
 
-// Gets a function's priority from its Scope's attributes. Errors if it sees any
+TraversalCode OuterIDBindingFinder::PostExpr(const Expr* expr)
-// problems.
-static int get_func_priotity(const attr_list& attrs)
 	{
-	int priority = 0;
+	if ( expr->Tag() == EXPR_LAMBDA )
-
-	for ( const auto& a : attrs )
 		{
-		if ( a->Tag() == ATTR_DEPRECATED )
+		--lambda_depth;
-			continue;
+		assert(lambda_depth >= 0);
-
-		if ( a->Tag() != ATTR_PRIORITY )
-			{
-			a->Error("illegal attribute for function body");
-			continue;
 		}
 
-		Val* v = a->AttrExpr()->Eval(0);
+	return TC_CONTINUE;
-		if ( ! v )
-			{
-			a->Error("cannot evaluate attribute expression");
-			continue;
-			}
-
-		if ( ! IsIntegral(v->Type()->Tag()) )
-			{
-			a->Error("expression is not of integral type");
-			continue;
-			}
-
-		priority = v->InternalInt();
-		}
-
-	return priority;
 	}
 
 void end_func(Stmt* body)
 	{
-	std::unique_ptr<function_ingredients> ingredients = gather_function_ingredients(pop_scope(), body);
+	auto ingredients = std::make_unique<function_ingredients>(pop_scope(), body);
 
 	if ( streq(ingredients->id->Name(), "anonymous-function") )
 		{
@@ -508,24 +500,10 @@ void end_func(Stmt* body)
 		}
 
 	ingredients->id->ID_Val()->AsFunc()->SetScope(ingredients->scope);
-	}
+	// Note: ideally, something would take ownership of this memory until the
-
+	// end of script execution, but that's essentially the same as the
-std::unique_ptr<function_ingredients> gather_function_ingredients(Scope* scope, Stmt* body)
+	// lifetime of the process at the moment, so ok to "leak" it.
-	{
+	ingredients.release();
-	auto ingredients = std::make_unique<function_ingredients>();
-
-	ingredients->frame_size = scope->Length();
-	ingredients->inits = scope->GetInits();
-
-	ingredients->scope = scope;
-	ingredients->id = scope->ScopeID();
-
-	auto attrs = scope->Attrs();
-
-	ingredients->priority = (attrs ? get_func_priotity(*attrs) : 0);
-	ingredients->body = body;
-
-	return ingredients;
 	}
 
 Val* internal_val(const char* name)
@@ -548,7 +526,14 @@ id_list gather_outer_ids(Scope* scope, Stmt* body)
 	id_list idl ( cb.outer_id_references.size() );
 
 	for ( size_t i = 0; i < cb.outer_id_references.size(); ++i )
-	        idl.append(cb.outer_id_references[i]->Id());
+		{
+		auto id = cb.outer_id_references[i]->Id();
+
+		if ( idl.is_member(id) )
+			continue;
+
+		idl.append(id);
+		}
 
 	return idl;
 	}

src/Var.h

@@ -26,11 +26,6 @@ extern void begin_func(ID* id, const char* module_name, function_flavor flavor,
 		       int is_redef, FuncType* t, attr_list* attrs = nullptr);
 extern void end_func(Stmt* body);
 
-// Gathers all of the information from a scope and a function body needed to
-// build a function and collects it into a function_ingredients struct.
-// Gathered elements are not refeed.
-extern std::unique_ptr<function_ingredients> gather_function_ingredients(Scope* scope, Stmt* body);
-
 // Gather all IDs referenced inside a body that aren't part of a given scope.
 extern id_list gather_outer_ids(Scope* scope, Stmt* body);
 

src/parse.y

@@ -1234,7 +1234,7 @@ anonymous_function:
 			// a lambda expression.
 
 			// Gather the ingredients for a BroFunc from the current scope
-			std::unique_ptr<function_ingredients> ingredients = gather_function_ingredients(current_scope(), $5);
+			auto ingredients = std::make_unique<function_ingredients>(current_scope(), $5);
 			id_list outer_ids = gather_outer_ids(pop_scope(), $5);
 
 			$$ = new LambdaExpr(std::move(ingredients), std::move(outer_ids));