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

…ray (#1209) 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. Separate tar to ext4 logic of Convert into a ConvertTarToExt4 function. Signed-off-by: Maksim An <maksiman@microsoft.com>
microsoft · Nov 11, 2021 · ddab09b · ddab09b
1 parent 3a8cd1e
commit ddab09b
Show file tree

Hide file tree

Showing 6 changed files with 191 additions and 161 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,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
 	},

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. 
+	// 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")
 		}
 	}

diff --git a/ext4/tar2ext4/tar2ext4.go b/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
+}
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.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,
 	}
 }