1 files changed, 17 insertions, 74 deletions

diff --git a/tools/syz-imagegen/imagegen.go b/tools/syz-imagegen/imagegen.go
index 5f120b29d..367e7d8b0 100644
--- a/tools/syz-imagegen/imagegen.go
+++ b/tools/syz-imagegen/imagegen.go
@@ -14,7 +14,6 @@ package main
 
 import (
 	"bytes"
-	"encoding/hex"
 	"errors"
 	"flag"
 	"fmt"
@@ -692,10 +691,11 @@ func (image *Image) generateSize() error {
 	}
 	image.hash = crc32.ChecksumIEEE(data)
 
-	segments := calculateSegments(data)
-
 	// Write out image *with* change of directory.
-	out := writeImage(image, len(data), segments, true)
+	out, err := writeImage(image, data, true)
+	if err != nil {
+		return fmt.Errorf("failed to write image: %v", err)
+	}
 
 	p, err := image.target.Deserialize(out, prog.Strict)
 	if err != nil {
@@ -767,85 +767,28 @@ func runCmd(cmd string, args ...string) ([]byte, error) {
 	return osutil.RunCmd(10*time.Minute, "", cmd, args...)
 }
 
-func writeImage(image *Image, length int, segments []Segment, chdir bool) []byte {
+func writeImage(image *Image, data []byte, chdir bool) ([]byte, error) {
 	buf := new(bytes.Buffer)
 	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", image)
-	fmt.Fprintf(buf, `syz_mount_image$%v(&(0x7f0000000000)='%v\x00', &(0x7f0000000100)='./file0\x00',`+
-		` 0x%x, 0x%x, &(0x7f0000000200)=[`,
-		image.fs.Name, image.fs.Name, length, len(segments))
-	addr := 0x7f0000010000
-	for i, seg := range segments {
-		if i != 0 {
-			fmt.Fprintf(buf, ", ")
-		}
-		fmt.Fprintf(buf, `{&(0x%x)="%v", 0x%x, 0x%x}`,
-			addr, hex.EncodeToString(seg.data), len(seg.data), seg.offset)
-		addr += len(seg.data)
-	}
+	compressedData := prog.Compress(data)
+	b64Data := prog.EncodeB64(compressedData)
 	chdirAsInt := 0
 	if chdir {
 		chdirAsInt = 1
 	}
-	fmt.Fprintf(buf, "], 0x0, &(0x%x), 0x%x)\n", addr, chdirAsInt)
-	return buf.Bytes()
-}
-
-type Segment struct {
-	offset int
-	data   []byte
-}
+	fmt.Fprintf(buf, `syz_mount_image$%v(&AUTO='%v\x00', &AUTO='./file0\x00',`+
+		` 0x%x, AUTO, 0x0, &AUTO, 0x%x, &AUTO="$`,
+		image.fs.Name, image.fs.Name, len(data), chdirAsInt)
+	buf.Write(b64Data)
+	fmt.Fprintf(buf, "\")\n")
 
-func calculateSegments(data []byte) []Segment {
-	const (
-		skip  = 32 // min zero bytes to skip
-		align = 32 // non-zero block alignment
-	)
-	data0 := data
-	zeros := make([]byte, skip+align)
-	var segs []Segment
-	offset := 0
-	for len(data) != 0 {
-		pos := bytes.Index(data, zeros)
-		if pos == -1 {
-			segs = append(segs, Segment{offset, data})
-			break
-		}
-		pos = (pos + align - 1) & ^(align - 1)
-		if pos != 0 {
-			segs = append(segs, Segment{offset, data[:pos]})
-		}
-		for pos < len(data) && data[pos] == 0 {
-			pos++
-		}
-		pos = pos & ^(align - 1)
-		offset += pos
-		data = data[pos:]
-	}
-	if false {
-		// self-test.
-		restored := make([]byte, len(data0))
-		for _, seg := range segs {
-			copy(restored[seg.offset:], seg.data)
-		}
-		if !bytes.Equal(data0, restored) {
-			panic("restored data differs!")
-		}
-	}
-	return segs
+	return buf.Bytes(), nil
 }
 
 // TODO: also generate syz_read_part_table tests:
-//	fmt.Printf(`syz_read_part_table(0x%x, 0x%x, &(0x7f0000000200)=[`,
-//		len(data0), len(segs))
-//	addr := 0x7f0000010000
-//	for i, seg := range segs {
-//		if i != 0 {
-//			fmt.Printf(", ")
-//		}
-//		fmt.Printf(`{&(0x%x)="%v", 0x%x, 0x%x}`,
-//			addr, hex.EncodeToString(seg.data), len(seg.data), seg.offset)
-//		addr = (addr + len(seg.data) + 0xff) & ^0xff
-//	}
-//	fmt.Printf("])\n")
+//	fmt.Fprintf(buf, `syz_read_part_table(0x%x, 0x%x, &AUTO="$`,
+//		len(data), len(compressedData))
+//  buf.Write(b64Data)
+//  fmt.Fprintf(buf, "\")\n")