Puncturable Pseudorandom Set-Based Private Information Retrieval

experiment/pir/pps/BUILD.bazel

@@ -0,0 +1,109 @@
+# Copyright 2023 Ant Group Co., Ltd.
+#
+# 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 = "ggm_pset",
+    srcs = ["ggm_pset.cc"],
+    hdrs = ["ggm_pset.h"],
+    deps = [
+        "@yacl//yacl/base:dynamic_bitset",
+        "@yacl//yacl/crypto/rand",
+        "@yacl//yacl/crypto/tools:prg",
+    ],
+)
+
+psi_cc_library(
+    name = "client",
+    srcs = ["client.cc"],
+    hdrs = ["client.h"],
+    deps = [
+        ":ggm_pset",
+        "@yacl//yacl/base:dynamic_bitset",
+        "@yacl//yacl/crypto/rand",
+        "@yacl//yacl/crypto/tools:prg",
+    ],
+)
+
+psi_cc_library(
+    name = "server",
+    srcs = ["server.cc"],
+    hdrs = ["server.h"],
+    deps = [
+        ":ggm_pset",
+        "@yacl//yacl/base:dynamic_bitset",
+        "@yacl//yacl/crypto/rand",
+        "@yacl//yacl/crypto/tools:prg",
+    ],
+)
+
+psi_cc_library(
+    name = "sender",
+    srcs = ["sender.cc"],
+    hdrs = ["sender.h"],
+    deps = [
+        ":ggm_pset",
+        "@yacl//yacl/base:dynamic_bitset",
+        "@yacl//yacl/link:context",
+    ],
+)
+
+psi_cc_library(
+    name = "receiver",
+    srcs = ["receiver.cc"],
+    hdrs = ["receiver.h"],
+    deps = [
+        ":ggm_pset",
+        "@yacl//yacl/base:dynamic_bitset",
+        "@yacl//yacl/link:context",
+    ],
+)
+
+psi_cc_test(
+    name = "pps_test",
+    srcs = ["ggm_pset_test.cc"],
+    deps = [
+        ":ggm_pset",
+    ],
+)
+
+psi_cc_test(
+    name = "pps_pir_test",
+    srcs = ["pps_pir_test.cc"],
+    deps = [
+        ":client",
+        ":receiver",
+        ":sender",
+        ":server",
+        "@yacl//yacl/base:dynamic_bitset",
+        "@yacl//yacl/link",
+    ],
+)
+
+psi_cc_binary(
+    name = "pps_pir_benchmark",
+    srcs = ["pps_pir_benchmark.cc"],
+    deps = [
+        ":client",
+        ":receiver",
+        ":sender",
+        ":server",
+        "@com_github_google_benchmark//:benchmark_main",
+        "@yacl//yacl/base:dynamic_bitset",
+        "@yacl//yacl/link",
+    ],
+)

experiment/pir/pps/README.md

@@ -0,0 +1,46 @@
+# PPS-Based PIR Implementation
+
+## Overview
+
+This repository contains the implementation of the Single Bits and Multi-Bits Private Information Retrieval (PIR) algorithms, which incorporate Puncturable Pseudorandom Sets (PPS) as described in the paper "Private Information Retrieval with Sublinear Online Time" (https://eprint.iacr.org/2019/1075.pdf) by Henry Corrigan-Gibbs and Dmitry Kogan. By utilizing PPS, these implementations achieve fast, sublinear-time database lookups without increasing the server-side storage requirements. The use of PPS enables the system to efficiently manage the complexity of data retrieval while ensuring the privacy of the client's queries, optimizing both Single Bits and Multi-Bits PIR functionalities.
+
+
+
+## Implementation Details
+
+
+This project only implements the dual-server model described in the paper. The implementation of Puncturable Pseudorandom Sets is based on Chapter 2 of the paper, "Construction 4 (Puncturable pseudorandom set from puncturable PRF)." The implementation of the Single-Bits PIR protocol follows "Construction 16 (Two-server PIR with sublinear online time)" from Chapter 3 of the original paper. The implementation of the Multi-Bits PIR protocol is based on Chapter 4 and Appendix D's "Construction 44 (Multi-query offline/online PIR)."
+
+**(Offline/online PIR)**. An offline/online PIR scheme is a tuple $\Pi$ = (**Setup**, **Hint**, **Query**, **Answer**, **Reconstruct**) of five efficient algorithms:
+
+* **Setup**$(1^{\lambda}, n) \to (ck, q_h)$,  a randomized algorithm that takes in security parameter $\lambda$ and database length $n$ and outputs a client key $ck$ and a hint request $q_h$.
+
+* **Hint**$(x, q_h) \to h$,  a deterministic algorithm that takes in a database $x \in \{0, 1\}^n$ and a hint request $q_h$ and outputs a hint $h$.
+
+* **Query**$(ck, i) \to h$,  a randomized algorithm that takes in the client’s key $ck$ and an index $i \in [n]$, and outputs a a query $q$.
+
+* **Answer**$(q) \to a$,  a deterministic algorithm that takes as input a query $q$ and gets access to an oracle that:
+    * takes as input an index $j \in [n]$, and
+    * returns the $j$-th bit of the database $x_j ∈ \{0, 1\}$
+
+    outputs an answer string $a$, and
+
+* **Reconstruct**$(h, a) \to x_i$,  , a deterministic algorithm that takes as a hint $h$ and an answer $a$, and outputs a bit $x_i$.
+
+
+**Puncturable pseudorandom sets**
+
+Puncturable pseudorandom sets are an extension of puncturable pseudorandom functions (PRFs). A typical pseudorandom function generates outputs that are indistinguishable from random by any efficient algorithm, given only the outputs and not the secret key. When a PRF is punctured at a particular point, it behaves like a normal PRF for all inputs except for the punctured point, for which the output or behavior is obscured or undefined.
+
+Our implementation utilizes the AES pseudorandom generator (PRG) to construct a GGM tree-based pseudorandom function (PRF), and then employs Construction 4 from the original paper to build the Puncturable Pseudorandom Set (PPS).
+
+**Single-Bits PIR**
+
+Construction 16 (Two-server PIR with sublinear online time).
+
+**Multi-Bits PIR**
+
+Construction 44 (Multi-query offline/online PIR).
+
+** Contact **
+email: yangw.ing@foxmail.com

experiment/pir/pps/client.cc

@@ -0,0 +1,152 @@
+#include "client.h"
+
+#include <spdlog/spdlog.h>
+
+namespace pir::pps {
+
+bool PpsPirClient::Bernoulli() {
+  std::random_device rand;
+  std::mt19937 gen(rand());
+  double p =
+      static_cast<double>(set_size_ - 1) / static_cast<double>(universe_size_);
+  std::bernoulli_distribution distribution(p);
+  return distribution(gen);
+}
+
+uint64_t PpsPirClient::GetRandomU64Less() {
+  return LemireTrick(yacl::crypto::RandU64(), set_size_);
+}
+
+// Generate sk and m random numbers \in [n]
+void PpsPirClient::Setup(PIRKey& sk, std::set<uint64_t>& deltas) {
+  sk = pps_.Gen(lambda_);
+  // The map.size() must be equal to SET_SIZE.
+  std::vector<uint64_t> rand =
+      yacl::crypto::PrgAesCtr<uint64_t>(yacl::crypto::RandU64(), M());
+  for (uint64_t i = 0; i < M(); i++) {
+    // The most expensive operation.
+    uint64_t r = LemireTrick(rand[i], universe_size_);
+    if (!deltas.insert(r).second) {
+      rand[i] = yacl::crypto::RandU64();
+      i--;
+    }
+  }
+}
+
+// Params:
+// For input index i_{pir} \in [n], ouput punc key.
+void PpsPirClient::Query(uint64_t i, PIRKey& sk, std::set<uint64_t>& deltas,
+                         PIRQueryParam& param, PIRPuncKey& sk_punc) {
+  std::set<uint64_t>::iterator iter = deltas.begin();
+  const PIREvalMap map = pps_.getMap();
+  for (param.j_ = 0; iter != deltas.end(); ++param.j_, ++iter) {
+    uint64_t r = MODULE_SUB(i, *iter, universe_size_);
+    if (map.find(r) != map.end()) {
+      sk_punc.delta_ = *iter;
+      break;
+    }
+  }
+  if (iter == deltas.end()) {
+    SPDLOG_INFO("Can't find a j \\in m such that i - \\delta_j \\in Eval(sk)");
+    param.j_ = PIR_ABORT;
+    auto map_iter = std::next(map.begin(), GetRandomU64Less());
+    sk_punc.delta_ = MODULE_SUB(i, map_iter->first, universe_size_);
+  }
+
+  uint64_t i_punc;
+  if ((param.b_ = Bernoulli())) {
+    auto map_iter = std::next(map.begin(), GetRandomU64Less());
+    i_punc = map_iter->first;
+  } else {
+    i_punc = MODULE_SUB(i, sk_punc.delta_, universe_size_);
+  }
+  pps_.Punc(i_punc, sk, sk_punc);
+}
+
+uint64_t PpsPirClient::Reconstruct(PIRQueryParam& param,
+                                   yacl::dynamic_bitset<>& h, bool a, bool& r) {
+  if (param.b_ || (param.j_ == PIR_ABORT)) {
+    SPDLOG_INFO("Reconstruct: Param b == 1 OR j == \\abort");
+    return PIR_ABORT;
+  }
+  r = a ^ h[param.j_];
+  return PIR_OK;
+}
+
+void PpsPirClient::Setup(std::vector<PIRKeyUnion>& ck,
+                         std::vector<std::unordered_set<uint64_t>>& v) {
+  ck.resize(MM());
+  v.resize(MM());
+  std::vector<uint128_t> rand =
+      yacl::crypto::PrgAesCtr<uint128_t>(yacl::crypto::RandU128(), MM());
+  for (uint64_t i = 0; i < MM(); ++i) {
+    pps_.Eval(rand[i], v[i]);
+    if (v[i].size() == set_size_) {
+      ck[i] = PIRKeyUnion(rand[i]);
+    } else {
+      v[i].clear();
+      rand[i] = yacl::crypto::RandU128();
+      --i;
+    }
+  }
+}
+
+void PpsPirClient::Query(uint64_t i, std::vector<PIRKeyUnion>& ck,
+                         std::vector<std::unordered_set<uint64_t>>& v,
+                         PIRQueryParam& param, PIRPuncKey& punc_l,
+                         PIRPuncKey& punc_r) {
+  // GenWith(1^\lambda, n, i)
+
+  PIRKey k_new = pps_.Gen(lambda_);
+  PIRKeyUnion k_right;
+  uint64_t rand_i = std::next(pps_.getMap().begin(), GetRandomU64Less())->first;
+  punc_l.delta_ = MODULE_SUB(i, rand_i, universe_size_);
+
+  uint64_t i_punc;
+  param.j_ = PIR_ABORT;
+  if ((param.b_ = Bernoulli())) {
+    rand_i = std::next(pps_.getMap().begin(), GetRandomU64Less())->first;
+    i_punc = MODULE_ADD(rand_i, punc_l.delta_, universe_size_);
+  } else {
+    for (param.j_ = 0; param.j_ < ck.size(); ++param.j_) {
+      rand_i = MODULE_SUB(i, ck[param.j_].delta_, universe_size_);
+      if (v[param.j_].find(rand_i) != v[param.j_].end()) {
+        k_right = PIRKeyUnion(ck[param.j_].k_, ck[param.j_].delta_);
+        ck[param.j_] = PIRKeyUnion(k_new, punc_l.delta_);
+        v[param.j_].clear();
+        pps_.Eval(k_new, v[param.j_]);
+        break;
+      }
+    }
+
+    if (param.j_ == ck.size()) {
+      SPDLOG_INFO("Can't find a j \\in [m] such that i \\in Eval(sk_j)");
+      param.j_ = PIR_ABORT;
+    }
+    i_punc = i;
+  }
+
+  uint64_t i_punc_l = MODULE_SUB(i_punc, punc_l.delta_, universe_size_);
+  pps_.Punc(i_punc_l, k_new, punc_l);
+  if (param.j_ != PIR_ABORT) {
+    pps_.EvalMap(k_right.k_);
+    uint64_t i_punc_r = MODULE_SUB(i_punc, k_right.delta_, universe_size_);
+    pps_.Punc(i_punc_r, k_right.k_, punc_r);
+    punc_r.delta_ = k_right.delta_;
+  } else {
+    punc_r = punc_l;
+  }
+}
+
+uint64_t PpsPirClient::Reconstruct(PIRQueryParam& param,
+                                   yacl::dynamic_bitset<>& h, bool a_left,
+                                   bool a_right, bool& r) {
+  if (param.j_ != PIR_ABORT) {
+    r = a_right ^ h[param.j_];
+    h[param.j_] = a_left ^ r;
+    return PIR_OK;
+  }
+  SPDLOG_INFO("Reconstruct: j == \\abort");
+  return PIR_ABORT;
+}
+}  // namespace pir::pps

experiment/pir/pps/client.h

@@ -0,0 +1,82 @@
+#include <random>
+#include <set>
+
+#include "ggm_pset.h"
+#include "yacl/crypto/rand/rand.h"
+
+namespace pir::pps {
+
+// j, b_ in Query
+struct PIRQueryParam {
+  uint64_t j_;
+  bool b_;
+};
+
+class PpsPirClient {
+ public:
+  PpsPirClient() : pps_(), lambda_(0), universe_size_(0), set_size_(0) {}
+
+  PpsPirClient(uint32_t lambda, uint64_t universe_size, uint64_t set_size)
+      : pps_(universe_size, set_size),
+        lambda_(lambda),
+        universe_size_(universe_size),
+        set_size_(set_size) {}
+
+  // Get m = (n / s(n)) * log(n)
+  uint64_t M() {
+    uint64_t q = universe_size_ / set_size_;
+    uint64_t r = universe_size_ % set_size_;
+    if (r) {
+      q += 1;
+    }
+    return q * static_cast<uint32_t>(std::ceil(std::log(universe_size_)));
+  }
+
+  // sample a bit b from Bernoulli((s - 1) / n)
+  bool Bernoulli();
+
+  // sample a random from [0, universe_size_]
+  uint64_t GetRandomU64Less();
+
+  // Setup(1^\lambda, universe_size_) -> ck, q_h
+  void Setup(PIRKey& sk, std::set<uint64_t>& deltas);
+
+  // Query(ck, i \in [n]) -> q \in K_p
+  void Query(uint64_t i, PIRKey& sk, std::set<uint64_t>& deltas,
+             PIRQueryParam& param, PIRPuncKey& sk_punc);
+
+  // Reconstruct(h ∈ {0, 1}^m, a ∈ {0, 1}) → x_i
+  uint64_t Reconstruct(PIRQueryParam& param, yacl::dynamic_bitset<>& h, bool a,
+                       bool& r);
+
+  // Get m = (2 * n / s(n)) * log(n)
+  uint64_t MM() {
+    uint64_t q = 2 * universe_size_ / set_size_;
+    uint64_t r = 2 * universe_size_ % set_size_;
+    if (r) {
+      q += 1;
+    }
+    return q * static_cast<uint32_t>(std::ceil(std::log(universe_size_)));
+  }
+
+  // Construction 44 (Multi-query offline/online PIR)
+  // Setup(1^λ, n) → (ck, q_h)
+  void Setup(std::vector<PIRKeyUnion>& ck,
+             std::vector<std::unordered_set<uint64_t>>& v);
+
+  // Query(ck, i) → (ck, q_left, q_right)
+  void Query(uint64_t i, std::vector<PIRKeyUnion>& ck,
+             std::vector<std::unordered_set<uint64_t>>& v, PIRQueryParam& param,
+             PIRPuncKey& punc_l, PIRPuncKey& punc_r);
+
+  // Reconstruct(h, a_left, a_right) → (h′, x_i)
+  uint64_t Reconstruct(PIRQueryParam& param, yacl::dynamic_bitset<>& h,
+                       bool a_left, bool a_right, bool& r);
+
+ private:
+  PPS pps_;
+  uint32_t lambda_;
+  uint64_t universe_size_;
+  uint64_t set_size_;
+};
+}  // namespace pir::pps