From f597e20569647d70f5eb4d3afde13c69d944a262 Mon Sep 17 00:00:00 2001
From: Oleksandr Brezhniev <oleksandr.brezhniev@gmail.com>
Date: Sun, 21 Nov 2021 02:58:03 +0200
Subject: [PATCH] Fix Poseidon Hash check for inputs being in Finite Field
 (#42)

---
 poseidon/poseidon.go      |  4 ++++
 poseidon/poseidon_test.go | 23 ++++++++++++++---------
 2 files changed, 18 insertions(+), 9 deletions(-)

diff --git a/poseidon/poseidon.go b/poseidon/poseidon.go
index cd708bb..8ee54b7 100644
--- a/poseidon/poseidon.go
+++ b/poseidon/poseidon.go
@@ -1,6 +1,7 @@
 package poseidon
 
 import (
+	"errors"
 	"fmt"
 	"math/big"
 
@@ -59,6 +60,9 @@ func Hash(inpBI []*big.Int) (*big.Int, error) {
 	if len(inpBI) == 0 || len(inpBI) > len(NROUNDSP) {
 		return nil, fmt.Errorf("invalid inputs length %d, max %d", len(inpBI), len(NROUNDSP)) //nolint:gomnd,lll
 	}
+	if !utils.CheckBigIntArrayInField(inpBI[:]) {
+		return nil, errors.New("inputs values not inside Finite Field")
+	}
 	inp := utils.BigIntArrayToElementArray(inpBI[:])
 
 	nRoundsF := NROUNDSF
diff --git a/poseidon/poseidon_test.go b/poseidon/poseidon_test.go
index e544a30..6e0e422 100644
--- a/poseidon/poseidon_test.go
+++ b/poseidon/poseidon_test.go
@@ -1,23 +1,14 @@
 package poseidon
 
 import (
-	"encoding/hex"
 	"math/big"
 	"testing"
 
 	"github.com/iden3/go-iden3-crypto/utils"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
-	"golang.org/x/crypto/blake2b"
 )
 
-func TestBlake2bVersion(t *testing.T) {
-	h := blake2b.Sum256([]byte("poseidon_constants"))
-	assert.Equal(t,
-		"e57ba154fb2c47811dc1a2369b27e25a44915b4e4ece4eb8ec74850cb78e01b1",
-		hex.EncodeToString(h[:]))
-}
-
 func TestPoseidonHash(t *testing.T) {
 	b0 := big.NewInt(0)
 	b1 := big.NewInt(1)
@@ -121,6 +112,20 @@ func TestErrorInputs(t *testing.T) {
 	assert.Equal(t, "invalid inputs length 18, max 16", err.Error())
 }
 
+func TestInputsNotInField(t *testing.T) {
+	var err error
+
+	// Very big number, should just return error and not go into endless loop
+	b1 := utils.NewIntFromString("12242166908188651009877250812424843524687801523336557272219921456462821518061999999999999999999999999999999999999999999999999999999999") //nolint:lll
+	_, err = Hash([]*big.Int{b1})
+	require.Error(t, err, "inputs values not inside Finite Field")
+
+	// Finite Field const Q, should return error
+	b2 := utils.NewIntFromString("21888242871839275222246405745257275088548364400416034343698204186575808495617") //nolint:lll
+	_, err = Hash([]*big.Int{b2})
+	require.Error(t, err, "inputs values not inside Finite Field")
+}
+
 func BenchmarkPoseidonHash(b *testing.B) {
 	b0 := big.NewInt(0)
 	b1 := utils.NewIntFromString("12242166908188651009877250812424843524687801523336557272219921456462821518061") //nolint:lll