1package poseidon
  2
  3import (
  4	"bytes"
  5	"errors"
  6	"math/big"
  7	"strconv"
  8
  9	"github.com/iden3/go-iden3-crypto/field"
 10	"golang.org/x/crypto/blake2b"
 11)
 12
 13const SEED = "poseidon"
 14const NROUNDSF = 8
 15const NROUNDSP = 57
 16const T = 6
 17
 18var constants = generateConstantsData()
 19
 20type constantsData struct {
 21	fqR field.Fq
 22	c   []*big.Int
 23	m   [][]*big.Int
 24}
 25
 26// checkBigIntInField checks if given big.Int fits in a Field R element
 27func checkBigIntInField(a *big.Int, q *big.Int) bool {
 28	if a.Cmp(q) != -1 {
 29		return false
 30	}
 31	return true
 32}
 33
 34// checkBigIntArrayInField checks if given big.Int fits in a Field R element
 35func checkBigIntArrayInField(arr []*big.Int, q *big.Int) bool {
 36	for _, a := range arr {
 37		if !checkBigIntInField(a, q) {
 38			return false
 39		}
 40	}
 41	return true
 42}
 43
 44func generateConstantsData() constantsData {
 45	var constants constantsData
 46
 47	r, ok := new(big.Int).SetString("21888242871839275222246405745257275088548364400416034343698204186575808495617", 10)
 48	if !ok {
 49
 50	}
 51	fqR := field.NewFq(r)
 52	constants.fqR = fqR
 53	constants.c = getPseudoRandom(fqR, SEED+"_constants", big.NewInt(int64(NROUNDSF+NROUNDSP)))
 54	constants.m = getMDS(fqR)
 55
 56	return constants
 57}
 58
 59func getPseudoRandom(fqR field.Fq, seed string, n *big.Int) []*big.Int {
 60	var res []*big.Int
 61	hash := blake2b.Sum256([]byte(seed))
 62	for big.NewInt(int64(len(res))).Cmp(n) == -1 { // res < n
 63		newN := fqR.Affine(leByteArrayToBigInt(fqR, hash[:]))
 64		res = append(res, newN)
 65		hash = blake2b.Sum256(hash[:])
 66	}
 67	return res
 68}
 69
 70func leByteArrayToBigInt(fqR field.Fq, b []byte) *big.Int {
 71	res := fqR.Zero()
 72	for i := 0; i < len(b); i++ {
 73		n := big.NewInt(int64(b[i]))
 74		res = new(big.Int).Add(res, new(big.Int).Lsh(n, uint(i*8)))
 75	}
 76	return res
 77}
 78
 79func nonceToString(n int) string {
 80	r := strconv.Itoa(n)
 81	for len(r) < 4 {
 82		r = "0" + r
 83	}
 84	return r
 85}
 86
 87// https://eprint.iacr.org/2019/458.pdf pag.8
 88func getMDS(fqR field.Fq) [][]*big.Int {
 89	nonce := 0
 90	cauchyMatrix := getPseudoRandom(fqR, SEED+"_matrix_"+nonceToString(nonce), big.NewInt(T*2))
 91	for !checkAllDifferent(cauchyMatrix) {
 92		nonce += 1
 93		cauchyMatrix = getPseudoRandom(fqR, SEED+"_matrix_"+nonceToString(nonce), big.NewInt(T*2))
 94	}
 95	var m [][]*big.Int
 96	for i := 0; i < T; i++ {
 97		var mi []*big.Int
 98		for j := 0; j < T; j++ {
 99			mi = append(mi, fqR.Inverse(fqR.Sub(cauchyMatrix[i], cauchyMatrix[T+j])))
100		}
101		m = append(m, mi)
102	}
103	return m
104}
105
106func checkAllDifferent(v []*big.Int) bool {
107	for i := 0; i < len(v); i++ {
108		if bytes.Equal(v[i].Bytes(), big.NewInt(int64(0)).Bytes()) {
109			return false
110		}
111		for j := i + 1; j < len(v); j++ {
112			if bytes.Equal(v[i].Bytes(), v[j].Bytes()) {
113				return false
114			}
115		}
116	}
117	return true
118}
119
120// ark computes Add-Round Key, from the paper https://eprint.iacr.org/2019/458.pdf
121func ark(state []*big.Int, c *big.Int) []*big.Int {
122	for i := 0; i < len(state); i++ {
123		state[i] = constants.fqR.Add(state[i], c)
124	}
125	return state
126}
127
128// cubic performs x^3 mod p
129func cubic(a *big.Int) *big.Int {
130	return constants.fqR.Mul(a, constants.fqR.Square(constants.fqR.Square(a)))
131}
132
133// sbox https://eprint.iacr.org/2019/458.pdf pag.6
134func sbox(state []*big.Int, i int) []*big.Int {
135	if (i < NROUNDSF/2) || (i >= NROUNDSF/2+NROUNDSP) {
136		for j := 0; j < T; j++ {
137			state[j] = cubic(state[j])
138		}
139	} else {
140		state[0] = cubic(state[0])
141	}
142	return state
143}
144
145// mix returns [[matrix]] * [vector]
146func mix(state []*big.Int, m [][]*big.Int) []*big.Int {
147	var newState []*big.Int
148	for i := 0; i < len(state); i++ {
149		newState = append(newState, constants.fqR.Zero())
150		for j := 0; j < len(state); j++ {
151			newState[i] = constants.fqR.Add(newState[i], constants.fqR.Mul(m[i][j], state[j]))
152		}
153	}
154	for i := 0; i < len(state); i++ {
155		state[i] = newState[i]
156	}
157	return state
158}
159
160// Hash computes the Poseidon hash for the given inputs
161func Hash(inp []*big.Int) (*big.Int, error) {
162	var state []*big.Int
163	if len(inp) < 0 || len(inp) > T {
164		return nil, errors.New("wrong inputs length")
165	}
166	if !checkBigIntArrayInField(inp, constants.fqR.Q) {
167		return nil, errors.New("inputs values not inside Finite Field")
168	}
169
170	for i := 0; i < len(inp); i++ {
171		state = append(state, inp[i])
172	}
173	for i := len(inp); i < T; i++ {
174		state = append(state, constants.fqR.Zero())
175	}
176
177	// ARK --> SBox --> M, https://eprint.iacr.org/2019/458.pdf pag.5
178	for i := 0; i < NROUNDSF+NROUNDSP; i++ {
179		state = ark(state, constants.c[i])
180		state = sbox(state, i)
181		state = mix(state, constants.m)
182	}
183	return state[0], nil
184}