Multi Nonce Implementation for Gnark ZK Proofs

.github/workflows/test.yaml

@@ -12,7 +12,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        node-version: [22]
+        node-version: [24]
         package-dir: [js]
 
     steps:

gnark/circuits/aesV2/aes128_test.go

@@ -23,14 +23,28 @@ func TestAES128(t *testing.T) {
 	Nonce := "006CB6DBC0543B59DA48D90B"
 	Counter := 1
 
-	// calculate ciphertext ourselves
+	// calculate ciphertext ourselves for each block
 	block, err := aes.NewCipher(mustHex(key))
 	if err != nil {
 		panic(err)
 	}
-	ctr := cipher.NewCTR(block, append(mustHex(Nonce), binary.BigEndian.AppendUint32(nil, uint32(Counter))...))
-	ciphertext := make([]byte, len(mustHex(plaintext)))
-	ctr.XORKeyStream(ciphertext, mustHex(plaintext))
+	plaintextBytes := mustHex(plaintext)
+	ciphertext := make([]byte, len(plaintextBytes))
+
+	// Process each block with its own counter
+	blockSize := 16
+	for b := 0; b < BLOCKS; b++ {
+		start := b * blockSize
+		end := start + blockSize
+		if end > len(plaintextBytes) {
+			end = len(plaintextBytes)
+		}
+
+		// Create CTR mode with the counter for this block
+		iv := append(mustHex(Nonce), binary.BigEndian.AppendUint32(nil, uint32(Counter+b))...)
+		ctr := cipher.NewCTR(block, iv)
+		ctr.XORKeyStream(ciphertext[start:end], plaintextBytes[start:end])
+	}
 
 	keyAssign := StrToIntSlice(key, true)
 	ptAssign := StrToIntSlice(plaintext, true)
@@ -41,8 +55,8 @@ func TestAES128(t *testing.T) {
 	assignment := AESCircuit{
 		AESBaseCircuit: AESBaseCircuit{
 			Key:     make([]frontend.Variable, 16),
-			Counter: Counter,
-			Nonce:   [12]frontend.Variable{},
+			Counter: [BLOCKS]frontend.Variable{},
+			Nonce:   [BLOCKS][12]frontend.Variable{},
 			In:      [BLOCKS * 16]frontend.Variable{},
 		},
 		Out: [BLOCKS * 16]frontend.Variable{},
@@ -59,15 +73,20 @@ func TestAES128(t *testing.T) {
 		assignment.Out[i] = ctAssign[i]
 	}
 
-	for i := 0; i < len(nonceAssign); i++ {
-		assignment.Nonce[i] = nonceAssign[i]
+	// Set the same nonce for all blocks in this test
+	for b := 0; b < BLOCKS; b++ {
+		for i := 0; i < len(nonceAssign); i++ {
+			assignment.Nonce[b][i] = nonceAssign[i]
+		}
+		// Set counter for each block (incrementing)
+		assignment.Counter[b] = Counter + b
 	}
 
 	assert.CheckCircuit(&AESCircuit{
 		AESBaseCircuit: AESBaseCircuit{
 			Key:     make([]frontend.Variable, 16),
-			Counter: Counter,
-			Nonce:   [12]frontend.Variable{},
+			Counter: [BLOCKS]frontend.Variable{},
+			Nonce:   [BLOCKS][12]frontend.Variable{},
 			In:      [BLOCKS * 16]frontend.Variable{},
 		},
 		Out: [BLOCKS * 16]frontend.Variable{},

gnark/circuits/aesV2/aes256_test.go

@@ -23,14 +23,28 @@ func TestAES256(t *testing.T) {
 	Nonce := "00FAAC24C1585EF15A43D875"
 	Counter := 1
 
-	// calculate ciphertext ourselves
+	// calculate ciphertext ourselves for each block
 	block, err := aes.NewCipher(mustHex(key))
 	if err != nil {
 		panic(err)
 	}
-	ctr := cipher.NewCTR(block, append(mustHex(Nonce), binary.BigEndian.AppendUint32(nil, uint32(Counter))...))
-	ciphertext := make([]byte, len(mustHex(plaintext)))
-	ctr.XORKeyStream(ciphertext, mustHex(plaintext))
+	plaintextBytes := mustHex(plaintext)
+	ciphertext := make([]byte, len(plaintextBytes))
+
+	// Process each block with its own counter
+	blockSize := 16
+	for b := 0; b < BLOCKS; b++ {
+		start := b * blockSize
+		end := start + blockSize
+		if end > len(plaintextBytes) {
+			end = len(plaintextBytes)
+		}
+
+		// Create CTR mode with the counter for this block
+		iv := append(mustHex(Nonce), binary.BigEndian.AppendUint32(nil, uint32(Counter+b))...)
+		ctr := cipher.NewCTR(block, iv)
+		ctr.XORKeyStream(ciphertext[start:end], plaintextBytes[start:end])
+	}
 
 	keyAssign := StrToIntSlice(key, true)
 	ptAssign := StrToIntSlice(plaintext, true)
@@ -41,8 +55,8 @@ func TestAES256(t *testing.T) {
 	assignment := AESCircuit{
 		AESBaseCircuit: AESBaseCircuit{
 			Key:     make([]frontend.Variable, 32),
-			Counter: Counter,
-			Nonce:   [12]frontend.Variable{},
+			Counter: [BLOCKS]frontend.Variable{},
+			Nonce:   [BLOCKS][12]frontend.Variable{},
 			In:      [BLOCKS * 16]frontend.Variable{},
 		},
 		Out: [BLOCKS * 16]frontend.Variable{},
@@ -59,15 +73,20 @@ func TestAES256(t *testing.T) {
 		assignment.Out[i] = ciphertext[i]
 	}
 
-	for i := 0; i < len(nonceAssign); i++ {
-		assignment.Nonce[i] = nonceAssign[i]
+	// Set the same nonce for all blocks in this test
+	for b := 0; b < BLOCKS; b++ {
+		for i := 0; i < len(nonceAssign); i++ {
+			assignment.Nonce[b][i] = nonceAssign[i]
+		}
+		// Set counter for each block (incrementing)
+		assignment.Counter[b] = Counter + b
 	}
 
 	assert.CheckCircuit(&AESCircuit{
 		AESBaseCircuit: AESBaseCircuit{
 			Key:     make([]frontend.Variable, 32),
-			Counter: Counter,
-			Nonce:   [12]frontend.Variable{},
+			Counter: [BLOCKS]frontend.Variable{},
+			Nonce:   [BLOCKS][12]frontend.Variable{},
 			In:      [BLOCKS * 16]frontend.Variable{},
 		},
 		Out: [BLOCKS * 16]frontend.Variable{},

gnark/circuits/aesV2/common.go

@@ -12,16 +12,16 @@ const NB = 4
 
 type AESBaseCircuit struct {
 	Key     []frontend.Variable
-	Nonce   [12]frontend.Variable          `gnark:",public"`
-	Counter frontend.Variable              `gnark:",public"`
+	Nonce   [BLOCKS][12]frontend.Variable  `gnark:",public"`
+	Counter [BLOCKS]frontend.Variable      `gnark:",public"`
 	In      [BLOCKS * 16]frontend.Variable `gnark:",public"`
 }
 
 type AESGadget struct {
 	api            frontend.API
-	sbox           *logderivlookup.Table
+	sbox           logderivlookup.Table
 	RCon           [11]frontend.Variable
-	t0, t1, t2, t3 *logderivlookup.Table
+	t0, t1, t2, t3 logderivlookup.Table
 	keySize        int
 }
 
@@ -53,25 +53,24 @@ func (aes *AESBaseCircuit) Define(api frontend.API, out [BLOCKS * 16]frontend.Va
 		return errors.New("key size must be 16 or 32")
 	}
 
-	counter := aes.Counter
 	var counterBlock [16]frontend.Variable
 
 	gAes := NewAESGadget(api, keySize)
 
-	for i := 0; i < 12; i++ {
-		counterBlock[i] = aes.Nonce[i]
-	}
 	for b := 0; b < BLOCKS; b++ {
-		gAes.createIV(counter, counterBlock[:])
+		// Use per-block nonce
+		for i := 0; i < 12; i++ {
+			counterBlock[i] = aes.Nonce[b][i]
+		}
+		// Use per-block counter
+		gAes.createIV(aes.Counter[b], counterBlock[:])
 		// encrypt counter under key
 		keystream := gAes.Encrypt(aes.Key, counterBlock)
 
 		for i := 0; i < 16; i++ {
 			api.AssertIsEqual(out[b*16+i], gAes.VariableXor(keystream[i], aes.In[b*16+i], 8))
 		}
-		counter = api.Add(counter, 1)
 	}
-	api.AssertIsEqual(counter, api.Add(aes.Counter, BLOCKS))
 
 	return nil
 }
@@ -135,7 +134,7 @@ func (aes *AESGadget) ShiftSub(state [16]frontend.Variable) []frontend.Variable
 }
 
 // substitute word with naive lookup of sbox
-func (aes *AESGadget) Subws(sbox *logderivlookup.Table, a ...frontend.Variable) []frontend.Variable {
+func (aes *AESGadget) Subws(sbox logderivlookup.Table, a ...frontend.Variable) []frontend.Variable {
 	return sbox.Lookup(a...)
 }
 

gnark/circuits/aesV2_oprf/aes128_test.go

@@ -32,14 +32,25 @@ func TestAES128(t *testing.T) {
 	plaintext := make([]byte, aes_v2.BLOCKS*16)
 	copy(plaintext[pos:], secretBytes)
 
-	// calculate ciphertext ourselves
+	// calculate ciphertext ourselves for each block
 	block, err := aes.NewCipher(key)
 	if err != nil {
 		panic(err)
 	}
-	ctr := cipher.NewCTR(block, append(mustHex(Nonce), binary.BigEndian.AppendUint32(nil, uint32(Counter))...))
 	ciphertext := make([]byte, len(plaintext))
-	ctr.XORKeyStream(ciphertext, plaintext)
+	blockSize := 16
+
+	for b := 0; b < aes_v2.BLOCKS; b++ {
+		start := b * blockSize
+		end := start + blockSize
+		if end > len(plaintext) {
+			end = len(plaintext)
+		}
+
+		iv := append(mustHex(Nonce), binary.BigEndian.AppendUint32(nil, uint32(Counter+b))...)
+		ctr := cipher.NewCTR(block, iv)
+		ctr.XORKeyStream(ciphertext[start:end], plaintext[start:end])
+	}
 
 	nonceAssign := mustHex(Nonce)
 
@@ -71,8 +82,8 @@ func createWitness(d *toprf.Params, bKey []uint8, bNonce []uint8, counter int, c
 	witness := AESTOPRFCircuit{
 		AESBaseCircuit: aes_v2.AESBaseCircuit{
 			Key:     make([]frontend.Variable, 16),
-			Counter: counter,
-			Nonce:   [12]frontend.Variable{},
+			Counter: [aes_v2.BLOCKS]frontend.Variable{},
+			Nonce:   [aes_v2.BLOCKS][12]frontend.Variable{},
 			In:      [aes_v2.BLOCKS * 16]frontend.Variable{},
 		},
 		Out: [aes_v2.BLOCKS * 16]frontend.Variable{},
@@ -99,8 +110,12 @@ func createWitness(d *toprf.Params, bKey []uint8, bNonce []uint8, counter int, c
 		witness.Out[i] = plaintext[i]
 	}
 
-	for i := 0; i < len(bNonce); i++ {
-		witness.Nonce[i] = bNonce[i]
+	// Set per-block nonce and counter
+	for b := 0; b < aes_v2.BLOCKS; b++ {
+		for i := 0; i < len(bNonce); i++ {
+			witness.Nonce[b][i] = bNonce[i]
+		}
+		witness.Counter[b] = counter + b
 	}
 	utils.SetBitmask(witness.Bitmask[:], uint32(pos), uint32(l))
 	return witness

gnark/circuits/aesV2_oprf/aes256_test.go

@@ -32,14 +32,25 @@ func TestAES256(t *testing.T) {
 	plaintext := make([]byte, aes_v2.BLOCKS*16)
 	copy(plaintext[pos:], secretBytes)
 
-	// calculate ciphertext ourselves
+	// calculate ciphertext ourselves for each block
 	block, err := aes.NewCipher(mustHex(key))
 	if err != nil {
 		panic(err)
 	}
-	ctr := cipher.NewCTR(block, append(mustHex(Nonce), binary.BigEndian.AppendUint32(nil, uint32(Counter))...))
 	ciphertext := make([]byte, len(plaintext))
-	ctr.XORKeyStream(ciphertext, plaintext)
+	blockSize := 16
+
+	for b := 0; b < aes_v2.BLOCKS; b++ {
+		start := b * blockSize
+		end := start + blockSize
+		if end > len(plaintext) {
+			end = len(plaintext)
+		}
+
+		iv := append(mustHex(Nonce), binary.BigEndian.AppendUint32(nil, uint32(Counter+b))...)
+		ctr := cipher.NewCTR(block, iv)
+		ctr.XORKeyStream(ciphertext[start:end], plaintext[start:end])
+	}
 
 	keyAssign := mustHex(key)
 	nonceAssign := mustHex(Nonce)
@@ -60,8 +71,8 @@ func createWitness256(d *toprf.Params, bKey []uint8, bNonce []uint8, counter int
 	witness := AESTOPRFCircuit{
 		AESBaseCircuit: aes_v2.AESBaseCircuit{
 			Key:     make([]frontend.Variable, 32),
-			Counter: counter,
-			Nonce:   [12]frontend.Variable{},
+			Counter: [aes_v2.BLOCKS]frontend.Variable{},
+			Nonce:   [aes_v2.BLOCKS][12]frontend.Variable{},
 			In:      [aes_v2.BLOCKS * 16]frontend.Variable{},
 		},
 		Out: [aes_v2.BLOCKS * 16]frontend.Variable{},
@@ -88,8 +99,12 @@ func createWitness256(d *toprf.Params, bKey []uint8, bNonce []uint8, counter int
 		witness.Out[i] = plaintext[i]
 	}
 
-	for i := 0; i < len(bNonce); i++ {
-		witness.Nonce[i] = bNonce[i]
+	// Set per-block nonce and counter
+	for b := 0; b < aes_v2.BLOCKS; b++ {
+		for i := 0; i < len(bNonce); i++ {
+			witness.Nonce[b][i] = bNonce[i]
+		}
+		witness.Counter[b] = counter + b
 	}
 	utils.SetBitmask(witness.Bitmask[:], uint32(pos), uint32(l))
 	return witness