From 9ea470c623711a8d06723569796df74e1ace0179 Mon Sep 17 00:00:00 2001
From: Tarek Elgamal <telgamal@meta.com>
Date: Thu, 30 Mar 2023 14:48:07 -0700
Subject: [PATCH] Adding tests to the primitive operations of HEAggGame (#2254)

Summary:
Pull Request resolved: https://github.com/facebookresearch/fbpcs/pull/2254

Adding two tests for the primitive operations of Homomorphic encryption. The main operations we use in the HE Aggregation Game are:

add_with_ciphertext(c1, c2): It adds two ciphertexts c1 & c2
add_with_plaintext(c1, p1): It adds a ciphertext c1 to a plaintext  p1 and the results goes to a ciphertext

In this diff we test that these two operations return the expected sum after decryption.

These tests is also intended to act as small example of how homomorphic encryption works which can be used in talks/tutorials.

Reviewed By: xyguo, zhangpuhan

Differential Revision: D44549251

fbshipit-source-id: 8b281da75ee54254c55a504b06c37f3c475453dd
---
 .../he_aggregation/test/HEAggGameTest.cpp     | 50 +++++++++++++++++++
 1 file changed, 50 insertions(+)

diff --git a/fbpcs/emp_games/he_aggregation/test/HEAggGameTest.cpp b/fbpcs/emp_games/he_aggregation/test/HEAggGameTest.cpp
index 37434d738..41614a52e 100644
--- a/fbpcs/emp_games/he_aggregation/test/HEAggGameTest.cpp
+++ b/fbpcs/emp_games/he_aggregation/test/HEAggGameTest.cpp
@@ -29,6 +29,10 @@
 #include "fbpcs/emp_games/he_aggregation/HEAggGame.h"
 #include "fbpcs/emp_games/he_aggregation/HEAggOptions.h"
 
+#include "privacy_infra/elgamal/ElGamal.h"
+
+namespace heschme = facebook::privacy_infra::elgamal;
+
 namespace pcf2_he {
 
 std::unordered_map<uint64_t, uint64_t> runGame(
@@ -56,6 +60,52 @@ void verifyOutput(
       folly::toJson(actualOutput), folly::toJson(expectedOutput));
 }
 
+TEST(HEAggGameTest, HECiphertextAdditionTest) {
+  const std::string baseDir_ =
+      private_measurement::test_util::getBaseDirFromPath(__FILE__);
+
+  // Generate private key, public key and decryption table
+  auto sk = heschme::PrivateKey::generate();
+  auto pk = sk.toPublicKey();
+  heschme::initializeElGamalDecryptionTable(FLAGS_decryption_table_size);
+
+  // Encrypt values
+  int x = 111;
+  int y = 222;
+  heschme::Ciphertext c1 = pk.encrypt(x);
+  heschme::Ciphertext c2 = pk.encrypt(y);
+
+  // Perform addition
+  heschme::Ciphertext c3 = heschme::Ciphertext::add_with_ciphertext(c1, c2);
+
+  // Decrypt
+  uint64_t decrypted = sk.decrypt(c3);
+
+  EXPECT_EQ(decrypted, x + y);
+}
+TEST(HEAggGameTest, HEPlaintextAdditionTest) {
+  const std::string baseDir_ =
+      private_measurement::test_util::getBaseDirFromPath(__FILE__);
+
+  // Generate private key, public key and decryption table
+  auto sk = heschme::PrivateKey::generate();
+  auto pk = sk.toPublicKey();
+  heschme::initializeElGamalDecryptionTable(FLAGS_decryption_table_size);
+
+  // Encrypt values
+  int x = 111;
+  int y = 444;
+  heschme::Ciphertext c1 = pk.encrypt(x);
+
+  // Perform addition
+  heschme::Ciphertext c2 = heschme::Ciphertext::add_with_plaintext(c1, y);
+
+  // Decrypt
+  uint64_t decrypted = sk.decrypt(c2);
+
+  EXPECT_EQ(decrypted, x + y);
+}
+
 TEST(HEAggGameTest, HEAggGameCorrectnessTest) {
   const std::string baseDir_ =
       private_measurement::test_util::getBaseDirFromPath(__FILE__);