pkg/image: factor out from prog

Move image compression-related function to a separate package.
In preperation for subsequent changes that make decompression
more complex. Prog package is already large and complex.
Also makes running compression tests/benchmarks much faster.

8 files changed, 26 insertions, 19 deletions

diff --git a/prog/compression.go b/pkg/image/compression.go
index 87820b5ee..9878b460d 100644
--- a/prog/compression.go
+++ b/pkg/image/compression.go
@@ -1,7 +1,7 @@
 // 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
+package image
 
 import (
 	"bytes"
diff --git a/prog/compression_test.go b/pkg/image/compression_test.go
index db5d95070..cf18ed340 100644
--- a/prog/compression_test.go
+++ b/pkg/image/compression_test.go
@@ -1,7 +1,7 @@
 // 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
+package image
 
 import (
 	"bytes"
diff --git a/pkg/ipc/ipc_test.go b/pkg/ipc/ipc_test.go
index c1d3af69a..a3ab140f1 100644
--- a/pkg/ipc/ipc_test.go
+++ b/pkg/ipc/ipc_test.go
@@ -13,6 +13,7 @@ import (
 	"time"
 
 	"github.com/google/syzkaller/pkg/csource"
+	"github.com/google/syzkaller/pkg/image"
 	. "github.com/google/syzkaller/pkg/ipc"
 	"github.com/google/syzkaller/pkg/ipc/ipcconfig"
 	"github.com/google/syzkaller/pkg/osutil"
@@ -210,9 +211,9 @@ func TestZlib(t *testing.T) {
 	r := rand.New(testutil.RandSource(t))
 	for i := 0; i < 10; i++ {
 		data := testutil.RandMountImage(r)
-		compressed := prog.Compress(data)
+		compressed := image.Compress(data)
 		text := fmt.Sprintf(`syz_compare_zlib(&(0x7f0000000000)="$%s", AUTO, &(0x7f0000800000)="$%s", AUTO)`,
-			prog.EncodeB64(data), prog.EncodeB64(compressed))
+			image.EncodeB64(data), image.EncodeB64(compressed))
 		p, err := target.Deserialize([]byte(text), prog.Strict)
 		if err != nil {
 			t.Fatalf("failed to deserialize empty program: %v", err)
diff --git a/prog/encoding.go b/prog/encoding.go
index 6b9fbe501..bfa80b983 100644
--- a/prog/encoding.go
+++ b/prog/encoding.go
@@ -11,6 +11,8 @@ import (
 	"reflect"
 	"strconv"
 	"strings"
+
+	"github.com/google/syzkaller/pkg/image"
 )
 
 // String generates a very compact program description (mostly for debug output).
@@ -604,10 +606,10 @@ func (p *parser) parseArgString(t Type, dir Dir) (Arg, error) {
 	}
 	// Check compressed data for validity.
 	if typ.IsCompressed() {
-		if err := DecompressWriter(ioutil.Discard, data); err != nil {
+		if err := image.DecompressWriter(ioutil.Discard, data); err != nil {
 			p.strictFailf("invalid compressed data in arg: %v", err)
 			// In non-strict mode, empty the data slice.
-			data = Compress(nil)
+			data = image.Compress(nil)
 		}
 	}
 	size := ^uint64(0)
@@ -887,8 +889,7 @@ func serializeData(buf *bytes.Buffer, data []byte, readable bool) {
 func serializeCompressedData(buf *bytes.Buffer, data []byte) {
 	buf.WriteByte('"')
 	buf.WriteByte('$')
-	encoded := EncodeB64(data)
-	buf.Write(encoded)
+	buf.Write(image.EncodeB64(data))
 	buf.WriteByte('"')
 }
 
@@ -983,7 +984,7 @@ func (p *parser) deserializeData() ([]byte, bool, error) {
 				rawData = append(rawData, v)
 			}
 			p.Parse('"')
-			decoded, err := DecodeB64(rawData)
+			decoded, err := image.DecodeB64(rawData)
 			if err != nil {
 				return nil, false, fmt.Errorf("data arg is corrupt: %v", err)
 			}
diff --git a/prog/heatmap_test.go b/prog/heatmap_test.go
index a981945ea..585bc095c 100644
--- a/prog/heatmap_test.go
+++ b/prog/heatmap_test.go
@@ -8,6 +8,7 @@ import (
 	"sort"
 	"testing"
 
+	"github.com/google/syzkaller/pkg/image"
 	"github.com/google/syzkaller/pkg/testutil"
 )
 
@@ -52,7 +53,7 @@ func TestGenericHeatmap(t *testing.T) {
 
 	r := rand.New(testutil.RandSource(t))
 	for _, test := range testData {
-		data, err := DecodeB64(test.data)
+		data, err := image.DecodeB64(test.data)
 		if err != nil {
 			t.Fatalf("bad decode: %v", err)
 		}
diff --git a/prog/mutation.go b/prog/mutation.go
index 96e11516c..100c47a91 100644
--- a/prog/mutation.go
+++ b/prog/mutation.go
@@ -10,6 +10,8 @@ import (
 	"math/rand"
 	"sort"
 	"sync"
+
+	"github.com/google/syzkaller/pkg/image"
 )
 
 // Maximum length of generated binary blobs inserted into the program.
@@ -385,21 +387,21 @@ func (t *BufferType) mutate(r *randGen, s *state, arg Arg, ctx ArgCtx) (calls []
 
 var imageMu sync.Mutex
 
-func (r *randGen) mutateImage(image []byte) (data []byte, retry bool) {
-	if len(image) == 0 {
-		return image, true
+func (r *randGen) mutateImage(compressed []byte) (data []byte, retry bool) {
+	if len(compressed) == 0 {
+		return compressed, true
 	}
 	// Don't decompress more than one image at a time
 	// since it can consume lots of memory.
 	// Reconsider when/if we move mutation to the host process.
 	imageMu.Lock()
 	defer imageMu.Unlock()
-	data, err := Decompress(image)
+	data, err := image.Decompress(compressed)
 	if err != nil {
 		panic(fmt.Sprintf("could not decompress data: %v", err))
 	}
 	if len(data) == 0 {
-		return image, true // Do not mutate empty data.
+		return compressed, true // Do not mutate empty data.
 	}
 	hm := MakeGenericHeatmap(data, r.Rand)
 	for i := hm.NumMutations(); i > 0; i-- {
@@ -410,7 +412,7 @@ func (r *randGen) mutateImage(image []byte) (data []byte, retry bool) {
 		}
 		storeInt(data[index:], r.randInt(uint64(width*8)), width)
 	}
-	return Compress(data), false
+	return image.Compress(data), false
 }
 
 func mutateBufferSize(r *randGen, arg *DataArg, minLen, maxLen uint64) {
diff --git a/sys/linux/init_images.go b/sys/linux/init_images.go
index 8b62455d2..c9a3035ef 100644
--- a/sys/linux/init_images.go
+++ b/sys/linux/init_images.go
@@ -8,6 +8,7 @@ import (
 	"fmt"
 	"io"
 
+	"github.com/google/syzkaller/pkg/image"
 	"github.com/google/syzkaller/prog"
 )
 
@@ -25,7 +26,7 @@ func (arch *arch) extractSyzMountImage(c *prog.Call) (io.Reader, error) {
 		return nil, fmt.Errorf("an empty image")
 	}
 	buf := new(bytes.Buffer)
-	if err := prog.DecompressWriter(buf, data); err != nil {
+	if err := image.DecompressWriter(buf, data); err != nil {
 		return nil, err
 	}
 	return buf, nil
diff --git a/tools/syz-imagegen/imagegen.go b/tools/syz-imagegen/imagegen.go
index f42cda778..4cabcede2 100644
--- a/tools/syz-imagegen/imagegen.go
+++ b/tools/syz-imagegen/imagegen.go
@@ -27,6 +27,7 @@ import (
 	"syscall"
 	"time"
 
+	"github.com/google/syzkaller/pkg/image"
 	"github.com/google/syzkaller/pkg/osutil"
 	"github.com/google/syzkaller/pkg/tool"
 	"github.com/google/syzkaller/prog"
@@ -877,8 +878,8 @@ func writeImage(img *Image, data []byte) ([]byte, error) {
 	fmt.Fprintf(buf, "# Code generated by tools/syz-imagegen. DO NOT EDIT.\n")
 	fmt.Fprintf(buf, "# requires: manual\n\n")
 	fmt.Fprintf(buf, "# %v\n\n", img)
-	compressedData := prog.Compress(data)
-	b64Data := prog.EncodeB64(compressedData)
+	compressedData := image.Compress(data)
+	b64Data := image.EncodeB64(compressedData)
 	if img.fs.Name == parttable {
 		fmt.Fprintf(buf, `%s(AUTO, &AUTO="$`, syzReadPartTable)
 	} else {