syz-fuzzer: add program priority in corpus

Update #534

3 files changed, 131 insertions, 8 deletions

diff --git a/syz-fuzzer/fuzzer.go b/syz-fuzzer/fuzzer.go
index 22387135a..b005f91b8 100644
--- a/syz-fuzzer/fuzzer.go
+++ b/syz-fuzzer/fuzzer.go
@@ -6,11 +6,13 @@ package main
 import (
 	"flag"
 	"fmt"
+	"math/rand"
 	"net/http"
 	_ "net/http/pprof"
 	"os"
 	"runtime"
 	"runtime/debug"
+	"sort"
 	"sync"
 	"sync/atomic"
 	"time"
@@ -49,6 +51,8 @@ type Fuzzer struct {
 	corpusMu     sync.RWMutex
 	corpus       []*prog.Prog
 	corpusHashes map[hash.Sig]struct{}
+	corpusPrios  []int64
+	sumPrios     int64
 
 	signalMu     sync.RWMutex
 	corpusSignal signal.Signal // signal of inputs in corpus
@@ -58,6 +62,12 @@ type Fuzzer struct {
 	logMu sync.Mutex
 }
 
+type FuzzerSnapshot struct {
+	corpus      []*prog.Prog
+	corpusPrios []int64
+	sumPrios    int64
+}
+
 type Stat int
 
 const (
@@ -375,11 +385,25 @@ func (fuzzer *Fuzzer) addInputFromAnotherFuzzer(inp rpctype.RPCInput) {
 	fuzzer.addInputToCorpus(p, sign, sig)
 }
 
+func (fuzzer *FuzzerSnapshot) chooseProgram(r *rand.Rand) *prog.Prog {
+	randVal := r.Int63n(fuzzer.sumPrios + 1)
+	idx := sort.Search(len(fuzzer.corpusPrios), func(i int) bool {
+		return fuzzer.corpusPrios[i] >= randVal
+	})
+	return fuzzer.corpus[idx]
+}
+
 func (fuzzer *Fuzzer) addInputToCorpus(p *prog.Prog, sign signal.Signal, sig hash.Sig) {
 	fuzzer.corpusMu.Lock()
 	if _, ok := fuzzer.corpusHashes[sig]; !ok {
 		fuzzer.corpus = append(fuzzer.corpus, p)
 		fuzzer.corpusHashes[sig] = struct{}{}
+		prio := int64(len(sign))
+		if sign.Empty() {
+			prio = 1
+		}
+		fuzzer.sumPrios += prio
+		fuzzer.corpusPrios = append(fuzzer.corpusPrios, fuzzer.sumPrios)
 	}
 	fuzzer.corpusMu.Unlock()
 
@@ -391,10 +415,10 @@ func (fuzzer *Fuzzer) addInputToCorpus(p *prog.Prog, sign signal.Signal, sig has
 	}
 }
 
-func (fuzzer *Fuzzer) corpusSnapshot() []*prog.Prog {
+func (fuzzer *Fuzzer) snapshot() FuzzerSnapshot {
 	fuzzer.corpusMu.RLock()
 	defer fuzzer.corpusMu.RUnlock()
-	return fuzzer.corpus
+	return FuzzerSnapshot{fuzzer.corpus, fuzzer.corpusPrios, fuzzer.sumPrios}
 }
 
 func (fuzzer *Fuzzer) addMaxSignal(sign signal.Signal) {
diff --git a/syz-fuzzer/fuzzer_test.go b/syz-fuzzer/fuzzer_test.go
new file mode 100644
index 000000000..41887896a
--- /dev/null
+++ b/syz-fuzzer/fuzzer_test.go
@@ -0,0 +1,99 @@
+// Copyright 2019 syzkaller project authors. All rights reserved.
+// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.
+
+package main
+
+import (
+	"math"
+	"math/rand"
+	"testing"
+
+	"github.com/google/syzkaller/pkg/hash"
+	"github.com/google/syzkaller/pkg/signal"
+	"github.com/google/syzkaller/prog"
+)
+
+type InputTest struct {
+	p    *prog.Prog
+	sign signal.Signal
+	sig  hash.Sig
+}
+
+func TestChooseProgram(t *testing.T) {
+	rs := rand.NewSource(0)
+	r := rand.New(rs)
+	target := getTarget(t, "test", "64")
+	fuzzer := &Fuzzer{corpusHashes: make(map[hash.Sig]struct{})}
+
+	const (
+		maxIters   = 1000
+		sizeCorpus = 1000
+		eps        = 0.01
+	)
+
+	priorities := make(map[*prog.Prog]int64)
+	for i := 0; i < sizeCorpus; i++ {
+		sizeSig := i + 1
+		if sizeSig%250 == 0 {
+			sizeSig = 0
+		}
+		inp := generateInput(target, rs, 10, sizeSig)
+		fuzzer.addInputToCorpus(inp.p, inp.sign, inp.sig)
+		priorities[inp.p] = int64(len(inp.sign))
+	}
+	snapshot := fuzzer.snapshot()
+	counters := make(map[*prog.Prog]int)
+	for it := 0; it < maxIters; it++ {
+		counters[snapshot.chooseProgram(r)]++
+	}
+	for p, prio := range priorities {
+		prob := float64(prio) / float64(fuzzer.sumPrios)
+		diff := math.Abs(prob*maxIters - float64(counters[p]))
+		if diff > eps*maxIters {
+			t.Fatalf("the difference (%f) is higher than %f%%", diff, eps*100)
+		}
+	}
+}
+
+func TestAddInputConcurrency(t *testing.T) {
+	target := getTarget(t, "test", "64")
+	fuzzer := &Fuzzer{corpusHashes: make(map[hash.Sig]struct{})}
+
+	const (
+		routines = 10
+		iters    = 100
+	)
+
+	for i := 0; i < routines; i++ {
+		go func() {
+			rs := rand.NewSource(0)
+			r := rand.New(rs)
+			for it := 0; it < iters; it++ {
+				inp := generateInput(target, rs, 10, it)
+				fuzzer.addInputToCorpus(inp.p, inp.sign, inp.sig)
+				snapshot := fuzzer.snapshot()
+				snapshot.chooseProgram(r).Clone()
+			}
+		}()
+	}
+}
+
+func generateInput(target *prog.Target, rs rand.Source, ncalls int, sizeSig int) (inp InputTest) {
+	inp.p = target.Generate(rs, ncalls, nil)
+	var raw []uint32
+	for i := 1; i <= sizeSig; i++ {
+		raw = append(raw, uint32(i))
+	}
+	inp.sign = signal.FromRaw(raw, 0)
+	inp.sig = hash.Hash(inp.p.Serialize())
+	return
+}
+
+func getTarget(t *testing.T, os, arch string) *prog.Target {
+	t.Parallel()
+	target, err := prog.GetTarget(os, arch)
+	if err != nil {
+		t.Fatal(err)
+	}
+	return target
+}
diff --git a/syz-fuzzer/proc.go b/syz-fuzzer/proc.go
index 3b884b2ac..7cf222c60 100644
--- a/syz-fuzzer/proc.go
+++ b/syz-fuzzer/proc.go
@@ -87,16 +87,16 @@ func (proc *Proc) loop() {
 		}
 
 		ct := proc.fuzzer.choiceTable
-		corpus := proc.fuzzer.corpusSnapshot()
-		if len(corpus) == 0 || i%generatePeriod == 0 {
+		fuzzerSnapshot := proc.fuzzer.snapshot()
+		if len(fuzzerSnapshot.corpus) == 0 || i%generatePeriod == 0 {
 			// Generate a new prog.
 			p := proc.fuzzer.target.Generate(proc.rnd, programLength, ct)
 			log.Logf(1, "#%v: generated", proc.pid)
 			proc.execute(proc.execOpts, p, ProgNormal, StatGenerate)
 		} else {
 			// Mutate an existing prog.
-			p := corpus[proc.rnd.Intn(len(corpus))].Clone()
-			p.Mutate(proc.rnd, programLength, ct, corpus)
+			p := fuzzerSnapshot.chooseProgram(proc.rnd).Clone()
+			p.Mutate(proc.rnd, programLength, ct, fuzzerSnapshot.corpus)
 			log.Logf(1, "#%v: mutated", proc.pid)
 			proc.execute(proc.execOpts, p, ProgNormal, StatFuzz)
 		}
@@ -211,10 +211,10 @@ func (proc *Proc) smashInput(item *WorkSmash) {
 	if proc.fuzzer.comparisonTracingEnabled && item.call != -1 {
 		proc.executeHintSeed(item.p, item.call)
 	}
-	corpus := proc.fuzzer.corpusSnapshot()
+	fuzzerSnapshot := proc.fuzzer.snapshot()
 	for i := 0; i < 100; i++ {
 		p := item.p.Clone()
-		p.Mutate(proc.rnd, programLength, proc.fuzzer.choiceTable, corpus)
+		p.Mutate(proc.rnd, programLength, proc.fuzzer.choiceTable, fuzzerSnapshot.corpus)
 		log.Logf(1, "#%v: smash mutated", proc.pid)
 		proc.execute(proc.execOpts, p, ProgNormal, StatSmash)
 	}