1package poseidon
  2
  3import (
  4	"errors"
  5	"math/big"
  6	"strconv"
  7
  8	"github.com/iden3/go-iden3-crypto/constants"
  9	"github.com/iden3/go-iden3-crypto/ff"
 10	"github.com/iden3/go-iden3-crypto/utils"
 11	"golang.org/x/crypto/blake2b"
 12)
 13
 14const SEED = "poseidon"
 15const NROUNDSF = 8
 16const NROUNDSP = 57
 17const T = 6
 18
 19var constC []*ff.Element
 20var constM [T][T]*ff.Element
 21
 22func Zero() *ff.Element {
 23	return ff.NewElement().SetZero()
 24}
 25
 26func modQ(v *big.Int) {
 27	v.Mod(v, constants.Q)
 28}
 29
 30func init() {
 31	constC = getPseudoRandom(SEED+"_constants", NROUNDSF+NROUNDSP)
 32	constM = getMDS()
 33}
 34
 35func getPseudoRandom(seed string, n int) []*ff.Element {
 36	res := make([]*ff.Element, n)
 37	hash := blake2b.Sum256([]byte(seed))
 38	for i := 0; i < n; i++ {
 39		hashBigInt := big.NewInt(int64(0))
 40		res[i] = ff.NewElement().SetBigInt(utils.SetBigIntFromLEBytes(hashBigInt, hash[:]))
 41		hash = blake2b.Sum256(hash[:])
 42	}
 43	return res
 44}
 45
 46func nonceToString(n int) string {
 47	r := strconv.Itoa(n)
 48	for len(r) < 4 {
 49		r = "0" + r
 50	}
 51	return r
 52}
 53
 54// https://eprint.iacr.org/2019/458.pdf pag.8
 55func getMDS() [T][T]*ff.Element {
 56	nonce := 0
 57	cauchyMatrix := getPseudoRandom(SEED+"_matrix_"+nonceToString(nonce), T*2)
 58	for !checkAllDifferent(cauchyMatrix) {
 59		nonce += 1
 60		cauchyMatrix = getPseudoRandom(SEED+"_matrix_"+nonceToString(nonce), T*2)
 61	}
 62	var m [T][T]*ff.Element
 63	for i := 0; i < T; i++ {
 64		for j := 0; j < T; j++ {
 65			m[i][j] = ff.NewElement().Sub(cauchyMatrix[i], cauchyMatrix[T+j])
 66			m[i][j].Inverse(m[i][j])
 67		}
 68	}
 69	return m
 70}
 71
 72func checkAllDifferent(v []*ff.Element) bool {
 73	for i := 0; i < len(v); i++ {
 74		if v[i].Equal(ff.NewElement().SetZero()) {
 75			return false
 76		}
 77		for j := i + 1; j < len(v); j++ {
 78			if v[i].Equal(v[j]) {
 79				return false
 80			}
 81		}
 82	}
 83	return true
 84}
 85
 86// ark computes Add-Round Key, from the paper https://eprint.iacr.org/2019/458.pdf
 87func ark(state [T]*ff.Element, c *ff.Element) {
 88	for i := 0; i < T; i++ {
 89		state[i].Add(state[i], c)
 90	}
 91}
 92
 93// cubic performs x^5 mod p
 94// https://eprint.iacr.org/2019/458.pdf page 8
 95// var five = big.NewInt(5)
 96
 97func cubic(a *ff.Element) {
 98	a.Exp(*a, 5)
 99}
100
101// sbox https://eprint.iacr.org/2019/458.pdf page 6
102func sbox(state [T]*ff.Element, i int) {
103	if (i < NROUNDSF/2) || (i >= NROUNDSF/2+NROUNDSP) {
104		for j := 0; j < T; j++ {
105			cubic(state[j])
106		}
107	} else {
108		cubic(state[0])
109	}
110}
111
112// mix returns [[matrix]] * [vector]
113func mix(state [T]*ff.Element, newState [T]*ff.Element, m [T][T]*ff.Element) {
114	mul := Zero()
115	for i := 0; i < T; i++ {
116		newState[i].SetUint64(0)
117		for j := 0; j < T; j++ {
118			mul.Mul(m[i][j], state[j])
119			newState[i].Add(newState[i], mul)
120		}
121	}
122}
123
124// PoseidonHash computes the Poseidon hash for the given inputs
125func PoseidonHash(inpBI [T]*big.Int) (*big.Int, error) {
126	if !utils.CheckBigIntArrayInField(inpBI[:]) {
127		return nil, errors.New("inputs values not inside Finite Field")
128	}
129	inp := utils.BigIntArrayToElementArray(inpBI[:])
130	state := [T]*ff.Element{}
131	for i := 0; i < T; i++ {
132		state[i] = ff.NewElement().Set(inp[i])
133	}
134
135	// ARK --> SBox --> M, https://eprint.iacr.org/2019/458.pdf pag.5
136	var newState [T]*ff.Element
137	for i := 0; i < T; i++ {
138		newState[i] = Zero()
139	}
140	for i := 0; i < NROUNDSF+NROUNDSP; i++ {
141		ark(state, constC[i])
142		sbox(state, i)
143		mix(state, newState, constM)
144		state, newState = newState, state
145	}
146	rE := state[0]
147	r := big.NewInt(0)
148	rE.ToBigIntRegular(r)
149	return r, nil
150}
151
152// Hash performs the Poseidon hash over a ff.Element array
153// in chunks of 5 elements
154func Hash(arr []*big.Int) (*big.Int, error) {
155	r := big.NewInt(int64(1))
156	for i := 0; i < len(arr); i = i + T - 1 {
157		var toHash [T]*big.Int
158		j := 0
159		for ; j < T-1; j++ {
160			if i+j >= len(arr) {
161				break
162			}
163			toHash[j] = arr[i+j]
164		}
165		toHash[j] = r
166		j++
167		for ; j < T; j++ {
168			toHash[j] = big.NewInt(0)
169		}
170
171		ph, err := PoseidonHash(toHash)
172		if err != nil {
173			return nil, err
174		}
175		modQ(r.Add(r, ph))
176	}
177
178	return r, nil
179}
180
181// HashBytes hashes a msg byte slice by blocks of 31 bytes encoded as
182// little-endian
183func HashBytes(b []byte) (*big.Int, error) {
184	n := 31
185	bElems := make([]*big.Int, 0, len(b)/n+1)
186	for i := 0; i < len(b)/n; i++ {
187		v := big.NewInt(int64(0))
188		utils.SetBigIntFromLEBytes(v, b[n*i:n*(i+1)])
189		bElems = append(bElems, v)
190
191	}
192	if len(b)%n != 0 {
193		v := big.NewInt(int64(0))
194		utils.SetBigIntFromLEBytes(v, b[(len(b)/n)*n:])
195		bElems = append(bElems, v)
196	}
197	return Hash(bElems)
198}