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