prog: move image extraction from sys/linux

Now that images are not linux-specific,
we can move all image-related logic directly into prog package
and significantly simplify the logic.

11 files changed, 123 insertions, 291 deletions

diff --git a/prog/analysis.go b/prog/analysis.go
index b4a30dbe8..fec43e1bd 100644
--- a/prog/analysis.go
+++ b/prog/analysis.go
@@ -9,8 +9,11 @@
 package prog
 
 import (
+	"bytes"
 	"fmt"
 	"io"
+
+	"github.com/google/syzkaller/pkg/image"
 )
 
 type state struct {
@@ -340,39 +343,27 @@ func checkMaxCallID(id int) {
 	}
 }
 
-type ExtractedAssetType int
+type AssetType int
 
 const (
-	MountInRepro ExtractedAssetType = iota
+	MountInRepro AssetType = iota
 )
 
-type ExtractedAsset struct {
-	Call   int
-	Type   ExtractedAssetType
-	Reader io.Reader
-	Error  error
-}
-
-func (p *Prog) ExtractAssets() []*ExtractedAsset {
-	handler := p.Target.ExtractMountedImage
-	if handler == nil {
-		// Such an operation is not supported by the target.
-		return nil
-	}
-	ret := []*ExtractedAsset{}
+func (p *Prog) ForEachAsset(cb func(name string, typ AssetType, r io.Reader)) {
 	for id, c := range p.Calls {
-		// So far we only support the MountInRepro asset.
-		reader, err := handler(c)
-		if reader == nil && err == nil {
-			// This is not the call that contains the mount image.
-			continue
-		}
-		ret = append(ret, &ExtractedAsset{
-			Type:   MountInRepro,
-			Call:   id,
-			Reader: reader,
-			Error:  err,
+		ForeachArg(c, func(arg Arg, _ *ArgCtx) {
+			a, ok := arg.(*DataArg)
+			if !ok || a.Type().(*BufferType).Kind != BufferCompressed {
+				return
+			}
+			data, err := image.Decompress(a.Data())
+			if err != nil {
+				panic(err)
+			}
+			if len(data) == 0 {
+				return
+			}
+			cb(fmt.Sprintf("mount_%v", id), MountInRepro, bytes.NewReader(data))
 		})
 	}
-	return ret
 }
diff --git a/prog/images_test.go b/prog/images_test.go
new file mode 100644
index 000000000..86914105f
--- /dev/null
+++ b/prog/images_test.go
@@ -0,0 +1,81 @@
+// Copyright 2022 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 prog_test
+
+import (
+	"flag"
+	"fmt"
+	"io"
+	"os"
+	"path/filepath"
+	"reflect"
+	"sort"
+	"strings"
+	"testing"
+
+	"github.com/google/go-cmp/cmp"
+	"github.com/google/syzkaller/pkg/osutil"
+	. "github.com/google/syzkaller/prog"
+	"github.com/google/syzkaller/sys/targets"
+)
+
+var flagUpdate = flag.Bool("update", false, "update test files accordingly to current results")
+
+func TestForEachAsset(t *testing.T) {
+	target, err := GetTarget(targets.Linux, targets.AMD64)
+	if err != nil {
+		t.Fatal(err)
+	}
+	files, err := filepath.Glob(filepath.Join("testdata", "fs_images", "*.in"))
+	if err != nil {
+		t.Fatalf("directory read failed: %v", err)
+	}
+	allOutFiles, err := filepath.Glob(filepath.Join("testdata", "fs_images", "*.out*"))
+	if err != nil {
+		t.Fatalf("directory read failed: %v", err)
+	}
+	testedOutFiles := []string{}
+	for _, file := range files {
+		sourceProg, err := os.ReadFile(file)
+		if err != nil {
+			t.Fatal(err)
+		}
+		p, err := target.Deserialize(sourceProg, NonStrict)
+		if err != nil {
+			t.Fatalf("failed to deserialize %s: %s", file, err)
+		}
+		base := strings.TrimSuffix(file, ".in")
+		p.ForEachAsset(func(name string, typ AssetType, r io.Reader) {
+			if typ != MountInRepro {
+				t.Fatalf("unknown asset type %v", typ)
+			}
+			testResult, err := io.ReadAll(r)
+			if err != nil {
+				t.Fatal(err)
+			}
+			outFilePath := fmt.Sprintf("%v.out_%v", base, name)
+			if *flagUpdate {
+				if err := osutil.WriteFile(outFilePath, testResult); err != nil {
+					t.Fatal(err)
+				}
+			}
+			if !osutil.IsExist(outFilePath) {
+				t.Fatalf("asset %v does not exist", outFilePath)
+			}
+			testedOutFiles = append(testedOutFiles, outFilePath)
+			outFile, err := os.ReadFile(outFilePath)
+			if err != nil {
+				t.Fatal(err)
+			}
+			if !reflect.DeepEqual(testResult, outFile) {
+				t.Fatalf("output not equal:\nWant: %x\nGot: %x", outFile, testResult)
+			}
+		})
+	}
+	sort.Strings(testedOutFiles)
+	sort.Strings(allOutFiles)
+	if diff := cmp.Diff(allOutFiles, testedOutFiles); diff != "" {
+		t.Fatalf("not all output files used: %v", diff)
+	}
+}
diff --git a/prog/target.go b/prog/target.go
index d50f1bf8f..36fde13de 100644
--- a/prog/target.go
+++ b/prog/target.go
@@ -5,7 +5,6 @@ package prog
 
 import (
 	"fmt"
-	"io"
 	"math/rand"
 	"sort"
 	"sync"
@@ -65,11 +64,6 @@ type Target struct {
 	SyscallMap map[string]*Syscall
 	ConstMap   map[string]uint64
 
-	// The extracted images will be then saved to the disk or uploaded to the asset storage.
-	// Returns nil if the call does not mount any image.
-	// We have to use io.Reader since such blobs can get pretty large.
-	ExtractMountedImage func(c *Call) (io.Reader, error)
-
 	init        sync.Once
 	initArch    func(target *Target)
 	types       []Type
diff --git a/sys/linux/testdata/fs_images/0.in b/prog/testdata/fs_images/0.in
index 0ea5bce18..0ea5bce18 100644
--- a/sys/linux/testdata/fs_images/0.in
+++ b/prog/testdata/fs_images/0.in
diff --git a/sys/linux/testdata/fs_images/0.out0 b/prog/testdata/fs_images/0.out_mount_0
index 4eb5233c9..4eb5233c9 100644
--- a/sys/linux/testdata/fs_images/0.out0
+++ b/prog/testdata/fs_images/0.out_mount_0
diff --git a/sys/linux/testdata/fs_images/1.in b/prog/testdata/fs_images/1.in
index 6e8992119..6e8992119 100644
--- a/sys/linux/testdata/fs_images/1.in
+++ b/prog/testdata/fs_images/1.in
diff --git a/sys/linux/testdata/fs_images/1.out0 b/prog/testdata/fs_images/1.out_mount_0
index c498f4d53..c498f4d53 100644
--- a/sys/linux/testdata/fs_images/1.out0
+++ b/prog/testdata/fs_images/1.out_mount_0
diff --git a/sys/linux/init.go b/sys/linux/init.go
index 7d537075c..faf4a98af 100644
--- a/sys/linux/init.go
+++ b/sys/linux/init.go
@@ -54,7 +54,6 @@ func InitTarget(target *prog.Target) {
 
 	target.MakeDataMmap = targets.MakePosixMmap(target, true, true)
 	target.Neutralize = arch.neutralize
-	target.ExtractMountedImage = arch.extractSyzMountImage
 	target.SpecialTypes = map[string]func(g *prog.Gen, typ prog.Type, dir prog.Dir, old prog.Arg) (
 		prog.Arg, []*prog.Call){
 		"timespec":                  arch.generateTimespec,
@@ -245,8 +244,6 @@ func (arch *arch) neutralize(c *prog.Call, fixStructure bool) error {
 	case "sched_setattr":
 		// Enabling a SCHED_FIFO or a SCHED_RR policy may lead to false positive stall-related crashes.
 		neutralizeSchedAttr(c.Args[1])
-	case "syz_mount_image":
-		return arch.fixUpSyzMountImage(c, fixStructure)
 	}
 
 	switch c.Meta.Name {
diff --git a/sys/linux/init_images.go b/sys/linux/init_images.go
deleted file mode 100644
index c9a3035ef..000000000
--- a/sys/linux/init_images.go
+++ /dev/null
@@ -1,88 +0,0 @@
-// Copyright 2022 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 linux
-
-import (
-	"bytes"
-	"fmt"
-	"io"
-
-	"github.com/google/syzkaller/pkg/image"
-	"github.com/google/syzkaller/prog"
-)
-
-func (arch *arch) extractSyzMountImage(c *prog.Call) (io.Reader, error) {
-	// In order to reduce the size of syzlang programs, disk images are compressed.
-	// Here we extract the compressed image.
-	if c.Meta.CallName != "syz_mount_image" {
-		return nil, nil
-	}
-	data, _, err := parseSyzMountImage(c)
-	if err != nil {
-		// Parsing failed --> do not try to recover, just ignore.
-		return nil, err
-	} else if len(data) == 0 {
-		return nil, fmt.Errorf("an empty image")
-	}
-	buf := new(bytes.Buffer)
-	if err := image.DecompressWriter(buf, data); err != nil {
-		return nil, err
-	}
-	return buf, nil
-}
-
-func (arch *arch) fixUpSyzMountImage(c *prog.Call, fixStructure bool) error {
-	// Previously we did such a sanitization right in the common_linux.h, but this was problematic
-	// for two reasons:
-	// 1) It further complicates the already complicated executor code.
-	// 2) We'd need to duplicate the logic in Go for raw image extraction.
-	// So now we do all the initialization in Go and let the C code only interpret the commands.
-	data, sizeArg, err := parseSyzMountImage(c)
-	sizeArg.Val = uint64(len(data))
-	if err != nil {
-		if fixStructure {
-			deactivateSyzMountImage(c)
-			return nil
-		}
-		return err
-	}
-	return nil
-}
-
-func deactivateSyzMountImage(c *prog.Call) {
-	dataPointer := c.Args[6]
-	newArg := dataPointer.Type().DefaultArg(dataPointer.Dir())
-	prog.RemoveArg(dataPointer)
-	c.Args[6] = newArg
-	// Also set the size field to 0.
-	c.Args[5].(*prog.ConstArg).Val = 0
-}
-
-// Returns the compressed disk image and the corresponding length argument.
-func parseSyzMountImage(c *prog.Call) ([]byte, *prog.ConstArg, error) {
-	if len(c.Args) < 7 {
-		panic("invalid number of arguments in syz_mount_image")
-	}
-
-	// Check `size` argument.
-	sizeArg, ok := c.Args[5].(*prog.ConstArg)
-	if !ok {
-		panic("syz_mount_image's size arg is not const")
-	}
-
-	dataPointer, ok := c.Args[6].(*prog.PointerArg)
-	if !ok {
-		panic("syz_mount_image's data pointer is invalid")
-	}
-
-	dataArg, ok := dataPointer.Res.(*prog.DataArg)
-	if !ok {
-		return nil, sizeArg, fmt.Errorf("could not find raw image data")
-	}
-	if dataArg == nil {
-		return nil, sizeArg, fmt.Errorf("image argument contains no data")
-	}
-
-	return dataArg.Data(), sizeArg, nil
-}
diff --git a/sys/linux/init_images_test.go b/sys/linux/init_images_test.go
deleted file mode 100644
index c6af2049c..000000000
--- a/sys/linux/init_images_test.go
+++ /dev/null
@@ -1,114 +0,0 @@
-// Copyright 2022 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 linux_test
-
-import (
-	"flag"
-	"fmt"
-	"io"
-	"os"
-	"path/filepath"
-	"reflect"
-	"sort"
-	"strings"
-	"testing"
-
-	"github.com/google/syzkaller/pkg/osutil"
-	"github.com/google/syzkaller/prog"
-	"github.com/google/syzkaller/sys/targets"
-)
-
-// nolint: lll
-
-func TestSyzMountImageNeutralize(t *testing.T) {
-	prog.TestDeserializeHelper(t, targets.Linux, targets.AMD64, nil, []prog.DeserializeTest{
-		{
-			// A valid call, nothing should change.
-			In: `syz_mount_image$bfs(&(0x7f0000000000)='bfs\x00', &(0x7f0000000100)='./file0\x00', ` +
-				`0x0, &(0x7f00000100a0)={[], [], 0x0}, 0x0, 0x15, ` +
-				`&(0x7f0000000200)="$eJwqrqzKTszJSS0CBAAA//8TyQPi")`,
-		},
-		{
-			// Invalid compressed size.
-			In: `syz_mount_image$bfs(&(0x7f0000000000)='bfs\x00', &(0x7f0000000100)='./file0\x00', ` +
-				`0x0, &(0x7f00000100a0)={[], [], 0x0}, 0x0, 0xdeadbeef, ` +
-				`&(0x7f0000000200)="$eJwqrqzKTszJSS0CBAAA//8TyQPi")`,
-			// It should be able to fix up the size.
-			Out: `syz_mount_image$bfs(&(0x7f0000000000)='bfs\x00', &(0x7f0000000100)='./file0\x00', ` +
-				`0x0, &(0x7f00000100a0)={[], [], 0x0}, 0x0, 0x15, ` +
-				`&(0x7f0000000200)="$eJwqrqzKTszJSS0CBAAA//8TyQPi")`,
-		},
-	})
-}
-
-var flagUpdate = flag.Bool("update", false, "update test files accordingly to current results")
-
-func TestExtractSyzMountImage(t *testing.T) {
-	target, err := prog.GetTarget(targets.Linux, targets.AMD64)
-	if err != nil {
-		t.Fatal(err)
-	}
-	files, err := filepath.Glob(filepath.Join("testdata", "fs_images", "*.in"))
-	if err != nil {
-		t.Fatalf("directory read failed: %v", err)
-	}
-	allOutFiles, err := filepath.Glob(filepath.Join("testdata", "fs_images", "*.out*"))
-	if err != nil {
-		t.Fatalf("directory read failed: %v", err)
-	}
-	testedOutFiles := []string{}
-	for _, file := range files {
-		if !strings.HasSuffix(file, ".in") {
-			continue
-		}
-		sourceProg, err := os.ReadFile(file)
-		if err != nil {
-			t.Fatal(err)
-		}
-		p, err := target.Deserialize(sourceProg, prog.NonStrict)
-		if err != nil {
-			t.Fatalf("failed to deserialize %s: %s", file, err)
-		}
-		base := strings.TrimSuffix(file, ".in")
-		for _, asset := range p.ExtractAssets() {
-			if asset.Type != prog.MountInRepro {
-				continue
-			}
-			outFilePath := fmt.Sprintf("%s.out%d", base, asset.Call)
-			var testResult []byte
-			if asset.Reader != nil {
-				var err error
-				testResult, err = io.ReadAll(asset.Reader)
-				if err != nil {
-					t.Fatal(err)
-				}
-			}
-			if *flagUpdate && asset.Reader != nil {
-				err := osutil.WriteFile(outFilePath, testResult)
-				if err != nil {
-					t.Fatal(err)
-				}
-			}
-			outExists := osutil.IsExist(outFilePath)
-			if !outExists && asset.Reader != nil {
-				t.Fatalf("#%d: mount found, but does not exist in the answer", asset.Call)
-			}
-			if testResult != nil {
-				testedOutFiles = append(testedOutFiles, outFilePath)
-				outFile, err := os.ReadFile(outFilePath)
-				if err != nil {
-					t.Fatal(err)
-				}
-				if !reflect.DeepEqual(testResult, outFile) {
-					t.Fatalf("output not equal:\nWant: %x\nGot: %x", outFile, testResult)
-				}
-			}
-		}
-	}
-	sort.Strings(testedOutFiles)
-	sort.Strings(allOutFiles)
-	if !reflect.DeepEqual(testedOutFiles, allOutFiles) {
-		t.Fatalf("all out files: %v\ntested files: %v", allOutFiles, testedOutFiles)
-	}
-}
diff --git a/syz-manager/manager.go b/syz-manager/manager.go
index 4db377f0c..dd7d5942f 100644
--- a/syz-manager/manager.go
+++ b/syz-manager/manager.go
@@ -8,6 +8,7 @@ import (
 	"encoding/json"
 	"flag"
 	"fmt"
+	"io"
 	"io/ioutil"
 	"math/rand"
 	"net"
@@ -1014,12 +1015,12 @@ func (mgr *Manager) saveFailedRepro(rep *report.Report, stats *repro.Stats) {
 func (mgr *Manager) saveRepro(res *ReproResult) {
 	repro := res.repro
 	opts := fmt.Sprintf("# %+v\n", repro.Opts)
-	prog := repro.Prog.Serialize()
+	progText := repro.Prog.Serialize()
 
 	// Append this repro to repro list to send to hub if it didn't come from hub originally.
 	if !res.hub {
 		progForHub := []byte(fmt.Sprintf("# %+v\n# %v\n# %v\n%s",
-			repro.Opts, repro.Report.Title, mgr.cfg.Tag, prog))
+			repro.Opts, repro.Report.Title, mgr.cfg.Tag, progText))
 		mgr.mu.Lock()
 		mgr.newRepros = append(mgr.newRepros, progForHub)
 		mgr.mu.Unlock()
@@ -1068,7 +1069,7 @@ func (mgr *Manager) saveRepro(res *ReproResult) {
 			Flags:      crashFlags,
 			Report:     report.Report,
 			ReproOpts:  repro.Opts.Serialize(),
-			ReproSyz:   repro.Prog.Serialize(),
+			ReproSyz:   progText,
 			ReproC:     cprogText,
 			Assets:     mgr.uploadReproAssets(repro),
 		}
@@ -1088,7 +1089,7 @@ func (mgr *Manager) saveRepro(res *ReproResult) {
 	if err := osutil.WriteFile(filepath.Join(dir, "description"), []byte(rep.Title+"\n")); err != nil {
 		log.Logf(0, "failed to write crash: %v", err)
 	}
-	osutil.WriteFile(filepath.Join(dir, "repro.prog"), append([]byte(opts), prog...))
+	osutil.WriteFile(filepath.Join(dir, "repro.prog"), append([]byte(opts), progText...))
 	if mgr.cfg.Tag != "" {
 		osutil.WriteFile(filepath.Join(dir, "repro.tag"), []byte(mgr.cfg.Tag))
 	}
@@ -1101,9 +1102,11 @@ func (mgr *Manager) saveRepro(res *ReproResult) {
 	if len(cprogText) > 0 {
 		osutil.WriteFile(filepath.Join(dir, "repro.cprog"), cprogText)
 	}
-	for _, asset := range repro.Prog.ExtractAssets() {
-		saveReproAsset(dir, asset)
-	}
+	repro.Prog.ForEachAsset(func(name string, typ prog.AssetType, r io.Reader) {
+		if err := osutil.WriteGzipStream(name+".gz", r); err != nil {
+			log.Logf(0, "failed to write crash asset: type %d, write error %v", typ, err)
+		}
+	})
 	if res.strace != nil {
 		// Unlike dashboard reporting, we save strace output separately from the original log.
 		if res.strace.Error != nil {
@@ -1118,60 +1121,28 @@ func (mgr *Manager) saveRepro(res *ReproResult) {
 }
 
 func (mgr *Manager) uploadReproAssets(repro *repro.Result) []dashapi.NewAsset {
-	ret := []dashapi.NewAsset{}
 	if mgr.assetStorage == nil {
-		return ret
+		return nil
 	}
 
-	for _, asset := range repro.Prog.ExtractAssets() {
-		if asset.Reader == nil {
-			// Skip empty assets.
-			continue
+	ret := []dashapi.NewAsset{}
+	repro.Prog.ForEachAsset(func(name string, typ prog.AssetType, r io.Reader) {
+		dashTyp, ok := map[prog.AssetType]dashapi.AssetType{
+			prog.MountInRepro: dashapi.MountInRepro,
+		}[typ]
+		if !ok {
+			panic("unknown extracted prog asset")
 		}
-		newAsset, err := mgr.uploadReproAsset(asset)
+		asset, err := mgr.assetStorage.UploadCrashAsset(r, name, dashTyp, nil)
 		if err != nil {
-			log.Logf(1, "processing of the asset #%d,%v failed: %s",
-				asset.Call, asset.Type, err)
-			continue
+			log.Logf(1, "processing of the asset %v (%v) failed: %v", name, typ, err)
+			return
 		}
-		ret = append(ret, newAsset)
-	}
+		ret = append(ret, asset)
+	})
 	return ret
 }
 
-func (mgr *Manager) uploadReproAsset(asset *prog.ExtractedAsset) (dashapi.NewAsset, error) {
-	if asset.Error != nil {
-		return dashapi.NewAsset{}, fmt.Errorf("failed to extract: %w", asset.Error)
-	}
-	switch asset.Type {
-	case prog.MountInRepro:
-		name := fmt.Sprintf("mount_%d", asset.Call)
-		return mgr.assetStorage.UploadCrashAsset(asset.Reader, name,
-			dashapi.MountInRepro, nil)
-	default:
-		panic("unknown extracted prog asset")
-	}
-}
-
-func saveReproAsset(dir string, asset *prog.ExtractedAsset) {
-	path := ""
-	switch asset.Type {
-	case prog.MountInRepro:
-		path = filepath.Join(dir, fmt.Sprintf("repro.mount%d", asset.Call))
-	default:
-		panic("unknown extracted prog asset")
-	}
-	var err error
-	if asset.Error != nil {
-		err = osutil.WriteFile(path+".error", []byte(asset.Error.Error()))
-	} else if asset.Reader != nil {
-		err = osutil.WriteGzipStream(path+".gz", asset.Reader)
-	}
-	if err != nil {
-		log.Logf(0, "failed to write crash asset: type %d, write error %v", asset.Type, err)
-	}
-}
-
 func saveReproStats(filename string, stats *repro.Stats) {
 	text := ""
 	if stats != nil {