Add YPIR algorithm implementation by sh1k4ku · Pull Request #380 · secretflow/psi

sh1k4ku · 2026-03-11T11:31:30Z

Add YPIR PIR scheme with SimplePIR and DoublePIR support
Add AES PRNG, HEXL NTT integration, and utility modules
Add Bazel build rules for YPIR and hexl dependency
Update Spiral algorithm (NTT, params, poly matrix, client/server)
Add integration tests and basic flow tests

- Add YPIR PIR scheme with SimplePIR and DoublePIR support - Add AES PRNG, HEXL NTT integration, and utility modules - Add Bazel build rules for YPIR and hexl dependency - Update Spiral algorithm (NTT, params, poly matrix, client/server) - Add integration tests and basic flow tests

github-actions · 2026-03-11T11:31:48Z

All contributors have signed the CLA ✍️ ✅
_{Posted by the CLA Assistant Lite bot.}

gemini-code-assist · 2026-03-11T12:14:10Z

Summary of Changes

Hello, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly expands the system's PIR capabilities by adding a YPIR scheme with SimplePIR and DoublePIR support. It incorporates several key improvements, including optimized NTT operations via HEXL, enhanced build rules, and updates to the Spiral algorithm. The addition of integration and flow tests ensures the stability and correctness of these new features.

Highlights

YPIR Algorithm Implementation: This pull request introduces a new YPIR PIR scheme, enhancing the system's private information retrieval capabilities.
HEXL NTT Integration: The integration of HEXL NTT provides optimized Number Theoretic Transform operations, crucial for efficient homomorphic computations.
Bazel Build Rules: Added Bazel build rules for YPIR and its hexl dependency, streamlining the build process and dependency management.
Spiral Algorithm Update: Updates to the Spiral algorithm include NTT improvements, parameter adjustments, and client/server modifications for enhanced performance.
Integration and Basic Flow Tests: Comprehensive integration and basic flow tests have been added to ensure the YPIR implementation functions correctly and reliably.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Activity

Added YPIR PIR scheme with SimplePIR and DoublePIR support
Integrated AES PRNG and HEXL NTT
Added Bazel build rules for YPIR and hexl dependency
Updated Spiral algorithm (NTT, params, poly matrix, client/server)
Added integration tests and basic flow tests

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	`/gemini review`	Performs a code review for the current pull request in its current state.
Pull Request Summary	`/gemini summary`	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	`/gemini help`	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

gemini-code-assist

Code Review

This pull request introduces a new PIR scheme, YPIR, with support for both SimplePIR and DoublePIR. It includes the core algorithm implementation, integration with HEXL for NTT, and new utility modules. The changes also update the existing Spiral algorithm, including parameters, NTT implementation, and client/server logic.

My review focuses on several key areas. I've identified critical security vulnerabilities related to the use of fixed seeds for random number generation in core cryptographic functions. These must be addressed by using a secure source of randomness. I also found a potential bug in a utility function that appears to discard computed results, and some maintainability issues with hardcoded parameters. Additionally, there are build-related concerns regarding portability due to a hardcoded dependency path and architecture-specific compiler flags. While the overall implementation of the new algorithm is substantial, these issues should be resolved to ensure the code is secure, correct, and maintainable.

_{Note: Security Review did not run due to the size of the PR.}

gemini-code-assist · 2026-03-11T12:18:54Z

 PolyMatrixRaw PolyMatrixRaw::Random(const Params& params, size_t rows,
                                    size_t cols) {
-  yacl::crypto::Prg<uint64_t> prg(yacl::crypto::SecureRandU128());
+  yacl::crypto::Prg<uint64_t> prg(yacl::MakeUint128(0, 20001));


The Random method is implemented using a fixed seed. This means it will produce the same "random" matrix every time, which is a major security risk if this function is used in a production context where unpredictability is required. For testing, a fixed seed is acceptable, but for production code, a source of true randomness like yacl::crypto::SecureRandU128() should be used. The function name is also misleading. If this is for testing, it should be renamed to something like RandomForTest. A similar issue exists for PolyMatrixNtt::Random on line 333.

Suggested change

yacl::crypto::Prg<uint64_t> prg(yacl::MakeUint128(0, 20001));

yacl::crypto::Prg<uint64_t> prg(yacl::crypto::SecureRandU128());

gemini-code-assist · 2026-03-11T12:18:54Z

+  uint128_t seed = yacl::MakeUint128(0, 10001);
  yacl::crypto::Prg<uint64_t> rng(seed);
-  uint128_t pub_seed = yacl::crypto::SecureRandU128();
+  uint128_t pub_seed = yacl::MakeUint128(0, 10002);
  yacl::crypto::Prg<uint64_t> rng_pub(pub_seed);


The GenPublicKeys function uses fixed seeds for its random number generators. This is a critical security vulnerability as it makes the public keys, and any encryption performed with them, deterministic and predictable. Cryptographic keys must be generated with a secure source of randomness. Please use yacl::crypto::SecureRandU128() to seed the PRNGs.

Suggested change

uint128_t seed = yacl::MakeUint128(0, 10001);

yacl::crypto::Prg<uint64_t> rng(seed);

uint128_t pub_seed = yacl::crypto::SecureRandU128();

uint128_t pub_seed = yacl::MakeUint128(0, 10002);

yacl::crypto::Prg<uint64_t> rng_pub(pub_seed);

uint128_t seed = yacl::crypto::SecureRandU128();

yacl::crypto::Prg<uint64_t> rng(seed);

uint128_t pub_seed = yacl::crypto::SecureRandU128();

yacl::crypto::Prg<uint64_t> rng_pub(pub_seed);

gemini-code-assist · 2026-03-11T12:18:54Z

+  yacl::crypto::Prg<uint64_t> rng(yacl::MakeUint128(0, 10003));
  // a empty query
  SpiralQuery query;
-  uint128_t query_seed;
-  rng.Fill(absl::MakeSpan(reinterpret_cast<uint8_t*>(&query_seed),
-                          sizeof(query_seed)));
+  uint128_t query_seed = yacl::MakeUint128(0, 10004);
  yacl::crypto::Prg<uint64_t> rng_pub(query_seed);
  query.seed_ = query_seed;


The GenQueryInternal function uses fixed seeds for its random number generators. This makes the query generation deterministic, which is a critical security flaw in a cryptographic protocol. An attacker could potentially predict queries. A secure random seed must be used.

Suggested change

yacl::crypto::Prg<uint64_t> rng(yacl::MakeUint128(0, 10003));

// a empty query

SpiralQuery query;

uint128_t query_seed;

rng.Fill(absl::MakeSpan(reinterpret_cast<uint8_t*>(&query_seed),

sizeof(query_seed)));

uint128_t query_seed = yacl::MakeUint128(0, 10004);

yacl::crypto::Prg<uint64_t> rng_pub(query_seed);

query.seed_ = query_seed;

yacl::crypto::Prg<uint64_t> rng(yacl::crypto::SecureRandU128());

// a empty query

SpiralQuery query;

uint128_t query_seed;

rng.Fill(absl::MakeSpan(reinterpret_cast<uint8_t*>(&query_seed),

sizeof(query_seed)));

yacl::crypto::Prg<uint64_t> rng_pub(query_seed);

query.seed_ = query_seed;

gemini-code-assist · 2026-03-11T12:18:54Z

+new_local_repository(
+    name = "hexl",
+    build_file = "//bazel:hexl.BUILD",
+    path = "/usr/local",
+)


The new_local_repository for hexl uses a hardcoded path /usr/local. This makes the build non-portable and dependent on the local environment setup. It will fail on systems where hexl is installed in a different location or not at all. Consider using http_archive to fetch and build hexl from a source repository to ensure a hermetic and portable build. If hexl must be a local dependency, its path should be configurable, for example, via an environment variable or a .bazelrc file.

gemini-code-assist · 2026-03-11T12:18:54Z

+  for (size_t r = 0; r < in.Rows(); ++r) {
+    size_t c = 0;
+
+    auto in_poly = in.Poly(r, c);
+    auto scratch_sub = scratch.subspan(0, poly_size);
+    std::copy(in_poly.begin(), in_poly.end(), scratch_sub.begin());
+
+    arith::NttInverse(params, scratch_sub);
+
+    if (r == 0) {
+      size_t raw_poly_idx = out.PolyStartIndex(r, c);
+      std::vector<uint64_t> temp_vec(scratch_sub.begin(), scratch_sub.end());
+      for (size_t i = 0; i < params.PolyLen(); ++i) {
+        out.Data()[raw_poly_idx + i] = params.CrtCompose(temp_vec, i);
+      }
+    }
+  }


In FromNttScratch, the loop iterates over the rows of the input matrix in, which is enforced to have 2 rows. However, the if (r == 0) block means that the output matrix out is only updated for the first row (r=0). For the second row (r=1), arith::NttInverse is called, but its result is not used to update out. This seems incorrect, as the second row of the result is never computed. The function should likely process and store the results for all rows.

gemini-code-assist · 2026-03-11T12:18:54Z

 DEBUG_FLAGS = ["-O0", "-g", "-DSPDLOG_ACTIVE_LEVEL=1"]
 RELEASE_FLAGS = ["-O2"]
-FAST_FLAGS = ["-O1"]
+FAST_FLAGS = ["-O3", "-march=native"]


The FAST_FLAGS now include -march=native. While this can improve performance by generating code optimized for the specific CPU of the build machine, it can lead to portability issues. The resulting binary may not run on machines with older CPUs that lack the instruction sets used. If the compiled artifacts are intended to be distributed and run on various machines, it's safer to target a more general architecture (e.g., -march=x86-64-v2) or remove this flag.

sh1k4ku · 2026-03-26T08:31:12Z

I have read the CLA Document and I hereby sign the CLA

sh1k4ku · 2026-03-26T08:31:17Z

recheck

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

gemini-code-assist Bot reviewed Mar 11, 2026

View reviewed changes

sh1k4ku added 3 commits March 25, 2026 13:05

Integrate ypir interface refactor

82cdca0

fix some issues

e834bdd

fix issues again

7f23043

sh1k4ku and others added 8 commits March 26, 2026 17:19

fix unittest and ci issues

9efab1d

fix clang format

6a0bd25

fix sth

f776f21

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

no avx implementation

47b7275

fix linux_ut-2xlarge

a1cc2ab

fix macos

62eecf9

fix macos yacl

c796183

fix macos brpc

7fa2ab9

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add YPIR algorithm implementation#380

Add YPIR algorithm implementation#380
sh1k4ku wants to merge 12 commits into
secretflow:mainfrom
changtong9:ypir-dev

sh1k4ku commented Mar 11, 2026

Uh oh!

github-actions Bot commented Mar 11, 2026 •

edited

Loading

Uh oh!

gemini-code-assist Bot commented Mar 11, 2026

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

gemini-code-assist Bot Mar 11, 2026

Uh oh!

gemini-code-assist Bot Mar 11, 2026

Uh oh!

gemini-code-assist Bot Mar 11, 2026

Uh oh!

gemini-code-assist Bot Mar 11, 2026

Uh oh!

gemini-code-assist Bot Mar 11, 2026

Uh oh!

gemini-code-assist Bot Mar 11, 2026

Uh oh!

Uh oh!

sh1k4ku commented Mar 26, 2026

Uh oh!

sh1k4ku commented Mar 26, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

	yacl::crypto::Prg<uint64_t> prg(yacl::MakeUint128(0, 20001));
	yacl::crypto::Prg<uint64_t> prg(yacl::crypto::SecureRandU128());

Conversation

sh1k4ku commented Mar 11, 2026

Uh oh!

github-actions Bot commented Mar 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

gemini-code-assist Bot commented Mar 11, 2026

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Mar 11, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Mar 11, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Mar 11, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Mar 11, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Mar 11, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Mar 11, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

sh1k4ku commented Mar 26, 2026

Uh oh!

sh1k4ku commented Mar 26, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

github-actions Bot commented Mar 11, 2026 •

edited

Loading