【OSCP】implement KMPRT17 multiparty PSI

WORKSPACE

@@ -23,6 +23,9 @@ psi_deps()
 # Warning: psi relies on yacl to bring in common 3p libraries.
 #          Please make sure yacl_deps are called right after psi_deps.
 #
+load("@bazel_tools//tools/build_defs/repo:git.bzl", "git_repository")
+
+
 load("@yacl//bazel:repositories.bzl", "yacl_deps")
 
 yacl_deps()

psi/legacy/BUILD.bazel

@@ -86,6 +86,30 @@ psi_cc_library(
     ],
 )
 
+# psi_cc_library(
+#     name = "bc22_2party_psi",
+#     srcs = ["bc22_2party_psi.cc"],
+#     hdrs = ["bc22_2party_psi.h"],
+#     deps = [
+#         ":base_operator",
+#         ":factory",
+#         "//psi/bc22:bc22_psi",
+#     ],
+#     alwayslink = True,
+# )
+
+psi_cc_library(
+    name = "kmprt17_mp_psi",
+    srcs = ["kmprt17_mp_psi.cc"],
+    hdrs = ["kmprt17_mp_psi.h"],
+    deps = [
+        ":base_operator",
+        ":factory",
+        "//psi/legacy/kmprt17_mp_psi",
+    ],
+    alwayslink = True,
+)
+
 psi_cc_library(
     name = "dp_2party_psi",
     srcs = ["dp_2party_psi.cc"],
@@ -117,6 +141,7 @@ psi_cc_library(
         ":dp_2party_psi",
         ":ecdh_3party_psi",
         ":kkrt_2party_psi",
+        ":kmprt17_mp_psi",
         ":nparty_psi",
         ":rr22_2party_psi",
     ],

psi/legacy/kmprt17_mp_psi.cc

@@ -0,0 +1,36 @@
+// Copyright 2024 zhangwfjh
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "psi/legacy/kmprt17_mp_psi.h"
+
+#include <memory>
+
+#include "psi/legacy/factory.h"
+
+namespace psi::psi {
+
+namespace {
+
+std::unique_ptr<PsiBaseOperator> CreateOperator(
+    const MemoryPsiConfig& config,
+    const std::shared_ptr<yacl::link::Context>& lctx) {
+  return std::make_unique<KmprtMpPsiOperator>(
+      KmprtMpPsiOperator::Options({lctx, config.receiver_rank()}));
+}
+
+REGISTER_OPERATOR(KMPRT_PSI_NPC, CreateOperator);
+
+}  // namespace
+
+}  // namespace psi::psi

psi/legacy/kmprt17_mp_psi.h

@@ -0,0 +1,42 @@
+// Copyright 2024 zhangwfjh
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include <string>
+#include <vector>
+
+#include "yacl/link/link.h"
+
+#include "psi/legacy/base_operator.h"
+#include "psi/legacy/kmprt17_mp_psi/kmprt17_mp_psi.h"
+
+namespace psi::psi {
+
+class KmprtMpPsiOperator : public PsiBaseOperator {
+ public:
+  using Options = KmprtParty::Options;
+
+  explicit KmprtMpPsiOperator(const Options& options)
+      : PsiBaseOperator(options.link_ctx), options_(options) {}
+
+  std::vector<std::string> OnRun(const std::vector<std::string>& inputs) final {
+    return KmprtParty{options_}.Run(inputs);
+  }
+
+ private:
+  Options options_;
+};
+
+}  // namespace psi::psi

psi/legacy/kmprt17_mp_psi/BUILD.bazel

@@ -0,0 +1,51 @@
+# Copyright 2024 zhangwfjh
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+load("//bazel:psi.bzl", "psi_cc_library", "psi_cc_test")
+
+package(default_visibility = ["//visibility:public"])
+
+psi_cc_library(
+    name = "kmprt17_mp_psi",
+    srcs = [
+        "kmprt17_hashing.cc",
+        "kmprt17_mp_psi.cc",
+        "kmprt17_opprf.cc",
+    ],
+    hdrs = [
+        "kmprt17_hashing.h",
+        "kmprt17_mp_psi.h",
+        "kmprt17_opprf.h",
+    ],
+    deps = [
+        "//psi/utils:communication",
+        "//psi/utils:sync",
+        "//psi/utils:test_utils",
+        "@com_google_absl//absl/types:span",
+        "@yacl//yacl/base:exception",
+        "@yacl//yacl/base:int128",
+        "@yacl//yacl/crypto/hash:hash_utils",
+        "@yacl//yacl/crypto/rand",
+        "@yacl//yacl/kernels/algorithms:base_ot",
+        "@yacl//yacl/kernels/algorithms:iknp_ote",
+        "@yacl//yacl/kernels/algorithms:kkrt_ote",
+        "@yacl//yacl/link",
+    ],
+)
+
+psi_cc_test(
+    name = "kmprt17_mp_psi_test",
+    srcs = ["kmprt17_mp_psi_test.cc"],
+    deps = [":kmprt17_mp_psi"],
+)

psi/legacy/kmprt17_mp_psi/kmprt17_hashing.cc

@@ -0,0 +1,69 @@
+// Copyright 2024 zhangwfjh
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "psi/legacy/kmprt17_mp_psi/kmprt17_hashing.h"
+
+#include <random>
+#include <utility>
+
+#include "yacl/base/exception.h"
+#include "yacl/crypto/rand/rand.h"
+
+namespace psi::psi {
+
+void KmprtCuckooHashing::Insert(uint128_t elem) {
+  auto insert_into = [this, &elem](uint8_t c) {
+    for (uint8_t retry{}; retry != 128 && elem != NONE; ++retry) {
+      uint8_t rand_idx = yacl::crypto::FastRandU64() % num_hashes_[c];
+      uint8_t idx = (rand_idx + 1) % num_hashes_[c];
+      size_t addr;
+      do {
+        addr = HashU128{}(elem, idx) % num_bins_[c];
+        if (auto &bin = bins_[c][addr]; bin == NONE || bin == elem) {
+          bin = std::exchange(elem, NONE);
+          return;
+        }
+        idx = (idx + 1) % num_hashes_[c];
+      } while (idx != rand_idx);
+      std::swap(bins_[c][addr], elem);
+    }
+  };
+  for (uint8_t c{}; c != 2; ++c) {
+    insert_into(c);
+  }
+  YACL_ENFORCE_EQ(elem, NONE, "Failed to insert element.");
+}
+
+auto KmprtCuckooHashing::Lookup(uint128_t elem) const
+    -> std::pair<uint8_t, size_t> {
+  for (uint8_t c{}; c != 2; ++c) {
+    for (uint8_t idx{}; idx != num_hashes_[c]; ++idx) {
+      if (size_t addr = HashU128{}(elem, idx) % num_bins_[c];
+          bins_[c][addr] == elem) {
+        return {c, addr};
+      }
+    }
+  }
+  return {-1, -1};
+}
+
+void KmprtSimpleHashing::Insert(std::pair<uint128_t, uint64_t> point) {
+  for (uint8_t c{}; c != 2; ++c) {
+    for (size_t idx{}; idx != num_hashes_[c]; ++idx) {
+      bins_[c][HashU128{}(point.first, idx) % num_bins_[c]].emplace(point);
+    }
+  }
+}
+
+}  // namespace psi::psi

psi/legacy/kmprt17_mp_psi/kmprt17_hashing.h

@@ -0,0 +1,65 @@
+// Copyright 2024 zhangwfjh
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include <tuple>
+#include <unordered_map>
+#include <vector>
+
+#include "absl/numeric/int128.h"
+#include "yacl/base/int128.h"
+
+namespace psi::psi {
+
+struct HashU128 {
+  size_t operator()(uint128_t x, uint8_t idx = 0) const {
+    return absl::Uint128High64(x) + idx * absl::Uint128Low64(x);
+  }
+};
+
+template <typename Bin>
+struct KmprtDoubleHashing {
+  KmprtDoubleHashing(size_t m1, size_t m2) : num_bins_{m1, m2} {
+    bins_[0].resize(m1);
+    bins_[1].resize(m2);
+  }
+
+  const Bin &GetBin(uint8_t c, size_t addr) const { return bins_[c][addr]; }
+
+  const uint8_t num_hashes_[2]{3, 2};
+  const size_t num_bins_[2];
+  std::vector<Bin> bins_[2];
+};
+
+class KmprtCuckooHashing : public KmprtDoubleHashing<uint128_t> {
+ public:
+  constexpr static uint128_t NONE{};
+
+  KmprtCuckooHashing(size_t m1, size_t m2) : KmprtDoubleHashing{m1, m2} {}
+
+  void Insert(uint128_t);
+  std::pair<uint8_t, size_t> Lookup(uint128_t) const;
+};
+
+class KmprtSimpleHashing
+    : public KmprtDoubleHashing<
+          std::unordered_map<uint128_t, uint64_t, HashU128>> {
+ public:
+  KmprtSimpleHashing(size_t m1, size_t m2) : KmprtDoubleHashing{m1, m2} {}
+
+  void Insert(std::pair<uint128_t, uint64_t>);
+};
+
+}  // namespace psi::psi