1package arbo
 2
 3import (
 4	"crypto/sha256"
 5	"math/big"
 6
 7	"github.com/iden3/go-iden3-crypto/poseidon"
 8)
 9
10var (
11	// TypeHashSha256 represents the label for the HashFunction of Sha256
12	TypeHashSha256 = []byte("sha256")
13	// TypeHashPoseidon represents the label for the HashFunction of
14	// Poseidon
15	TypeHashPoseidon = []byte("poseidon")
16
17	// HashFunctionSha256 contains the HashSha256 struct which implements
18	// the HashFunction interface
19	HashFunctionSha256 HashSha256
20	// HashFunctionPoseidon contains the HashPoseidon struct which implements
21	// the HashFunction interface
22	HashFunctionPoseidon HashPoseidon
23)
24
25// Once Generics are at Go, this will be updated (August 2021
26// https://blog.golang.org/generics-next-step)
27
28// HashFunction defines the interface that is expected for a hash function to be
29// used in a generic way in the Tree.
30type HashFunction interface {
31	Type() []byte
32	Len() int
33	Hash(...[]byte) ([]byte, error)
34	// CheckInputs checks if the inputs are valid without computing the hash
35	// CheckInputs(...[]byte) error
36}
37
38// HashSha256 implements the HashFunction interface for the Sha256 hash
39type HashSha256 struct{}
40
41// Type returns the type of HashFunction for the HashSha256
42func (f HashSha256) Type() []byte {
43	return TypeHashSha256
44}
45
46// Len returns the length of the Hash output
47func (f HashSha256) Len() int {
48	return 32 //nolint:gomnd
49}
50
51// Hash implements the hash method for the HashFunction HashSha256
52func (f HashSha256) Hash(b ...[]byte) ([]byte, error) {
53	var toHash []byte
54	for i := 0; i < len(b); i++ {
55		toHash = append(toHash, b[i]...)
56	}
57	h := sha256.Sum256(toHash)
58	return h[:], nil
59}
60
61// HashPoseidon implements the HashFunction interface for the Poseidon hash
62type HashPoseidon struct{}
63
64// Type returns the type of HashFunction for the HashPoseidon
65func (f HashPoseidon) Type() []byte {
66	return TypeHashPoseidon
67}
68
69// Len returns the length of the Hash output
70func (f HashPoseidon) Len() int {
71	return 32 //nolint:gomnd
72}
73
74// Hash implements the hash method for the HashFunction HashPoseidon
75func (f HashPoseidon) Hash(b ...[]byte) ([]byte, error) {
76	var toHash []*big.Int
77	for i := 0; i < len(b); i++ {
78		bi := BytesToBigInt(b[i])
79		toHash = append(toHash, bi)
80	}
81	h, err := poseidon.Hash(toHash)
82	if err != nil {
83		return nil, err
84	}
85	hB := BigIntToBytes(h)
86	return hB, nil
87}