Fix bugs in Zero Balance Verification

pkg/zkproofs/ciphertext_ciphertext_equality_test.go

@@ -452,3 +452,40 @@ func TestVerifyCiphertextCiphertextEquality_InvalidInputs(t *testing.T) {
 		require.False(t, valid, "Proof verification should fail for proof params with zero value")
 	})
 }
+
+// Test that the proof is still valid for cases where Ciphertext.D is the identity point.
+func TestCiphertextCiphertextEqualityProof_IdentityD(t *testing.T) {
+	// Key generation
+	sourcePrivateKey, _ := testutils.GenerateKey()
+	destPrivateKey, _ := testutils.GenerateKey()
+	eg := elgamal.NewTwistedElgamal()
+	sourceKeypair, _ := eg.KeyGen(*sourcePrivateKey, TestDenom)
+	destinationKeypair, _ := eg.KeyGen(*destPrivateKey, TestDenom)
+
+	// Encrypt the source amount
+	sourceCiphertext, _, err := eg.Encrypt(sourceKeypair.PublicKey, big.NewInt(100))
+	assert.NoError(t, err, "Encryption should succeed for sourceCiphertext")
+
+	destCiphertextZero, destZeroRandomness, _ := eg.Encrypt(destinationKeypair.PublicKey, big.NewInt(0))
+	zeroCiphertext, _ := elgamal.SubtractCiphertext(sourceCiphertext, sourceCiphertext)
+	zeroScalar := curves.ED25519().Scalar.Zero()
+	// Generate the proof
+	proof, err := NewCiphertextCiphertextEqualityProof(
+		sourceKeypair,
+		&destinationKeypair.PublicKey,
+		zeroCiphertext,
+		&destZeroRandomness,
+		&zeroScalar,
+	)
+	assert.NoError(t, err, "Proof generation should not fail")
+
+	// VerifyCipherCipherEquality the proof
+	valid := VerifyCiphertextCiphertextEquality(
+		proof,
+		&sourceKeypair.PublicKey,
+		&destinationKeypair.PublicKey,
+		zeroCiphertext,
+		destCiphertextZero,
+	)
+	assert.True(t, valid, "Proof verification should not fail")
+}

pkg/zkproofs/ciphertext_commitment_equality_test.go

@@ -327,3 +327,36 @@ func TestVerifyCiphertextCommitmentEquality_InvalidInput(t *testing.T) {
 		require.False(t, valid, "Proof verification should fail for nil Pedersen commitment")
 	})
 }
+
+// Test that the proof is still valid for cases where Ciphertext.D is the identity point.
+func TestCiphertextCommitmentEqualityProof_IdentityD(t *testing.T) {
+	// Key generation
+	sourcePrivateKey, _ := testutils.GenerateKey()
+	eg := elgamal.NewTwistedElgamal()
+	sourceKeypair, _ := eg.KeyGen(*sourcePrivateKey, TestDenom)
+
+	// Encrypt the source amount
+	sourceCiphertext, _, err := eg.Encrypt(sourceKeypair.PublicKey, big.NewInt(100))
+	assert.NoError(t, err, "Encryption should succeed for sourceCiphertext")
+
+	zeroCommitment, zeroRandomness, _ := eg.Encrypt(sourceKeypair.PublicKey, big.NewInt(0))
+	zeroCiphertext, _ := elgamal.SubtractCiphertext(sourceCiphertext, sourceCiphertext)
+	zeroScalar := curves.ED25519().Scalar.Zero()
+	// Generate the proof
+	proof, err := NewCiphertextCommitmentEqualityProof(
+		sourceKeypair,
+		zeroCiphertext,
+		&zeroRandomness,
+		&zeroScalar,
+	)
+	assert.NoError(t, err, "Proof generation should not fail")
+
+	// VerifyCipherCipherEquality the proof
+	valid := VerifyCiphertextCommitmentEquality(
+		proof,
+		&sourceKeypair.PublicKey,
+		zeroCiphertext,
+		&zeroCommitment.C,
+	)
+	assert.True(t, valid, "Proof verification should not fail")
+}

pkg/zkproofs/zero_balance.go

@@ -120,7 +120,8 @@ func (p *ZeroBalanceProof) validateContents() bool {
 		return false
 	}
 
-	if p.Yp.IsIdentity() || p.Yd.IsIdentity() || p.Z.IsZero() {
+	// We leave out p.Yd since it is valid for Yd to be zero (Ciphertext.D can be Identity point)
+	if p.Yp.IsIdentity() || p.Z.IsZero() {
 		return false
 	}
 

pkg/zkproofs/zero_balance_test.go

@@ -295,3 +295,24 @@ func TestVerifyZeroProof_InvalidInput(t *testing.T) {
 	valid = VerifyZeroBalance(&invalidParamsProof, &keypair.PublicKey, ciphertext)
 	require.False(t, valid, "Verification should fail with invalid proof params")
 }
+
+func TestZeroBalanceProof_IdentityD(t *testing.T) {
+	// Setup keypair
+	privateKey, _ := testutils.GenerateKey()
+
+	eg := elgamal.NewTwistedElgamal()
+	keypair, _ := eg.KeyGen(*privateKey, TestDenom)
+
+	ciphertext, _, err := eg.Encrypt(keypair.PublicKey, big.NewInt(100))
+	require.NoError(t, err, "Failed to encrypt amount")
+
+	ciphertextZero, _ := elgamal.SubtractCiphertext(ciphertext, ciphertext)
+
+	// Generate ZeroBalanceProof
+	proof, err := NewZeroBalanceProof(keypair, ciphertextZero)
+	require.NoError(t, err, "Failed to generate proof")
+
+	// Verify the proof
+	valid := VerifyZeroBalance(proof, &keypair.PublicKey, ciphertextZero)
+	require.True(t, valid, "Proof should be valid")
+}