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