Add helper functions for file read/write (#510)

Summary:
Pull Request resolved: #510

Helper functions to read from/write to files for UDP related types

Reviewed By: robotal

Differential Revision: D43752780

fbshipit-source-id: 30e452a7f223fb457ad5e1f8162baf420bc15511

docker/CMakeLists.txt

@@ -17,6 +17,8 @@ set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -mno-avx512f")
 
 include(cmake/fbpcf.cmake)
 install(DIRECTORY cmake/ DESTINATION cmake/)
+find_package(Boost REQUIRED COMPONENTS serialization)
+INCLUDE_DIRECTORIES(${Boost_INCLUDE_DIR})
 
 # Build fbpcf static library
 file(GLOB_RECURSE fbpcf_src
@@ -35,6 +37,8 @@ target_link_libraries(
   ${EMP-OT_LIBRARIES}
   google-cloud-cpp::storage
   Folly::folly
+  Boost::serialization
+  ${Boost_LIBRARIES}
   re2)
 
 # Build Edit Distance Example Library

docker/cmake/fbpcf.cmake

@@ -34,6 +34,11 @@ find_package(google_cloud_cpp_storage REQUIRED)
 
 find_library(re2 libre2.so)
 
+# Find boost
+find_package(Boost REQUIRED)
+include_directories(${Boost_INCLUDE_DIRS})
+add_compile_options(-lboost_serialization)
+
 # since emp-tool is compiled with cc++11 and our games needs c++17 overwrite the
 # compile option to c++17
 add_compile_options(-std=c++17)

fbpcf/mpc_std_lib/unified_data_process/data_processor/UdpUtil.cpp

@@ -6,8 +6,12 @@
  */
 
 #include "fbpcf/mpc_std_lib/unified_data_process/data_processor/UdpUtil.h"
+#include <boost/archive/text_iarchive.hpp>
+#include <boost/archive/text_oarchive.hpp>
+#include <boost/serialization/vector.hpp>
 #include <cstdint>
 #include <cstring>
+#include "fbpcf/io/api/FileIOWrappers.h"
 
 namespace fbpcf::mpc_std_lib::unified_data_process::data_processor {
 
@@ -79,4 +83,82 @@ UdpUtil::localEncryption(
   return {ciphertextByte, s2vVec};
 }
 
+static std::vector<__m128i> convertCharVecToM128i(
+    const std::vector<unsigned char>& src) {
+  if (src.size() % sizeof(__m128i) != 0) {
+    throw std::invalid_argument("Unable to convert to std::vector<__m128i>");
+  }
+  std::vector<__m128i> rst(src.size() / sizeof(__m128i));
+  for (size_t i = 0; i < rst.size(); i++) {
+    rst.at(i) = _mm_set_epi8(
+        src.at(16 * i + 15),
+        src.at(16 * i + 14),
+        src.at(16 * i + 13),
+        src.at(16 * i + 12),
+        src.at(16 * i + 11),
+        src.at(16 * i + 10),
+        src.at(16 * i + 9),
+        src.at(16 * i + 8),
+        src.at(16 * i + 7),
+        src.at(16 * i + 6),
+        src.at(16 * i + 5),
+        src.at(16 * i + 4),
+        src.at(16 * i + 3),
+        src.at(16 * i + 2),
+        src.at(16 * i + 1),
+        src.at(16 * i));
+  }
+  return rst;
+}
+
+static std::vector<unsigned char> convertM128iToCharVec(
+    const std::vector<__m128i>& src) {
+  std::vector<unsigned char> rst(src.size() * sizeof(__m128i));
+  memcpy(rst.data(), src.data(), rst.size());
+  return rst;
+}
+
+void writeEncryptionResultsToFile(
+    const IUdpEncryption::EncryptionResuts& EncryptionResuts,
+    const std::string& file) {
+  std::ostringstream s;
+  boost::archive::text_oarchive oa(s);
+  oa << EncryptionResuts.ciphertexts << EncryptionResuts.indexes
+     << convertM128iToCharVec(EncryptionResuts.nonces);
+
+  fbpcf::io::FileIOWrappers::writeFile(file, s.str());
+}
+
+void writeExpandedKeyToFile(
+    const std::vector<__m128i>& expandedKey,
+    const std::string& file) {
+  std::ostringstream s;
+  boost::archive::text_oarchive oa(s);
+  oa << convertM128iToCharVec(expandedKey);
+  fbpcf::io::FileIOWrappers::writeFile(file, s.str());
+}
+
+IUdpEncryption::EncryptionResuts readEncryptionResultsFromFile(
+    const std::string& file) {
+  std::istringstream s(fbpcf::io::FileIOWrappers::readFile(file));
+
+  IUdpEncryption::EncryptionResuts data;
+  boost::archive::text_iarchive ia(s);
+  ia >> data.ciphertexts;
+  ia >> data.indexes;
+  std::vector<unsigned char> tmp;
+  ia >> tmp;
+
+  data.nonces = convertCharVecToM128i(tmp);
+  return data;
+}
+
+std::vector<__m128i> readExpandedKeyFromFile(const std::string& file) {
+  std::vector<unsigned char> data(sizeof(__m128i) * 11);
+  std::istringstream s(fbpcf::io::FileIOWrappers::readFile(file));
+  boost::archive::text_iarchive ia(s);
+  ia >> data;
+  return convertCharVecToM128i(data);
+}
+
 } // namespace fbpcf::mpc_std_lib::unified_data_process::data_processor

fbpcf/mpc_std_lib/unified_data_process/data_processor/UdpUtil.h

@@ -10,6 +10,7 @@
 #include <emmintrin.h>
 #include <sys/types.h>
 #include "fbpcf/frontend/BitString.h"
+#include "fbpcf/mpc_std_lib/unified_data_process/data_processor/IUdpEncryption.h"
 #include "fbpcf/primitive/mac/S2v.h"
 #include "fbpcf/primitive/mac/S2vFactory.h"
 
@@ -48,6 +49,19 @@ class UdpUtil {
       int inputPartyID);
 };
 
+void writeEncryptionResultsToFile(
+    const IUdpEncryption::EncryptionResuts& encryptionResults,
+    const std::string& file);
+
+void writeExpandedKeyToFile(
+    const std::vector<__m128i>& expandedKey,
+    const std::string& file);
+
+std::vector<__m128i> readExpandedKeyFromFile(const std::string& file);
+
+IUdpEncryption::EncryptionResuts readEncryptionResultsFromFile(
+    const std::string& file);
+
 } // namespace fbpcf::mpc_std_lib::unified_data_process::data_processor
 
 #include "fbpcf/mpc_std_lib/unified_data_process/data_processor/UdpUtil_impl.h"

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

@@ -21,6 +21,8 @@
 #include "fbpcf/mpc_std_lib/unified_data_process/data_processor/IDataProcessor.h"
 #include "fbpcf/scheduler/SchedulerHelper.h"
 #include "fbpcf/test/TestHelper.h"
+#include "folly/Format.h"
+#include "folly/Random.h"
 
 namespace fbpcf::mpc_std_lib::unified_data_process::data_processor {
 
@@ -306,4 +308,55 @@ TEST(DataProcessorSubObjects, testUdpEncryptionAndDecryptionObjects) {
       agentFactories.at(1)->create(0, "receiver"));
 }
 
+TEST(TestFileReadAndWrite, testExpandedKeyReadAndWrite) {
+  const int batchSize = 11;
+  std::vector<__m128i> key(batchSize);
+  for (size_t i = 0; i < batchSize; i++) {
+    key.at(i) = _mm_set_epi32(
+        folly::Random::rand32(),
+        folly::Random::rand32(),
+        folly::Random::rand32(),
+        folly::Random::rand32());
+  }
+  const std::string expandedKeyFile =
+      folly::sformat("./expanded_key_{}", folly::Random::rand32());
+
+  writeExpandedKeyToFile(key, expandedKeyFile);
+  auto key1 = readExpandedKeyFromFile(expandedKeyFile);
+  std::remove(expandedKeyFile.c_str());
+  fbpcf::testEq(key, key1);
+}
+
+TEST(TestFileReadAndWrite, testEncryptionResultReadAndWrite) {
+  IUdpEncryption::EncryptionResuts results;
+
+  const int batchSize = 11;
+  results.ciphertexts = std::vector<std::vector<unsigned char>>(batchSize);
+  results.indexes = std::vector<int32_t>(batchSize);
+  results.nonces = std::vector<__m128i>(batchSize);
+  for (size_t i = 0; i < batchSize; i++) {
+    results.nonces.at(i) = _mm_set_epi32(
+        folly::Random::rand32(),
+        folly::Random::rand32(),
+        folly::Random::rand32(),
+        folly::Random::rand32());
+    results.indexes.at(i) = folly::Random::rand32();
+    results.ciphertexts.at(i) = std::vector<unsigned char>(60);
+    for (size_t j = 0; j < results.ciphertexts.size(); j++) {
+      results.ciphertexts.at(i).at(j) = i + j;
+    }
+  }
+  const std::string resultFile =
+      folly::sformat("./encryptionResult_{}", folly::Random::rand32());
+
+  writeEncryptionResultsToFile(results, resultFile);
+  auto results1 = readEncryptionResultsFromFile(resultFile);
+  std::remove(resultFile.c_str());
+  fbpcf::testEq(results.nonces, results1.nonces);
+  fbpcf::testVectorEq(results.indexes, results1.indexes);
+  for (size_t i = 0; i < batchSize; i++) {
+    fbpcf::testVectorEq(results.ciphertexts.at(i), results1.ciphertexts.at(i));
+  }
+}
+
 } // namespace fbpcf::mpc_std_lib::unified_data_process::data_processor