From a31c5056f5d3ad0d7a1ef613ca5e43665acb99ee Mon Sep 17 00:00:00 2001
From: Maksim An <maksiman@microsoft.com>
Date: Tue, 26 Oct 2021 23:56:26 -0700
Subject: [PATCH 1/4] 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>
---
 cmd/dmverity-vhd/main.go                      | 34 +--------
 ext4/dmverity/dmverity.go                     | 44 +++++++-----
 ext4/tar2ext4/tar2ext4.go                     | 69 +++++++++++++------
 internal/tools/securitypolicy/main.go         | 44 +++---------
 .../hcsshim/ext4/dmverity/dmverity.go         | 44 +++++++-----
 .../hcsshim/ext4/tar2ext4/tar2ext4.go         | 69 +++++++++++++------
 6 files changed, 166 insertions(+), 138 deletions(-)

diff --git a/cmd/dmverity-vhd/main.go b/cmd/dmverity-vhd/main.go
index 74cf3f6679..7e54ff6b53 100644
--- 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,10 @@ 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.ConvertAndRootDigest(rc)
 			if err != nil {
-				return errors.Wrap(err, "failed to create merkle tree")
+				return errors.Wrapf(err, "failed to compute root hash")
 			}
-			hash := dmverity.RootHash(tree)
 			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
index 8948421172..e824a92fc3 100644
--- 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,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, 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")
 		}
 	}
@@ -107,6 +114,11 @@ func MerkleTree(data []byte) ([]byte, error) {
 	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
index d9ba3131c7..0f94bcd2cf 100644
--- a/ext4/tar2ext4/tar2ext4.go
+++ b/ext4/tar2ext4/tar2ext4.go
@@ -5,6 +5,7 @@ import (
 	"bufio"
 	"bytes"
 	"encoding/binary"
+	"fmt"
 	"io"
 	"io/ioutil"
 	"os"
@@ -188,42 +189,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)
+		merkleTree, 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(merkleTree); err != nil {
 			return err
 		}
 	}
@@ -273,3 +268,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
index 58100455c9..c5fddd9713 100644
--- 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,19 +219,19 @@ 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{
@@ -263,11 +239,11 @@ func convertEnvironmentVariableRules(toml []EnvironmentVariableRule) securitypol
 			Rule:     rule.Rule,
 		}
 
-		json[strconv.Itoa(i)] = jsonRule
+		jsn[strconv.Itoa(i)] = jsonRule
 	}
 
 	return securitypolicy.EnvRules{
-		Elements: json,
+		Elements: jsn,
 	}
 }
 
diff --git a/test/vendor/github.com/Microsoft/hcsshim/ext4/dmverity/dmverity.go b/test/vendor/github.com/Microsoft/hcsshim/ext4/dmverity/dmverity.go
index 8948421172..e824a92fc3 100644
--- a/test/vendor/github.com/Microsoft/hcsshim/ext4/dmverity/dmverity.go
+++ b/test/vendor/github.com/Microsoft/hcsshim/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,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, 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")
 		}
 	}
@@ -107,6 +114,11 @@ func MerkleTree(data []byte) ([]byte, error) {
 	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/test/vendor/github.com/Microsoft/hcsshim/ext4/tar2ext4/tar2ext4.go b/test/vendor/github.com/Microsoft/hcsshim/ext4/tar2ext4/tar2ext4.go
index d9ba3131c7..0f94bcd2cf 100644
--- a/test/vendor/github.com/Microsoft/hcsshim/ext4/tar2ext4/tar2ext4.go
+++ b/test/vendor/github.com/Microsoft/hcsshim/ext4/tar2ext4/tar2ext4.go
@@ -5,6 +5,7 @@ import (
 	"bufio"
 	"bytes"
 	"encoding/binary"
+	"fmt"
 	"io"
 	"io/ioutil"
 	"os"
@@ -188,42 +189,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)
+		merkleTree, 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(merkleTree); err != nil {
 			return err
 		}
 	}
@@ -273,3 +268,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
+}

From 68707319fbe7ddda1768ad3d74295d5a3109c890 Mon Sep 17 00:00:00 2001
From: Maksim An <maksiman@microsoft.com>
Date: Thu, 28 Oct 2021 10:56:43 -0700
Subject: [PATCH 2/4] pr feedback: use bufio.Reader, Read vs io.ReadFull,
 renaming

Signed-off-by: Maksim An <maksiman@microsoft.com>
---
 cmd/dmverity-vhd/main.go                      |  4 ++--
 ext4/dmverity/dmverity.go                     | 19 +++++++------------
 ext4/tar2ext4/tar2ext4.go                     |  8 ++++----
 internal/tools/securitypolicy/main.go         |  2 +-
 .../hcsshim/ext4/dmverity/dmverity.go         | 19 +++++++------------
 .../hcsshim/ext4/tar2ext4/tar2ext4.go         |  8 ++++----
 6 files changed, 25 insertions(+), 35 deletions(-)

diff --git a/cmd/dmverity-vhd/main.go b/cmd/dmverity-vhd/main.go
index 7e54ff6b53..9e8ef9f54e 100644
--- a/cmd/dmverity-vhd/main.go
+++ b/cmd/dmverity-vhd/main.go
@@ -213,9 +213,9 @@ var rootHashVHDCommand = cli.Command{
 				return errors.Wrapf(err, "failed to uncompress layer %s", diffID.String())
 			}
 
-			hash, err := tar2ext4.ConvertAndRootDigest(rc)
+			hash, err := tar2ext4.ConvertAndComputeRootDigest(rc)
 			if err != nil {
-				return errors.Wrapf(err, "failed to compute root hash")
+				return errors.Wrap(err, "failed to compute root hash")
 			}
 			fmt.Fprintf(os.Stdout, "Layer %d\nroot hash: %x\n", layerNumber, hash)
 		}
diff --git a/ext4/dmverity/dmverity.go b/ext4/dmverity/dmverity.go
index e824a92fc3..a9f4b92460 100644
--- a/ext4/dmverity/dmverity.go
+++ b/ext4/dmverity/dmverity.go
@@ -17,8 +17,8 @@ import (
 
 const (
 	blockSize = compactext4.BlockSize
-	// bufioSize is a default buffer size to use with bufio.Reader
-	bufioSize = 1024 * 1024 // 1MB
+	// 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
 )
@@ -73,16 +73,16 @@ type VerityInfo struct {
 	Version       uint32
 }
 
-// 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) {
+// MerkleTree constructs dm-verity hash-tree for a given bufio.Reader with a fixed salt (0-byte) and algorithm (sha256).
+func MerkleTree(br *bufio.Reader) ([]byte, error) {
 	layers := make([][]byte, 0)
-	currentLevel := bufio.NewReaderSize(r, bufioSize)
+	var currentLevel io.Reader = br
 
 	for {
 		nextLevel := bytes.NewBuffer(make([]byte, 0))
 		for {
 			block := make([]byte, blockSize)
-			if _, err := currentLevel.Read(block); err != nil {
+			if _, err := io.ReadFull(currentLevel, block); err != nil {
 				if err == io.EOF {
 					break
 				}
@@ -96,7 +96,7 @@ func MerkleTreeWithReader(r io.Reader) ([]byte, error) {
 		nextLevel.Write(padding)
 
 		layers = append(layers, nextLevel.Bytes())
-		currentLevel = bufio.NewReaderSize(nextLevel, bufioSize)
+		currentLevel = bufio.NewReaderSize(nextLevel, MerkleTreeBufioSize)
 
 		// This means that only root hash remains and our job is done
 		if nextLevel.Len() == blockSize {
@@ -114,11 +114,6 @@ func MerkleTreeWithReader(r io.Reader) ([]byte, error) {
 	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
index 0f94bcd2cf..809f4b3eef 100644
--- a/ext4/tar2ext4/tar2ext4.go
+++ b/ext4/tar2ext4/tar2ext4.go
@@ -194,7 +194,7 @@ func Convert(r io.Reader, w io.ReadWriteSeeker, options ...Option) error {
 			return err
 		}
 
-		merkleTree, err := dmverity.MerkleTreeWithReader(w)
+		merkleTree, err := dmverity.MerkleTree(bufio.NewReaderSize(w, dmverity.MerkleTreeBufioSize))
 		if err != nil {
 			return errors.Wrap(err, "failed to build merkle tree")
 		}
@@ -269,10 +269,10 @@ func ReadExt4SuperBlock(vhdPath string) (*format.SuperBlock, error) {
 	return &sb, nil
 }
 
-// ConvertAndRootDigest writes a compact ext4 file system image that contains the files in the
+// ConvertAndComputeRootDigest 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) {
+func ConvertAndComputeRootDigest(r io.Reader) (string, error) {
 	out, err := ioutil.TempFile("", "")
 	if err != nil {
 		return "", fmt.Errorf("failed to create temporary file: %s", err)
@@ -294,7 +294,7 @@ func ConvertAndRootDigest(r io.Reader) (string, error) {
 		return "", fmt.Errorf("failed to seek start on temp file when creating merkle tree: %s", err)
 	}
 
-	tree, err := dmverity.MerkleTreeWithReader(r)
+	tree, err := dmverity.MerkleTree(bufio.NewReaderSize(out, dmverity.MerkleTreeBufioSize))
 	if err != nil {
 		return "", fmt.Errorf("failed to create merkle tree: %s", err)
 	}
diff --git a/internal/tools/securitypolicy/main.go b/internal/tools/securitypolicy/main.go
index c5fddd9713..ad7e6d4498 100644
--- a/internal/tools/securitypolicy/main.go
+++ b/internal/tools/securitypolicy/main.go
@@ -167,7 +167,7 @@ func createPolicyFromConfig(config Config) (securitypolicy.SecurityPolicy, error
 				return p, err
 			}
 
-			hashString, err := tar2ext4.ConvertAndRootDigest(r)
+			hashString, err := tar2ext4.ConvertAndComputeRootDigest(r)
 			if err != nil {
 				return p, err
 			}
diff --git a/test/vendor/github.com/Microsoft/hcsshim/ext4/dmverity/dmverity.go b/test/vendor/github.com/Microsoft/hcsshim/ext4/dmverity/dmverity.go
index e824a92fc3..a9f4b92460 100644
--- a/test/vendor/github.com/Microsoft/hcsshim/ext4/dmverity/dmverity.go
+++ b/test/vendor/github.com/Microsoft/hcsshim/ext4/dmverity/dmverity.go
@@ -17,8 +17,8 @@ import (
 
 const (
 	blockSize = compactext4.BlockSize
-	// bufioSize is a default buffer size to use with bufio.Reader
-	bufioSize = 1024 * 1024 // 1MB
+	// 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
 )
@@ -73,16 +73,16 @@ type VerityInfo struct {
 	Version       uint32
 }
 
-// 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) {
+// MerkleTree constructs dm-verity hash-tree for a given bufio.Reader with a fixed salt (0-byte) and algorithm (sha256).
+func MerkleTree(br *bufio.Reader) ([]byte, error) {
 	layers := make([][]byte, 0)
-	currentLevel := bufio.NewReaderSize(r, bufioSize)
+	var currentLevel io.Reader = br
 
 	for {
 		nextLevel := bytes.NewBuffer(make([]byte, 0))
 		for {
 			block := make([]byte, blockSize)
-			if _, err := currentLevel.Read(block); err != nil {
+			if _, err := io.ReadFull(currentLevel, block); err != nil {
 				if err == io.EOF {
 					break
 				}
@@ -96,7 +96,7 @@ func MerkleTreeWithReader(r io.Reader) ([]byte, error) {
 		nextLevel.Write(padding)
 
 		layers = append(layers, nextLevel.Bytes())
-		currentLevel = bufio.NewReaderSize(nextLevel, bufioSize)
+		currentLevel = bufio.NewReaderSize(nextLevel, MerkleTreeBufioSize)
 
 		// This means that only root hash remains and our job is done
 		if nextLevel.Len() == blockSize {
@@ -114,11 +114,6 @@ func MerkleTreeWithReader(r io.Reader) ([]byte, error) {
 	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/test/vendor/github.com/Microsoft/hcsshim/ext4/tar2ext4/tar2ext4.go b/test/vendor/github.com/Microsoft/hcsshim/ext4/tar2ext4/tar2ext4.go
index 0f94bcd2cf..809f4b3eef 100644
--- a/test/vendor/github.com/Microsoft/hcsshim/ext4/tar2ext4/tar2ext4.go
+++ b/test/vendor/github.com/Microsoft/hcsshim/ext4/tar2ext4/tar2ext4.go
@@ -194,7 +194,7 @@ func Convert(r io.Reader, w io.ReadWriteSeeker, options ...Option) error {
 			return err
 		}
 
-		merkleTree, err := dmverity.MerkleTreeWithReader(w)
+		merkleTree, err := dmverity.MerkleTree(bufio.NewReaderSize(w, dmverity.MerkleTreeBufioSize))
 		if err != nil {
 			return errors.Wrap(err, "failed to build merkle tree")
 		}
@@ -269,10 +269,10 @@ func ReadExt4SuperBlock(vhdPath string) (*format.SuperBlock, error) {
 	return &sb, nil
 }
 
-// ConvertAndRootDigest writes a compact ext4 file system image that contains the files in the
+// ConvertAndComputeRootDigest 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) {
+func ConvertAndComputeRootDigest(r io.Reader) (string, error) {
 	out, err := ioutil.TempFile("", "")
 	if err != nil {
 		return "", fmt.Errorf("failed to create temporary file: %s", err)
@@ -294,7 +294,7 @@ func ConvertAndRootDigest(r io.Reader) (string, error) {
 		return "", fmt.Errorf("failed to seek start on temp file when creating merkle tree: %s", err)
 	}
 
-	tree, err := dmverity.MerkleTreeWithReader(r)
+	tree, err := dmverity.MerkleTree(bufio.NewReaderSize(out, dmverity.MerkleTreeBufioSize))
 	if err != nil {
 		return "", fmt.Errorf("failed to create merkle tree: %s", err)
 	}

From 90aa8cfee76a4e22b62a9a67ee8e918973f2b15e Mon Sep 17 00:00:00 2001
From: Maksim An <maksiman@microsoft.com>
Date: Wed, 3 Nov 2021 21:20:39 -0700
Subject: [PATCH 3/4] pr feedback: add ConvertTarToExt4

Separate tar to ext4 logic of Convert into a ConvertTarToExt4
function.

Signed-off-by: Maksim An <maksiman@microsoft.com>
---
 ext4/tar2ext4/tar2ext4.go                     | 46 +++++++++++--------
 .../hcsshim/ext4/tar2ext4/tar2ext4.go         | 46 +++++++++++--------
 2 files changed, 52 insertions(+), 40 deletions(-)

diff --git a/ext4/tar2ext4/tar2ext4.go b/ext4/tar2ext4/tar2ext4.go
index 809f4b3eef..1ffca8ecfc 100644
--- a/ext4/tar2ext4/tar2ext4.go
+++ b/ext4/tar2ext4/tar2ext4.go
@@ -67,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 {
@@ -177,20 +178,24 @@ 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
-		// dm-verity processing
-		if _, err = w.Seek(0, io.SeekStart); err != nil {
+		// Rewind the stream for dm-verity processing
+		if _, err := w.Seek(0, io.SeekStart); err != nil {
 			return err
 		}
 
@@ -201,7 +206,8 @@ func Convert(r io.Reader, w io.ReadWriteSeeker, options ...Option) error {
 
 		// Write dm-verity super-block and then the merkle tree after the end of the
 		// ext4 filesystem
-		if _, err = w.Seek(0, io.SeekEnd); err != nil {
+		ext4size, err := w.Seek(0, io.SeekEnd)
+		if err != nil {
 			return err
 		}
 
@@ -212,7 +218,7 @@ func Convert(r io.Reader, w io.ReadWriteSeeker, options ...Option) error {
 
 		// pad the super-block
 		sbsize := int(unsafe.Sizeof(*superBlock))
-		padding := bytes.Repeat([]byte{0}, ext4blocksize-(sbsize%ext4blocksize))
+		padding := bytes.Repeat([]byte{0}, ext4BlockSize-(sbsize%ext4BlockSize))
 		if _, err = w.Write(padding); err != nil {
 			return err
 		}
@@ -270,8 +276,9 @@ func ReadExt4SuperBlock(vhdPath string) (*format.SuperBlock, error) {
 }
 
 // ConvertAndComputeRootDigest 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.
+// 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 {
@@ -281,12 +288,11 @@ func ConvertAndComputeRootDigest(r io.Reader) (string, error) {
 		_ = os.Remove(out.Name())
 	}()
 
-	opts := []Option{
+	options := []Option{
 		ConvertWhiteout,
 		MaximumDiskSize(dmverity.RecommendedVHDSizeGB),
 	}
-
-	if err := Convert(r, out, opts...); err != nil {
+	if err := ConvertTarToExt4(r, out, options...); err != nil {
 		return "", fmt.Errorf("failed to convert tar to ext4: %s", err)
 	}
 
diff --git a/test/vendor/github.com/Microsoft/hcsshim/ext4/tar2ext4/tar2ext4.go b/test/vendor/github.com/Microsoft/hcsshim/ext4/tar2ext4/tar2ext4.go
index 809f4b3eef..1ffca8ecfc 100644
--- a/test/vendor/github.com/Microsoft/hcsshim/ext4/tar2ext4/tar2ext4.go
+++ b/test/vendor/github.com/Microsoft/hcsshim/ext4/tar2ext4/tar2ext4.go
@@ -67,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 {
@@ -177,20 +178,24 @@ 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
-		// dm-verity processing
-		if _, err = w.Seek(0, io.SeekStart); err != nil {
+		// Rewind the stream for dm-verity processing
+		if _, err := w.Seek(0, io.SeekStart); err != nil {
 			return err
 		}
 
@@ -201,7 +206,8 @@ func Convert(r io.Reader, w io.ReadWriteSeeker, options ...Option) error {
 
 		// Write dm-verity super-block and then the merkle tree after the end of the
 		// ext4 filesystem
-		if _, err = w.Seek(0, io.SeekEnd); err != nil {
+		ext4size, err := w.Seek(0, io.SeekEnd)
+		if err != nil {
 			return err
 		}
 
@@ -212,7 +218,7 @@ func Convert(r io.Reader, w io.ReadWriteSeeker, options ...Option) error {
 
 		// pad the super-block
 		sbsize := int(unsafe.Sizeof(*superBlock))
-		padding := bytes.Repeat([]byte{0}, ext4blocksize-(sbsize%ext4blocksize))
+		padding := bytes.Repeat([]byte{0}, ext4BlockSize-(sbsize%ext4BlockSize))
 		if _, err = w.Write(padding); err != nil {
 			return err
 		}
@@ -270,8 +276,9 @@ func ReadExt4SuperBlock(vhdPath string) (*format.SuperBlock, error) {
 }
 
 // ConvertAndComputeRootDigest 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.
+// 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 {
@@ -281,12 +288,11 @@ func ConvertAndComputeRootDigest(r io.Reader) (string, error) {
 		_ = os.Remove(out.Name())
 	}()
 
-	opts := []Option{
+	options := []Option{
 		ConvertWhiteout,
 		MaximumDiskSize(dmverity.RecommendedVHDSizeGB),
 	}
-
-	if err := Convert(r, out, opts...); err != nil {
+	if err := ConvertTarToExt4(r, out, options...); err != nil {
 		return "", fmt.Errorf("failed to convert tar to ext4: %s", err)
 	}
 

From a02bdf5167a08857c12db952f58fae1932a72b49 Mon Sep 17 00:00:00 2001
From: Maksim An <maksiman@microsoft.com>
Date: Tue, 9 Nov 2021 12:44:19 -0800
Subject: [PATCH 4/4] pr feedback: switch from bufio.Reader to io.Reader

Signed-off-by: Maksim An <maksiman@microsoft.com>
---
 cmd/dmverity-vhd/main.go                                    | 2 +-
 ext4/dmverity/dmverity.go                                   | 6 +++---
 .../github.com/Microsoft/hcsshim/ext4/dmverity/dmverity.go  | 6 +++---
 3 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/cmd/dmverity-vhd/main.go b/cmd/dmverity-vhd/main.go
index 9e8ef9f54e..d8520bb5ca 100644
--- a/cmd/dmverity-vhd/main.go
+++ b/cmd/dmverity-vhd/main.go
@@ -217,7 +217,7 @@ var rootHashVHDCommand = cli.Command{
 			if err != nil {
 				return errors.Wrap(err, "failed to compute root hash")
 			}
-			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
 	},
diff --git a/ext4/dmverity/dmverity.go b/ext4/dmverity/dmverity.go
index a9f4b92460..86da6a9968 100644
--- a/ext4/dmverity/dmverity.go
+++ b/ext4/dmverity/dmverity.go
@@ -73,10 +73,10 @@ type VerityInfo struct {
 	Version       uint32
 }
 
-// MerkleTree constructs dm-verity hash-tree for a given bufio.Reader with a fixed salt (0-byte) and algorithm (sha256).
-func MerkleTree(br *bufio.Reader) ([]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)
-	var currentLevel io.Reader = br
+	currentLevel := r
 
 	for {
 		nextLevel := bytes.NewBuffer(make([]byte, 0))
diff --git a/test/vendor/github.com/Microsoft/hcsshim/ext4/dmverity/dmverity.go b/test/vendor/github.com/Microsoft/hcsshim/ext4/dmverity/dmverity.go
index a9f4b92460..86da6a9968 100644
--- a/test/vendor/github.com/Microsoft/hcsshim/ext4/dmverity/dmverity.go
+++ b/test/vendor/github.com/Microsoft/hcsshim/ext4/dmverity/dmverity.go
@@ -73,10 +73,10 @@ type VerityInfo struct {
 	Version       uint32
 }
 
-// MerkleTree constructs dm-verity hash-tree for a given bufio.Reader with a fixed salt (0-byte) and algorithm (sha256).
-func MerkleTree(br *bufio.Reader) ([]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)
-	var currentLevel io.Reader = br
+	currentLevel := r
 
 	for {
 		nextLevel := bytes.NewBuffer(make([]byte, 0))