add IUdpEncryption API for test (#504)

Summary:
Pull Request resolved: #504

Adding an abstract interface for UDP encryption. This allows adding a mock later for this interface

Reviewed By: adshastri

Differential Revision:
D43576519

Privacy Context Container: L416713

fbshipit-source-id: 32191957d42eeb3cd346db61dd31f535cb008173

fbpcf/mpc_std_lib/unified_data_process/data_processor/IUdpEncryption.h

@@ -0,0 +1,57 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <emmintrin.h>
+#include <cstddef>
+#include <vector>
+
+namespace fbpcf::mpc_std_lib::unified_data_process::data_processor {
+
+/**
+ * This is merely an interface to accommodate mocking for test.
+ **/
+class IUdpEncryption {
+ public:
+  virtual ~IUdpEncryption() = default;
+
+  virtual void prepareToProcessMyData(size_t myDataWidth) = 0;
+
+  /**
+   * Process my data via UDP encryption. This API should be called in coordinate
+   * with "ProcessPeerData" on peer's side. If this API is ever called, calling
+   * "getExpandedKey" to retrieve the expanded key for decryption later.
+   */
+  virtual void processMyData(
+      const std::vector<std::vector<unsigned char>>& plaintextData) = 0;
+
+  virtual std::vector<__m128i> getExpandedKey() = 0;
+
+  virtual void prepareToProcessPeerData(
+      size_t peerDataWidth,
+      const std::vector<int32_t>& indexes) = 0;
+
+  /*
+   * process peer data via UDP encryption. This API should be called in
+   * coordinate with "ProcessMyData" on peer's side. This API is ever
+   * called, calling "getProcessedData" to retrive the cherry-picked
+   * encryption later.
+   */
+  virtual void processPeerData(size_t dataSize) = 0;
+
+  struct EncryptionResuts {
+    std::vector<std::vector<unsigned char>> ciphertexts;
+    std::vector<__m128i> nonces;
+    std::vector<int32_t> indexes;
+  };
+
+  // returning the ciphertext, nonce, and index of cherry-picked rows
+  virtual EncryptionResuts getProcessedData() = 0;
+};
+
+} // namespace fbpcf::mpc_std_lib::unified_data_process::data_processor

fbpcf/mpc_std_lib/unified_data_process/data_processor/UdpEncryption.cpp

@@ -14,7 +14,6 @@
 namespace fbpcf::mpc_std_lib::unified_data_process::data_processor {
 
 UdpEncryption::UdpEncryption(
-
     std::unique_ptr<fbpcf::engine::communication::IPartyCommunicationAgent>
         agent)
     : agent_(std::move(agent)),

fbpcf/mpc_std_lib/unified_data_process/data_processor/UdpEncryption.h

@@ -10,6 +10,7 @@
 #include <emmintrin.h>
 #include "fbpcf/engine/communication/IPartyCommunicationAgent.h"
 #include "fbpcf/engine/util/util.h"
+#include "fbpcf/mpc_std_lib/unified_data_process/data_processor/IUdpEncryption.h"
 #include "fbpcf/mpc_std_lib/unified_data_process/data_processor/UdpUtil.h"
 
 namespace fbpcf::mpc_std_lib::unified_data_process::data_processor {
@@ -22,7 +23,7 @@ class DataProcessor;
  *be passed into this object in batches. This object is not thread-safe but
  *it will spin up multiple threads internally.
  **/
-class UdpEncryption {
+class UdpEncryption final : public IUdpEncryption {
   template <int schedulerId>
   friend class DataProcessor;
 
@@ -31,15 +32,15 @@ class UdpEncryption {
       std::unique_ptr<fbpcf::engine::communication::IPartyCommunicationAgent>
           agent);
 
-  void prepareToProcessMyData(size_t myDataWidth);
+  void prepareToProcessMyData(size_t myDataWidth) override;
 
   // process my data via UDP encryption. This API should be called in coordinate
   // 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);
+      const std::vector<std::vector<unsigned char>>& plaintextData) override;
 
-  std::vector<__m128i> getExpandedKey() {
+  std::vector<__m128i> getExpandedKey() override {
     if (statusOfProcessingMyData_ != Status::inProgress) {
       throw std::runtime_error(
           "Can't call get ExapndedKey before preparation!");
@@ -52,19 +53,15 @@ class UdpEncryption {
 
   void prepareToProcessPeerData(
       size_t peerDataWidth,
-      const std::vector<int32_t>& indexes);
+      const std::vector<int32_t>& indexes) override;
 
   // process peer data via UDP encryption. This API should be called in
   // coordinate with "ProcessMyData" on peer's side. This API is ever called,
   // calling "getProcessedData" to retrive the cherry-picked encryption later.
-  void processPeerData(size_t dataSize);
+  void processPeerData(size_t dataSize) override;
 
   // returning the ciphertext, nonce, and index of cherry-picked rows
-  std::tuple<
-      std::vector<std::vector<unsigned char>>,
-      std::vector<__m128i>,
-      std::vector<int32_t>>
-  getProcessedData() {
+  EncryptionResuts getProcessedData() override {
     if (statusOfProcessingPeerData_ != Status::inProgress) {
       throw std::runtime_error(
           "Can't call getProcessedData before preparation!");
@@ -77,8 +74,7 @@ class UdpEncryption {
   }
 
  private:
-  std::unique_ptr<fbpcf::engine::communication::IPartyCommunicationAgent>
-      agent_;
+  std::unique_ptr<engine::communication::IPartyCommunicationAgent> agent_;
 
   uint64_t myDataIndexOffset_;
   uint64_t peerDataIndexOffset_;