Rework merkle tree implementation to use io.Reader instead of byte array

MerkleTree implementation requires the entire content of ext4 file system to be read into a byte array when computing cryptographic digest. This PR reworks the existing implementation to work with io.Reader interface instead. Additionally update the existing usages of MerkleTree with the new MerkleTreeWithReader implementation. Signed-off-by: Maksim An <maksiman@microsoft.com>
microsoft · Oct 27, 2021 · b3a2950 · b3a2950
1 parent f174aa8
commit b3a2950
Show file tree

Hide file tree

Showing 6 changed files with 165 additions and 140 deletions.
diff --git a/cmd/dmverity-vhd/main.go b/cmd/dmverity-vhd/main.go
@@ -2,8 +2,6 @@ package main
 
 import (
 	"fmt"
-	"io"
-	"io/ioutil"
 	"os"
 	"path/filepath"
 
@@ -203,12 +201,6 @@ var rootHashVHDCommand = cli.Command{
 		}
 		log.Debugf("%d layers found", len(layers))
 
-		tmpFile, err := ioutil.TempFile("", "")
-		if err != nil {
-			return errors.Wrap(err, "failed to create temporary file")
-		}
-		defer os.Remove(tmpFile.Name())
-
 		for layerNumber, layer := range layers {
 			diffID, err := layer.DiffID()
 			if err != nil {
@@ -221,32 +213,7 @@ var rootHashVHDCommand = cli.Command{
 				return errors.Wrapf(err, "failed to uncompress layer %s", diffID.String())
 			}
 
-			opts := []tar2ext4.Option{
-				tar2ext4.ConvertWhiteout,
-				tar2ext4.MaximumDiskSize(maxVHDSize),
-			}
-
-			if _, err := tmpFile.Seek(0, io.SeekStart); err != nil {
-				return errors.Wrapf(err, "failed seek start on temp file when processing layer %d", layerNumber)
-			}
-			if err := tmpFile.Truncate(0); err != nil {
-				return errors.Wrapf(err, "failed truncate temp file when processing layer %d", layerNumber)
-			}
-
-			if err := tar2ext4.Convert(rc, tmpFile, opts...); err != nil {
-				return errors.Wrap(err, "failed to convert tar to ext4")
-			}
-
-			data, err := ioutil.ReadFile(tmpFile.Name())
-			if err != nil {
-				return errors.Wrap(err, "failed to read temporary VHD file")
-			}
-
-			tree, err := dmverity.MerkleTree(data)
-			if err != nil {
-				return errors.Wrap(err, "failed to create merkle tree")
-			}
-			hash := dmverity.RootHash(tree)
+			hash, err := tar2ext4.ConvertAndRootDigest(rc)
 			fmt.Fprintf(os.Stdout, "Layer %d\nroot hash: %x\n", layerNumber, hash)
 		}
 		return nil

diff --git a/ext4/dmverity/dmverity.go b/ext4/dmverity/dmverity.go
@@ -1,6 +1,7 @@
 package dmverity
 
 import (
+	"bufio"
 	"bytes"
 	"crypto/rand"
 	"crypto/sha256"
@@ -16,9 +17,12 @@ import (
 
 const (
 	blockSize = compactext4.BlockSize
-	// RecommendedVHDSizeGB is the recommended size in GB for VHDs, which is not a hard limit. 
+	// bufioSize is a default buffer size to use with bufio.Reader
+	bufioSize = 1024 * 1024 // 1MB
+	// RecommendedVHDSizeGB is the recommended size in GB for VHDs, which is not a hard limit.
 	RecommendedVHDSizeGB = 128 * 1024 * 1024 * 1024
 )
+
 var salt = bytes.Repeat([]byte{0}, 32)
 
 var (
@@ -69,20 +73,19 @@ type VerityInfo struct {
 	Version       uint32
 }
 
-// MerkleTree constructs dm-verity hash-tree for a given byte array with a fixed salt (0-byte) and algorithm (sha256).
-func MerkleTree(data []byte) ([]byte, error) {
+// MerkleTreeWithReader constructs dm-verity hash-tree for a given io.Reader with a fixed salt (0-byte) and algorithm (sha256).
+func MerkleTreeWithReader(r io.Reader) ([]byte, error) {
 	layers := make([][]byte, 0)
+	currentLevel := bufio.NewReaderSize(r, bufioSize)
 
-	currentLevel := bytes.NewBuffer(data)
-
-	for currentLevel.Len() != blockSize {
-		blocks := currentLevel.Len() / blockSize
+	for {
 		nextLevel := bytes.NewBuffer(make([]byte, 0))
-
-		for i := 0; i < blocks; i++ {
+		for {
 			block := make([]byte, blockSize)
-			_, err := currentLevel.Read(block)
-			if err != nil {
+			if _, err := currentLevel.Read(block); err != nil {
+				if err == io.EOF {
+					break
+				}
 				return nil, errors.Wrap(err, "failed to read data block")
 			}
 			h := hash2(salt, block)
@@ -92,21 +95,30 @@ func MerkleTree(data []byte) ([]byte, error) {
 		padding := bytes.Repeat([]byte{0}, blockSize-(nextLevel.Len()%blockSize))
 		nextLevel.Write(padding)
 
-		currentLevel = nextLevel
-		layers = append(layers, currentLevel.Bytes())
+		layers = append(layers, nextLevel.Bytes())
+		currentLevel = bufio.NewReaderSize(nextLevel, bufioSize)
+
+		// This means that only root hash remains and our job is done
+		if nextLevel.Len() == blockSize {
+			break
+		}
 	}
 
-	var tree = bytes.NewBuffer(make([]byte, 0))
+	tree := bytes.NewBuffer(make([]byte, 0))
 	for i := len(layers) - 1; i >= 0; i-- {
-		_, err := tree.Write(layers[i])
-		if err != nil {
+		if _, err := tree.Write(layers[i]); err != nil {
 			return nil, errors.Wrap(err, "failed to write merkle tree")
 		}
 	}
 
 	return tree.Bytes(), nil
 }
 
+// MerkleTree constructs dm-verity hash-tree for a given byte array with a fixed salt (0-byte) and algorithm (sha256).
+func MerkleTree(data []byte) ([]byte, error) {
+	return MerkleTreeWithReader(bytes.NewBuffer(data))
+}
+
 // RootHash computes root hash of dm-verity hash-tree
 func RootHash(tree []byte) []byte {
 	return hash2(salt, tree[:blockSize])

diff --git a/ext4/tar2ext4/tar2ext4.go b/ext4/tar2ext4/tar2ext4.go
@@ -5,6 +5,7 @@ import (
 	"bufio"
 	"bytes"
 	"encoding/binary"
+	"fmt"
 	"github.com/pkg/errors"
 	"io"
 	"io/ioutil"
@@ -183,42 +184,36 @@ func Convert(r io.Reader, w io.ReadWriteSeeker, options ...Option) error {
 		}
 
 		// Rewind the stream and then read it all into a []byte for
-		// dmverity processing
-		_, err = w.Seek(0, io.SeekStart)
-		if err != nil {
-			return err
-		}
-		data, err := ioutil.ReadAll(w)
-		if err != nil {
+		// dm-verity processing
+		if _, err = w.Seek(0, io.SeekStart); err != nil {
 			return err
 		}
 
-		mtree, err := dmverity.MerkleTree(data)
+		merkelTree, err := dmverity.MerkleTreeWithReader(w)
 		if err != nil {
 			return errors.Wrap(err, "failed to build merkle tree")
 		}
 
-		// Write dmverity superblock and then the merkle tree after the end of the
+		// Write dm-verity super-block and then the merkle tree after the end of the
 		// ext4 filesystem
-		_, err = w.Seek(0, io.SeekEnd)
-		if err != nil {
+		if _, err = w.Seek(0, io.SeekEnd); err != nil {
 			return err
 		}
-		superblock := dmverity.NewDMVeritySuperblock(uint64(ext4size))
-		err = binary.Write(w, binary.LittleEndian, superblock)
-		if err != nil {
+
+		superBlock := dmverity.NewDMVeritySuperblock(uint64(ext4size))
+		if err = binary.Write(w, binary.LittleEndian, superBlock); err != nil {
 			return err
 		}
-		// pad the superblock
-		sbsize := int(unsafe.Sizeof(*superblock))
+
+		// pad the super-block
+		sbsize := int(unsafe.Sizeof(*superBlock))
 		padding := bytes.Repeat([]byte{0}, ext4blocksize-(sbsize%ext4blocksize))
-		_, err = w.Write(padding)
-		if err != nil {
+		if _, err = w.Write(padding); err != nil {
 			return err
 		}
+
 		// write the tree
-		_, err = w.Write(mtree)
-		if err != nil {
+		if _, err = w.Write(merkelTree); err != nil {
 			return err
 		}
 	}
@@ -268,3 +263,37 @@ func ReadExt4SuperBlock(vhdPath string) (*format.SuperBlock, error) {
 	}
 	return &sb, nil
 }
+
+// ConvertAndRootDigest writes a compact ext4 file system image that contains the files in the
+// input tar stream, computes and returns its cryptographic digest. Convert is called with minimal
+// options: ConvertWhiteout and MaximumDiskSize set to dmverity.RecommendedVHDSizeGB.
+func ConvertAndRootDigest(r io.Reader) (string, error) {
+	out, err := ioutil.TempFile("", "")
+	if err != nil {
+		return "", fmt.Errorf("failed to create temporary file: %s", err)
+	}
+	defer func() {
+		_ = os.Remove(out.Name())
+	}()
+
+	opts := []Option{
+		ConvertWhiteout,
+		MaximumDiskSize(dmverity.RecommendedVHDSizeGB),
+	}
+
+	if err := Convert(r, out, opts...); err != nil {
+		return "", fmt.Errorf("failed to convert tar to ext4: %s", err)
+	}
+
+	if _, err := out.Seek(0, io.SeekStart); err != nil {
+		return "", fmt.Errorf("failed to seek start on temp file when creating merkle tree: %s", err)
+	}
+
+	tree, err := dmverity.MerkleTreeWithReader(r)
+	if err != nil {
+		return "", fmt.Errorf("failed to create merkle tree: %s", err)
+	}
+
+	hash := dmverity.RootHash(tree)
+	return fmt.Sprintf("%x", hash), nil
+}
diff --git a/internal/tools/securitypolicy/main.go b/internal/tools/securitypolicy/main.go
@@ -11,7 +11,6 @@ import (
 	"strconv"
 
 	"github.com/BurntSushi/toml"
-	"github.com/Microsoft/hcsshim/ext4/dmverity"
 	"github.com/Microsoft/hcsshim/ext4/tar2ext4"
 	"github.com/Microsoft/hcsshim/pkg/securitypolicy"
 	"github.com/google/go-containerregistry/pkg/authn"
@@ -168,43 +167,20 @@ func createPolicyFromConfig(config Config) (securitypolicy.SecurityPolicy, error
 				return p, err
 			}
 
-			out, err := ioutil.TempFile("", "")
+			hashString, err := tar2ext4.ConvertAndRootDigest(r)
 			if err != nil {
 				return p, err
 			}
-			defer os.Remove(out.Name())
-
-			opts := []tar2ext4.Option{
-				tar2ext4.ConvertWhiteout,
-				tar2ext4.MaximumDiskSize(dmverity.RecommendedVHDSizeGB),
-			}
-
-			err = tar2ext4.Convert(r, out, opts...)
-			if err != nil {
-				return p, err
-			}
-
-			data, err := ioutil.ReadFile(out.Name())
-			if err != nil {
-				return p, err
-			}
-
-			tree, err := dmverity.MerkleTree(data)
-			if err != nil {
-				return p, err
-			}
-			hash := dmverity.RootHash(tree)
-			hashString := fmt.Sprintf("%x", hash)
 			addLayer(&container.Layers, hashString)
 		}
 
 		// add rules for all known environment variables from the configuration
 		// these are in addition to "other rules" from the policy definition file
-		config, err := img.ConfigFile()
+		imgConfig, err := img.ConfigFile()
 		if err != nil {
 			return p, err
 		}
-		for _, env := range config.Config.Env {
+		for _, env := range imgConfig.Config.Env {
 			rule := securitypolicy.EnvRule{
 				Strategy: securitypolicy.EnvVarRuleString,
 				Rule:     env,
@@ -214,7 +190,7 @@ func createPolicyFromConfig(config Config) (securitypolicy.SecurityPolicy, error
 		}
 
 		// cri adds TERM=xterm for all workload containers. we add to all containers
-		// to prevent any possble erroring
+		// to prevent any possible error
 		rule := securitypolicy.EnvRule{
 			Strategy: securitypolicy.EnvVarRuleString,
 			Rule:     "TERM=xterm",
@@ -243,31 +219,31 @@ func validateEnvRules(rules []EnvironmentVariableRule) error {
 }
 
 func convertCommand(toml []string) securitypolicy.CommandArgs {
-	json := map[string]string{}
+	jsn := map[string]string{}
 
 	for i, arg := range toml {
-		json[strconv.Itoa(i)] = arg
+		jsn[strconv.Itoa(i)] = arg
 	}
 
 	return securitypolicy.CommandArgs{
-		Elements: json,
+		Elements: jsn,
 	}
 }
 
 func convertEnvironmentVariableRules(toml []EnvironmentVariableRule) securitypolicy.EnvRules {
-	json := map[string]securitypolicy.EnvRule{}
+	jsn := map[string]securitypolicy.EnvRule{}
 
 	for i, rule := range toml {
 		jsonRule := securitypolicy.EnvRule{
 			Strategy: rule.Strategy,
 			Rule:     rule.Rule,
 		}
 
-		json[strconv.Itoa(i)] = jsonRule
+		jsn[strconv.Itoa(i)] = jsonRule
 	}
 
 	return securitypolicy.EnvRules{
-		Elements: json,
+		Elements: jsn,
 	}
 }