From e1d2173dc9d4c5ddd2c7b8fc647bd6fa6905678f Mon Sep 17 00:00:00 2001
From: arnaucube <root@arnaucube.com>
Date: Tue, 30 Mar 2021 21:59:24 +0200
Subject: [PATCH] Add HashFunction interface (w/ Poseidon & Sha256)

---
 hash.go      | 87 ++++++++++++++++++++++++++++++++++++++++++++++++++++
 hash_test.go | 38 +++++++++++++++++++++++
 utils.go     | 25 +++++++++++++++
 3 files changed, 150 insertions(+)
 create mode 100644 hash.go
 create mode 100644 hash_test.go
 create mode 100644 utils.go

diff --git a/hash.go b/hash.go
new file mode 100644
index 0000000..c1ee64b
--- /dev/null
+++ b/hash.go
@@ -0,0 +1,87 @@
+package arbo
+
+import (
+	"crypto/sha256"
+	"math/big"
+
+	"github.com/iden3/go-iden3-crypto/poseidon"
+)
+
+var (
+	// TypeHashSha256 represents the label for the HashFunction of Sha256
+	TypeHashSha256 = []byte("sha256")
+	// TypeHashPoseidon represents the label for the HashFunction of
+	// Poseidon
+	TypeHashPoseidon = []byte("poseidon")
+
+	// HashFunctionSha256 contains the HashSha256 struct which implements
+	// the HashFunction interface
+	HashFunctionSha256 HashSha256
+	// HashFunctionPoseidon contains the HashPoseidon struct which implements
+	// the HashFunction interface
+	HashFunctionPoseidon HashPoseidon
+)
+
+// Once Generics are at Go, this will be updated (August 2021
+// https://blog.golang.org/generics-next-step)
+
+// HashFunction defines the interface that is expected for a hash function to be
+// used in a generic way in the Tree.
+type HashFunction interface {
+	Type() []byte
+	Len() int
+	Hash(...[]byte) ([]byte, error)
+	// CheckInputs checks if the inputs are valid without computing the hash
+	// CheckInputs(...[]byte) error
+}
+
+// HashSha256 implements the HashFunction interface for the Sha256 hash
+type HashSha256 struct{}
+
+// Type returns the type of HashFunction for the HashSha256
+func (f HashSha256) Type() []byte {
+	return TypeHashSha256
+}
+
+// Len returns the length of the Hash output
+func (f HashSha256) Len() int {
+	return 32 //nolint:gomnd
+}
+
+// Hash implements the hash method for the HashFunction HashSha256
+func (f HashSha256) Hash(b ...[]byte) ([]byte, error) {
+	var toHash []byte
+	for i := 0; i < len(b); i++ {
+		toHash = append(toHash, b[i]...)
+	}
+	h := sha256.Sum256(toHash)
+	return h[:], nil
+}
+
+// HashPoseidon implements the HashFunction interface for the Poseidon hash
+type HashPoseidon struct{}
+
+// Type returns the type of HashFunction for the HashPoseidon
+func (f HashPoseidon) Type() []byte {
+	return TypeHashPoseidon
+}
+
+// Len returns the length of the Hash output
+func (f HashPoseidon) Len() int {
+	return 32 //nolint:gomnd
+}
+
+// Hash implements the hash method for the HashFunction HashPoseidon
+func (f HashPoseidon) Hash(b ...[]byte) ([]byte, error) {
+	var toHash []*big.Int
+	for i := 0; i < len(b); i++ {
+		bi := BytesToBigInt(b[i])
+		toHash = append(toHash, bi)
+	}
+	h, err := poseidon.Hash(toHash)
+	if err != nil {
+		return nil, err
+	}
+	hB := BigIntToBytes(h)
+	return hB, nil
+}
diff --git a/hash_test.go b/hash_test.go
new file mode 100644
index 0000000..dd17ac3
--- /dev/null
+++ b/hash_test.go
@@ -0,0 +1,38 @@
+package arbo
+
+import (
+	"encoding/hex"
+	"math/big"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestHashSha256(t *testing.T) {
+	// Sha256 hash
+	hashFunc := &HashSha256{}
+	b := []byte("test")
+	h, err := hashFunc.Hash(b)
+	if err != nil {
+		t.Fatal(err)
+	}
+	assert.Equal(t,
+		"9f86d081884c7d659a2feaa0c55ad015a3bf4f1b2b0b822cd15d6c15b0f00a08",
+		hex.EncodeToString(h))
+}
+
+func TestHashPoseidon(t *testing.T) {
+	// Poseidon hash
+	hashFunc := &HashPoseidon{}
+	h, err := hashFunc.Hash(
+		BigIntToBytes(big.NewInt(1)),
+		BigIntToBytes(big.NewInt(2)))
+	if err != nil {
+		t.Fatal(err)
+	}
+	hBI := BytesToBigInt(h)
+	// value checked with circomlib
+	assert.Equal(t,
+		"7853200120776062878684798364095072458815029376092732009249414926327459813530",
+		hBI.String())
+}
diff --git a/utils.go b/utils.go
new file mode 100644
index 0000000..47358fd
--- /dev/null
+++ b/utils.go
@@ -0,0 +1,25 @@
+package arbo
+
+import "math/big"
+
+// SwapEndianness swaps the order of the bytes in the byte slice.
+func SwapEndianness(b []byte) []byte {
+	o := make([]byte, len(b))
+	for i := range b {
+		o[len(b)-1-i] = b[i]
+	}
+	return o
+}
+
+// BigIntToBytes converts a *big.Int into a byte array in Little-Endian
+func BigIntToBytes(bi *big.Int) []byte {
+	var b [32]byte
+	copy(b[:], SwapEndianness(bi.Bytes()))
+	return b[:]
+}
+
+// BytesToBigInt converts a byte array in Little-Endian representation into
+// *big.Int
+func BytesToBigInt(b []byte) *big.Int {
+	return new(big.Int).SetBytes(SwapEndianness(b))
+}