pkg/corpus: increase the priority of programs that were rarely selectedcorpusfix

2 files changed, 31 insertions, 1 deletions

diff --git a/pkg/corpus/corpus.go b/pkg/corpus/corpus.go
index e6569dbc6..d57b5fd25 100644
--- a/pkg/corpus/corpus.go
+++ b/pkg/corpus/corpus.go
@@ -98,6 +98,7 @@ type Item struct {
 	Call         int
 	Prog         *prog.Prog
 	HasAny       bool // whether the prog contains squashed arguments
+	ExecLast     int64
 	Signal       signal.Signal
 	StableSignal signal.Signal
 	Cover        []uint64
@@ -106,6 +107,8 @@ type Item struct {
 	areas map[*focusAreaState]struct{}
 }
 
+const maxExecs int64 = 5
+
 func (item Item) StringCall() string {
 	return item.Prog.CallName(item.Call)
 }
@@ -155,6 +158,7 @@ func (corpus *Corpus) Save(inp NewInput) {
 			Prog:         old.Prog,
 			Call:         old.Call,
 			HasAny:       old.HasAny,
+			ExecLast:     old.ExecLast,
 			Signal:       newSignal,
 			StableSignal: newStableSignal,
 			Cover:        newCover.Serialize(),
@@ -171,6 +175,7 @@ func (corpus *Corpus) Save(inp NewInput) {
 			Call:         inp.Call,
 			Prog:         inp.Prog,
 			HasAny:       inp.Prog.ContainsAny(),
+			ExecLast:     maxExecs,
 			Signal:       inp.Signal,
 			StableSignal: inp.StableSignal,
 			Cover:        inp.Cover,
@@ -181,6 +186,8 @@ func (corpus *Corpus) Save(inp NewInput) {
 	corpus.progsMap[sig] = newItem
 	corpus.applyFocusAreas(newItem, newItem.Cover)
 	corpus.saveProgram(inp.Prog, newItem.StableSignal)
+	corpus.modifyPriority(inp.Prog, newItem.Signal, newItem.ExecLast)
+	// TODO do same thing with Focus Areas
 
 	corpus.signal.Merge(inp.Signal)
 	newCover := corpus.cover.MergeDiff(inp.Cover)
diff --git a/pkg/corpus/prio.go b/pkg/corpus/prio.go
index 4d0053850..db2a4c28d 100644
--- a/pkg/corpus/prio.go
+++ b/pkg/corpus/prio.go
@@ -7,6 +7,7 @@ import (
 	"math/rand"
 	"sort"
 
+	"github.com/google/syzkaller/pkg/hash"
 	"github.com/google/syzkaller/pkg/signal"
 	"github.com/google/syzkaller/prog"
 )
@@ -49,6 +50,17 @@ func (pl *ProgramsList) saveProgram(p *prog.Prog, signal signal.Signal) {
 	}
 }
 
+func (pl *ProgramsList) modifyPriority(p *prog.Prog, probSignal signal.Signal, koef int64) {
+	prio := int64(len(probSignal))
+	if prio == 0 {
+		prio = 1
+	}
+
+	if ind, ok := pl.progToInd[p]; ok {
+		pl.prios.add(ind, koef*prio)
+	}
+}
+
 func (corpus *Corpus) ChooseProgram(r *rand.Rand) *prog.Prog {
 	corpus.mu.RLock()
 	defer corpus.mu.RUnlock()
@@ -80,7 +92,18 @@ func (corpus *Corpus) ChooseProgram(r *rand.Rand) *prog.Prog {
 	if randArea != nil {
 		return randArea.chooseProgram(r)
 	}
-	return corpus.chooseProgram(r)
+
+	prog := corpus.chooseProgram(r)
+	progData := prog.Serialize()
+	sig := hash.String(progData)
+
+	item := corpus.progsMap[sig]
+	if item.ExecLast > 0 {
+		corpus.modifyPriority(prog, item.Signal, -1)
+		item.ExecLast--
+	}
+
+	return prog
 }
 
 func (corpus *Corpus) Programs() []*prog.Prog {