sys/linux: extract raw images from syz_mount_image

To simplify the extraction code, let's make segments non-overlapping
even before execution.

1 files changed, 83 insertions, 12 deletions

diff --git a/sys/linux/init_images.go b/sys/linux/init_images.go
index e914b8b4d..eb1d7c7ad 100644
--- a/sys/linux/init_images.go
+++ b/sys/linux/init_images.go
@@ -4,12 +4,45 @@
 package linux
 
 import (
+	"bytes"
 	"fmt"
 	"io"
+	"sort"
 
 	"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 imags are de facto compressed.
+	// Here we do the uncompression.
+	if c.Meta.CallName != "syz_mount_image" {
+		return nil, nil
+	}
+	ret, err := parseSyzMountImage(c)
+	if err != nil {
+		// Parsing failed --> do not try to recover, just ignore.
+		return nil, err
+	} else if len(ret.segments) == 0 {
+		return nil, fmt.Errorf("an empty image")
+	}
+	readers := []io.Reader{}
+	nextPos := 0
+	// Add fake zero readers between segments, so that we can combine them to read the whole image.
+	for _, segment := range ret.segments {
+		offset := int(segment.offset.Val)
+		if offset > nextPos {
+			readers = append(readers, &zeroReader{left: offset - nextPos})
+		}
+		size := int(segment.size.Val)
+		readers = append(readers, bytes.NewReader(segment.data.Data()[0:size]))
+		nextPos = offset + size
+	}
+	if int(ret.size.Val) > nextPos {
+		readers = append(readers, &zeroReader{left: int(ret.size.Val) - nextPos})
+	}
+	return io.MultiReader(readers...), nil
+}
+
 type zeroReader struct {
 	left int
 }
@@ -25,39 +58,67 @@ func (zr *zeroReader) Read(p []byte) (n int, err error) {
 	for i := 0; i < toRead; i++ {
 		p[i] = 0
 	}
+	zr.left -= toRead
 	return toRead, nil
 }
 
-func newZeroReader(size int) io.Reader {
-	return &zeroReader{left: size}
-}
-
 const imageMaxSize = 129 << 20
 
 func fixUpImageSegments(parsed *mountImageArgs) {
+	const maxImageSegments = 4096
 	parsed.filterSegments(func(i int, segment *mountImageSegment) bool {
+		if i >= maxImageSegments {
+			// We don't want more than maxImageSegments segments.
+			return false
+		}
 		if segment.parseError != nil {
+			// Delete mangled segment structures.
 			return false
 		}
 		return segment.offset.Val < imageMaxSize && segment.size.Val < imageMaxSize
 	})
+	// Overwriting of the image multiple times is not efficient and complicates image extraction in Go.
+	// So let's make segments non-overlapping.
+	sort.Stable(parsed)
+	resizeSegment := func(s *mountImageSegment, size uint64) {
+		s.size.Val = size
+		s.data.SetData(s.data.Data()[0:size])
+	}
+
 	newSize := parsed.size.Val
-	for _, segment := range parsed.segments {
+	if newSize > imageMaxSize {
+		newSize = imageMaxSize
+	}
+
+	for idx, segment := range parsed.segments {
 		actualSize := uint64(len(segment.data.Data()))
-		if segment.size.Val > actualSize {
+		if segment.size.Val != actualSize {
 			segment.size.Val = actualSize
 		}
+		if idx > 0 {
+			// Adjust the end of the previous segment.
+			prevSegment := parsed.segments[idx-1]
+			if prevSegment.offset.Val+prevSegment.size.Val > segment.offset.Val {
+				resizeSegment(prevSegment, segment.offset.Val-prevSegment.offset.Val)
+			}
+		}
 		if segment.offset.Val+segment.size.Val > imageMaxSize {
-			segment.offset.Val = imageMaxSize - segment.size.Val
+			resizeSegment(segment, imageMaxSize-segment.offset.Val)
 		}
 		if segment.offset.Val+segment.size.Val > newSize {
 			newSize = segment.offset.Val + segment.size.Val
 		}
 	}
 	if newSize > imageMaxSize {
-		newSize = imageMaxSize
+		// Assert that the logic above is not broken.
+		panic("newSize > imageMaxSize")
 	}
 	parsed.size.Val = newSize
+
+	// Drop 0-size segments.
+	parsed.filterSegments(func(i int, segment *mountImageSegment) bool {
+		return segment.size.Val > 0
+	})
 }
 
 func (arch *arch) fixUpSyzMountImage(c *prog.Call) {
@@ -71,10 +132,6 @@ func (arch *arch) fixUpSyzMountImage(c *prog.Call) {
 		deactivateSyzMountImage(c)
 		return
 	}
-	const maxImageSegments = 4096
-	ret.filterSegments(func(i int, _ *mountImageSegment) bool {
-		return i < maxImageSegments
-	})
 	fixUpImageSegments(ret)
 }
 
@@ -99,6 +156,20 @@ func (m *mountImageArgs) filterSegments(filter func(int, *mountImageSegment) boo
 	m.segmentsCount.Val = uint64(len(newArgs))
 }
 
+// Methods for segment sorting.
+func (m *mountImageArgs) Len() int { return len(m.segments) }
+func (m *mountImageArgs) Swap(i, j int) {
+	inner := m.segmentsGroup.Inner
+	inner[i], inner[j] = inner[j], inner[i]
+	m.segments[i], m.segments[j] = m.segments[j], m.segments[i]
+}
+func (m *mountImageArgs) Less(i, j int) bool {
+	if m.segments[i].offset.Val != m.segments[j].offset.Val {
+		return m.segments[i].offset.Val < m.segments[j].offset.Val
+	}
+	return m.segments[i].size.Val < m.segments[j].size.Val
+}
+
 type mountImageSegment struct {
 	data       *prog.DataArg
 	offset     *prog.ConstArg