update udp encryption API to take in customized indexes in fbpcf (#518)

Summary:
Pull Request resolved: #518

We will use 3 diffs to change the udp encryption API to support customized indexes.
This diff makes the change in fbpcf, while maintaining a temporary API for backward compatibility.
In the next diff, we will make the change in fbpcs, which complete the change.
In the last diff, we will clean up the temporary APIs in fbpcf.

Reviewed By: robotal

Differential Revision: D44242200

fbshipit-source-id: a36346f269a7d3a6727a214cc7f5e181e8f1474d

fbpcf/mpc_std_lib/unified_data_process/data_processor/DataProcessor_impl.h

@@ -9,6 +9,7 @@
 
 #include <emmintrin.h>
 #include <fbpcf/engine/util/util.h>
+#include <numeric>
 #include "fbpcf/engine/util/aes.h"
 #include "fbpcf/mpc_std_lib/aes_circuit/AesCircuitCtr.h"
 #include "fbpcf/mpc_std_lib/unified_data_process/data_processor/DataProcessor.h"
@@ -24,9 +25,12 @@ DataProcessor<schedulerId>::processMyData(
     size_t outputSize) {
   size_t dataSize = plaintextData.size();
   size_t dataWidth = plaintextData.at(0).size();
+  std::vector<uint64_t> indexes(plaintextData.size());
+  // generate 0 to n-1 vector
+  std::iota(indexes.begin(), indexes.end(), 0);
 
   encrypter_.prepareToProcessMyData(dataWidth);
-  encrypter_.processMyData(plaintextData);
+  encrypter_.processMyData(plaintextData, indexes);
 
   // 1b. (peer)receive encryted data from peer
   // 2b. (peer)pick desired ciphertext blocks

fbpcf/mpc_std_lib/unified_data_process/data_processor/IUdpEncryption.h

@@ -30,7 +30,15 @@ class IUdpEncryption {
    * "getExpandedKey" to retrieve the expanded key for decryption later.
    */
   virtual void processMyData(
-      const std::vector<std::vector<unsigned char>>& plaintextData) = 0;
+      const std::vector<std::vector<unsigned char>>& plaintextData,
+      const std::vector<uint64_t>& indexes) = 0;
+
+  // a temp API to maintain backward compatibility
+  virtual void processMyData(
+      const std::vector<std::vector<unsigned char>>& plaintextData) {
+    processMyData(
+        plaintextData, std::vector<uint64_t>(plaintextData.size(), 0));
+  }
 
   virtual std::vector<__m128i> getExpandedKey() = 0;
 

fbpcf/mpc_std_lib/unified_data_process/data_processor/UdpEncryption.cpp

@@ -31,7 +31,8 @@ void UdpEncryption::prepareToProcessMyData(size_t myDataWidth) {
 }
 
 void UdpEncryption::processMyData(
-    const std::vector<std::vector<unsigned char>>& plaintextData) {
+    const std::vector<std::vector<unsigned char>>& plaintextData,
+    const std::vector<uint64_t>& indexes) {
   if (statusOfProcessingMyData_ != Status::inProgress) {
     throw std::runtime_error("Can't call procesMyData before preparation!");
   }
@@ -43,10 +44,13 @@ void UdpEncryption::processMyData(
         "Inconsistent data width, expecting " + std::to_string(myDataWidth_) +
         " but get " + std::to_string(plaintextData.at(0).size()));
   }
+  /* will use indexes to replace myDataIndexOffset_ once the underlying change
+   * is done*/
   auto [ciphertext, nonce] =
       UdpUtil::localEncryption(plaintextData, prgKey_, myDataIndexOffset_);
   agent_->send(nonce);
   for (size_t i = 0; i < ciphertext.size(); i++) {
+    agent_->sendSingleT<uint64_t>(indexes.at(i));
     agent_->send(ciphertext.at(i));
   }
   myDataIndexOffset_ += plaintextData.size();
@@ -85,8 +89,10 @@ void UdpEncryption::processPeerData(size_t dataSize) {
   }
 
   for (size_t i = 0; i < dataSize; i++) {
+    auto index = agent_->receiveSingleT<uint64_t>();
     auto ciphertext = agent_->receive(peerDataWidth_);
-    auto pos = indexToOrderMap_.find(i + peerDataIndexOffset_);
+    auto pos = indexToOrderMap_.find(
+        i + peerDataIndexOffset_ /* will use index instead in next diff*/);
     if (pos != indexToOrderMap_.end()) {
       // this ciphertext should be picked up
       cherryPickedEncryption_.at(pos->second) = std::move(ciphertext);

fbpcf/mpc_std_lib/unified_data_process/data_processor/UdpEncryption.h

@@ -38,7 +38,8 @@ class UdpEncryption final : public IUdpEncryption {
   // with "ProcessPeerData" on peer's side. If this API is ever called, calling
   // "getExpandedKey" to retrive the expanded key for decryption later.
   void processMyData(
-      const std::vector<std::vector<unsigned char>>& plaintextData) override;
+      const std::vector<std::vector<unsigned char>>& plaintextData,
+      const std::vector<uint64_t>& indexes) override;
 
   std::vector<__m128i> getExpandedKey() override {
     if (statusOfProcessingMyData_ != Status::inProgress) {

fbpcf/mpc_std_lib/unified_data_process/data_processor/test/DataProcessorTest.cpp

@@ -191,8 +191,13 @@ void testUdpEncryptionAndDecryptionObjects(
                   const std::vector<uint64_t>& indexes,
                   const std::vector<size_t>& sizes) {
     udpEnc->prepareToProcessMyData(dataWidth);
+    size_t myDataIndexOffset = 0;
     for (size_t i = 0; i < plaintextDataInShards.size(); i++) {
-      udpEnc->processMyData(plaintextDataInShards.at(i));
+      std::vector<uint64_t> u64indexes(plaintextDataInShards.at(i).size());
+      // generate 0 to n-1 vector
+      std::iota(u64indexes.begin(), u64indexes.end(), myDataIndexOffset);
+      myDataIndexOffset += plaintextDataInShards.at(i).size();
+      udpEnc->processMyData(plaintextDataInShards.at(i), u64indexes);
     };
     udpEnc->prepareToProcessPeerData(dataWidth, indexes);
     for (size_t i = 0; i < sizes.size(); i++) {
@@ -249,8 +254,13 @@ void testUdpEncryptionAndDecryptionObjects(
       udpEnc->processPeerData(sizes.at(i));
     }
     udpEnc->prepareToProcessMyData(dataWidth);
+    size_t myDataIndexOffset = 0;
     for (size_t i = 0; i < plaintextDataInShards.size(); i++) {
-      udpEnc->processMyData(plaintextDataInShards.at(i));
+      std::vector<uint64_t> u64indexes(plaintextDataInShards.at(i).size());
+      // generate 0 to n-1 vector
+      std::iota(u64indexes.begin(), u64indexes.end(), myDataIndexOffset);
+      myDataIndexOffset += plaintextDataInShards.at(i).size();
+      udpEnc->processMyData(plaintextDataInShards.at(i), u64indexes);
     };
 
     auto [intersection, nonces, pickedIndexes] = udpEnc->getProcessedData();

fbpcf/mpc_std_lib/unified_data_process/data_processor/test/UdpEncryptionMock.h

@@ -19,6 +19,12 @@ class UdpEncryptionMock final : public IUdpEncryption {
  public:
   MOCK_METHOD(void, prepareToProcessMyData, (size_t));
 
+  MOCK_METHOD(
+      void,
+      processMyData,
+      (const std::vector<std::vector<unsigned char>>&,
+       const std::vector<uint64_t>& indexes));
+
   MOCK_METHOD(
       void,
       processMyData,