From 93f329fc4b7324177ea2e822de3d155c936ff8c7 Mon Sep 17 00:00:00 2001
From: Sylvain Pelissier <sylvain.pelissier@gmail.com>
Date: Fri, 10 Nov 2023 11:25:40 +0100
Subject: [PATCH] Encryption optimization

---
 paillier.go      | 19 ++++++++++---------
 paillier_test.go | 29 +++++++++++++++++++++++++++++
 2 files changed, 39 insertions(+), 9 deletions(-)

diff --git a/paillier.go b/paillier.go
index a7153e8..bbcc7c8 100644
--- a/paillier.go
+++ b/paillier.go
@@ -39,10 +39,11 @@ func (pk *PublicKey) EncryptWithR(m *big.Int, r *big.Int) (*Cypher, error) {
 	nSquare := pk.GetNSquare()
 
 	// g is _always_ equal n+1
+	// g^m = (n+1)^m = 1 + mn
 	// Threshold encryption is safe only for g=n+1 choice.
 	// See [DJN 10], section 5.1
-	g := new(big.Int).Add(pk.N, big.NewInt(1))
-	gm := new(big.Int).Exp(g, m, nSquare)
+	gm := new(big.Int).Mul(pk.N, m)
+	gm = gm.Mod(new(big.Int).Add(gm, big.NewInt(1)), nSquare)
 	rn := new(big.Int).Exp(r, pk.N, nSquare)
 	return &Cypher{new(big.Int).Mod(new(big.Int).Mul(rn, gm), nSquare)}, nil
 }
@@ -151,14 +152,14 @@ func computePhi(p, q *big.Int) *big.Int {
 // except that instead of generating Lambda private key component from LCM
 // of p and q we use Euler's totient function as suggested in [KL 08].
 //
-//     [KL 08]:  Jonathan Katz, Yehuda Lindell, (2008)
-//               Introduction to Modern Cryptography: Principles and Protocols,
-//               Chapman & Hall/CRC
+//	[KL 08]:  Jonathan Katz, Yehuda Lindell, (2008)
+//	          Introduction to Modern Cryptography: Principles and Protocols,
+//	          Chapman & Hall/CRC
 //
-//     [DJN 10]: Ivan Damgard, Mads Jurik, Jesper Buus Nielsen, (2010)
-//               A Generalization of Paillier’s Public-Key System
-//               with Applications to Electronic Voting
-//               Aarhus University, Dept. of Computer Science, BRICS
+//	[DJN 10]: Ivan Damgard, Mads Jurik, Jesper Buus Nielsen, (2010)
+//	          A Generalization of Paillier’s Public-Key System
+//	          with Applications to Electronic Voting
+//	          Aarhus University, Dept. of Computer Science, BRICS
 func CreatePrivateKey(p, q *big.Int) *PrivateKey {
 	n := new(big.Int).Mul(p, q)
 	lambda := computePhi(p, q)
diff --git a/paillier_test.go b/paillier_test.go
index 5df9256..232fd2d 100644
--- a/paillier_test.go
+++ b/paillier_test.go
@@ -65,6 +65,35 @@ func TestEncryptDecryptSmall(t *testing.T) {
 	}
 }
 
+func TestEncrypt512(t *testing.T) {
+	p, _ := rand.Prime(rand.Reader, 512)
+	q, _ := rand.Prime(rand.Reader, 512)
+	privateKey := CreatePrivateKey(p, q)
+
+	for i := 1; i < 10; i++ {
+		initialValue, _ := rand.Int(rand.Reader, privateKey.PublicKey.N)
+		cypher, err := privateKey.Encrypt(initialValue, rand.Reader)
+		if err != nil {
+			t.Error(err)
+		}
+		returnedValue := privateKey.Decrypt(cypher)
+		if initialValue.Cmp(returnedValue) != 0 {
+			t.Error("wrong decryption ", returnedValue, " is not ", initialValue)
+		}
+	}
+}
+
+func BenchmarkEncrypt(b *testing.B) {
+	p, _ := rand.Prime(rand.Reader, 2048)
+	q, _ := rand.Prime(rand.Reader, 2048)
+	privateKey := CreatePrivateKey(p, q)
+	// run the Encrypt function b.N times
+	for n := 0; n < b.N; n++ {
+		initialValue, _ := rand.Int(rand.Reader, privateKey.PublicKey.N)
+		privateKey.Encrypt(initialValue, rand.Reader)
+	}
+}
+
 func TestCheckPlaintextSpace(t *testing.T) {
 	p := big.NewInt(13)
 	q := big.NewInt(11)