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