From 8482d3c1035095c89d112c75bfcc2e4095b486bf Mon Sep 17 00:00:00 2001 From: Dmitry Vyukov Date: Sat, 17 Dec 2022 11:59:24 +0100 Subject: 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. --- pkg/image/compression.go | 71 ++++++++++++++++++++++++++++++++++++++++++ pkg/image/compression_test.go | 44 ++++++++++++++++++++++++++ pkg/ipc/ipc_test.go | 5 +-- prog/compression.go | 71 ------------------------------------------ prog/compression_test.go | 44 -------------------------- prog/encoding.go | 11 ++++--- prog/heatmap_test.go | 3 +- prog/mutation.go | 14 +++++---- sys/linux/init_images.go | 3 +- tools/syz-imagegen/imagegen.go | 5 +-- 10 files changed, 139 insertions(+), 132 deletions(-) create mode 100644 pkg/image/compression.go create mode 100644 pkg/image/compression_test.go delete mode 100644 prog/compression.go delete mode 100644 prog/compression_test.go diff --git a/pkg/image/compression.go b/pkg/image/compression.go new file mode 100644 index 000000000..9878b460d --- /dev/null +++ b/pkg/image/compression.go @@ -0,0 +1,71 @@ +// 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 image + +import ( + "bytes" + "compress/zlib" + "encoding/base64" + "fmt" + "io" +) + +func Compress(rawData []byte) []byte { + var buffer bytes.Buffer + zlibWriter := zlib.NewWriter(&buffer) + + _, err := zlibWriter.Write(rawData) + if err != nil { + panic(fmt.Sprintf("could not compress with zlib: %v", err)) + } + + err = zlibWriter.Close() + if err != nil { + panic(fmt.Sprintf("could not finalize compression with zlib: %v", err)) + } + + return buffer.Bytes() +} + +func Decompress(compressedData []byte) ([]byte, error) { + buf := new(bytes.Buffer) + err := DecompressWriter(buf, compressedData) + return buf.Bytes(), err +} + +func DecompressWriter(w io.Writer, compressedData []byte) error { + zlibReader, err := zlib.NewReader(bytes.NewReader(compressedData)) + if err != nil { + return fmt.Errorf("could not initialise zlib: %v", err) + } + + if _, err := io.Copy(w, zlibReader); err != nil { + return fmt.Errorf("could not read data with zlib: %v", err) + } + + return zlibReader.Close() +} + +func DecodeB64(b64Data []byte) ([]byte, error) { + decoder := base64.NewDecoder(base64.StdEncoding, bytes.NewReader(b64Data)) + rawData, err := io.ReadAll(decoder) + if err != nil { + return nil, fmt.Errorf("could not decode Base64: %v", err) + } + return rawData, nil +} + +func EncodeB64(rawData []byte) []byte { + var buf bytes.Buffer + encoder := base64.NewEncoder(base64.StdEncoding, &buf) + _, err := encoder.Write(rawData) + if err != nil { + panic(fmt.Sprintf("could not encode Base64: %v", err)) + } + err = encoder.Close() + if err != nil { + panic(fmt.Sprintf("could not finalize encoding to Base64: %v", err)) + } + return buf.Bytes() +} diff --git a/pkg/image/compression_test.go b/pkg/image/compression_test.go new file mode 100644 index 000000000..cf18ed340 --- /dev/null +++ b/pkg/image/compression_test.go @@ -0,0 +1,44 @@ +// 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 image + +import ( + "bytes" + "fmt" + "math/rand" + "testing" + + "github.com/google/syzkaller/pkg/testutil" +) + +func TestCompress(t *testing.T) { + r := rand.New(testutil.RandSource(t)) + err := testRoundTrip(r, Compress, Decompress) + if err != nil { + t.Fatalf("compress/decompress %v", err) + } +} + +func TestEncode(t *testing.T) { + r := rand.New(testutil.RandSource(t)) + err := testRoundTrip(r, EncodeB64, DecodeB64) + if err != nil { + t.Fatalf("encode/decode Base64 %v", err) + } +} + +func testRoundTrip(r *rand.Rand, transform func([]byte) []byte, inverse func([]byte) ([]byte, error)) error { + for i := 0; i < testutil.IterCount(); i++ { + randBytes := testutil.RandMountImage(r) + resultBytes := transform(randBytes) + resultBytes, err := inverse(resultBytes) + if err != nil { + return err + } + if !bytes.Equal(randBytes, resultBytes) { + return fmt.Errorf("roundtrip changes data (original length %d)", len(randBytes)) + } + } + return nil +} 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/compression.go b/prog/compression.go deleted file mode 100644 index 87820b5ee..000000000 --- a/prog/compression.go +++ /dev/null @@ -1,71 +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 prog - -import ( - "bytes" - "compress/zlib" - "encoding/base64" - "fmt" - "io" -) - -func Compress(rawData []byte) []byte { - var buffer bytes.Buffer - zlibWriter := zlib.NewWriter(&buffer) - - _, err := zlibWriter.Write(rawData) - if err != nil { - panic(fmt.Sprintf("could not compress with zlib: %v", err)) - } - - err = zlibWriter.Close() - if err != nil { - panic(fmt.Sprintf("could not finalize compression with zlib: %v", err)) - } - - return buffer.Bytes() -} - -func Decompress(compressedData []byte) ([]byte, error) { - buf := new(bytes.Buffer) - err := DecompressWriter(buf, compressedData) - return buf.Bytes(), err -} - -func DecompressWriter(w io.Writer, compressedData []byte) error { - zlibReader, err := zlib.NewReader(bytes.NewReader(compressedData)) - if err != nil { - return fmt.Errorf("could not initialise zlib: %v", err) - } - - if _, err := io.Copy(w, zlibReader); err != nil { - return fmt.Errorf("could not read data with zlib: %v", err) - } - - return zlibReader.Close() -} - -func DecodeB64(b64Data []byte) ([]byte, error) { - decoder := base64.NewDecoder(base64.StdEncoding, bytes.NewReader(b64Data)) - rawData, err := io.ReadAll(decoder) - if err != nil { - return nil, fmt.Errorf("could not decode Base64: %v", err) - } - return rawData, nil -} - -func EncodeB64(rawData []byte) []byte { - var buf bytes.Buffer - encoder := base64.NewEncoder(base64.StdEncoding, &buf) - _, err := encoder.Write(rawData) - if err != nil { - panic(fmt.Sprintf("could not encode Base64: %v", err)) - } - err = encoder.Close() - if err != nil { - panic(fmt.Sprintf("could not finalize encoding to Base64: %v", err)) - } - return buf.Bytes() -} diff --git a/prog/compression_test.go b/prog/compression_test.go deleted file mode 100644 index db5d95070..000000000 --- a/prog/compression_test.go +++ /dev/null @@ -1,44 +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 prog - -import ( - "bytes" - "fmt" - "math/rand" - "testing" - - "github.com/google/syzkaller/pkg/testutil" -) - -func TestCompress(t *testing.T) { - r := rand.New(testutil.RandSource(t)) - err := testRoundTrip(r, Compress, Decompress) - if err != nil { - t.Fatalf("compress/decompress %v", err) - } -} - -func TestEncode(t *testing.T) { - r := rand.New(testutil.RandSource(t)) - err := testRoundTrip(r, EncodeB64, DecodeB64) - if err != nil { - t.Fatalf("encode/decode Base64 %v", err) - } -} - -func testRoundTrip(r *rand.Rand, transform func([]byte) []byte, inverse func([]byte) ([]byte, error)) error { - for i := 0; i < testutil.IterCount(); i++ { - randBytes := testutil.RandMountImage(r) - resultBytes := transform(randBytes) - resultBytes, err := inverse(resultBytes) - if err != nil { - return err - } - if !bytes.Equal(randBytes, resultBytes) { - return fmt.Errorf("roundtrip changes data (original length %d)", len(randBytes)) - } - } - return nil -} 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 { -- cgit mrf-deployment