Add feature for threshold ecdh psi

examples/psi/config/threshold_ecdh_psi_receiver.json

@@ -0,0 +1,36 @@
+{
+  "psi_config": {
+    "protocol_config": {
+      "protocol": "PROTOCOL_ECDH",
+      "ecdh_config": {
+        "curve": "CURVE_FOURQ"
+      },
+      "role": "ROLE_RECEIVER",
+      "broadcast_result": true
+    },
+    "input_config": {
+      "type": "IO_TYPE_FILE_CSV",
+      "path": "/tmp/receiver_input.csv"
+    },
+    "output_config": {
+      "type": "IO_TYPE_FILE_CSV",
+      "path": "/tmp/receiver_output.csv"
+    },
+    "keys": ["id_0", "id_1"],
+    "intersection_threshold": 50,
+    "count_path": "/tmp/receiver_count.csv"
+  },
+  "link_config": {
+    "parties": [
+      {
+        "id": "receiver",
+        "host": "127.0.0.1:5300"
+      },
+      {
+        "id": "sender",
+        "host": "127.0.0.1:5400"
+      }
+    ]
+  },
+  "self_link_party": "receiver"
+}

examples/psi/config/threshold_ecdh_psi_sender.json

@@ -0,0 +1,36 @@
+{
+  "psi_config": {
+    "protocol_config": {
+      "protocol": "PROTOCOL_ECDH",
+      "ecdh_config": {
+        "curve": "CURVE_FOURQ"
+      },
+      "role": "ROLE_SENDER",
+      "broadcast_result": true
+    },
+    "input_config": {
+      "type": "IO_TYPE_FILE_CSV",
+      "path": "/tmp/sender_input.csv"
+    },
+    "output_config": {
+      "type": "IO_TYPE_FILE_CSV",
+      "path": "/tmp/sender_output.csv"
+    },
+    "keys": ["id_0", "id_1"],
+    "intersection_threshold": 50,
+    "count_path": "/tmp/sender_count.csv"
+  },
+  "link_config": {
+    "parties": [
+      {
+        "id": "receiver",
+        "host": "127.0.0.1:5300"
+      },
+      {
+        "id": "sender",
+        "host": "127.0.0.1:5400"
+      }
+    ]
+  },
+  "self_link_party": "sender"
+}

experiment/psi/threshold_ecdh_psi/BUILD.bazel

@@ -0,0 +1,103 @@
+# Copyright 2025
+#
+# 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_binary", "psi_cc_library", "psi_cc_test")
+
+package(default_visibility = ["//visibility:public"])
+
+psi_cc_library(
+    name = "common",
+    srcs = ["common.cc"],
+    hdrs = ["common.h"],
+    deps = [
+        "//psi/proto:psi_cc_proto",
+        "//psi/proto:psi_v2_cc_proto",
+        "//psi/utils:io",
+        "//psi/utils:test_utils",
+    ],
+)
+
+psi_cc_library(
+    name = "receiver",
+    srcs = ["receiver.cc"],
+    hdrs = ["receiver.h"],
+    deps = [
+        ":common",
+        ":threshold_ecdh_psi",
+        "//psi/algorithm/ecdh:receiver",
+        "//psi/utils:serialize",
+        "//psi/utils:sync",
+    ],
+)
+
+psi_cc_library(
+    name = "sender",
+    srcs = ["sender.cc"],
+    hdrs = ["sender.h"],
+    deps = [
+        ":common",
+        ":threshold_ecdh_psi",
+        ":threshold_ecdh_psi_cache",
+        "//psi/algorithm/ecdh:sender",
+        "//psi/utils:serialize",
+        "//psi/utils:sync",
+    ],
+)
+
+psi_cc_library(
+    name = "threshold_ecdh_psi_cache",
+    srcs = ["threshold_ecdh_psi_cache.cc"],
+    hdrs = ["threshold_ecdh_psi_cache.h"],
+    deps = [
+        "//psi/utils:batch_provider",
+        "//psi/utils:io",
+    ],
+)
+
+psi_cc_library(
+    name = "threshold_ecdh_psi",
+    srcs = ["threshold_ecdh_psi.cc"],
+    hdrs = ["threshold_ecdh_psi.h"],
+    deps = [
+        ":threshold_ecdh_psi_cache",
+        "//psi/algorithm/ecdh:ecdh_psi",
+        "@yacl//yacl/crypto/rand",
+    ],
+)
+
+psi_cc_binary(
+    name = "threshold_ecdh_psi_benchmark",
+    srcs = ["threshold_ecdh_psi_benchmark.cc"],
+    deps = ["@com_github_google_benchmark//:benchmark_main"] + [
+        ":receiver",
+        ":common",
+        ":sender",
+        "//psi/utils:test_utils",
+        "//psi/utils:random_str",
+        "@yacl//yacl/link",
+    ],
+)
+
+psi_cc_test(
+    name = "threshold_ecdh_psi_test",
+    srcs = ["threshold_ecdh_psi_test.cc"],
+    deps = [
+        ":common",
+        ":receiver",
+        ":sender",
+        "//psi/utils:random_str",
+        "//psi/utils:test_utils",
+        "@yacl//yacl/link",
+    ],
+)

experiment/psi/threshold_ecdh_psi/README.md

@@ -0,0 +1,49 @@
+# Threshold ECDH Private Set Intersection
+
+## Introduction
+
+Threshold Private Set Intersection (Threshold PSI) allows users to set an intersection threshold to limit the size of the intersection. Specifically, the final intersection size is the smaller count between the threshold and the actual intersection size. Additionally, both parties are restricted from accessing any part of the intersection beyond the defined threshold. This protocol is a functional extension based on ECDH PSI.
+
+## Threshold ECDH PSI
+
+$P_1$ and $P_2$ are the participants in the protocol, where $P_1$ is the receiver and $P_2$ is the sender. $X$ and $Y$ are the respective input sets of $P_1$ and $P_2$ , with size $n_1 = \left\vert X \right\vert$ and $n_2 = \left\vert Y \right\vert$.
+
+1. For $1 \le i \le n_1$ , $P_1$ randomly chooses $\alpha \gets Z_q$ , computes $a_i = H(x_i)^{\alpha}$ , and sends the set $\\{a_1,a_2,...,a_{n_1}\\}$ to $P_2$ .
+2. For $1 \le i \le n_1$ , $P_2$ randomly chooses $\beta \gets Z_q$ , computes $a_i' = (H(x_i)^{\alpha})^{\beta}$ . Then it samples a random permutation $\pi$ and sends the permuted set $TX = \pi(\\{a_1',a_2',...,a'_{n_1}\\})$ to $P_1$.
+3. For $1 \le j \le n_2$ , $P_2$ computes $b_j = H(y_j)^{\beta}$ , and sends the set $\\{b_1,b_2,...,b_{n_2}\\}$ to $P_1$ .
+4. For $1 \le j \le n_2$ , $P_1$ computes $b_j' = (H(y_j)^{\beta})^{\alpha}$ and gets the set $TY = \\{b_1',b_2',...,b'_{n_2}\\}$ .
+5. $P_1$ truncates the intersection $TX \cap TY$ based on the threshold and sends the truncated index set to $P_2$ .
+6. $P_2$ restores the index set based on the permutation $\pi$ from the second step and sends the restored index set to $P_1$ .
+7. $P_1$ obtains the intersection through the restored index set.
+
+## Test
+
+1. Compile the binary
+
+    ```bash
+    bazel build //psi/apps/psi_launcher:main
+    ```
+
+2. Generate test data
+
+    ```bash
+    $ python examples/psi/generate_psi_data.py --receiver_item_cnt 1e3 \
+        --sender_item_cnt 1e3 --intersection_cnt 1e2 --id_cnt 2 \
+        --receiver_path /tmp/receiver_input.csv \
+        --sender_path /tmp/sender_input.csv \
+        --intersection_path /tmp/intersection.csv
+    ```
+
+3. Launch Threshold ECDH PSI
+
+    For sender terminal,
+
+    ```bash
+    ./bazel-bin/psi/apps/psi_launcher/main --config $(pwd)/examples/psi/config/threshold_ecdh_psi_sender.json
+    ```
+
+    For receiver terminal,
+
+    ```bash
+    ./bazel-bin/psi/apps/psi_launcher/main --config $(pwd)/examples/psi/config/threshold_ecdh_psi_receiver.json
+    ```

experiment/psi/threshold_ecdh_psi/common.cc

@@ -0,0 +1,79 @@
+// Copyright 2025
+//
+// 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 "experiment/psi/threshold_ecdh_psi/common.h"
+
+#include "spdlog/spdlog.h"
+
+#include "psi/utils/io.h"
+#include "psi/utils/test_utils.h"
+
+#include "psi/proto/psi.pb.h"
+
+namespace psi::ecdh {
+void SaveIntersectionCount(const std::string &count_path, uint32_t real_count,
+                           uint32_t final_count) {
+  auto ofs = io::BuildOutputStream(io::FileIoOptions(count_path));
+
+  ofs->Write("real_count,final_count\n");
+  ofs->Write(fmt::format("{},{}\n", real_count, final_count));
+
+  ofs->Close();
+  SPDLOG_INFO("Save intersection count to {}", count_path);
+}
+
+void GeneratePsiConfig(const std::filesystem::path &tmp_folder,
+                       const std::vector<std::string> &items_sender,
+                       const std::vector<std::string> &items_receiver,
+                       uint32_t threshold, v2::PsiConfig &sender_config,
+                       v2::PsiConfig &receiver_config) {
+  std::string sender_input_path = tmp_folder / "sender_input.csv";
+  std::string receiver_input_path = tmp_folder / "receiver_input.csv";
+  std::string sender_output_path = tmp_folder / "sender_output.csv";
+  std::string receiver_output_path = tmp_folder / "receiver_output.csv";
+  std::string sender_cache_path = tmp_folder / "sender_cache";
+  std::string receiver_cache_path = tmp_folder / "receiver_cache";
+
+  test::WriteCsvFile(sender_input_path, items_sender);
+  test::WriteCsvFile(receiver_input_path, items_receiver);
+
+  sender_config.mutable_protocol_config()->set_protocol(
+      v2::Protocol::PROTOCOL_ECDH);
+  sender_config.mutable_protocol_config()->mutable_ecdh_config()->set_curve(
+      CurveType::CURVE_FOURQ);
+  sender_config.mutable_protocol_config()->set_role(v2::Role::ROLE_SENDER);
+  sender_config.mutable_protocol_config()->set_broadcast_result(true);
+  sender_config.mutable_input_config()->set_type(v2::IoType::IO_TYPE_FILE_CSV);
+  sender_config.mutable_input_config()->set_path(sender_input_path);
+  sender_config.mutable_output_config()->set_type(v2::IoType::IO_TYPE_FILE_CSV);
+  sender_config.mutable_output_config()->set_path(sender_output_path);
+  sender_config.add_keys("id");
+  sender_config.set_intersection_threshold(threshold);
+
+  receiver_config.mutable_protocol_config()->set_protocol(
+      v2::Protocol::PROTOCOL_ECDH);
+  receiver_config.mutable_protocol_config()->mutable_ecdh_config()->set_curve(
+      CurveType::CURVE_FOURQ);
+  receiver_config.mutable_protocol_config()->set_role(v2::Role::ROLE_RECEIVER);
+  receiver_config.mutable_protocol_config()->set_broadcast_result(true);
+  receiver_config.mutable_input_config()->set_type(
+      v2::IoType::IO_TYPE_FILE_CSV);
+  receiver_config.mutable_input_config()->set_path(receiver_input_path);
+  receiver_config.mutable_output_config()->set_type(
+      v2::IoType::IO_TYPE_FILE_CSV);
+  receiver_config.mutable_output_config()->set_path(receiver_output_path);
+  receiver_config.add_keys("id");
+  receiver_config.set_intersection_threshold(threshold);
+}
+}  // namespace psi::ecdh

experiment/psi/threshold_ecdh_psi/common.h

@@ -0,0 +1,36 @@
+// Copyright 2025
+//
+// 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 <filesystem>
+#include <fstream>
+#include <string>
+#include <vector>
+
+#include "psi/proto/psi_v2.pb.h"
+
+namespace psi::ecdh {
+// The receiver or sender may need to know the size of the real intersection and
+// the size of the intersection after being restricted by the threshold.
+void SaveIntersectionCount(const std::string &count_path, uint32_t real_count,
+                           uint32_t final_count);
+
+// Generate PsiConfig for test and benchmark.
+void GeneratePsiConfig(const std::filesystem::path &tmp_folder,
+                       const std::vector<std::string> &items_sender,
+                       const std::vector<std::string> &items_receiver,
+                       uint32_t threshold, v2::PsiConfig &sender_config,
+                       v2::PsiConfig &receiver_config);
+}  // namespace psi::ecdh