treehash.go 1.5 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071
  1. package glacier
  2. import (
  3. "crypto/sha256"
  4. "io"
  5. )
  6. const bufsize = 1024 * 1024
  7. // Hash contains information about the tree-hash and linear hash of a
  8. // Glacier payload. This structure is generated by ComputeHashes().
  9. type Hash struct {
  10. TreeHash []byte
  11. LinearHash []byte
  12. }
  13. // ComputeHashes computes the tree-hash and linear hash of a seekable reader r.
  14. func ComputeHashes(r io.ReadSeeker) Hash {
  15. r.Seek(0, 0) // Read the whole stream
  16. defer r.Seek(0, 0) // Rewind stream at end
  17. buf := make([]byte, bufsize)
  18. hashes := [][]byte{}
  19. hsh := sha256.New()
  20. for {
  21. // Build leaf nodes in 1MB chunks
  22. n, err := io.ReadAtLeast(r, buf, bufsize)
  23. if n == 0 {
  24. break
  25. }
  26. tmpHash := sha256.Sum256(buf[:n])
  27. hashes = append(hashes, tmpHash[:])
  28. hsh.Write(buf[:n]) // Track linear hash while we're at it
  29. if err != nil {
  30. break // This is the last chunk
  31. }
  32. }
  33. return Hash{
  34. LinearHash: hsh.Sum(nil),
  35. TreeHash: buildHashTree(hashes),
  36. }
  37. }
  38. // buildHashTree builds a hash tree root node given a set of hashes.
  39. func buildHashTree(hashes [][]byte) []byte {
  40. if hashes == nil || len(hashes) == 0 {
  41. return nil
  42. }
  43. for len(hashes) > 1 {
  44. tmpHashes := [][]byte{}
  45. for i := 0; i < len(hashes); i += 2 {
  46. if i+1 <= len(hashes)-1 {
  47. tmpHash := append(append([]byte{}, hashes[i]...), hashes[i+1]...)
  48. tmpSum := sha256.Sum256(tmpHash)
  49. tmpHashes = append(tmpHashes, tmpSum[:])
  50. } else {
  51. tmpHashes = append(tmpHashes, hashes[i])
  52. }
  53. }
  54. hashes = tmpHashes
  55. }
  56. return hashes[0]
  57. }