From 053ae45b59bef06f1ec0f92c6d46945470c863e1 Mon Sep 17 00:00:00 2001
From: Piotr Dyraga <piotr.dyraga@akena.engineering>
Date: Thu, 28 Jun 2018 14:28:05 +0200
Subject: [PATCH] Let to pass an arbitrary number of cyphers to PublicKey.Add

`PublicKey.Add` accepts now an arbitrary number of `*Cypher`s.
Previously, it was possible to add only two `*Cypher*` instances and it
was quite unwieldy for a lot of use cases.

The current implementation is backward compatible. It's still possible
to pass just two `*Cypher` arguments, so every project referencing the
library will observe no compilation issues.
---
 paillier.go      | 20 +++++++++++++++-----
 paillier_test.go | 20 ++++++++++++--------
 2 files changed, 27 insertions(+), 13 deletions(-)

diff --git a/paillier.go b/paillier.go
index 9b6fba2..62d3e35 100644
--- a/paillier.go
+++ b/paillier.go
@@ -40,14 +40,24 @@ func (pub *PublicKey) Encrypt(m *big.Int, random io.Reader) (*Cypher, error) {
 	return &Cypher{new(big.Int).Mod(new(big.Int).Mul(rn, gm), nSquare)}, nil
 }
 
-// Takes two cypher texts and returns a 3rd one that encode
-// the sum of the two plain texts.
+// Add takes an arbitrary number of cyphertexts and returns one that encodes
+// their sum.
 //
 // It's possible because Paillier is a homomorphic encryption scheme, where
 // E(m1) * E(m2) = E(m1 + m2)
-func (this *PublicKey) Add(cypher1, cypher2 *Cypher) *Cypher {
-	m := new(big.Int).Mul(cypher1.C, cypher2.C)
-	return &Cypher{new(big.Int).Mod(m, this.GetNSquare())}
+func (pub *PublicKey) Add(cypher ...*Cypher) *Cypher {
+	accumulator := big.NewInt(1)
+
+	for _, c := range cypher {
+		accumulator = new(big.Int).Mod(
+			new(big.Int).Mul(accumulator, c.C),
+			pub.GetNSquare(),
+		)
+	}
+
+	return &Cypher{
+		C: accumulator,
+	}
 }
 
 type PrivateKey struct {
diff --git a/paillier_test.go b/paillier_test.go
index 5e9fe9b..8918e27 100644
--- a/paillier_test.go
+++ b/paillier_test.go
@@ -64,13 +64,17 @@ func TestEncryptDecryptSmall(t *testing.T) {
 	}
 }
 
-func TestAddCypher(t *testing.T) {
-	privateKey := CreatePrivateKey(big.NewInt(13), big.NewInt(11))
-	cypher1, _ := privateKey.Encrypt(big.NewInt(12), rand.Reader)
-	cypher2, _ := privateKey.Encrypt(big.NewInt(13), rand.Reader)
-	cypher3 := privateKey.Add(cypher1, cypher2)
-	m := privateKey.Decrypt(cypher3)
-	if !reflect.DeepEqual(m, big.NewInt(25)) {
-		t.Error(m)
+func TestAddCyphers(t *testing.T) {
+	privateKey := CreatePrivateKey(big.NewInt(17), big.NewInt(13))
+
+	cypher1, _ := privateKey.Encrypt(big.NewInt(5), rand.Reader)
+	cypher2, _ := privateKey.Encrypt(big.NewInt(6), rand.Reader)
+	cypher3, _ := privateKey.Encrypt(big.NewInt(7), rand.Reader)
+	cypher4, _ := privateKey.Encrypt(big.NewInt(8), rand.Reader)
+	cypher5 := privateKey.Add(cypher1, cypher2, cypher3, cypher4)
+
+	m := privateKey.Decrypt(cypher5)
+	if !reflect.DeepEqual(m, big.NewInt(26)) {
+		t.Errorf("Unexpected decrypted value [%v]", m)
 	}
 }