treehash.go

package glacier

import (
	"crypto/sha256"
	"io"

	"github.com/aws/aws-sdk-go/internal/sdkio"
)

const bufsize = 1024 * 1024

// Hash contains information about the tree-hash and linear hash of a
// Glacier payload. This structure is generated by ComputeHashes().
type Hash struct {
	TreeHash   []byte
	LinearHash []byte
}

// ComputeHashes computes the tree-hash and linear hash of a seekable reader r.
//
// See http://docs.aws.amazon.com/amazonglacier/latest/dev/checksum-calculations.html for more information.
func ComputeHashes(r io.ReadSeeker) Hash {
	start, _ := r.Seek(0, sdkio.SeekCurrent) // Read the whole stream
	defer r.Seek(start, sdkio.SeekStart)     // Rewind stream at end

	buf := make([]byte, bufsize)
	hashes := [][]byte{}
	hsh := sha256.New()

	for {
		// Build leaf nodes in 1MB chunks
		n, err := io.ReadAtLeast(r, buf, bufsize)
		if n == 0 {
			break
		}

		tmpHash := sha256.Sum256(buf[:n])
		hashes = append(hashes, tmpHash[:])
		hsh.Write(buf[:n]) // Track linear hash while we're at it

		if err != nil {
			break // This is the last chunk
		}
	}

	return Hash{
		LinearHash: hsh.Sum(nil),
		TreeHash:   ComputeTreeHash(hashes),
	}
}

// ComputeTreeHash builds a tree hash root node given a slice of
// hashes. Glacier tree hash to be derived from SHA256 hashes of 1MB
// chucks of the data.
//
// See http://docs.aws.amazon.com/amazonglacier/latest/dev/checksum-calculations.html for more information.
func ComputeTreeHash(hashes [][]byte) []byte {
	hashCount := len(hashes)
	switch hashCount {
	case 0:
		return nil
	case 1:
		return hashes[0]
	}
	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()
		}
		leaves = queue
		queue = queue[:0]
	}
	return leaves[0][:]
}