From 76713834e4ab721bc6265c89203ee171152cc017 Mon Sep 17 00:00:00 2001
From: Aleksandr Nogikh <nogikh@google.com>
Date: Tue, 5 Mar 2024 15:14:07 +0100
Subject: prog: fix resource leak during replaceArg() of union fields

The replaced union field may contain resource references that must also
be cleaned up.

The bug was triggered via methods that patch conditional fields, so
let's add stress tests for the conditional fields + resources
combination.

Reported-by: Paul Chaignon <paul.chaignon@gmail.com>
---
 prog/expr_test.go | 17 +++++++++++++++++
 prog/mutation.go  |  1 -
 prog/prog.go      |  3 +++
 3 files changed, 20 insertions(+), 1 deletion(-)

(limited to 'prog')

diff --git a/prog/expr_test.go b/prog/expr_test.go
index f6c8c8603..adff50c6a 100644
--- a/prog/expr_test.go
+++ b/prog/expr_test.go
@@ -42,6 +42,23 @@ func TestGenerateConditionalFields(t *testing.T) {
 	}
 }
 
+func TestConditionalResources(t *testing.T) {
+	// Let's stress test the code and rely on various internal checks.
+	target, rs, _ := initRandomTargetTest(t, "test", "64")
+	ct := target.BuildChoiceTable(nil, map[*Syscall]bool{
+		target.SyscallMap["test$create_cond_resource"]: true,
+		target.SyscallMap["test$use_cond_resource"]:    true,
+	})
+	iters := 500
+	if testing.Short() {
+		iters /= 10
+	}
+	for i := 0; i < iters; i++ {
+		p := target.Generate(rs, 10, ct)
+		p.Mutate(rs, 10, ct, nil, nil)
+	}
+}
+
 func TestMutateConditionalFields(t *testing.T) {
 	target, rs, _ := initRandomTargetTest(t, "test", "64")
 	ct := target.DefaultChoiceTable()
diff --git a/prog/mutation.go b/prog/mutation.go
index cf93ce4eb..54f88c175 100644
--- a/prog/mutation.go
+++ b/prog/mutation.go
@@ -505,7 +505,6 @@ func (t *UnionType) mutate(r *randGen, s *state, arg Arg, ctx ArgCtx) (calls []*
 		index++
 	}
 	optType, optDir := t.Fields[index].Type, t.Fields[index].Dir(a.Dir())
-	removeArg(a.Option)
 	var newOpt Arg
 	newOpt, calls = r.generateArg(s, optType, optDir)
 	replaceArg(arg, MakeUnionArg(t, a.Dir(), newOpt, index))
diff --git a/prog/prog.go b/prog/prog.go
index 34fcfa6e8..f93293b64 100644
--- a/prog/prog.go
+++ b/prog/prog.go
@@ -362,6 +362,9 @@ func replaceArg(arg, arg1 Arg) {
 	case *PointerArg:
 		*a = *arg1.(*PointerArg)
 	case *UnionArg:
+		if a.Option != nil {
+			removeArg(a.Option)
+		}
 		*a = *arg1.(*UnionArg)
 	case *DataArg:
 		*a = *arg1.(*DataArg)
-- 
cgit mrf-deployment