Fix ReadDMVeritySuperBlock function

ext4/dmverity/dmverity.go

@@ -30,6 +30,7 @@ var (
 	ErrSuperBlockReadFailure  = errors.New("failed to read dm-verity super block")
 	ErrSuperBlockParseFailure = errors.New("failed to parse dm-verity super block")
 	ErrRootHashReadFailure    = errors.New("failed to read dm-verity root hash")
+	ErrNotVeritySuperBlock    = errors.New("not dm-verity super-block")
 )
 
 type dmveritySuperblock struct {
@@ -183,7 +184,9 @@ func ReadDMVerityInfo(vhdPath string, offsetInBytes int64) (*VerityInfo, error)
 	if err := binary.Read(b, binary.LittleEndian, dmvSB); err != nil {
 		return nil, errors.Wrapf(ErrSuperBlockParseFailure, "%s", err)
 	}
-
+	if string(bytes.Trim(dmvSB.Signature[:], "\x00")[:]) != "verity" {
+		return nil, ErrNotVeritySuperBlock
+	}
 	// read the merkle tree root
 	if s, err := vhd.Read(block); err != nil || s != blockSize {
 		if err != nil {

ext4/dmverity/dmverity_test.go

@@ -0,0 +1,102 @@
+package dmverity
+
+import (
+	"bytes"
+	"encoding/binary"
+	"io"
+	"io/ioutil"
+	"math/rand"
+	"os"
+	"strings"
+	"testing"
+	"unsafe"
+
+	"github.com/pkg/errors"
+)
+
+func tempFileWithContentLength(t *testing.T, length int) *os.File {
+	tmpFile, err := ioutil.TempFile("", "")
+	if err != nil {
+		t.Fatalf("failed to create temp file")
+	}
+	defer tmpFile.Close()
+
+	content := make([]byte, length)
+	if _, err := rand.Read(content); err != nil {
+		t.Fatalf("failed to write random bytes to buffer")
+	}
+	if _, err := tmpFile.Write(content); err != nil {
+		t.Fatalf("failed to write random bytes to temp file")
+	}
+	return tmpFile
+}
+
+func writeDMVeritySuperBlock(filename string) (*os.File, error) {
+	out, err := os.OpenFile(filename, os.O_RDWR, 0777)
+	if err != nil {
+		return nil, err
+	}
+	defer out.Close()
+	fsSize, err := out.Seek(0, io.SeekEnd)
+	if err != nil {
+		return nil, err
+	}
+	sb := NewDMVeritySuperblock(uint64(fsSize))
+	if err := binary.Write(out, binary.LittleEndian, sb); err != nil {
+		return nil, err
+	}
+	sbSize := int(unsafe.Sizeof(*sb))
+	padding := bytes.Repeat([]byte{0}, blockSize-(sbSize%blockSize))
+	if _, err = out.Write(padding); err != nil {
+		return nil, err
+	}
+	return out, nil
+}
+
+func TestInvalidReadEOF(t *testing.T) {
+	tmpFile := tempFileWithContentLength(t, blockSize)
+	_, err := ReadDMVerityInfo(tmpFile.Name(), blockSize)
+	if err == nil {
+		t.Fatalf("no error returned")
+	}
+	if errors.Cause(err) != ErrSuperBlockReadFailure || !strings.Contains(err.Error(), "EOF") {
+		t.Fatalf("unexpected error: %s", err)
+	}
+}
+
+func TestInvalidReadNotEnoughBytes(t *testing.T) {
+	tmpFile := tempFileWithContentLength(t, blockSize+blockSize/2)
+	_, err := ReadDMVerityInfo(tmpFile.Name(), blockSize)
+	if err == nil {
+		t.Fatalf("no error returned")
+	}
+	if errors.Cause(err) != ErrSuperBlockReadFailure || !strings.Contains(err.Error(), "unexpected bytes read") {
+		t.Fatalf("unexpected error: %s", err)
+	}
+}
+
+func TestNotVeritySuperBlock(t *testing.T) {
+	tmpFile := tempFileWithContentLength(t, 2*blockSize)
+	_, err := ReadDMVerityInfo(tmpFile.Name(), blockSize)
+	if err == nil {
+		t.Fatalf("no error returned")
+	}
+	if err != ErrNotVeritySuperBlock {
+		t.Fatalf("expected %q, got %q", ErrNotVeritySuperBlock, err)
+	}
+}
+
+func TestNoMerkleTree(t *testing.T) {
+	tmpFile := tempFileWithContentLength(t, blockSize)
+	targetFile, err := writeDMVeritySuperBlock(tmpFile.Name())
+	if err != nil {
+		t.Fatalf("failed to write dm-verity super-block: %s", err)
+	}
+	_, err = ReadDMVerityInfo(targetFile.Name(), blockSize)
+	if err == nil {
+		t.Fatalf("no error returned")
+	}
+	if errors.Cause(err) != ErrRootHashReadFailure {
+		t.Fatalf("expected %q, got %q", ErrRootHashReadFailure, err)
+	}
+}