1package arbo
  2
  3import (
  4	"crypto/sha256"
  5	"math/big"
  6
  7	"github.com/iden3/go-iden3-crypto/poseidon"
  8	"golang.org/x/crypto/blake2b"
  9)
 10
 11var (
 12	// TypeHashSha256 represents the label for the HashFunction of Sha256
 13	TypeHashSha256 = []byte("sha256")
 14	// TypeHashPoseidon represents the label for the HashFunction of
 15	// Poseidon
 16	TypeHashPoseidon = []byte("poseidon")
 17	// TypeHashBlake2b represents the label for the HashFunction of Blake2b
 18	TypeHashBlake2b = []byte("blake2b")
 19
 20	// HashFunctionSha256 contains the HashSha256 struct which implements
 21	// the HashFunction interface
 22	HashFunctionSha256 HashSha256
 23	// HashFunctionPoseidon contains the HashPoseidon struct which implements
 24	// the HashFunction interface
 25	HashFunctionPoseidon HashPoseidon
 26	// HashFunctionBlake2b contains the HashBlake2b struct which implements
 27	// the HashFunction interface
 28	HashFunctionBlake2b HashBlake2b
 29)
 30
 31// Once Generics are at Go, this will be updated (August 2021
 32// https://blog.golang.org/generics-next-step)
 33
 34// HashFunction defines the interface that is expected for a hash function to be
 35// used in a generic way in the Tree.
 36type HashFunction interface {
 37	Type() []byte
 38	Len() int
 39	Hash(...[]byte) ([]byte, error)
 40	// CheckInput checks if the input is valid without computing the hash
 41	// CheckInput(...[]byte) error
 42}
 43
 44// HashSha256 implements the HashFunction interface for the Sha256 hash
 45type HashSha256 struct{}
 46
 47// Type returns the type of HashFunction for the HashSha256
 48func (f HashSha256) Type() []byte {
 49	return TypeHashSha256
 50}
 51
 52// Len returns the length of the Hash output
 53func (f HashSha256) Len() int {
 54	return 32 //nolint:gomnd
 55}
 56
 57// Hash implements the hash method for the HashFunction HashSha256
 58func (f HashSha256) Hash(b ...[]byte) ([]byte, error) {
 59	var toHash []byte
 60	for i := 0; i < len(b); i++ {
 61		toHash = append(toHash, b[i]...)
 62	}
 63	h := sha256.Sum256(toHash)
 64	return h[:], nil
 65}
 66
 67// HashPoseidon implements the HashFunction interface for the Poseidon hash
 68type HashPoseidon struct{}
 69
 70// Type returns the type of HashFunction for the HashPoseidon
 71func (f HashPoseidon) Type() []byte {
 72	return TypeHashPoseidon
 73}
 74
 75// Len returns the length of the Hash output
 76func (f HashPoseidon) Len() int {
 77	return 32 //nolint:gomnd
 78}
 79
 80// Hash implements the hash method for the HashFunction HashPoseidon. It
 81// expects the byte arrays to be little-endian representations of big.Int
 82// values.
 83func (f HashPoseidon) Hash(b ...[]byte) ([]byte, error) {
 84	var toHash []*big.Int
 85	for i := 0; i < len(b); i++ {
 86		bi := BytesToBigInt(b[i])
 87		toHash = append(toHash, bi)
 88	}
 89	h, err := poseidon.Hash(toHash)
 90	if err != nil {
 91		return nil, err
 92	}
 93	hB := BigIntToBytes(f.Len(), h)
 94	return hB, nil
 95}
 96
 97// HashBlake2b implements the HashFunction interface for the Blake2b hash
 98type HashBlake2b struct{}
 99
100// Type returns the type of HashFunction for the HashBlake2b
101func (f HashBlake2b) Type() []byte {
102	return TypeHashBlake2b
103}
104
105// Len returns the length of the Hash output
106func (f HashBlake2b) Len() int {
107	return 32 //nolint:gomnd
108}
109
110// Hash implements the hash method for the HashFunction HashBlake2b
111func (f HashBlake2b) Hash(b ...[]byte) ([]byte, error) {
112	hasher, err := blake2b.New256(nil)
113	if err != nil {
114		return nil, err
115	}
116	for i := 0; i < len(b); i++ {
117		if _, err = hasher.Write(b[i]); err != nil {
118			return nil, err
119		}
120	}
121	return hasher.Sum(nil), nil
122}