prog: add implementation for resource centric

8 files changed, 122 insertions, 17 deletions

diff --git a/prog/analysis.go b/prog/analysis.go
index 383ba15d1..f8330ed1c 100644
--- a/prog/analysis.go
+++ b/prog/analysis.go
@@ -15,6 +15,7 @@ import (
 type state struct {
 	target    *Target
 	ct        *ChoiceTable
+	corpus    []*Prog
 	files     map[string]bool
 	resources map[string][]*ResultArg
 	strings   map[string]bool
@@ -23,8 +24,8 @@ type state struct {
 }
 
 // analyze analyzes the program p up to but not including call c.
-func analyze(ct *ChoiceTable, p *Prog, c *Call) *state {
-	s := newState(p.Target, ct)
+func analyze(ct *ChoiceTable, corpus []*Prog, p *Prog, c *Call) *state {
+	s := newState(p.Target, ct, corpus)
 	resources := true
 	for _, c1 := range p.Calls {
 		if c1 == c {
@@ -35,10 +36,11 @@ func analyze(ct *ChoiceTable, p *Prog, c *Call) *state {
 	return s
 }
 
-func newState(target *Target, ct *ChoiceTable) *state {
+func newState(target *Target, ct *ChoiceTable, corpus []*Prog) *state {
 	s := &state{
 		target:    target,
 		ct:        ct,
+		corpus:    corpus,
 		files:     make(map[string]bool),
 		resources: make(map[string][]*ResultArg),
 		strings:   make(map[string]bool),
diff --git a/prog/any_test.go b/prog/any_test.go
index 19a32d060..4e162258c 100644
--- a/prog/any_test.go
+++ b/prog/any_test.go
@@ -20,7 +20,7 @@ func TestIsComplexPtr(t *testing.T) {
 	compl := make(map[string]bool)
 	for _, meta := range target.Syscalls {
 		for i := 0; i < iters; i++ {
-			s := newState(target, nil)
+			s := newState(target, nil, nil)
 			calls := r.generateParticularCall(s, meta)
 			p := &Prog{Target: target, Calls: calls}
 			for _, arg := range p.complexPtrs() {
diff --git a/prog/encoding.go b/prog/encoding.go
index eb0fb79a9..904a3e11b 100644
--- a/prog/encoding.go
+++ b/prog/encoding.go
@@ -961,7 +961,7 @@ func (p *parser) auto(arg Arg) Arg {
 }
 
 func (p *parser) fixupAutos(prog *Prog) {
-	s := analyze(nil, prog, nil)
+	s := analyze(nil, nil, prog, nil)
 	for _, c := range prog.Calls {
 		p.target.assignSizesArray(c.Args, p.autos)
 		ForeachArg(c, func(arg Arg, _ *ArgCtx) {
diff --git a/prog/generation.go b/prog/generation.go
index df31f0271..dcb8dc2ac 100644
--- a/prog/generation.go
+++ b/prog/generation.go
@@ -14,7 +14,7 @@ func (target *Target) Generate(rs rand.Source, ncalls int, ct *ChoiceTable) *Pro
 		Target: target,
 	}
 	r := newRand(target, rs)
-	s := newState(target, ct)
+	s := newState(target, ct, nil)
 	for len(p.Calls) < ncalls {
 		calls := r.generateCall(s, p)
 		for _, c := range calls {
diff --git a/prog/mutation.go b/prog/mutation.go
index a3ab17637..8a33c3e2e 100644
--- a/prog/mutation.go
+++ b/prog/mutation.go
@@ -112,7 +112,7 @@ func (ctx *mutator) squashAny() bool {
 	arg.data = mutateData(r, arg.Data(), 0, maxBlobLen)
 	// Update base pointer if size has increased.
 	if baseSize < base.Res.Size() {
-		s := analyze(ctx.ct, p, p.Calls[0])
+		s := analyze(ctx.ct, ctx.corpus, p, p.Calls[0])
 		newArg := r.allocAddr(s, base.Type(), base.Res.Size(), base.Res)
 		*base = *newArg
 	}
@@ -131,8 +131,9 @@ func (ctx *mutator) insertCall() bool {
 	if idx < len(p.Calls) {
 		c = p.Calls[idx]
 	}
-	s := analyze(ctx.ct, p, c)
+	s := analyze(ctx.ct, ctx.corpus, p, c)
 	calls := r.generateCall(s, p)
+	// TODO: the program might have more than ncalls
 	p.insertBefore(c, calls)
 	return true
 }
@@ -158,7 +159,7 @@ func (ctx *mutator) mutateArg() bool {
 	if len(c.Args) == 0 {
 		return false
 	}
-	s := analyze(ctx.ct, p, c)
+	s := analyze(ctx.ct, ctx.corpus, p, c)
 	updateSizes := true
 	for stop, ok := false, false; !stop; stop = ok && r.oneOf(3) {
 		ok = true
diff --git a/prog/prog_test.go b/prog/prog_test.go
index 3300eb4da..1605f7991 100644
--- a/prog/prog_test.go
+++ b/prog/prog_test.go
@@ -74,7 +74,7 @@ func TestVmaType(t *testing.T) {
 	r := newRand(target, rs)
 	pageSize := target.PageSize
 	for i := 0; i < iters; i++ {
-		s := newState(target, nil)
+		s := newState(target, nil, nil)
 		calls := r.generateParticularCall(s, meta)
 		c := calls[len(calls)-1]
 		if c.Meta.Name != "test$vma0" {
@@ -198,7 +198,7 @@ func TestSpecialStructs(t *testing.T) {
 				if typ == nil {
 					t.Fatal("can't find struct description")
 				}
-				g := &Gen{newRand(target, rs), newState(target, nil)}
+				g := &Gen{newRand(target, rs), newState(target, nil, nil)}
 				for i := 0; i < iters/len(target.SpecialTypes); i++ {
 					arg, _ := gen(g, typ, nil)
 					gen(g, typ, arg)
diff --git a/prog/rand.go b/prog/rand.go
index be57c3e64..819b1104f 100644
--- a/prog/rand.go
+++ b/prog/rand.go
@@ -302,7 +302,7 @@ func (r *randGen) createResource(s *state, res *ResourceType) (arg Arg, calls []
 		// Generate one of them.
 		meta := metas[r.Intn(len(metas))]
 		calls := r.generateParticularCall(s, meta)
-		s1 := newState(r.target, s.ct)
+		s1 := newState(r.target, s.ct, nil)
 		s1.analyze(calls[len(calls)-1])
 		// Now see if we have what we want.
 		var allres []*ResultArg
@@ -480,7 +480,7 @@ func (target *Target) GenerateAllSyzProg(rs rand.Source) *Prog {
 		Target: target,
 	}
 	r := newRand(target, rs)
-	s := newState(target, nil)
+	s := newState(target, nil, nil)
 	handled := make(map[string]bool)
 	for _, meta := range target.Syscalls {
 		if !strings.HasPrefix(meta.CallName, "syz_") || handled[meta.CallName] {
@@ -589,7 +589,14 @@ func (r *randGen) generateArgImpl(s *state, typ Type, ignoreSpecial bool) (arg A
 
 func (a *ResourceType) generate(r *randGen, s *state) (arg Arg, calls []*Call) {
 	switch {
-	case r.nOutOf(1000, 1011):
+	case r.nOutOf(2, 5):
+		var res *ResultArg
+		res, calls = resourceCentric(a, s, r)
+		if res == nil {
+			return r.createResource(s, a)
+		}
+		arg = MakeResultArg(a, res, 0)
+	case r.nOutOf(1, 2):
 		// Get an existing resource.
 		alltypes := make([][]*ResultArg, 0, len(s.resources))
 		for _, res1 := range s.resources {
@@ -611,7 +618,7 @@ func (a *ResourceType) generate(r *randGen, s *state) (arg Arg, calls []*Call) {
 		} else {
 			arg, calls = r.createResource(s, a)
 		}
-	case r.nOutOf(10, 11):
+	case r.nOutOf(2, 3):
 		// Create a new resource.
 		arg, calls = r.createResource(s, a)
 	default:
@@ -754,3 +761,75 @@ func (a *CsumType) generate(r *randGen, s *state) (arg Arg, calls []*Call) {
 	// Filled at runtime by executor.
 	return MakeConstArg(a, 0), nil
 }
+
+// Finds a compatible resource with the type `t` and the calls that initialize that resource.
+func resourceCentric(t *ResourceType, s *state, r *randGen) (resource *ResultArg, calls []*Call) {
+	var p *Prog
+	for idx := range r.Perm(len(s.corpus)) {
+		p = s.corpus[idx].Clone()
+		resources := getCompatibleResources(p, t.TypeName, r)
+		if len(resources) > 0 {
+			resource = resources[r.Intn(len(resources))]
+			break
+		}
+	}
+
+	// No compatible resource was found.
+	if resource == nil {
+		return nil, nil
+	}
+
+	// Set that stores the resources that appear in the same calls with the selected resource.
+	relatedRes := map[*ResultArg]bool{resource: true}
+
+	// Remove unrelated calls from the program.
+	for idx := len(p.Calls) - 1; idx >= 0; idx-- {
+		includeCall := false
+		var newResources []*ResultArg
+		ForeachArg(p.Calls[idx], func(arg Arg, _ *ArgCtx) {
+			if a, ok := arg.(*ResultArg); ok {
+				if a.Res != nil && !relatedRes[a.Res] {
+					newResources = append(newResources, a.Res)
+				}
+				if relatedRes[a] || relatedRes[a.Res] {
+					includeCall = true
+				}
+			}
+		})
+		if !includeCall {
+			p.removeCall(idx)
+		} else {
+			for _, res := range newResources {
+				relatedRes[res] = true
+			}
+		}
+	}
+
+	// Selects a biased random length of the returned calls (more calls could offer more
+	// interesting programs). The values returned (n = len(calls): n, n-1, ..., 2.
+	biasedLen := 2 + r.biasedRand(len(calls)-1, 10)
+
+	// Removes the references that are not used anymore.
+	for i := biasedLen; i < len(calls); i++ {
+		p.removeCall(i)
+	}
+
+	return resource, p.Calls
+}
+
+func getCompatibleResources(p *Prog, resourceType string, r *randGen) (resources []*ResultArg) {
+	for _, c := range p.Calls {
+		ForeachArg(c, func(arg Arg, _ *ArgCtx) {
+			// Collect only initialized resources (the ones that are already used in other calls).
+			a, ok := arg.(*ResultArg)
+			if !ok || len(a.uses) == 0 || a.typ.Dir() != DirOut {
+				return
+			}
+			if !r.target.isCompatibleResource(resourceType, a.typ.Name()) {
+				return
+			}
+			resources = append(resources, a)
+		})
+	}
+	return resources
+}
diff --git a/tools/syz-mutate/mutate.go b/tools/syz-mutate/mutate.go
index b0361a20b..56aab8ee1 100644
--- a/tools/syz-mutate/mutate.go
+++ b/tools/syz-mutate/mutate.go
@@ -14,6 +14,7 @@ import (
 	"strings"
 	"time"
 
+	"github.com/google/syzkaller/pkg/db"
 	"github.com/google/syzkaller/pkg/mgrconfig"
 	"github.com/google/syzkaller/prog"
 	_ "github.com/google/syzkaller/sys"
@@ -25,6 +26,7 @@ var (
 	flagSeed   = flag.Int("seed", -1, "prng seed")
 	flagLen    = flag.Int("len", 30, "number of calls in programs")
 	flagEnable = flag.String("enable", "", "comma-separated list of enabled syscalls")
+	flagCorpus = flag.String("corpus", "", "name of the corpus file")
 )
 
 func main() {
@@ -56,8 +58,12 @@ func main() {
 	if *flagSeed != -1 {
 		seed = int64(*flagSeed)
 	}
+	var corpus []*prog.Prog
+	if *flagCorpus != "" {
+		corpus = readCorpus(*flagCorpus, target)
+	}
 	rs := rand.NewSource(seed)
-	prios := target.CalculatePriorities(nil)
+	prios := target.CalculatePriorities(corpus)
 	ct := target.BuildChoiceTable(prios, syscalls)
 	var p *prog.Prog
 	if flag.NArg() == 0 {
@@ -73,7 +79,24 @@ func main() {
 			fmt.Fprintf(os.Stderr, "failed to deserialize the program: %v\n", err)
 			os.Exit(1)
 		}
-		p.Mutate(rs, *flagLen, ct, nil)
+		p.Mutate(rs, *flagLen, ct, corpus)
 	}
 	fmt.Printf("%s\n", p.Serialize())
 }
+
+func readCorpus(filename string, target *prog.Target) (corpus []*prog.Prog) {
+	dbObj, err := db.Open(filename)
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "failed to open the corpus file: %v\n", err)
+		os.Exit(1)
+	}
+	for _, v := range dbObj.Records {
+		p, err := target.Deserialize(v.Val, prog.NonStrict)
+		if err != nil {
+			fmt.Fprintf(os.Stderr, "failed to deserialize the program: %v\n", err)
+			os.Exit(1)
+		}
+		corpus = append(corpus, p)
+	}
+	return corpus
+}