From 3f1d02b23f99beaf2bf3b06c11642e56578b12ee Mon Sep 17 00:00:00 2001
From: Jiaheng Hu <jiahengh@google.com>
Date: Sat, 15 Aug 2020 00:24:07 +0000
Subject: syz-manager: add test file as corpus

This commit enables the syz-manager to add unit test files as corpus to
accelerate fuzzing. The syz-ci would copy unit tests into the
worker/seeds folder for each manager process, and the manager would add
those tests as seed into the corpus.
---
 pkg/db/db.go           | 36 ++++++++++++++++++++++++++++++++++++
 pkg/osutil/osutil.go   | 15 +++++++++++++++
 syz-ci/manager.go      |  8 +++++++-
 syz-ci/syz-ci.go       |  3 ++-
 syz-manager/manager.go |  3 +++
 5 files changed, 63 insertions(+), 2 deletions(-)

diff --git a/pkg/db/db.go b/pkg/db/db.go
index 17723b224..7961c0797 100644
--- a/pkg/db/db.go
+++ b/pkg/db/db.go
@@ -17,6 +17,7 @@ import (
 	"io"
 	"io/ioutil"
 	"os"
+	"path/filepath"
 
 	"github.com/google/syzkaller/pkg/hash"
 	"github.com/google/syzkaller/pkg/log"
@@ -68,6 +69,41 @@ func (db *DB) Save(key string, val []byte, seq uint64) {
 	db.uncompacted++
 }
 
+// Load the test progs in the given directory store them inside db.
+func (db *DB) LoadTestAsSeed(target *prog.Target, seedsDir string) {
+	var files []string
+	err := filepath.Walk(seedsDir, func(path string, info os.FileInfo, err error) error {
+		if !info.IsDir() {
+			files = append(files, path)
+		}
+		return nil
+	})
+	if err != nil {
+		panic(err)
+	}
+	progEntries := loadPrograms(target, files)
+	if len(progEntries) == 0 {
+		return
+	}
+	for _, progEntry := range progEntries {
+		prog := progEntry.P.Serialize()
+		sig := hash.String(prog)
+		db.Save(sig, prog, 0)
+	}
+}
+
+func loadPrograms(target *prog.Target, files []string) []*prog.LogEntry {
+	var entries []*prog.LogEntry
+	for _, fn := range files {
+		data, err := ioutil.ReadFile(fn)
+		if err != nil {
+			log.Fatalf("failed to read log file: %v", err)
+		}
+		entries = append(entries, target.ParseLog(data)...)
+	}
+	return entries
+}
+
 func (db *DB) Delete(key string) {
 	if _, ok := db.Records[key]; !ok {
 		return
diff --git a/pkg/osutil/osutil.go b/pkg/osutil/osutil.go
index 47fcc8b7b..1eafafcb9 100644
--- a/pkg/osutil/osutil.go
+++ b/pkg/osutil/osutil.go
@@ -222,6 +222,21 @@ func LinkFiles(srcDir, dstDir string, files map[string]bool) error {
 	return nil
 }
 
+// LinkDir creates soft links from dstDir to srcDir.
+// All other files in dstDir are removed.
+func LinkDir(srcDir, dstDir string) error {
+	if err := os.RemoveAll(dstDir); err != nil {
+		return err
+	}
+	if err := MkdirAll(dstDir); err != nil {
+		return err
+	}
+	if err := os.Symlink(srcDir, dstDir); err != nil {
+		return err
+	}
+	return nil
+}
+
 func MkdirAll(dir string) error {
 	return os.MkdirAll(dir, DefaultDirPerm)
 }
diff --git a/syz-ci/manager.go b/syz-ci/manager.go
index 2d7d6e8c7..f58857873 100644
--- a/syz-ci/manager.go
+++ b/syz-ci/manager.go
@@ -79,7 +79,7 @@ type Manager struct {
 	stop       chan struct{}
 }
 
-func createManager(cfg *Config, mgrcfg *ManagerConfig, stop chan struct{}) (*Manager, error) {
+func createManager(cfg *Config, mgrcfg *ManagerConfig, stop chan struct{}, syzLinuxTestDir string) (*Manager, error) {
 	dir := osutil.Abs(filepath.Join("managers", mgrcfg.Name))
 	if err := osutil.MkdirAll(dir); err != nil {
 		log.Fatal(err)
@@ -126,6 +126,12 @@ func createManager(cfg *Config, mgrcfg *ManagerConfig, stop chan struct{}) (*Man
 		dash:       dash,
 		stop:       stop,
 	}
+
+	// Copy test files into the current manager workdir.
+	if err := osutil.LinkDir(syzLinuxTestDir, filepath.Join(dir, "workdir", "seeds")); err != nil {
+		log.Fatal(err)
+	}
+
 	os.RemoveAll(mgr.currentDir)
 	return mgr, nil
 }
diff --git a/syz-ci/syz-ci.go b/syz-ci/syz-ci.go
index faa51f4fc..7b4019375 100644
--- a/syz-ci/syz-ci.go
+++ b/syz-ci/syz-ci.go
@@ -170,9 +170,10 @@ func main() {
 		wg.Done()
 	}()
 
+	syzLinuxTestDir := filepath.Join(updater.syzkallerDir, "sys", "linux", "test")
 	var managers []*Manager
 	for _, mgrcfg := range cfg.Managers {
-		mgr, err := createManager(cfg, mgrcfg, stop)
+		mgr, err := createManager(cfg, mgrcfg, stop, syzLinuxTestDir)
 		if err != nil {
 			log.Logf(0, "failed to create manager %v: %v", mgrcfg.Name, err)
 			continue
diff --git a/syz-manager/manager.go b/syz-manager/manager.go
index 3a65a7300..cffb73817 100644
--- a/syz-manager/manager.go
+++ b/syz-manager/manager.go
@@ -188,6 +188,9 @@ func RunManager(cfg *mgrconfig.Config, target *prog.Target, sysTarget *targets.T
 		log.Fatalf("failed to open corpus database: %v", err)
 	}
 
+	// Load unit test program into corpusDB.
+	mgr.corpusDB.LoadTestAsSeed(target, filepath.Join(cfg.Workdir, "seeds/"))
+
 	// Create HTTP server.
 	mgr.initHTTP()
 	mgr.collectUsedFiles()
-- 
cgit mrf-deployment