Fix keccak tests and padding (#57)

* prover: keccak: fix tests Reduce the expected value by BN254 group order, because the circuit does the same. The expected value had just byte order adjusted, but it wasn't reduced modulo that big number, so the results didn't match and all tests failed. Signed-off-by: Wojciech Zmuda <[email protected]> * prover: keccak: account for domain separator size The length of the padded message did not consider size of the domain separation sequence. For input messages length close to the block size (i.e. 1080-1088 bits) the padding size was calculated to be 1088. Writing domain separator resulted in out of bounds write. Adjust padding size, so that input of size close to block size results in extending the message by the length of another block size, leaving room for the domain separator. Rename variable paddingSize to paddedSize, as this variable reflects the size of the input after size adjustments, not just the size of the padding. Signed-off-by: Wojciech Zmuda <[email protected]> --------- Signed-off-by: Wojciech Zmuda <[email protected]>
worldcoin · Mar 21, 2024 · 964cded · 964cded
1 parent 7f6099c
commit 964cded
Show file tree

Hide file tree

Showing 2 changed files with 51 additions and 9 deletions.
diff --git a/prover/keccak/keccak.go b/prover/keccak/keccak.go
@@ -12,14 +12,16 @@ import (
 
 const laneSize = 64
 const stateSize = 5
+const blockSize = 1088
+const domainSeparatorSize = 8
 
 func NewKeccak256(api frontend.API, inputSize int, data ...frontend.Variable) []frontend.Variable {
 	hash := abstractor.Call1(api, KeccakGadget{
 		InputSize:       inputSize,
 		InputData:       data,
 		OutputSize:      256,
 		Rounds:          24,
-		BlockSize:       1088,
+		BlockSize:       blockSize,
 		RotationOffsets: R,
 		RoundConstants:  RC,
 		Domain:          0x01,
@@ -33,7 +35,7 @@ func NewSHA3_256(api frontend.API, inputSize int, data ...frontend.Variable) []f
 		InputData:       data,
 		OutputSize:      256,
 		Rounds:          24,
-		BlockSize:       1088,
+		BlockSize:       blockSize,
 		RotationOffsets: R,
 		RoundConstants:  RC,
 		Domain:          0x06,
@@ -145,23 +147,23 @@ type KeccakGadget struct {
 
 func (g KeccakGadget) DefineGadget(api frontend.API) interface{} {
 	// Padding
-	paddingSize := int(math.Ceil(float64(g.InputSize)/float64(g.BlockSize))) * g.BlockSize
+	paddedSize := int(math.Ceil(float64(g.InputSize+domainSeparatorSize)/float64(g.BlockSize))) * g.BlockSize
 	if len(g.InputData) == 0 {
-		paddingSize = g.BlockSize
+		paddedSize = g.BlockSize
 	}
 
-	P := make([]frontend.Variable, paddingSize)
+	P := make([]frontend.Variable, paddedSize)
 	for i := 0; i < len(g.InputData); i += 1 {
 		P[i] = g.InputData[i]
 	}
 
 	// write domain separator
-	for i := 0; i < 8; i += 1 {
+	for i := 0; i < domainSeparatorSize; i += 1 {
 		P[i+len(g.InputData)] = (g.Domain >> i) & 1
 	}
 
 	// fill with zero bytes
-	for i := len(g.InputData) + 8; i < len(P); i += 1 {
+	for i := len(g.InputData) + domainSeparatorSize; i < len(P); i += 1 {
 		P[i] = 0
 	}
 

diff --git a/prover/keccak/keccak_test.go b/prover/keccak/keccak_test.go
@@ -34,6 +34,18 @@ func (circuit *TestKeccakCircuit2) Define(api frontend.API) error {
 	return nil
 }
 
+type TestKeccakCircuitBlockSize struct {
+	Input [blockSize]frontend.Variable `gnark:"input"`
+	Hash  frontend.Variable            `gnark:",public"`
+}
+
+func (circuit *TestKeccakCircuitBlockSize) Define(api frontend.API) error {
+	hash := NewKeccak256(api, len(circuit.Input), circuit.Input[:]...)
+	sum := api.FromBinary(hash...)
+	api.AssertIsEqual(circuit.Hash, sum)
+	return nil
+}
+
 type TestSHACircuit struct {
 	Input [0]frontend.Variable `gnark:"input"`
 	Hash  frontend.Variable    `gnark:",public"`
@@ -63,9 +75,20 @@ func TestKeccak(t *testing.T) {
 		Hash:  bigIntLE("0xc5d2460186f7233c927e7db2dcc703c0e500b653ca82273b7bfad8045d85a470"),
 	}, test.WithBackends(backend.GROTH16), test.WithCurves(ecc.BN254))
 
+	// Keccak: input equal block size
+	var circuit3 TestKeccakCircuitBlockSize
+	var inputArray [1088]frontend.Variable
+	fillArray(1, inputArray[:])
+	assert.ProverSucceeded(
+		&circuit3, &TestKeccakCircuitBlockSize{
+			Input: inputArray,
+			Hash:  bigIntLE("0x2d417340362cd4144efbf52adc1bfb7a4b40254f55f3b0f09efa6a1ef299b51a"),
+		}, test.WithBackends(backend.GROTH16), test.WithCurves(ecc.BN254),
+	)
+
 	// SHA3: hash empty input
-	var circuit3 TestSHACircuit
-	assert.ProverSucceeded(&circuit3, &TestSHACircuit{
+	var circuit4 TestSHACircuit
+	assert.ProverSucceeded(&circuit4, &TestSHACircuit{
 		Input: [0]frontend.Variable{},
 		Hash:  bigIntLE("0xa7ffc6f8bf1ed76651c14756a061d662f580ff4de43b49fa82d80a4b80f8434a"),
 	}, test.WithBackends(backend.GROTH16), test.WithCurves(ecc.BN254))
@@ -83,5 +106,22 @@ func bigIntLE(s string) big.Int {
 	}
 
 	bi.SetBytes(b)
+
+	// Reduce the number by BN254 group order, because the circuit does the same
+	modulus, ok := new(big.Int).SetString(
+		"21888242871839275222246405745257275088548364400416034343698204186575808495617", 10,
+	)
+	if !ok {
+		panic("can't set big int to BN254 group order")
+	}
+	bi.Mod(&bi, modulus)
+
 	return bi
 }
+
+// Fill an array with a specific value.
+func fillArray(value int, inputArray []frontend.Variable) {
+	for i := range inputArray {
+		inputArray[i] = value
+	}
+}