Reduce memory allocations when computing tree hash (aws#2279)

service/glacier/internal/customizations/treehash.go

@@ -118,25 +118,36 @@ func computeHashes(r io.Reader) Hash {
 //
 // See http://docs.aws.amazon.com/amazonglacier/latest/dev/checksum-calculations.html for more information.
 func computeTreeHash(hashes [][]byte) []byte {
-	if hashes == nil || len(hashes) == 0 {
+	hashCount := len(hashes)
+	switch hashCount {
+	case 0:
 		return nil
+	case 1:
+		return hashes[0]
 	}
-
-	for len(hashes) > 1 {
-		tmpHashes := [][]byte{}
-
-		for i := 0; i < len(hashes); i += 2 {
-			if i+1 <= len(hashes)-1 {
-				tmpHash := append(append([]byte{}, hashes[i]...), hashes[i+1]...)
-				tmpSum := sha256.Sum256(tmpHash)
-				tmpHashes = append(tmpHashes, tmpSum[:])
-			} else {
-				tmpHashes = append(tmpHashes, hashes[i])
+	leaves := make([][32]byte, hashCount)
+	for i := range leaves {
+		copy(leaves[i][:], hashes[i])
+	}
+	var (
+		queue = leaves[:0]
+		h256  = sha256.New()
+		buf   [32]byte
+	)
+	for len(leaves) > 1 {
+		for i := 0; i < len(leaves); i += 2 {
+			if i+1 == len(leaves) {
+				queue = append(queue, leaves[i])
+				break
 			}
+			h256.Write(leaves[i][:])
+			h256.Write(leaves[i+1][:])
+			h256.Sum(buf[:0])
+			queue = append(queue, buf)
+			h256.Reset()
 		}
-
-		hashes = tmpHashes
+		leaves = queue
+		queue = queue[:0]
 	}
-
-	return hashes[0]
+	return leaves[0][:]
 }

service/glacier/internal/customizations/treehash_test.go

@@ -3,6 +3,8 @@ package customizations
 import (
 	"bytes"
 	"context"
+	"crypto/sha256"
+	"encoding/hex"
 	"github.com/aws/smithy-go/middleware"
 	smithyhttp "github.com/aws/smithy-go/transport/http"
 	"io"
@@ -125,3 +127,57 @@ func TestTreeHashMiddleware(t *testing.T) {
 		})
 	}
 }
+
+func TestComputeHashes(t *testing.T) {
+
+	t.Run("no hash", func(t *testing.T) {
+		var hashes [][]byte
+		treeHash := computeTreeHash(hashes)
+		if treeHash != nil {
+			t.Fatalf("expected []byte(nil), got %v", treeHash)
+		}
+	})
+
+	t.Run("one hash", func(t *testing.T) {
+		hash := sha256.Sum256([]byte("hash"))
+		treeHash := computeTreeHash([][]byte{hash[:]})
+
+		expected, actual := hex.EncodeToString(hash[:]), hex.EncodeToString(treeHash)
+		if expected != actual {
+			t.Fatalf("expected %v, got %v", expected, actual)
+		}
+	})
+
+	t.Run("even hashes", func(t *testing.T) {
+		h1 := sha256.Sum256([]byte("h1"))
+		h2 := sha256.Sum256([]byte("h2"))
+
+		hash := sha256.Sum256(append(h1[:], h2[:]...))
+		expected := hex.EncodeToString(hash[:])
+
+		treeHash := computeTreeHash([][]byte{h1[:], h2[:]})
+		actual := hex.EncodeToString(treeHash)
+
+		if expected != actual {
+			t.Fatalf("expected %v, got %v", expected, actual)
+		}
+	})
+
+	t.Run("odd hashes", func(t *testing.T) {
+		h1 := sha256.Sum256([]byte("h1"))
+		h2 := sha256.Sum256([]byte("h2"))
+		h3 := sha256.Sum256([]byte("h3"))
+
+		h12 := sha256.Sum256(append(h1[:], h2[:]...))
+		hash := sha256.Sum256(append(h12[:], h3[:]...))
+		expected := hex.EncodeToString(hash[:])
+
+		treeHash := computeTreeHash([][]byte{h1[:], h2[:], h3[:]})
+		actual := hex.EncodeToString(treeHash)
+
+		if expected != actual {
+			t.Fatalf("expected %v, got %v", expected, actual)
+		}
+	})
+
+}