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
 96func cubic(a *ff.Element) {
 97	a.Exp(*a, 5)
 98}
 99
100// sbox https://eprint.iacr.org/2019/458.pdf page 6
101func sbox(state [T]*ff.Element, i int) {
102	if (i < NROUNDSF/2) || (i >= NROUNDSF/2+NROUNDSP) {
103		for j := 0; j < T; j++ {
104			cubic(state[j])
105		}
106	} else {
107		cubic(state[0])
108	}
109}
110
111// mix returns [[matrix]] * [vector]
112func mix(state [T]*ff.Element, newState [T]*ff.Element, m [T][T]*ff.Element) {
113	mul := Zero()
114	for i := 0; i < T; i++ {
115		newState[i].SetUint64(0)
116		for j := 0; j < T; j++ {
117			mul.Mul(m[i][j], state[j])
118			newState[i].Add(newState[i], mul)
119		}
120	}
121}
122
123// PoseidonHash computes the Poseidon hash for the given inputs
124func PoseidonHash(inpBI [T]*big.Int) (*big.Int, error) {
125	if !utils.CheckBigIntArrayInField(inpBI[:]) {
126		return nil, errors.New("inputs values not inside Finite Field")
127	}
128	inp := utils.BigIntArrayToElementArray(inpBI[:])
129	state := [T]*ff.Element{}
130	for i := 0; i < T; i++ {
131		state[i] = ff.NewElement().Set(inp[i])
132	}
133
134	// ARK --> SBox --> M, https://eprint.iacr.org/2019/458.pdf pag.5
135	var newState [T]*ff.Element
136	for i := 0; i < T; i++ {
137		newState[i] = Zero()
138	}
139	for i := 0; i < NROUNDSF+NROUNDSP; i++ {
140		ark(state, constC[i])
141		sbox(state, i)
142		mix(state, newState, constM)
143		state, newState = newState, state
144	}
145	rE := state[0]
146	r := big.NewInt(0)
147	rE.ToBigIntRegular(r)
148	return r, nil
149}
150
151// Hash performs the Poseidon hash over a ff.Element array
152// in chunks of 5 elements
153func Hash(arr []*big.Int) (*big.Int, error) {
154	r := big.NewInt(int64(1))
155	for i := 0; i < len(arr); i = i + T - 1 {
156		var toHash [T]*big.Int
157		j := 0
158		for ; j < T-1; j++ {
159			if i+j >= len(arr) {
160				break
161			}
162			toHash[j] = arr[i+j]
163		}
164		toHash[j] = r
165		j++
166		for ; j < T; j++ {
167			toHash[j] = big.NewInt(0)
168		}
169
170		ph, err := PoseidonHash(toHash)
171		if err != nil {
172			return nil, err
173		}
174		modQ(r.Add(r, ph))
175	}
176
177	return r, nil
178}
179
180// HashBytes hashes a msg byte slice by blocks of 31 bytes encoded as
181// little-endian
182func HashBytes(b []byte) (*big.Int, error) {
183	n := 31
184	bElems := make([]*big.Int, 0, len(b)/n+1)
185	for i := 0; i < len(b)/n; i++ {
186		v := big.NewInt(int64(0))
187		utils.SetBigIntFromLEBytes(v, b[n*i:n*(i+1)])
188		bElems = append(bElems, v)
189
190	}
191	if len(b)%n != 0 {
192		v := big.NewInt(int64(0))
193		utils.SetBigIntFromLEBytes(v, b[(len(b)/n)*n:])
194		bElems = append(bElems, v)
195	}
196	return Hash(bElems)
197}