Repository: secretflow/psi (language: python)

Issue:
选择并在 SecretFlow 中实现多方 PSI 协议
> 此 ISSUE 为 [隐语开源共建计划（SecretFlow Open Source Contribution Plan，简称 SF OSCP）第三期](https://github.com/orgs/secretflow/discussions/1007)任务 ISSUE，欢迎社区开发者参与共建～
> 若有感兴趣想要认领的任务，但还未报名，辛苦先完成[报名](https://www.wjx.top/vm/QKio0dq.aspx# )进行哈～


## 任务介绍
- 任务名称：选择并在 SecretFlow 中实现多方 PSI 协议
- 技术方向：PSI
- 任务难度：挑战🌟🌟🌟


## 详细要求
请选择并实现隐语的多方 PSI 协议，具体要求如下：
- 功能性：至少一个 大于 2 个参与方的半诚实 PSI 方案
- 安全性（尽量少 reveal）
- 代码规范：C++ 代码需要使用 Google C++ 代码规范进行格式化（流水线包含代码规范检查卡点）
- 提交说明：关联该 issue 并提交代码至 https://github.com/secretflow/spu/tree/main/libspu/psi/core

## 能力要求
- 熟悉 PSI 原理
- 了解多方 PSI 近年来的进展
- 熟悉隐语 PSI 接口

## 操作说明
- 参考内容：
  - [psi/core](https://github.com/secretflow/spu/tree/main/libspu/psi/core)
  - [psi/operator](https://github.com/secretflow/spu/tree/main/libspu/psi/operator)
- 范例：
  - [kkrt_psi.h](https://github.com/secretflow/spu/blob/main/libspu/psi/core/kkrt_psi.h)
  - [nparty_psi.h](https://github.com/secretflow/spu/blob/main/libspu/psi/operator/nparty_psi.h)

## 认领说明
本任务可有多种实现方式，故支持有多位开发者进行认领，请在认领任务后，在该 issue 下 comment 你的具体设计思路。

**设计思路说明**：简单说明计划使用哪个半诚实 PSI 方案以及简单的实现思路。

> 当同一任务有多位开发者提交设计思路时：
> - 若多位开发者的设计思路类似，则将按照 comment 的时间 assign 给第一位 comment 的开发者；
> - 若多位开发者的设计思路不同，则该 issue 将会拆分为以不同设计思路进行实现的子 issue，并 assign 给对应开发者。


Discussion / hints:
Stale issue message. Please comment to remove stale tag. Otherwise this issue will be closed soon.
zhangwfjh give it to me
scheme: FNP04
> zhangwfjh give it to me scheme: FNP04

感谢认领，可以在这个子 issue 下再认领一下～将把该实现方式 assign 给你哈～https://github.com/secretflow/psi/issues/9
zhangwfjh give it to me
scheme: KMPRT17

Fix patch:
diff --git a/WORKSPACE b/WORKSPACE
index 52192b0..3771892 100644
--- a/WORKSPACE
+++ b/WORKSPACE
@@ -23,6 +23,7 @@ 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()
diff --git a/psi/legacy/BUILD.bazel b/psi/legacy/BUILD.bazel
index 23a8707..118dd8c 100644
--- a/psi/legacy/BUILD.bazel
+++ b/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",
     ],
diff --git a/psi/legacy/kmprt17_mp_psi.cc b/psi/legacy/kmprt17_mp_psi.cc
new file mode 100644
index 0000000..f69f656
--- /dev/null
+++ b/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
diff --git a/psi/legacy/kmprt17_mp_psi.h b/psi/legacy/kmprt17_mp_psi.h
new file mode 100644
index 0000000..ce51a84
--- /dev/null
+++ b/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
diff --git a/psi/legacy/kmprt17_mp_psi/BUILD.bazel b/psi/legacy/kmprt17_mp_psi/BUILD.bazel
new file mode 100644
index 0000000..a548394
--- /dev/null
+++ b/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/kernel/algorithms:base_ot",
+        "@yacl//yacl/kernel/algorithms:iknp_ote",
+        "@yacl//yacl/kernel/algorithms:kkrt_ote",
+        "@yacl//yacl/link",
+    ],
+)
+
+psi_cc_test(
+    name = "kmprt17_mp_psi_test",
+    srcs = ["kmprt17_mp_psi_test.cc"],
+    deps = [":kmprt17_mp_psi"],
+)
diff --git a/psi/legacy/kmprt17_mp_psi/kmprt17_hashing.cc b/psi/legacy/kmprt17_mp_psi/kmprt17_hashing.cc
new file mode 100644
index 0000000..ebd537d
--- /dev/null
+++ b/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
diff --git a/psi/legacy/kmprt17_mp_psi/kmprt17_hashing.h b/psi/legacy/kmprt17_mp_psi/kmprt17_hashing.h
new file mode 100644
index 0000000..be80df9
--- /dev/null
+++ b/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
diff --git a/psi/legacy/kmprt17_mp_psi/kmprt17_mp_psi.cc b/psi/legacy/kmprt17_mp_psi/kmprt17_mp_psi.cc
new file mode 100644
index 0000000..8048881
--- /dev/null
+++ b/psi/legacy/kmprt17_mp_psi/kmprt17_mp_psi.cc
@@ -0,0 +1,163 @@
+// 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_mp_psi.h"
+
+#include <future>
+
+#include "yacl/crypto/hash/hash_utils.h"
+#include "yacl/crypto/rand/rand.h"
+#include "yacl/utils/serialize.h"
+
+#include "psi/legacy/kmprt17_mp_psi/kmprt17_opprf.h"
+#include "psi/utils/communication.h"
+#include "psi/utils/sync.h"
+
+namespace psi::psi {
+
+namespace {
+
+constexpr uint32_t kLinkRecvTimeout = 60 * 60 * 1000;
+
+}  // namespace
+
+KmprtParty::KmprtParty(const Options& options) : options_{options} {
+  auto [ctx, wsize, me, leader] = CollectContext();
+  ctx->SetRecvTimeout(kLinkRecvTimeout);
+  p2p_.resize(wsize);
+  for (size_t dst{}; dst != wsize; ++dst) {
+    if (me != dst) {
+      p2p_[dst] = CreateP2PLinkCtx("kmprt17_mp_psi", ctx, dst);
+    }
+  }
+}
+
+std::vector<std::string> KmprtParty::Run(
+    const std::vector<std::string>& inputs) {
+  auto [ctx, wsize, me, leader] = CollectContext();
+  auto counts = AllGatherItemsSize(ctx, inputs.size());
+  size_t count{};
+  for (auto cnt : counts) {
+    if (cnt == 0) {
+      return {};
+    }
+    count = std::max(cnt, count);
+  }
+  auto items = EncodeInputs(inputs, count);
+  auto shares = ZeroSharing(count);
+  auto recv_share = SwapShares(items, shares);
+  auto recons = Reconstruct(items, recv_share);
+  std::vector<std::string> intersection;
+  for (size_t k{}; k != count; ++k) {
+    if (recons[k] == 0) {
+      intersection.emplace_back(inputs[k]);
+    }
+  }
+  return intersection;
+}
+
+std::vector<uint128_t> KmprtParty::EncodeInputs(
+    const std::vector<std::string>& inputs, size_t count) const {
+  std::vector<uint128_t> items;
+  items.reserve(count);
+  std::transform(
+      inputs.begin(), inputs.end(), std::back_inserter(items),
+      [](std::string_view input) { return yacl::crypto::Blake3_128(input); });
+  //  Add random dummy elements
+  std::generate_n(std::back_inserter(items), count - inputs.size(),
+                  yacl::crypto::FastRandU128);
+  return items;
+}
+
+auto KmprtParty::ZeroSharing(size_t count) const -> std::vector<Share> {
+  auto [ctx, wsize, me, leader] = CollectContext();
+  std::vector<Share> shares(wsize, Share(count));
+  for (size_t k{}; k != count; ++k) {
+    uint64_t sum{};
+    for (size_t dst{1}; dst != wsize; ++dst) {
+      sum ^= shares[dst][k] = yacl::crypto::FastRandU64();
+    }
+    shares[0][k] = sum;
+  }
+  return shares;
+}
+
+auto KmprtParty::SwapShares(const std::vector<uint128_t>& items,
+                            const std::vector<Share>& shares) const -> Share {
+  auto [ctx, wsize, me, leader] = CollectContext();
+  auto count = shares.front().size();
+  std::vector<Share> recv_shares(count);
+  std::vector<std::future<Share>> futures(wsize);
+  // NOTE: First Send Then Receive for peers of smaller ranks
+  for (size_t id{}; id != me; ++id) {
+    futures[id] = std::async(
+        [&](size_t id) {
+          KmprtOpprfSend(p2p_[id], items, shares[id]);
+          return KmprtOpprfRecv(p2p_[id], items);
+        },
+        id);
+  }
+  // NOTE: First Receive Then Send for peers of larger ranks
+  for (size_t id{me + 1}; id != wsize; ++id) {
+    futures[id] = std::async(
+        [&](size_t id) {
+          auto ret = KmprtOpprfRecv(p2p_[id], items);
+          KmprtOpprfSend(p2p_[id], items, shares[id]);
+          return ret;
+        },
+        id);
+  }
+  for (size_t id{}; id != wsize; ++id) {
+    recv_shares[id] = (me == id ? shares[id] : futures[id].get());
+  }
+
+  Share share(count);  // S(x_k)
+  for (size_t k{}; k != count; ++k) {
+    for (size_t src{}; src != wsize; ++src) {
+      share[k] ^= recv_shares[src][k];
+    }
+  }
+  return share;
+}
+
+auto KmprtParty::Reconstruct(const std::vector<uint128_t>& items,
+                             const Share& share) const -> Share {
+  auto [ctx, wsize, me, leader] = CollectContext();
+  auto count = items.size();
+  if (me == leader) {
+    std::vector<Share> recv_shares(count);
+    std::vector<std::future<Share>> futures(wsize);
+    for (size_t src{}; src != wsize; ++src) {
+      if (me != src) {
+        futures[src] = std::async(
+            [&](size_t src) { return KmprtOpprfRecv(p2p_[src], items); }, src);
+      }
+    }
+    for (size_t src{}; src != wsize; ++src) {
+      recv_shares[src] = (me == src ? share : futures[src].get());
+    }
+    Share recons(count);  // sum of S_i(x_k) over i
+    for (size_t k{}; k != count; ++k) {
+      for (size_t src{}; src != wsize; ++src) {
+        recons[k] ^= recv_shares[src][k];
+      }
+    }
+    return recons;
+  } else {
+    KmprtOpprfSend(p2p_[leader], items, share);
+    return share;
+  }
+}
+
+}  // namespace psi::psi
diff --git a/psi/legacy/kmprt17_mp_psi/kmprt17_mp_psi.h b/psi/legacy/kmprt17_mp_psi/kmprt17_mp_psi.h
new file mode 100644
index 0000000..f642107
--- /dev/null
+++ b/psi/legacy/kmprt17_mp_psi/kmprt17_mp_psi.h
@@ -0,0 +1,60 @@
+// 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 <memory>
+#include <string>
+#include <vector>
+
+#include "yacl/base/int128.h"
+#include "yacl/link/link.h"
+
+namespace psi::psi {
+
+// Practical Multi-party Private Set Intersection from Symmetric-Key Techniques
+// https://eprint.iacr.org/2017/799.pdf
+
+class KmprtParty {
+ public:
+  struct Options {
+    std::shared_ptr<yacl::link::Context> link_ctx;
+    size_t leader_rank;
+  };
+
+  KmprtParty(const Options& options);
+  virtual std::vector<std::string> Run(const std::vector<std::string>& inputs);
+
+ private:
+  using Share = std::vector<uint64_t>;
+
+  std::vector<uint128_t> EncodeInputs(const std::vector<std::string>& inputs,
+                                      size_t count) const;
+  std::vector<Share> ZeroSharing(size_t count) const;
+  Share SwapShares(const std::vector<uint128_t>& items,
+                   const std::vector<Share>& shares) const;
+  Share Reconstruct(const std::vector<uint128_t>& items,
+                    const Share& share) const;
+
+  // (ctx, world_size, my_rank, leader_rank)
+  auto CollectContext() const {
+    return std::make_tuple(options_.link_ctx, options_.link_ctx->WorldSize(),
+                           options_.link_ctx->Rank(), options_.leader_rank);
+  }
+
+  Options options_;
+  std::vector<std::shared_ptr<yacl::link::Context>> p2p_;
+};
+
+}  // namespace psi::psi
diff --git a/psi/legacy/kmprt17_mp_psi/kmprt17_opprf.cc b/psi/legacy/kmprt17_mp_psi/kmprt17_opprf.cc
new file mode 100644
index 0000000..f5c0ae5
--- /dev/null
+++ b/psi/legacy/kmprt17_mp_psi/kmprt17_opprf.cc
@@ -0,0 +1,189 @@
+// 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_opprf.h"
+
+#include <algorithm>
+#include <cmath>
+#include <memory>
+
+#include "yacl/crypto/rand/rand.h"
+#include "yacl/crypto/tools/ro.h"
+#include "yacl/kernel/algorithms/base_ot.h"
+#include "yacl/kernel/algorithms/iknp_ote.h"
+#include "yacl/kernel/algorithms/kkrt_ote.h"
+#include "yacl/link/link.h"
+#include "yacl/utils/serialize.h"
+
+#include "psi/legacy/kmprt17_mp_psi/kmprt17_hashing.h"
+
+namespace psi::psi {
+
+namespace yc = yacl::crypto;
+
+namespace {
+
+// PSI-related constants
+// Ref https://eprint.iacr.org/2017/799.pdf (Table 2)
+constexpr float ZETA[]{1.12f, 0.17f};
+// constexpr size_t BETA[]{31, 63};
+constexpr size_t TABLE_SIZE[]{32, 64};  // next power of BETAs
+
+// OTe-related constants
+constexpr size_t NUM_BASE_OT{128};
+constexpr size_t NUM_INKP_OT{512};
+constexpr size_t BATCH_SIZE{896};
+
+static auto ro = yc::RandomOracle::GetDefault();
+
+}  // namespace
+
+// Ref https://eprint.iacr.org/2017/799.pdf (Figure 6, 7)
+
+std::vector<uint64_t> KmprtOpprfRecv(
+    const std::shared_ptr<yacl::link::Context>& ctx,
+    const std::vector<uint128_t>& queries) {
+  const size_t size{queries.size()};
+  const size_t bin_sizes[]{static_cast<size_t>(std::ceil(size * ZETA[0])),
+                           static_cast<size_t>(std::ceil(size * ZETA[1]))};
+  // Step 0. Prepares OPRF
+  yc::KkrtOtExtReceiver receiver;
+  size_t num_ot{bin_sizes[0] + bin_sizes[1]};
+  auto choice = yc::RandBits(NUM_BASE_OT);
+  auto base_ot = yc::BaseOtRecv(ctx, choice, NUM_BASE_OT);
+  auto store = yc::IknpOtExtSend(ctx, base_ot, NUM_INKP_OT);
+  receiver.Init(ctx, store, num_ot);
+  receiver.SetBatchSize(BATCH_SIZE);
+
+  // Step 1. Hashes queries into Cuckoo hashing
+  KmprtCuckooHashing hashing{bin_sizes[0], bin_sizes[1]};
+  for (size_t i{}; i != size; ++i) {
+    hashing.Insert(queries[i]);
+  }
+
+  std::vector<uint128_t> evals;
+  evals.reserve(num_ot);
+  size_t ot_idx{}, b{};
+  std::array<uint64_t, BATCH_SIZE> batch_evals;
+  // Step 2. For each bin, invokes single-query OPPRF
+  for (uint8_t c{}; c != 2; ++c) {
+    size_t ot_begin{c == uint8_t{0} ? 0 : bin_sizes[0]};
+    size_t ot_end{c == uint8_t{0} ? bin_sizes[0] : num_ot};
+    for (size_t addr{}; addr != hashing.num_bins_[c]; ++addr, ++ot_idx) {
+      auto elem = hashing.GetBin(c, addr);
+      elem == KmprtCuckooHashing::NONE && (elem = yc::FastRandU128());
+      receiver.Encode(
+          ot_idx, elem,
+          {reinterpret_cast<uint8_t*>(&batch_evals[b++]), sizeof(uint64_t)});
+      if (auto batch_size = (ot_idx - ot_begin) % BATCH_SIZE + 1;
+          batch_size == BATCH_SIZE || ot_idx + 1 == ot_end) {
+        b = 0;
+        receiver.SendCorrection(ctx, batch_size);
+        // For each query in a batch
+        for (size_t i{}; i != batch_size; ++i) {
+          uint128_t nonce = yacl::DeserializeUint128(
+              ctx->Recv(ctx->NextRank(), "Receive OPPRF nonce"));
+          std::vector<uint64_t> table(TABLE_SIZE[c]);
+          auto buf =
+              ctx->Recv(ctx->NextRank(), "Receive OPPRF EncryptionTable");
+          std::memcpy(table.data(), buf.data(),
+                      TABLE_SIZE[c] * sizeof(uint64_t));
+          uint64_t eval = batch_evals[i];
+          auto index =
+              ro.Gen<size_t>(absl::MakeSpan(reinterpret_cast<uint8_t*>(&eval),
+                                            sizeof eval),
+                             nonce) %
+              table.size();
+          evals.emplace_back(eval ^ table[index]);
+        }
+      }
+    }
+  }
+
+  // Step 3. Filters and obtains the results
+  std::vector<uint64_t> results(size);
+  std::transform(queries.cbegin(), queries.cend(), results.begin(),
+                 [&](auto q) {
+                   auto [c, addr] = hashing.Lookup(q);
+                   return evals[c * bin_sizes[0] + addr];
+                 });
+  return results;
+}
+
+void KmprtOpprfSend(const std::shared_ptr<yacl::link::Context>& ctx,
+                    const std::vector<uint128_t>& xs,
+                    const std::vector<uint64_t>& ys) {
+  YACL_ENFORCE_EQ(xs.size(), ys.size(), "Sizes mismatch.");
+  const size_t size{xs.size()};
+  const size_t bin_sizes[]{static_cast<size_t>(std::ceil(size * ZETA[0])),
+                           static_cast<size_t>(std::ceil(size * ZETA[1]))};
+  // Step 0. Prepares OPRF
+  yc::KkrtOtExtSender sender;
+  size_t num_ot{bin_sizes[0] + bin_sizes[1]};
+  auto base_ot = yc::BaseOtSend(ctx, NUM_BASE_OT);
+  auto choice = yc::RandBits(NUM_INKP_OT);
+  auto store = yc::IknpOtExtRecv(ctx, base_ot, choice, NUM_INKP_OT);
+  sender.Init(ctx, store, num_ot);
+  sender.SetBatchSize(BATCH_SIZE);
+
+  // Step 1. Hashes points into Simple hashing
+  KmprtSimpleHashing hashing{bin_sizes[0], bin_sizes[1]};
+  for (size_t i{}; i != size; ++i) {
+    hashing.Insert({xs[i], ys[i]});
+  }
+  size_t ot_idx{};
+  auto evaluator = sender.GetOprf();
+  // Step 2. For each bin, invokes single-query OPPRF
+  for (uint8_t c{}; c != 2; ++c) {
+    size_t ot_begin{c == uint8_t{0} ? 0 : bin_sizes[0]};
+    size_t ot_end{c == uint8_t{0} ? bin_sizes[0] : num_ot};
+    // For each programmable point in a batch
+    for (size_t addr{}; addr != hashing.num_bins_[c]; ++addr, ++ot_idx) {
+      if ((ot_idx - ot_begin) % BATCH_SIZE == 0) {
+        sender.RecvCorrection(
+            ctx, ot_idx + BATCH_SIZE >= ot_end ? ot_end - ot_idx : BATCH_SIZE);
+      }
+      auto bin = hashing.GetBin(c, addr);
+      uint128_t nonce;
+      std::vector<uint64_t> table;
+      bool separable;
+      do {
+        separable = true;
+        nonce = yc::FastRandSeed();
+        table.assign(TABLE_SIZE[c], uint64_t{0});
+        for (auto it = bin.cbegin(); it != bin.cend(); ++it) {
+          uint64_t eval = evaluator->Eval(ot_idx, it->first);
+          auto index =
+              ro.Gen<size_t>({reinterpret_cast<uint8_t*>(&eval), sizeof eval},
+                             nonce) %
+              table.size();
+          if (table[index] != uint64_t{0}) {
+            separable = false;
+            break;
+          }
+          table[index] = eval ^ it->second;
+        }
+      } while (!separable);
+      for (size_t i{}; i != TABLE_SIZE[c]; ++i) {
+        table[i] == uint64_t{0} && (table[i] = yc::FastRandU128());
+      }
+      ctx->SendAsync(ctx->NextRank(), yacl::SerializeUint128(nonce),
+                     fmt::format("OPPRF:Nonce={}", nonce));
+      yacl::Buffer buf(table.data(), table.size() * sizeof(uint64_t));
+      ctx->SendAsync(ctx->NextRank(), buf, "OPPRF:EncryptionTable");
+    }
+  }
+}
+
+}  // namespace psi::psi
diff --git a/psi/legacy/kmprt17_mp_psi/kmprt17_opprf.h b/psi/legacy/kmprt17_mp_psi/kmprt17_opprf.h
new file mode 100644
index 0000000..1d11fbc
--- /dev/null
+++ b/psi/legacy/kmprt17_mp_psi/kmprt17_opprf.h
@@ -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.
+
+#pragma once
+
+#include <memory>
+#include <tuple>
+#include <vector>
+
+#include "yacl/base/int128.h"
+#include "yacl/link/link.h"
+
+namespace psi::psi {
+
+// Table-based OPPRF, see https://eprint.iacr.org/2017/799.pdf (Figure 6)
+
+std::vector<uint64_t> KmprtOpprfRecv(
+    const std::shared_ptr<yacl::link::Context>&,
+    const std::vector<uint128_t>& queries);
+
+void KmprtOpprfSend(const std::shared_ptr<yacl::link::Context>&,
+                    const std::vector<uint128_t>& xs,
+                    const std::vector<uint64_t>& ys);
+
+}  // namespace psi::psi
diff --git a/psi/proto/psi.proto b/psi/proto/psi.proto
index 47e7a63..4c3b0de 100644
--- a/psi/proto/psi.proto
+++ b/psi/proto/psi.proto
@@ -68,6 +68,9 @@ enum PsiType {
   RR22_FAST_PSI_2PC = 13;
   RR22_LOWCOMM_PSI_2PC = 14;
   RR22_MALICIOUS_PSI_2PC = 15;
+
+  // KMPRT17 PSI https://eprint.iacr.org/2017/799.pdf
+  KMPRT17_MP_PSI = 16;
 }
 
 // The specified elliptic curve cryptography used in psi.


Test patch:
diff --git a/psi/legacy/kmprt17_mp_psi/kmprt17_mp_psi_test.cc b/psi/legacy/kmprt17_mp_psi/kmprt17_mp_psi_test.cc
new file mode 100644
index 0000000..f07858e
--- /dev/null
+++ b/psi/legacy/kmprt17_mp_psi/kmprt17_mp_psi_test.cc
@@ -0,0 +1,117 @@
+// 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_mp_psi.h"
+
+#include <random>
+#include <set>
+#include <vector>
+
+#include "gtest/gtest.h"
+#include "yacl/link/test_util.h"
+
+#include "psi/utils/test_utils.h"
+
+namespace psi::psi {
+
+namespace {
+
+struct KMPRTMpTestParams {
+  std::vector<size_t> item_size;
+  size_t intersection_size;
+};
+
+std::vector<std::vector<std::string>> CreateNPartyItems(
+    const KMPRTMpTestParams& params) {
+  std::vector<std::vector<std::string>> ret(params.item_size.size() + 1);
+  ret[params.item_size.size()] =
+      test::CreateRangeItems(1, params.intersection_size);
+
+  for (size_t idx = 0; idx < params.item_size.size(); ++idx) {
+    ret[idx] =
+        test::CreateRangeItems((idx + 1) * 1000000, params.item_size[idx]);
+  }
+
+  for (size_t idx = 0; idx < params.item_size.size(); ++idx) {
+    std::set<size_t> idx_set;
+    std::random_device rd;
+    std::mt19937 gen(rd());
+    std::uniform_int_distribution<> dis(0, params.item_size[idx] - 1);
+
+    while (idx_set.size() < params.intersection_size) {
+      idx_set.insert(dis(gen));
+    }
+    size_t j = 0;
+    for (const auto& iter : idx_set) {
+      ret[idx][iter] = ret[params.item_size.size()][j++];
+    }
+  }
+  return ret;
+}
+
+}  // namespace
+
+class KMPRTMpPsiTest : public testing::TestWithParam<KMPRTMpTestParams> {};
+
+TEST_P(KMPRTMpPsiTest, Works) {
+  std::vector<std::vector<std::string>> items;
+
+  auto params = GetParam();
+  items = CreateNPartyItems(params);
+  size_t leader_rank = 0;
+
+  auto ctxs = yacl::link::test::SetupWorld(params.item_size.size());
+
+  auto proc = [&](int idx) -> std::vector<std::string> {
+    KmprtParty::Options opts;
+    opts.link_ctx = ctxs[idx];
+    opts.leader_rank = leader_rank;
+    KmprtParty op(opts);
+    return op.Run(items[idx]);
+  };
+
+  size_t world_size = ctxs.size();
+  std::vector<std::future<std::vector<std::string>>> f_links(world_size);
+  for (size_t i = 0; i < world_size; i++) {
+    f_links[i] = std::async(proc, i);
+  }
+
+  std::vector<std::string> intersection = items[params.item_size.size()];
+  std::sort(intersection.begin(), intersection.end());
+
+  std::vector<std::vector<std::string>> results(world_size);
+  for (size_t i = 0; i < world_size; i++) {
+    results[i] = f_links[i].get();
+    if (i == leader_rank) {
+      std::sort(results[i].begin(), results[i].end());
+      EXPECT_EQ(results[i].size(), intersection.size());
+      EXPECT_EQ(results[i], intersection);
+    } else {
+      EXPECT_EQ(results[i].size(), 0);
+    }
+  }
+}
+
+INSTANTIATE_TEST_SUITE_P(
+    Works_Instances, KMPRTMpPsiTest,
+    testing::Values(KMPRTMpTestParams{{0, 3}, 0},                 //
+                    KMPRTMpTestParams{{3, 0}, 0},                 //
+                    KMPRTMpTestParams{{0, 0}, 0},                 //
+                    KMPRTMpTestParams{{4, 3}, 2},                 //
+                    KMPRTMpTestParams{{20, 17, 14}, 10},          //
+                    KMPRTMpTestParams{{20, 17, 14, 30}, 10},      //
+                    KMPRTMpTestParams{{20, 17, 14, 30, 35}, 11},  //
+                    KMPRTMpTestParams{{20, 17, 14, 30, 35}, 0}));
+
+}  // namespace psi::psi