pkg/corpus: a separate package for the corpus functionality

pkg/fuzzer and syz-manager have a common corpus functionality that can
be well be unified.

Create a separate pkg/corpus package that would be used by both of them.
It will simplify further work of moving pkg/fuzzer to the host.

6 files changed, 98 insertions, 243 deletions

diff --git a/pkg/fuzzer/corpus.go b/pkg/fuzzer/corpus.go
deleted file mode 100644
index b92ab1c64..000000000
--- a/pkg/fuzzer/corpus.go
+++ /dev/null
@@ -1,114 +0,0 @@
-// Copyright 2024 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 fuzzer
-
-import (
-	"math/rand"
-	"sort"
-	"sync"
-
-	"github.com/google/syzkaller/pkg/hash"
-	"github.com/google/syzkaller/pkg/signal"
-	"github.com/google/syzkaller/prog"
-)
-
-type Corpus struct {
-	mu        sync.RWMutex
-	progs     []*prog.Prog
-	hashes    map[hash.Sig]struct{}
-	sumPrios  int64
-	accPrios  []int64
-	signal    signal.Signal // signal of inputs in corpus
-	maxSignal signal.Signal // max signal ever observed (including flakes)
-	newSignal signal.Signal
-}
-
-// CorpusStat is a snapshot of the relevant current state figures.
-type CorpusStat struct {
-	Progs     int
-	Signal    int
-	MaxSignal int
-}
-
-func newCorpus() *Corpus {
-	return &Corpus{
-		hashes: make(map[hash.Sig]struct{}),
-	}
-}
-
-// TODO: maybe we want to treat progs from other fuzzers exactly like
-// candidates? And even triage them?
-func (corpus *Corpus) Save(p *prog.Prog, signal signal.Signal, sig hash.Sig) {
-	corpus.mu.Lock()
-	defer corpus.mu.Unlock()
-	if _, ok := corpus.hashes[sig]; !ok {
-		corpus.progs = append(corpus.progs, p)
-		corpus.hashes[sig] = struct{}{}
-		prio := int64(len(signal))
-		if prio == 0 {
-			prio = 1
-		}
-		corpus.sumPrios += prio
-		corpus.accPrios = append(corpus.accPrios, corpus.sumPrios)
-	}
-	corpus.signal.Merge(signal)
-	corpus.maxSignal.Merge(signal)
-}
-
-// Signal that should no longer be chased after.
-func (corpus *Corpus) AddMaxSignal(sign signal.Signal) {
-	// TODO: how do we ensure occasional drop of this max cover?
-	corpus.mu.Lock()
-	defer corpus.mu.Unlock()
-	corpus.maxSignal.Merge(sign)
-}
-
-func (corpus *Corpus) AddRawMaxSignal(signal []uint32, prio uint8) signal.Signal {
-	corpus.mu.Lock()
-	defer corpus.mu.Unlock()
-	diff := corpus.maxSignal.DiffRaw(signal, prio)
-	if diff.Empty() {
-		return diff
-	}
-	corpus.maxSignal.Merge(diff)
-	corpus.newSignal.Merge(diff)
-	return diff
-}
-
-func (corpus *Corpus) chooseProgram(r *rand.Rand) *prog.Prog {
-	corpus.mu.RLock()
-	defer corpus.mu.RUnlock()
-	if len(corpus.progs) == 0 {
-		return nil
-	}
-	randVal := r.Int63n(corpus.sumPrios + 1)
-	idx := sort.Search(len(corpus.accPrios), func(i int) bool {
-		return corpus.accPrios[i] >= randVal
-	})
-	return corpus.progs[idx]
-}
-
-func (corpus *Corpus) Programs() []*prog.Prog {
-	corpus.mu.RLock()
-	defer corpus.mu.RUnlock()
-	return corpus.progs
-}
-
-func (corpus *Corpus) GrabNewSignal() signal.Signal {
-	corpus.mu.Lock()
-	defer corpus.mu.Unlock()
-	sign := corpus.newSignal
-	corpus.newSignal = nil
-	return sign
-}
-
-func (corpus *Corpus) Stat() CorpusStat {
-	corpus.mu.RLock()
-	defer corpus.mu.RUnlock()
-	return CorpusStat{
-		Progs:     len(corpus.progs),
-		Signal:    len(corpus.signal),
-		MaxSignal: len(corpus.maxSignal),
-	}
-}
diff --git a/pkg/fuzzer/corpus_test.go b/pkg/fuzzer/corpus_test.go
deleted file mode 100644
index 0b62d8c5a..000000000
--- a/pkg/fuzzer/corpus_test.go
+++ /dev/null
@@ -1,98 +0,0 @@
-// Copyright 2024 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 fuzzer
-
-import (
-	"math"
-	"math/rand"
-	"testing"
-
-	"github.com/google/syzkaller/pkg/hash"
-	"github.com/google/syzkaller/pkg/signal"
-	"github.com/google/syzkaller/prog"
-	"github.com/google/syzkaller/sys/targets"
-)
-
-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, targets.TestOS, targets.TestArch64)
-	corpus := newCorpus()
-
-	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)
-		corpus.Save(inp.p, inp.sign, inp.sig)
-		priorities[inp.p] = int64(len(inp.sign))
-	}
-	counters := make(map[*prog.Prog]int)
-	for it := 0; it < maxIters; it++ {
-		counters[corpus.chooseProgram(r)]++
-	}
-	for p, prio := range priorities {
-		prob := float64(prio) / float64(corpus.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 TestCorpusSaveConcurrency(t *testing.T) {
-	target := getTarget(t, targets.TestOS, targets.TestArch64)
-	corpus := newCorpus()
-
-	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)
-				corpus.Save(inp.p, inp.sign, inp.sig)
-				corpus.chooseProgram(r).Clone()
-			}
-		}()
-	}
-}
-
-func generateInput(target *prog.Target, rs rand.Source, ncalls, sizeSig int) (inp InputTest) {
-	inp.p = target.Generate(rs, ncalls, target.DefaultChoiceTable())
-	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/pkg/fuzzer/cover.go b/pkg/fuzzer/cover.go
new file mode 100644
index 000000000..cc18862bd
--- /dev/null
+++ b/pkg/fuzzer/cover.go
@@ -0,0 +1,56 @@
+// Copyright 2024 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 fuzzer
+
+import (
+	"sync"
+
+	"github.com/google/syzkaller/pkg/signal"
+)
+
+// Cover keeps track of the signal known to the fuzzer.
+type Cover struct {
+	mu        sync.RWMutex
+	maxSignal signal.Signal // max signal ever observed (including flakes)
+	newSignal signal.Signal // newly identified max signal
+}
+
+// Signal that should no longer be chased after.
+func (cover *Cover) AddMaxSignal(sign signal.Signal) {
+	cover.mu.Lock()
+	defer cover.mu.Unlock()
+	cover.maxSignal.Merge(sign)
+}
+
+func (cover *Cover) addRawMaxSignal(signal []uint32, prio uint8) signal.Signal {
+	cover.mu.Lock()
+	defer cover.mu.Unlock()
+	diff := cover.maxSignal.DiffRaw(signal, prio)
+	if diff.Empty() {
+		return diff
+	}
+	cover.maxSignal.Merge(diff)
+	cover.newSignal.Merge(diff)
+	return diff
+}
+
+func (cover *Cover) GrabNewSignal() signal.Signal {
+	cover.mu.Lock()
+	defer cover.mu.Unlock()
+	sign := cover.newSignal
+	cover.newSignal = nil
+	return sign
+}
+
+type CoverStat struct {
+	MaxSignal int
+}
+
+func (cover *Cover) Stat() CoverStat {
+	cover.mu.RLock()
+	defer cover.mu.RUnlock()
+	return CoverStat{
+		MaxSignal: len(cover.maxSignal),
+	}
+}
diff --git a/pkg/fuzzer/fuzzer.go b/pkg/fuzzer/fuzzer.go
index 64be8325c..17513716a 100644
--- a/pkg/fuzzer/fuzzer.go
+++ b/pkg/fuzzer/fuzzer.go
@@ -12,15 +12,15 @@ import (
 	"sync/atomic"
 	"time"
 
+	"github.com/google/syzkaller/pkg/corpus"
 	"github.com/google/syzkaller/pkg/hash"
 	"github.com/google/syzkaller/pkg/ipc"
-	"github.com/google/syzkaller/pkg/rpctype"
 	"github.com/google/syzkaller/prog"
 )
 
 type Fuzzer struct {
 	Config         *Config
-	Corpus         *Corpus
+	Cover          *Cover
 	NeedCandidates chan struct{}
 
 	ctx    context.Context
@@ -49,7 +49,7 @@ func NewFuzzer(ctx context.Context, cfg *Config, rnd *rand.Rand,
 	target *prog.Target) *Fuzzer {
 	f := &Fuzzer{
 		Config:         cfg,
-		Corpus:         newCorpus(),
+		Cover:          &Cover{},
 		NeedCandidates: make(chan struct{}, 1),
 
 		ctx:    ctx,
@@ -75,6 +75,7 @@ func NewFuzzer(ctx context.Context, cfg *Config, rnd *rand.Rand,
 
 type Config struct {
 	Debug          bool
+	Corpus         *corpus.Corpus
 	Logf           func(level int, msg string, args ...interface{})
 	Coverage       bool
 	FaultInjection bool
@@ -87,7 +88,7 @@ type Config struct {
 	// If the number of queued candidates is less than MinCandidates,
 	// NeedCandidates is triggered.
 	MinCandidates uint
-	NewInputs     chan rpctype.Input
+	NewInputs     chan corpus.NewInput
 }
 
 type Request struct {
@@ -133,7 +134,7 @@ func (fuzzer *Fuzzer) Done(req *Request, res *Result) {
 func (fuzzer *Fuzzer) triageProgCall(p *prog.Prog, info *ipc.CallInfo, call int,
 	flags ProgTypes) {
 	prio := signalPrio(p, info, call)
-	newMaxSignal := fuzzer.Corpus.AddRawMaxSignal(info.Signal, prio)
+	newMaxSignal := fuzzer.Cover.addRawMaxSignal(info.Signal, prio)
 	if newMaxSignal.Empty() {
 		return
 	}
@@ -304,12 +305,12 @@ func (fuzzer *Fuzzer) choiceTableUpdater() {
 			return
 		case <-fuzzer.ctRegenerate:
 		}
-		fuzzer.updateChoiceTable(fuzzer.Corpus.Programs())
+		fuzzer.updateChoiceTable(fuzzer.Config.Corpus.Programs())
 	}
 }
 
 func (fuzzer *Fuzzer) ChoiceTable() *prog.ChoiceTable {
-	progs := fuzzer.Corpus.Programs()
+	progs := fuzzer.Config.Corpus.Programs()
 
 	fuzzer.ctMu.Lock()
 	defer fuzzer.ctMu.Unlock()
@@ -348,3 +349,15 @@ func (fuzzer *Fuzzer) logCurrentStats() {
 		fuzzer.Logf(0, "%s", str)
 	}
 }
+
+type Stat struct {
+	CoverStat
+	corpus.Stat
+}
+
+func (fuzzer *Fuzzer) Stat() Stat {
+	return Stat{
+		CoverStat: fuzzer.Cover.Stat(),
+		Stat:      fuzzer.Config.Corpus.Stat(),
+	}
+}
diff --git a/pkg/fuzzer/fuzzer_test.go b/pkg/fuzzer/fuzzer_test.go
index 1896b2b84..7000bb062 100644
--- a/pkg/fuzzer/fuzzer_test.go
+++ b/pkg/fuzzer/fuzzer_test.go
@@ -18,10 +18,10 @@ import (
 	"testing"
 	"time"
 
+	"github.com/google/syzkaller/pkg/corpus"
 	"github.com/google/syzkaller/pkg/csource"
 	"github.com/google/syzkaller/pkg/ipc"
 	"github.com/google/syzkaller/pkg/ipc/ipcconfig"
-	"github.com/google/syzkaller/pkg/rpctype"
 	"github.com/google/syzkaller/pkg/testutil"
 	"github.com/google/syzkaller/prog"
 	"github.com/google/syzkaller/sys/targets"
@@ -41,7 +41,8 @@ func TestFuzz(t *testing.T) {
 	defer cancel()
 
 	fuzzer := NewFuzzer(ctx, &Config{
-		Debug: true,
+		Debug:  true,
+		Corpus: corpus.NewCorpus(ctx),
 		Logf: func(level int, msg string, args ...interface{}) {
 			if level > 1 {
 				return
@@ -52,7 +53,7 @@ func TestFuzz(t *testing.T) {
 		EnabledCalls: map[*prog.Syscall]bool{
 			target.SyscallMap["syz_test_fuzzer1"]: true,
 		},
-		NewInputs: make(chan rpctype.Input),
+		NewInputs: make(chan corpus.NewInput),
 	}, rand.New(testutil.RandSource(t)), target)
 
 	go func() {
@@ -77,7 +78,7 @@ func TestFuzz(t *testing.T) {
 	tf.wait()
 
 	t.Logf("resulting corpus:")
-	for _, p := range fuzzer.Corpus.Programs() {
+	for _, p := range fuzzer.Config.Corpus.Programs() {
 		t.Logf("-----")
 		t.Logf("%s", p.Serialize())
 	}
@@ -99,6 +100,7 @@ func BenchmarkFuzzer(b *testing.B) {
 		calls[c] = true
 	}
 	fuzzer := NewFuzzer(ctx, &Config{
+		Corpus:       corpus.NewCorpus(ctx),
 		Coverage:     true,
 		EnabledCalls: calls,
 	}, rand.New(rand.NewSource(time.Now().UnixNano())), target)
@@ -160,7 +162,7 @@ func (f *testFuzzer) oneMore() bool {
 	defer f.mu.Unlock()
 	f.iter++
 	if f.iter%100 == 0 {
-		stat := f.fuzzer.Corpus.Stat()
+		stat := f.fuzzer.Stat()
 		f.t.Logf("<iter %d>: corpus %d, signal %d, max signal %d, crash types %d",
 			f.iter, stat.Progs, stat.Signal, stat.MaxSignal, len(f.crashes))
 	}
diff --git a/pkg/fuzzer/job.go b/pkg/fuzzer/job.go
index a0295affb..63493c914 100644
--- a/pkg/fuzzer/job.go
+++ b/pkg/fuzzer/job.go
@@ -7,10 +7,9 @@ import (
 	"fmt"
 	"math/rand"
 
+	"github.com/google/syzkaller/pkg/corpus"
 	"github.com/google/syzkaller/pkg/cover"
-	"github.com/google/syzkaller/pkg/hash"
 	"github.com/google/syzkaller/pkg/ipc"
-	"github.com/google/syzkaller/pkg/rpctype"
 	"github.com/google/syzkaller/pkg/signal"
 	"github.com/google/syzkaller/prog"
 )
@@ -71,7 +70,7 @@ func genProgRequest(fuzzer *Fuzzer, rnd *rand.Rand) *Request {
 }
 
 func mutateProgRequest(fuzzer *Fuzzer, rnd *rand.Rand) *Request {
-	p := fuzzer.Corpus.chooseProgram(rnd)
+	p := fuzzer.Config.Corpus.ChooseProgram(rnd)
 	if p == nil {
 		return nil
 	}
@@ -80,7 +79,7 @@ func mutateProgRequest(fuzzer *Fuzzer, rnd *rand.Rand) *Request {
 		prog.RecommendedCalls,
 		fuzzer.ChoiceTable(),
 		fuzzer.Config.NoMutateCalls,
-		fuzzer.Corpus.Programs(),
+		fuzzer.Config.Corpus.Programs(),
 	)
 	return &Request{
 		Prog:       newP,
@@ -126,10 +125,9 @@ func triageJobPrio(flags ProgTypes) jobPriority {
 }
 
 func (job *triageJob) run(fuzzer *Fuzzer) {
-	callName := ".extra"
 	logCallName := "extra"
 	if job.call != -1 {
-		callName = job.p.Calls[job.call].Meta.Name
+		callName := job.p.Calls[job.call].Meta.Name
 		logCallName = fmt.Sprintf("call #%v %v", job.call, callName)
 	}
 	fuzzer.Logf(3, "triaging input for %v (new signal=%v)", logCallName, job.newSignal.Len())
@@ -144,9 +142,7 @@ func (job *triageJob) run(fuzzer *Fuzzer) {
 			return
 		}
 	}
-	data := job.p.Serialize()
-	fuzzer.Logf(2, "added new input for %q to the corpus:\n%s",
-		logCallName, string(data))
+	fuzzer.Logf(2, "added new input for %q to the corpus:\n%s", logCallName, job.p.String())
 	if job.flags&ProgSmashed == 0 {
 		fuzzer.startJob(&smashJob{
 			p:           job.p.Clone(),
@@ -154,18 +150,18 @@ func (job *triageJob) run(fuzzer *Fuzzer) {
 			jobPriority: newJobPriority(smashPrio),
 		})
 	}
-	fuzzer.Corpus.Save(job.p, info.stableSignal, hash.Hash(data))
+	input := corpus.NewInput{
+		Prog:     job.p,
+		Call:     job.call,
+		Signal:   info.stableSignal,
+		Cover:    info.cover.Serialize(),
+		RawCover: info.rawCover,
+	}
+	fuzzer.Config.Corpus.Save(input)
 	if fuzzer.Config.NewInputs != nil {
 		select {
 		case <-fuzzer.ctx.Done():
-		case fuzzer.Config.NewInputs <- rpctype.Input{
-			Call:     callName,
-			CallID:   job.call,
-			Prog:     data,
-			Signal:   info.stableSignal.Serialize(),
-			Cover:    info.cover.Serialize(),
-			RawCover: info.rawCover,
-		}:
+		case fuzzer.Config.NewInputs <- input:
 		}
 	}
 }
@@ -298,7 +294,7 @@ func (job *smashJob) run(fuzzer *Fuzzer) {
 		p.Mutate(rnd, prog.RecommendedCalls,
 			fuzzer.ChoiceTable(),
 			fuzzer.Config.NoMutateCalls,
-			fuzzer.Corpus.Programs())
+			fuzzer.Config.Corpus.Programs())
 		result := fuzzer.exec(job, &Request{
 			Prog:       p,
 			NeedSignal: true,