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

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,33 +213,11 @@ 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)
+			hash, err := tar2ext4.ConvertAndComputeRootDigest(rc)
 			if err != nil {
-				return errors.Wrap(err, "failed to create merkle tree")
+				return errors.Wrap(err, "failed to compute root hash")
 			}
-			hash := dmverity.RootHash(tree)
-			fmt.Fprintf(os.Stdout, "Layer %d\nroot hash: %x\n", layerNumber, hash)
+			fmt.Fprintf(os.Stdout, "Layer %d\nroot hash: %s\n", layerNumber, hash)
 		}
 		return nil
 	},

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. 
+	// MerkleTreeBufioSize is a default buffer size to use with bufio.Reader
+	MerkleTreeBufioSize = 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) {
+// MerkleTree constructs dm-verity hash-tree for a given io.Reader with a fixed salt (0-byte) and algorithm (sha256).
+func MerkleTree(r io.Reader) ([]byte, error) {
 	layers := make([][]byte, 0)
+	currentLevel := r
 
-	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 := io.ReadFull(currentLevel, block); err != nil {
+				if err == io.EOF {
+					break
+				}
 				return nil, errors.Wrap(err, "failed to read data block")
 			}
 			h := hash2(salt, block)
@@ -92,14 +95,18 @@ 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, MerkleTreeBufioSize)
+
+		// 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")
 		}
 	}

ext4/tar2ext4/tar2ext4.go

@@ -5,6 +5,7 @@ import (
 	"bufio"
 	"bytes"
 	"encoding/binary"
+	"fmt"
 	"io"
 	"io/ioutil"
 	"os"
@@ -66,16 +67,17 @@ func MaximumDiskSize(size int64) Option {
 const (
 	whiteoutPrefix = ".wh."
 	opaqueWhiteout = ".wh..wh..opq"
-	ext4blocksize  = compactext4.BlockSize
+	ext4BlockSize  = compactext4.BlockSize
 )
 
-// Convert writes a compact ext4 file system image that contains the files in the
+// ConvertTarToExt4 writes a compact ext4 file system image that contains the files in the
 // input tar stream.
-func Convert(r io.Reader, w io.ReadWriteSeeker, options ...Option) error {
+func ConvertTarToExt4(r io.Reader, w io.ReadWriteSeeker, options ...Option) error {
 	var p params
 	for _, opt := range options {
 		opt(&p)
 	}
+
 	t := tar.NewReader(bufio.NewReader(r))
 	fs := compactext4.NewWriter(w, p.ext4opts...)
 	for {
@@ -176,54 +178,53 @@ func Convert(r io.Reader, w io.ReadWriteSeeker, options ...Option) error {
 			}
 		}
 	}
-	err := fs.Close()
-	if err != nil {
+	return fs.Close()
+}
+
+// Convert wraps ConvertTarToExt4 and conditionally computes (and appends) the file image's cryptographic
+// hashes (merkle tree) or/and appends a VHD footer.
+func Convert(r io.Reader, w io.ReadWriteSeeker, options ...Option) error {
+	var p params
+	for _, opt := range options {
+		opt(&p)
+	}
+
+	if err := ConvertTarToExt4(r, w, options...); err != nil {
 		return err
 	}
 
 	if p.appendDMVerity {
-		ext4size, err := w.Seek(0, io.SeekEnd)
-		if err != nil {
-			return err
-		}
-
-		// 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 {
+		// Rewind the stream for dm-verity processing
+		if _, err := w.Seek(0, io.SeekStart); err != nil {
 			return err
 		}
 
-		mtree, err := dmverity.MerkleTree(data)
+		merkleTree, err := dmverity.MerkleTree(bufio.NewReaderSize(w, dmverity.MerkleTreeBufioSize))
 		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)
+		ext4size, err := w.Seek(0, io.SeekEnd)
 		if 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))
-		padding := bytes.Repeat([]byte{0}, ext4blocksize-(sbsize%ext4blocksize))
-		_, err = w.Write(padding)
-		if err != nil {
+
+		// pad the super-block
+		sbsize := int(unsafe.Sizeof(*superBlock))
+		padding := bytes.Repeat([]byte{0}, ext4BlockSize-(sbsize%ext4BlockSize))
+		if _, err = w.Write(padding); err != nil {
 			return err
 		}
+
 		// write the tree
-		_, err = w.Write(mtree)
-		if err != nil {
+		if _, err = w.Write(merkleTree); err != nil {
 			return err
 		}
 	}
@@ -273,3 +274,37 @@ func ReadExt4SuperBlock(vhdPath string) (*format.SuperBlock, error) {
 	}
 	return &sb, nil
 }
+
+// ConvertAndComputeRootDigest writes a compact ext4 file system image that contains the files in the
+// input tar stream, computes the resulting file image's cryptographic hashes (merkle tree) and returns
+// merkle tree root digest. Convert is called with minimal options: ConvertWhiteout and MaximumDiskSize
+// set to dmverity.RecommendedVHDSizeGB.
+func ConvertAndComputeRootDigest(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())
+	}()
+
+	options := []Option{
+		ConvertWhiteout,
+		MaximumDiskSize(dmverity.RecommendedVHDSizeGB),
+	}
+	if err := ConvertTarToExt4(r, out, options...); 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.MerkleTree(bufio.NewReaderSize(out, dmverity.MerkleTreeBufioSize))
+	if err != nil {
+		return "", fmt.Errorf("failed to create merkle tree: %s", err)
+	}
+
+	hash := dmverity.RootHash(tree)
+	return fmt.Sprintf("%x", hash), nil
+}

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.ConvertAndComputeRootDigest(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,
 	}
 }