From 82b475e8ce4bc1610194d7fa14340a5031ddd54b Mon Sep 17 00:00:00 2001 From: Ruiyu Zhu Date: Fri, 10 Mar 2023 13:01:28 -0800 Subject: [PATCH] add IUdpEncryption API for test (#504) Summary: Pull Request resolved: https://github.com/facebookresearch/fbpcf/pull/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 --- .../data_processor/IUdpEncryption.h | 57 +++++++++++++++++++ .../data_processor/UdpEncryption.cpp | 1 - .../data_processor/UdpEncryption.h | 22 +++---- 3 files changed, 66 insertions(+), 14 deletions(-) create mode 100644 fbpcf/mpc_std_lib/unified_data_process/data_processor/IUdpEncryption.h diff --git a/fbpcf/mpc_std_lib/unified_data_process/data_processor/IUdpEncryption.h b/fbpcf/mpc_std_lib/unified_data_process/data_processor/IUdpEncryption.h new file mode 100644 index 00000000..7488399e --- /dev/null +++ b/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 +#include +#include + +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>& plaintextData) = 0; + + virtual std::vector<__m128i> getExpandedKey() = 0; + + virtual void prepareToProcessPeerData( + size_t peerDataWidth, + const std::vector& 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> ciphertexts; + std::vector<__m128i> nonces; + std::vector 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 diff --git a/fbpcf/mpc_std_lib/unified_data_process/data_processor/UdpEncryption.cpp b/fbpcf/mpc_std_lib/unified_data_process/data_processor/UdpEncryption.cpp index 78cefe09..344798c9 100644 --- a/fbpcf/mpc_std_lib/unified_data_process/data_processor/UdpEncryption.cpp +++ b/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 agent) : agent_(std::move(agent)), diff --git a/fbpcf/mpc_std_lib/unified_data_process/data_processor/UdpEncryption.h b/fbpcf/mpc_std_lib/unified_data_process/data_processor/UdpEncryption.h index 5094e338..49b7a9a3 100644 --- a/fbpcf/mpc_std_lib/unified_data_process/data_processor/UdpEncryption.h +++ b/fbpcf/mpc_std_lib/unified_data_process/data_processor/UdpEncryption.h @@ -10,6 +10,7 @@ #include #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 friend class DataProcessor; @@ -31,15 +32,15 @@ class UdpEncryption { std::unique_ptr 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>& plaintextData); + const std::vector>& 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& indexes); + const std::vector& 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<__m128i>, - std::vector> - 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 - agent_; + std::unique_ptr agent_; uint64_t myDataIndexOffset_; uint64_t peerDataIndexOffset_;