prog: implemented argument and call priorities

1 files changed, 200 insertions, 40 deletions

diff --git a/prog/mutation_test.go b/prog/mutation_test.go
index b041f02d2..9ae807e0a 100644
--- a/prog/mutation_test.go
+++ b/prog/mutation_test.go
@@ -6,6 +6,7 @@ package prog
 import (
 	"bytes"
 	"fmt"
+	"math"
 	"math/rand"
 	"sync"
 	"testing"
@@ -22,7 +23,6 @@ func TestMutationFlags(t *testing.T) {
 			`r0 = mutate$flags2(&(0x7f0000000000)="2e2f66696c653000", 0x0)`,
 			`r0 = mutate$flags2(&(0x7f0000000000)="2e2f66696c653000", 0xd9)`,
 		},
-
 		// Mutate flags (bitmask = false).
 		{
 			`r0 = mutate$flags3(&(0x7f0000000000)="2e2f66696c653000", 0x0)`,
@@ -33,10 +33,157 @@ func TestMutationFlags(t *testing.T) {
 			`r0 = mutate$flags3(&(0x7f0000000000)="2e2f66696c653000", 0xaaaaaaaaaaaaaaaa)`,
 		},
 	}
+	runMutationTests(t, tests)
+}
 
+func TestChooseCall(t *testing.T) {
+	tests := [][2]string{
+		// The call with many arguments has a higher mutation probability.
+		{
+			`mutate0()
+mutate$integer(0x0, 0x1, 0x1, 0x1, 0x0, 0x1, 0x0, 0x0, 0x1)
+mutate$integer2(0x00, 0x00, 0x20, 0x00, 0x01)`,
+			`mutate0()
+mutate$integer(0x0, 0x1, 0x1, 0x1, 0x0, 0x1, 0x0, 0x0, 0xffffffff)
+mutate$integer2(0x00, 0x00, 0x20, 0x00, 0x01)`,
+		},
+		// Calls with the same probability.
+		{
+			`mutate$integer(0x0, 0x1, 0x1, 0x1, 0x0, 0x1, 0x0, 0x0, 0x1)
+mutate$integer(0x0, 0x1, 0x1, 0x1, 0x0, 0x1, 0x0, 0x0, 0x1)
+mutate$integer(0x0, 0x1, 0x1, 0x1, 0x0, 0x1, 0x0, 0x0, 0x1)
+mutate$integer(0x0, 0x1, 0x1, 0x1, 0x0, 0x1, 0x0, 0x0, 0x1)
+mutate$integer(0x0, 0x1, 0x1, 0x1, 0x0, 0x1, 0x0, 0x0, 0x1)`,
+			`mutate$integer(0x0, 0x1, 0x1, 0x1, 0x0, 0x1, 0x0, 0x0, 0x1)
+mutate$integer(0x0, 0x1, 0x1, 0x1, 0x0, 0x1, 0x0, 0x0, 0x1)
+mutate$integer(0x0, 0x1, 0x1, 0x1, 0x0, 0x1, 0x0, 0x0, 0xff)
+mutate$integer(0x0, 0x1, 0x1, 0x1, 0x0, 0x1, 0x0, 0x0, 0x1)
+mutate$integer(0x0, 0x1, 0x1, 0x1, 0x0, 0x1, 0x0, 0x0, 0x1)`,
+		},
+		// The call with a lower probability can be mutated.
+		{
+			`mutate7(&(0x7f0000000000)='123', 0x3)
+mutate$integer(0x0, 0x1, 0x1, 0x1, 0x0, 0x1, 0x0, 0x0, 0x1)
+r0 = mutate$flags(&(0x7f0000000000)="2e2f66696c653000", 0x0, 0x1, 0x1)`,
+			`mutate7(&(0x7f0000000000)='123', 0x2)
+mutate$integer(0x0, 0x1, 0x1, 0x1, 0x0, 0x1, 0x0, 0x0, 0x1)
+r0 = mutate$flags(&(0x7f0000000000)="2e2f66696c653000", 0x0, 0x1, 0x1)`,
+		},
+		// Complex arguments.
+		{
+			`test$struct(&(0x7f0000000000)={0x0, {0x0}})
+test$array0(&(0x7f0000001000)={0x1, [@f0=0x2, @f1=0x3], 0x4})
+mutate$integer(0x0, 0x1, 0x1, 0x1, 0x0, 0x1, 0x0, 0x0, 0x1)`,
+			`test$struct(&(0x7f0000000000)={0xff, {0x0}})
+test$array0(&(0x7f0000001000)={0x1, [@f0=0x2, @f1=0x3], 0x4})
+mutate$integer(0x0, 0x1, 0x1, 0x1, 0x0, 0x1, 0x0, 0x0, 0x1)`,
+		},
+	}
 	runMutationTests(t, tests)
 }
 
+func TestMutateArgument(t *testing.T) {
+	tests := [][2]string{
+		// Mutate an integer with a higher priority than the boolean arguments.
+		{
+			`mutate$integer(0x0, 0x1, 0x1, 0x1, 0x0, 0x1, 0x0, 0x0, 0x1)`,
+			`mutate$integer(0x0, 0x1, 0x1, 0x1, 0x0, 0x1, 0x0, 0x0, 0xffffffff)`,
+		},
+		// Mutate a boolean.
+		{
+			`mutate$integer(0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0)`,
+			`mutate$integer(0x0, 0x0, 0x0, 0x1, 0x0, 0x0, 0x0, 0x0, 0x0)`,
+		},
+		// Mutate flags (bitmask = true).
+		{
+			`r0 = mutate$flags(&(0x7f0000000000)="2e2f66696c653000", 0x0, 0x1, 0x1)`,
+			`r0 = mutate$flags(&(0x7f0000000000)="2e2f66696c653000", 0x20, 0x1, 0x9)`,
+		},
+		// Mutate an int8 from a set of other arguments with higher priority.
+		{
+			`mutate$integer2(0x00, 0x00, 0x20, 0x00, 0x01)`,
+			`mutate$integer2(0x00, 0x00, 0x20, 0x00, 0x07)`,
+		},
+		// Mutate an array of structs
+		{
+			`mutate$array2(&(0x7f0000000000)=[{0x0}, {0x0}, {0x0}, {0x0}, {0x0}])`,
+			`mutate$array2(&(0x7f0000000000)=[{0x0}, {0x0}, {0x3}, {0x0}, {0x0}])`,
+		},
+		// Mutate a non-special union that have more than 1 option
+		{
+			`mutate$union(&(0x7f0000000000)=@f1=[0x0, 0x1, 0x2, 0x3, 0x0, 0x1, 0x2, 0x3, 0x0, 0x0])`,
+			`mutate$union(&(0x7f0000000000)=@f0=0x2)`,
+		},
+		// Mutate the value of the current option in union
+		{
+			`mutate$union(&(0x7f0000000000)=@f1=[0x0, 0x1, 0x2, 0x3, 0x0, 0x1, 0x2, 0x3, 0x0, 0x0])`,
+			`mutate$union(&(0x7f0000000000)=@f1=[0x0, 0x1, 0xff, 0x3, 0x0, 0x1, 0x2, 0x3, 0x0, 0x0])`,
+		},
+	}
+
+	target := initTargetTest(t, "test", "64")
+	for ti, test := range tests {
+		test := test
+		t.Run(fmt.Sprint(ti), func(t *testing.T) {
+			t.Parallel()
+			rs, ct, p, goal, err := buildTestContext(test, target)
+			if err != nil {
+				t.Fatalf("failed to deserialize the program: %v", err)
+			}
+			want := goal.Serialize()
+			for i := 0; i < 1e5; i++ {
+				p1 := p.Clone()
+				ctx := &mutator{
+					p:      p1,
+					r:      newRand(p1.Target, rs),
+					ncalls: 0,
+					ct:     ct,
+					corpus: nil,
+				}
+				ctx.mutateArg()
+				data1 := p1.Serialize()
+				if bytes.Equal(want, data1) {
+					t.Logf("success on iter %v", i)
+					return
+				}
+			}
+			t.Fatalf("failed to achieve goal, original:%s\ngoal:%s", test[0], test[1])
+		})
+	}
+}
+
+func TestRandomChoice(t *testing.T) {
+	t.Parallel()
+	target, err := GetTarget("test", "64")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	r := newRand(target, randSource(t))
+	priorities := []float64{1, 1, 1, 1, 1, 1, 1, 1, 2}
+
+	const (
+		maxIters    = 100000
+		searchedIdx = 8
+		prob        = 0.2
+		eps         = 0.01
+	)
+
+	var index, count int
+	for i := 0; i < maxIters; i++ {
+		index = randomChoice(priorities, r)
+
+		if index == searchedIdx {
+			count++
+		}
+	}
+
+	diff := math.Abs(prob*maxIters - float64(count))
+	if diff > eps*maxIters {
+		t.Fatalf("The difference (%f) is higher than %f%%", diff, eps*100)
+	}
+}
+
 func TestClone(t *testing.T) {
 	target, rs, iters := initTest(t)
 	for i := 0; i < iters; i++ {
@@ -157,6 +304,12 @@ mutate5(&(0x7f0000001000)="2e2f66696c653000", 0x22c0)
 mutate5(&(0x7f0000001000)="2e2f66696c653000", 0x22c0)
 mutate5(&(0x7f0000001000)="2e2f66696c653100", 0x22c0)
 `},
+		// Mutate the array.
+		{`
+mutate$array(0x1, 0x30, &(0x7f0000000000)=[0x1, 0x1, 0x1, 0x1, 0x1])
+`, `
+mutate$array(0x1, 0x30, &(0x7f0000000000)=[0x1, 0x1, 0x1, 0x1])
+`},
 		// Extend an array.
 		{`
 mutate3(&(0x7f0000000000)=[0x1, 0x1], 0x2)
@@ -180,45 +333,6 @@ mutate8(0xffffffffffffffff)
 	runMutationTests(t, tests)
 }
 
-func runMutationTests(t *testing.T, tests [][2]string) {
-	target := initTargetTest(t, "test", "64")
-
-	for ti, test := range tests {
-		test := test
-		t.Run(fmt.Sprint(ti), func(t *testing.T) {
-			t.Parallel()
-			p, err := target.Deserialize([]byte(test[0]), Strict)
-			if err != nil {
-				t.Fatalf("failed to deserialize original program: %v", err)
-			}
-			goal, err := target.Deserialize([]byte(test[1]), Strict)
-			if err != nil {
-				t.Fatalf("failed to deserialize goal program: %v", err)
-			}
-			want := goal.Serialize()
-			enabled := make(map[*Syscall]bool)
-			for _, c := range p.Calls {
-				enabled[c.Meta] = true
-			}
-			for _, c := range goal.Calls {
-				enabled[c.Meta] = true
-			}
-			ct := target.BuildChoiceTable(nil, enabled)
-			rs := rand.NewSource(0)
-			for i := 0; i < 1e5; i++ {
-				p1 := p.Clone()
-				p1.Mutate(rs, len(goal.Calls), ct, nil)
-				data1 := p1.Serialize()
-				if bytes.Equal(want, data1) {
-					t.Logf("success on iter %v", i)
-					return
-				}
-			}
-			t.Fatalf("failed to achieve goal, original:%s\ngoal:%s", test[0], test[1])
-		})
-	}
-}
-
 func BenchmarkMutate(b *testing.B) {
 	target, cleanup := initBench(b)
 	defer cleanup()
@@ -259,3 +373,49 @@ func linuxAmd64ChoiceTable(target *Target) *ChoiceTable {
 	})
 	return linuxCT
 }
+
+func runMutationTests(t *testing.T, tests [][2]string) {
+	target := initTargetTest(t, "test", "64")
+	for ti, test := range tests {
+		test := test
+		t.Run(fmt.Sprint(ti), func(t *testing.T) {
+			t.Parallel()
+			rs, ct, p, goal, err := buildTestContext(test, target)
+			if err != nil {
+				t.Fatalf("failed to deserialize the program: %v", err)
+			}
+			want := goal.Serialize()
+			for i := 0; i < 1e5; i++ {
+				p1 := p.Clone()
+				p1.Mutate(rs, len(goal.Calls), ct, nil)
+				data1 := p1.Serialize()
+				if bytes.Equal(want, data1) {
+					t.Logf("success on iter %v", i)
+					return
+				}
+			}
+			t.Fatalf("failed to achieve goal, original:%s\ngoal:%s", test[0], test[1])
+		})
+	}
+}
+
+func buildTestContext(test [2]string, target *Target) (rs rand.Source, ct *ChoiceTable, p, goal *Prog, err error) {
+	p, err = target.Deserialize([]byte(test[0]), Strict)
+	if err != nil {
+		return
+	}
+	goal, err = target.Deserialize([]byte(test[1]), Strict)
+	if err != nil {
+		return
+	}
+	enabled := make(map[*Syscall]bool)
+	for _, c := range p.Calls {
+		enabled[c.Meta] = true
+	}
+	for _, c := range goal.Calls {
+		enabled[c.Meta] = true
+	}
+	ct = target.BuildChoiceTable(nil, enabled)
+	rs = rand.NewSource(0)
+	return
+}