diff --git a/heu/library/algorithms/ashe/BUILD.bazel b/heu/library/algorithms/ashe/BUILD.bazel new file mode 100644 index 00000000..43bd1769 --- /dev/null +++ b/heu/library/algorithms/ashe/BUILD.bazel @@ -0,0 +1,113 @@ +# Copyright 2024 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("@yacl//bazel:yacl.bzl", "yacl_cc_library", "yacl_cc_test") + +package(default_visibility = ["//visibility:public"]) + +test_suite( + name = "ashe_tests", +) + +yacl_cc_library( + name = "ashe", + hdrs = [ + "ashe.h", + ], + deps = [ + ":decryptor", + ":encryptor", + ":evaluator", + ":key_generator", + ], +) + +yacl_cc_library( + name = "ciphertext", + hdrs = ["ciphertext.h"], + deps = [ + "//heu/library/algorithms/util", + "@msgpack-c//:msgpack", + ], +) + +yacl_cc_library( + name = "secret_key", + hdrs = ["secret_key.h"], + deps = [ + "//heu/library/algorithms/util", + "@msgpack-c//:msgpack", + ], +) + +yacl_cc_library( + name = "public_parameters", + srcs = ["public_parameters.cc"], + hdrs = ["public_parameters.h"], + deps = [ + "//heu/library/algorithms/util", + "@msgpack-c//:msgpack", + ], +) + +yacl_cc_library( + name = "key_generator", + srcs = ["key_generator.cc"], + hdrs = ["key_generator.h"], + deps = [ + ":encryptor", + ":public_parameters", + ":secret_key", + ], +) + +yacl_cc_library( + name = "encryptor", + srcs = ["encryptor.cc"], + hdrs = ["encryptor.h"], + deps = [ + ":ciphertext", + ":public_parameters", + ":secret_key", + ], +) + +yacl_cc_library( + name = "decryptor", + srcs = ["decryptor.cc"], + hdrs = ["decryptor.h"], + deps = [ + ":ciphertext", + ":public_parameters", + ":secret_key", + ], +) + +yacl_cc_library( + name = "evaluator", + srcs = ["evaluator.cc"], + hdrs = ["evaluator.h"], + deps = [ + ":ciphertext", + ":public_parameters", + ], +) + +yacl_cc_test( + name = "ashe_test", + srcs = ["ashe_tests.cc"], + deps = [ + ":ashe", + ], +) diff --git a/heu/library/algorithms/ashe/README.md b/heu/library/algorithms/ashe/README.md new file mode 100644 index 00000000..00b4882b --- /dev/null +++ b/heu/library/algorithms/ashe/README.md @@ -0,0 +1,33 @@ +Implementation of a Variant of Improved Symmetric Homomorphic Encryption(iSHE) + +KeyGen(key_size): + +- Sample a $k_p$-bit (usually $k_p=1536$) prime $p$ +- Sample a $k_q$-bit (usually $k_q=512$) primes $q$ +- Set $k_{r1}, k_{r2}$, the length of random numbers +- Set the message space parameter $k_m=64$ +- The private key is $(p, q)$ and the public parameter is all the parameters $(k_p, k_q, k_{r1}, k_{r2}, k_m)$ and + pre-computed ciphertext of $0$ + +Encryption(sk, pp, m): + +- Sample $r_1 \gets \{0,1\}^{k_{r1}}, r_2 \gets \{0,1\}^{k_{r2}}$ +- The ciphertext is $c = r_1 \cdot p + r_2 \cdot q + (m\ \text{mod}\ (2^{k_m-1}-1))$ + +Decryption(sk, c): + +- Compute $m' = c\ (\text{mod}\ p\ (\text{mod}\ q\ (\ \text{mod}\ (2^{k_m-1}-1)))$ +- If $m' \le 2^{k_m-1}-1$ return m' else return $m'- 2^{k_m-1}-1$ + +Additive homomorphisms: + +- Add $(c_1, c_2) = c_1 + c_2$ +- AddPlain $(c, m) = c + (m\ (\ \text{mod}\ 2^{k_m-1}-1))$ +- Negate $\text{MulPlain}(a, 2^{k_m-1}-2)$ + +Multiplicative homomorphism: + +- MulPlain $(c, m) = c \cdot m$ + +This implement focus on 2048bits key size and 112bit security. The argument keysize in key_generator is actually the +scale of dataset. diff --git a/heu/library/algorithms/ashe/ashe.h b/heu/library/algorithms/ashe/ashe.h new file mode 100644 index 00000000..052b6b58 --- /dev/null +++ b/heu/library/algorithms/ashe/ashe.h @@ -0,0 +1,22 @@ +// Copyright 2022 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. + +#pragma once + +#include "heu/library/algorithms/ashe/decryptor.h" +#include "heu/library/algorithms/ashe/encryptor.h" +#include "heu/library/algorithms/ashe/evaluator.h" +#include "heu/library/algorithms/ashe/key_generator.h" +#include "heu/library/algorithms/ashe/public_parameters.h" +#include "heu/library/algorithms/ashe/secret_key.h" diff --git a/heu/library/algorithms/ashe/ashe_tests.cc b/heu/library/algorithms/ashe/ashe_tests.cc new file mode 100644 index 00000000..d21dd90f --- /dev/null +++ b/heu/library/algorithms/ashe/ashe_tests.cc @@ -0,0 +1,184 @@ +// Copyright 2022 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. + +#include + +#include "gtest/gtest.h" + +#include "heu/library/algorithms/ashe/ashe.h" + +namespace heu::lib::algorithms::ashe::test { + +class asheTest : public testing::Test { + protected: + static void SetUpTestSuite() { KeyGenerator::Generate(2048, &sk_, &pp_); } + + static SecretKey sk_; + static PublicKey pp_; +}; + +SecretKey asheTest::sk_; +PublicKey asheTest::pp_; + +TEST_F(asheTest, SerializeTest) { + auto pp_buffer = pp_.Serialize(); + PublicKey pp2; + pp2.Deserialize(pp_buffer); + ASSERT_EQ(pp_.k_r1, pp2.k_r1); + ASSERT_EQ(pp_.k_r2, pp2.k_r2); + ASSERT_EQ(pp_.k_q, pp2.k_q); + ASSERT_EQ(pp_.k_p, pp2.k_p); + ASSERT_EQ(pp_.randomZeros, pp2.randomZeros); + + auto sk_buffer = sk_.Serialize(); + SecretKey sk2; + sk2.Deserialize(sk_buffer); + ASSERT_EQ(sk_.p_, sk2.p_); + ASSERT_EQ(sk_.q_, sk2.q_); + + Encryptor encryptor(pp2, sk2); + + Evaluator evaluator(pp2); + + BigInt m0(-12345); + Ciphertext ct = encryptor.Encrypt(m0); + + BigInt dc; + Decryptor decryptor(pp_, sk_); + decryptor.Decrypt(ct, &dc); + EXPECT_EQ(dc, m0); +} + +TEST_F(asheTest, OperationEvaluate) { + Encryptor encryptor_(pp_, sk_); + Evaluator evaluator_(pp_); + Decryptor decryptor_(pp_, sk_); + + Plaintext m0 = Plaintext(12345); + Plaintext m1 = Plaintext(-20000); + Plaintext m3 = Plaintext(0); + Ciphertext c0 = encryptor_.Encrypt(m0); + Ciphertext c1 = encryptor_.Encrypt(m1); + Ciphertext c2 = encryptor_.Encrypt(-m0); + Ciphertext c3 = encryptor_.Encrypt(m3); + EXPECT_EQ(m0, Plaintext(12345)); + + Plaintext plain; + Ciphertext res; + + // evaluate add + res = evaluator_.Add(c0, c0); + decryptor_.Decrypt(res, &plain); + EXPECT_EQ(plain, Plaintext(12345 * 2)); + res = evaluator_.Add(c1, c1); + decryptor_.Decrypt(res, &plain); + EXPECT_EQ(plain, Plaintext(-20000 * 2)); + res = evaluator_.Add(c0, c1); + decryptor_.Decrypt(res, &plain); + EXPECT_EQ(plain, Plaintext(12345 - 20000)); + res = evaluator_.Add(c1, m3); + decryptor_.Decrypt(res, &plain); + EXPECT_EQ(plain, Plaintext(-20000)); + res = evaluator_.Add(c0, m1); + decryptor_.Decrypt(res, &plain); + EXPECT_EQ(plain, Plaintext(12345 - 20000)); + res = evaluator_.Add(c1, m0); + decryptor_.Decrypt(res, &plain); + EXPECT_EQ(plain, Plaintext(12345 - 20000)); + res = evaluator_.Add(c2, c0); + decryptor_.Decrypt(res, &plain); + EXPECT_EQ(plain, Plaintext(0)); + + res = evaluator_.Mul(c0, m0); + decryptor_.Decrypt(res, &plain); + EXPECT_EQ(plain, Plaintext(12345 * 12345)); + res = evaluator_.Mul(c1, m0); + decryptor_.Decrypt(res, &plain); + + Ciphertext Zero = encryptor_.EncryptZero(); + decryptor_.Decrypt(Zero, &plain); + EXPECT_EQ(plain, BigInt(0)); + decryptor_.Decrypt(c1, &plain); + EXPECT_EQ(plain, BigInt(-20000)); + + Plaintext pt0 = Plaintext(12345); + Plaintext pt1 = Plaintext(20000); + Ciphertext ct0 = encryptor_.Encrypt(pt0); + Ciphertext ct1 = encryptor_.Encrypt(pt1); + evaluator_.AddInplace(&ct0, pt1); + decryptor_.Decrypt(ct0, &plain); + EXPECT_EQ(plain, BigInt(20000 + 12345)); + evaluator_.AddInplace(&ct0, ct1); + decryptor_.Decrypt(ct0, &plain); + EXPECT_EQ(plain, BigInt(20000 + 12345 + 20000)); + evaluator_.Randomize(&ct0); + decryptor_.Decrypt(ct0, &plain); + EXPECT_EQ(plain, BigInt(20000 + 12345 + 20000)); + Plaintext pt_min = Plaintext(pp_.MessageSpace().first); + Plaintext pt_max = Plaintext(pp_.MessageSpace().second - BigInt(1)); + Ciphertext ct_max = encryptor_.Encrypt(pt_max); + Ciphertext ct_min = encryptor_.Encrypt(pt_min); + Plaintext tmp = decryptor_.Decrypt(ct_min); + EXPECT_EQ(tmp, pt_min); + tmp = decryptor_.Decrypt(ct_max); + EXPECT_EQ(tmp, pt_max); +} + +TEST_F(asheTest, NegateEvalutate) { + Encryptor encryptor_(pp_, sk_); + Evaluator evaluator_(pp_); + Decryptor decryptor_(pp_, sk_); + Plaintext p1 = Plaintext(123456); + Plaintext p2 = Plaintext(23456); + Ciphertext c1 = encryptor_.Encrypt(p1); + Ciphertext c2 = encryptor_.Encrypt(p2); + Ciphertext c3 = evaluator_.Sub(c1, c2); + c2 = evaluator_.Negate(c2); + decryptor_.Decrypt(c2, &p1); + EXPECT_EQ(p1, -23456); + decryptor_.Decrypt(c3, &p1); + EXPECT_EQ(BigInt(123456 - 23456), p1); +} + +TEST_F(asheTest, RuntimeEfficientTest) { + Encryptor encryptor_(pp_, sk_); + Evaluator evaluator_(pp_); + Decryptor decryptor_(pp_, sk_); + Ciphertext c1, c2; + std::chrono::time_point t1, t2; + c1 = encryptor_.Encrypt(pp_.MessageSpace().second); + c2 = encryptor_.Encrypt(pp_.MessageSpace().second); + t1 = std::chrono::high_resolution_clock::now(); + for (int i = 0; i < 100000; i++) { + evaluator_.AddInplace(&c1, c2); + // Plaintext m = decryptor_.Decrypt(c2); + // std::cout << m << std::endl; + // EXPECT_EQ(m, BigInt(-1)*BigInt(i + 2)); + } + t2 = std::chrono::high_resolution_clock::now(); + auto duration = + std::chrono::duration_cast(t2 - t1); + std::cout << "add 1w times used " << duration.count() << std::endl; + auto p = decryptor_.Decrypt(c1); + std::cout << p << std::endl; + EXPECT_EQ(p, pp_.MessageSpace().second * 100001); + t1 = std::chrono::high_resolution_clock::now(); + for (int i = 0; i < 10000; i++) { + Plaintext m = decryptor_.Decrypt(c2); + } + t2 = std::chrono::high_resolution_clock::now(); + duration = std::chrono::duration_cast(t2 - t1); + std::cout << "decrypt 1w times used " << duration.count() << std::endl; +} +} // namespace heu::lib::algorithms::ashe::test diff --git a/heu/library/algorithms/ashe/ciphertext.h b/heu/library/algorithms/ashe/ciphertext.h new file mode 100644 index 00000000..96684036 --- /dev/null +++ b/heu/library/algorithms/ashe/ciphertext.h @@ -0,0 +1,45 @@ +// Copyright 2022 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. + +#pragma once + +#include + +#include "heu/library/algorithms/util/big_int.h" +#include "heu/library/algorithms/util/he_object.h" + +namespace heu::lib::algorithms::ashe { +using Plaintext = BigInt; + +class Ciphertext : public HeObject { + public: + Ciphertext() = default; + + explicit Ciphertext(BigInt n) : n_(std::move(n)) {} + + [[nodiscard]] std::string ToString() const override { + return fmt::format("CT: {}", n_); + } + + bool operator==(const Ciphertext &other) const { return n_ == other.n_; } + + bool operator!=(const Ciphertext &other) const { + return !this->operator==(other); + } + + MSGPACK_DEFINE(n_); + + BigInt n_; +}; +} // namespace heu::lib::algorithms::ashe diff --git a/heu/library/algorithms/ashe/decryptor.cc b/heu/library/algorithms/ashe/decryptor.cc new file mode 100644 index 00000000..0b65e4cd --- /dev/null +++ b/heu/library/algorithms/ashe/decryptor.cc @@ -0,0 +1,28 @@ +// Copyright 2022 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. + +#include "heu/library/algorithms/ashe/decryptor.h" + +namespace heu::lib::algorithms::ashe { +void Decryptor::Decrypt(const Ciphertext &ct, Plaintext *out) const { + *out = Decrypt(ct); +} + +Plaintext Decryptor::Decrypt(const Ciphertext &ct) const { + BigInt m = ct.n_ % p % q; + m = m % MOD_; + if (m >= half) m = m - MOD_; + return m; +} +} // namespace heu::lib::algorithms::ashe diff --git a/heu/library/algorithms/ashe/decryptor.h b/heu/library/algorithms/ashe/decryptor.h new file mode 100644 index 00000000..43295040 --- /dev/null +++ b/heu/library/algorithms/ashe/decryptor.h @@ -0,0 +1,46 @@ +// Copyright 2022 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. + +#pragma once + +#include + +#include "heu/library/algorithms/ashe/ciphertext.h" +#include "heu/library/algorithms/ashe/public_parameters.h" +#include "heu/library/algorithms/ashe/secret_key.h" + +namespace heu::lib::algorithms::ashe { +class Decryptor { + public: + explicit Decryptor(PublicKey pp, SecretKey sk) + : pp_(std::move(pp)), sk_(std::move(sk)) { + p = sk_.p_; + q = sk_.q_; + MOD_ = BigInt(1) << 192; + half = BigInt(1) << 191; + } + + void Decrypt(const Ciphertext &ct, Plaintext *out) const; + + [[nodiscard]] Plaintext Decrypt(const Ciphertext &ct) const; + + private: + PublicKey pp_; + SecretKey sk_; + BigInt half; + BigInt MOD_; + BigInt p; + BigInt q; +}; +} // namespace heu::lib::algorithms::ashe diff --git a/heu/library/algorithms/ashe/encryptor.cc b/heu/library/algorithms/ashe/encryptor.cc new file mode 100644 index 00000000..2977a3cf --- /dev/null +++ b/heu/library/algorithms/ashe/encryptor.cc @@ -0,0 +1,54 @@ +// Copyright 2022 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. + +#include "heu/library/algorithms/ashe/encryptor.h" + +namespace heu::lib::algorithms::ashe { +Ciphertext Encryptor::EncryptZero() const { return Encrypt(BigInt(0)); } + +Ciphertext Encryptor::Encrypt(const Plaintext &m) const { + return EncryptImpl(m, nullptr); +} + +void Encryptor::Encrypt(const Plaintext &m, Ciphertext *out) const { + *out = Encrypt(m); +} + +std::pair Encryptor::EncryptWithAudit( + const Plaintext &m) const { + std::string audit_out; + Ciphertext ct_out = EncryptImpl(m, &audit_out); + audit_out.append( + fmt::format("pt:{}\n ct:{}", m.ToString(), ct_out.n_.ToString())); + return std::make_pair(ct_out, audit_out); +} + +template +Ciphertext Encryptor::EncryptImpl(const Plaintext &m, + std::string *audit_str) const { + YACL_ENFORCE(m <= pp_.MessageSpace().second && m >= pp_.MessageSpace().first, + "Plaintext {} is too large, cannot encrypt.", m); + BigInt r, r1; + BigInt::RandomExactBits(pp_.k_r1, &r); + BigInt::RandomExactBits(pp_.k_r2, &r1); + const BigInt m1 = r * sk_.p_ + r1 * sk_.q_ + m % MOD_; + + if constexpr (audit) { + YACL_ENFORCE(audit_str != nullptr); + *audit_str = + fmt::format("r:{}\n r':{}\n", r.ToHexString(), r1.ToHexString()); + } + return Ciphertext(m1); +} +} // namespace heu::lib::algorithms::ashe diff --git a/heu/library/algorithms/ashe/encryptor.h b/heu/library/algorithms/ashe/encryptor.h new file mode 100644 index 00000000..fee29158 --- /dev/null +++ b/heu/library/algorithms/ashe/encryptor.h @@ -0,0 +1,46 @@ +// Copyright 2022 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. + +#pragma once + +#include + +#include "heu/library/algorithms/ashe/ciphertext.h" +#include "heu/library/algorithms/ashe/public_parameters.h" +#include "heu/library/algorithms/ashe/secret_key.h" + +namespace heu::lib::algorithms::ashe { +class Encryptor { + public: + explicit Encryptor(PublicKey pk, SecretKey sk) + : pp_(std::move(pk)), sk_(std::move(sk)) { + MOD_ = BigInt(1) << 192; + } + + [[nodiscard]] Ciphertext EncryptZero() const; + [[nodiscard]] Ciphertext Encrypt(const Plaintext &m) const; + + void Encrypt(const Plaintext &m, Ciphertext *out) const; + + [[nodiscard]] std::pair EncryptWithAudit( + const Plaintext &m) const; + + private: + template + Ciphertext EncryptImpl(const Plaintext &m, std::string *audit_str) const; + PublicKey pp_; + SecretKey sk_; + BigInt MOD_; +}; +} // namespace heu::lib::algorithms::ashe diff --git a/heu/library/algorithms/ashe/evaluator.cc b/heu/library/algorithms/ashe/evaluator.cc new file mode 100644 index 00000000..ef945a0b --- /dev/null +++ b/heu/library/algorithms/ashe/evaluator.cc @@ -0,0 +1,117 @@ +// Copyright 2022 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. + +#include "heu/library/algorithms/ashe/evaluator.h" + +namespace heu::lib::algorithms::ashe { +void Evaluator::Randomize(Ciphertext *ct) const { + BigInt r; + BigInt::RandomLtN(BigInt(pp_.randomZeros.size()), &r); + AddInplace(ct, Ciphertext(pp_.randomZeros[r.Get()])); +} + +Ciphertext Evaluator::Add(const Ciphertext &a, const Ciphertext &b) const { + return Ciphertext(a.n_ + b.n_); +} + +Ciphertext Evaluator::Add(const Ciphertext &a, const Plaintext &b) const { + return Ciphertext(a.n_ + b % MOD_); +} + +Ciphertext Evaluator::Add(const Plaintext &a, const Ciphertext &b) const { + return Add(b, a); +} + +Plaintext Evaluator::Add(const Plaintext &a, const Plaintext &b) const { + return a + b; +} + +void Evaluator::AddInplace(Ciphertext *a, const Ciphertext &b) const { + a->n_ += b.n_; +} + +void Evaluator::AddInplace(Ciphertext *a, const Plaintext &b) const { + if (b.IsNegative()) { + a->n_ += b + MOD_; + } else { + a->n_ += b; + } +} + +void Evaluator::AddInplace(Plaintext *a, const Plaintext &b) const { + *a = Add(*a, b); +} + +Ciphertext Evaluator::Sub(const Ciphertext &a, const Ciphertext &b) const { + return Ciphertext(a.n_ - b.n_); +} + +Ciphertext Evaluator::Sub(const Ciphertext &a, const Plaintext &b) const { + return Add(a, -b); +} + +Ciphertext Evaluator::Sub(const Plaintext &a, const Ciphertext &b) const { + return Add(Negate(b), a); +} + +Plaintext Evaluator::Sub(const Plaintext &a, const Plaintext &b) const { + return a - b; +} + +void Evaluator::SubInplace(Ciphertext *a, const Ciphertext &b) const { + *a = Sub(*a, b); +} + +void Evaluator::SubInplace(Ciphertext *a, const Plaintext &p) const { + *a = Sub(*a, p); +} + +void Evaluator::SubInplace(Plaintext *a, const Plaintext &b) const { + *a = Sub(*a, b); +} + +Ciphertext Evaluator::Mul(const Ciphertext &a, const Plaintext &b) const { + if (!b.IsNegative()) { + return Ciphertext(b * a.n_); + } else { + Ciphertext neg_a = Negate(a); + return Ciphertext((-b) * neg_a.n_); + } +} + +Ciphertext Evaluator::Mul(const Plaintext &a, const Ciphertext &b) const { + return Mul(b, a); +} + +Plaintext Evaluator::Mul(const Plaintext &a, const Plaintext &b) const { + return a * b; +} + +void Evaluator::MulInplace(Ciphertext *a, const Plaintext &b) const { + *a = Mul(*a, b); +} + +void Evaluator::MulInplace(Plaintext *a, const Plaintext &b) const { + *a = Mul(*a, b); +} + +Ciphertext Evaluator::Negate(const Ciphertext &a) const { + BigInt r; + BigInt::RandomLtN(BigInt(pp_.randomZeros.size()), &r); + BigInt random0 = pp_.randomZeros[r.Get()]; + return Ciphertext(MOD_ - a.n_ + random0); +} + +void Evaluator::NegateInplace(Ciphertext *a) const { *a = Negate(*a); } +} // namespace heu::lib::algorithms::ashe diff --git a/heu/library/algorithms/ashe/evaluator.h b/heu/library/algorithms/ashe/evaluator.h new file mode 100644 index 00000000..4cb23827 --- /dev/null +++ b/heu/library/algorithms/ashe/evaluator.h @@ -0,0 +1,63 @@ +// Copyright 2022 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. + +#pragma once + +#include + +#include "heu/library/algorithms/ashe/ciphertext.h" +#include "heu/library/algorithms/ashe/public_parameters.h" + +namespace heu::lib::algorithms::ashe { +class Evaluator { + public: + explicit Evaluator(PublicKey pp) : pp_(std::move(pp)) { + MOD_ = BigInt(1) << 192; + } + + void Randomize(Ciphertext *ct) const; + + [[nodiscard]] Ciphertext Add(const Ciphertext &a, const Ciphertext &b) const; + [[nodiscard]] Ciphertext Add(const Ciphertext &a, const Plaintext &b) const; + [[nodiscard]] Ciphertext Add(const Plaintext &a, const Ciphertext &b) const; + [[nodiscard]] Plaintext Add(const Plaintext &a, const Plaintext &b) const; + + void AddInplace(Ciphertext *a, const Ciphertext &b) const; + void AddInplace(Ciphertext *a, const Plaintext &b) const; + void AddInplace(Plaintext *a, const Plaintext &b) const; + + [[nodiscard]] Ciphertext Sub(const Ciphertext &a, const Ciphertext &b) const; + [[nodiscard]] Ciphertext Sub(const Ciphertext &a, const Plaintext &b) const; + [[nodiscard]] Ciphertext Sub(const Plaintext &a, const Ciphertext &b) const; + [[nodiscard]] Plaintext Sub(const Plaintext &a, const Plaintext &b) const; + + void SubInplace(Ciphertext *a, const Ciphertext &b) const; + void SubInplace(Ciphertext *a, const Plaintext &p) const; + void SubInplace(Plaintext *a, const Plaintext &b) const; + + [[nodiscard]] Ciphertext Mul(const Ciphertext &a, const Plaintext &b) const; + [[nodiscard]] Ciphertext Mul(const Plaintext &a, const Ciphertext &b) const; + [[nodiscard]] Plaintext Mul(const Plaintext &a, const Plaintext &b) const; + + void MulInplace(Ciphertext *a, const Plaintext &b) const; + void MulInplace(Plaintext *a, const Plaintext &b) const; + + [[nodiscard]] Ciphertext Negate(const Ciphertext &a) const; + void NegateInplace(Ciphertext *a) const; + + private: + PublicKey pp_; + BigInt MOD_; +}; +} // namespace heu::lib::algorithms::ashe diff --git a/heu/library/algorithms/ashe/key_generator.cc b/heu/library/algorithms/ashe/key_generator.cc new file mode 100644 index 00000000..2c306f3d --- /dev/null +++ b/heu/library/algorithms/ashe/key_generator.cc @@ -0,0 +1,73 @@ +// Copyright 2022 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. + +#include "heu/library/algorithms/ashe/key_generator.h" + +#include + +namespace heu::lib::algorithms::ashe { +void KeyGenerator::Generate(int key_size, SecretKey *sk, PublicKey *pk) { + int64_t k_r1, k_p, k_q, k_r2, k_m; + if (key_size == 2048) { + // < 10w + k_r1 = 3978; + k_p = 777; + k_q = 256; + k_r2 = 504; + k_m = 128; + } else if (key_size == 219) { + // < 50w + k_r1 = 5160; + k_p = 779; + k_q = 256; + k_r2 = 504; + k_m = 128; + } else if (key_size == 220) { + // < 100w + k_r1 = 5800; + k_p = 780; + k_q = 256; + k_r2 = 504; + k_m = 128; + } else { + // <300w + k_r1 = 7180; + k_p = 782; + k_q = 256; + k_r2 = 504; + k_m = 128; + } + std::vector zeros; + BigInt p = BigInt::RandPrimeOver(k_p); + const BigInt q = BigInt::RandPrimeOver(k_q); + *sk = SecretKey(p, q); + + InitZeros(k_r1, k_p, k_q, k_r2, k_m, *sk, &zeros); + *pk = PublicKey(k_r1, k_p, k_q, k_r2, k_m, zeros); +} + +void KeyGenerator::Generate(SecretKey *sk, PublicKey *pk) { + Generate(2048, sk, pk); +} + +void KeyGenerator::InitZeros(int64_t k_r1, int64_t k_p, int64_t k_q, + int64_t k_r2, int64_t k_m, SecretKey sk_, + std::vector *zeros) { + auto tmp = PublicKey(k_r1 + 8, k_p, k_q, k_r2 + 8, k_m); + auto et = Encryptor(tmp, std::move(sk_)); + for (int i = 1; i <= 100; ++i) { + zeros->emplace_back(et.Encrypt(BigInt(0)).n_); + } +} +} // namespace heu::lib::algorithms::ashe diff --git a/heu/library/algorithms/ashe/key_generator.h b/heu/library/algorithms/ashe/key_generator.h new file mode 100644 index 00000000..9081b35a --- /dev/null +++ b/heu/library/algorithms/ashe/key_generator.h @@ -0,0 +1,31 @@ +// Copyright 2022 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. + +#pragma once + +#include "heu/library/algorithms/ashe/encryptor.h" +#include "heu/library/algorithms/ashe/public_parameters.h" +#include "heu/library/algorithms/ashe/secret_key.h" + +namespace heu::lib::algorithms::ashe { +class KeyGenerator { + public: + static void Generate(int key_size, SecretKey *sk, PublicKey *pk); + static void Generate(SecretKey *sk, PublicKey *pk); + + private: + static void InitZeros(int64_t k_r1, int64_t k_p, int64_t k_q, int64_t k_r2, + int64_t k_m, SecretKey sk_, std::vector *zeros); +}; +} // namespace heu::lib::algorithms::ashe diff --git a/heu/library/algorithms/ashe/public_parameters.cc b/heu/library/algorithms/ashe/public_parameters.cc new file mode 100644 index 00000000..b12de39d --- /dev/null +++ b/heu/library/algorithms/ashe/public_parameters.cc @@ -0,0 +1,48 @@ +// Copyright 2022 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. + +#include "heu/library/algorithms/ashe/public_parameters.h" + +namespace heu::lib::algorithms::ashe { +PublicKey::PublicKey(int64_t k_r1, int64_t k_p, int64_t k_q, + int64_t k_r2, int64_t k_m) { + this->k_r1 = k_r1; + this->k_p = k_p; + this->k_q = k_q; + this->k_r2 = k_r2; + this->k_m = k_m; + Init(); +} + +PublicKey::PublicKey(int64_t k_r1, int64_t k_p, int64_t k_q, + int64_t k_r2, int64_t k_m, + const std::vector &zeros) + : PublicKey(k_r1, k_p, k_q, k_r2, k_m) { + this->randomZeros = zeros; +} + +std::string PublicKey::ToString() const { + return fmt::format( + "ashe PP: k_r1={}, k_p={}, k_q={}, " + "k_r2={}, k_m={}, randomZeros={}[size:{}]", + std::to_string(k_r1), std::to_string(k_p), std::to_string(k_q), + std::to_string(k_r2), std::to_string(k_m), ToHexString(randomZeros), + randomZeros.size()); +} + +void PublicKey::Init() { + this->M[1] = BigInt(2).Pow(k_m - 1) - BigInt(1); + this->M[0] = -BigInt(2).Pow(k_m - 1); +} +} // namespace heu::lib::algorithms::ashe diff --git a/heu/library/algorithms/ashe/public_parameters.h b/heu/library/algorithms/ashe/public_parameters.h new file mode 100644 index 00000000..a0a37bf5 --- /dev/null +++ b/heu/library/algorithms/ashe/public_parameters.h @@ -0,0 +1,77 @@ +// Copyright 2022 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. + +#pragma once + +#include + +#include "heu/library/algorithms/util/big_int.h" +#include "heu/library/algorithms/util/he_object.h" + +namespace heu::lib::algorithms::ashe { +class PublicKey : public HeObject { + private: + BigInt plaintextBound; + + static std::string ToHexString(const std::vector &vec) { + std::ostringstream oss; + oss << "["; + for (size_t i = 0; i < vec.size(); ++i) { + if (i != 0) { + oss << ", "; + } + oss << "0x" << vec[i].ToHexString(); + } + oss << "]"; + return oss.str(); + } + + public: + int64_t k_r1 = 2196; + int64_t k_p = 768; + int64_t k_q = 256; + int64_t k_r2 = 504; + int64_t k_m = 128; + std::vector randomZeros; + BigInt M[2]; + + PublicKey() = default; + + PublicKey(int64_t k_r1, int64_t k_p, int64_t k_q, int64_t k_r2, int64_t k_m); + + PublicKey(int64_t k_r1, int64_t k_p, int64_t k_q, int64_t k_r2, int64_t k_m, + const std::vector &zeros); + + bool operator==(const PublicKey &other) const { + return k_r1 == other.k_r1 && k_p == other.k_p && k_q == other.k_q && + k_r2 == other.k_r2 && k_m == other.k_m; + } + + bool operator!=(const PublicKey &other) const { + return !this->operator==(other); + } + + [[nodiscard]] std::string ToString() const override; + + [[nodiscard]] const BigInt &PlaintextBound() const & { return M[1]; } + + void Init(); + + [[nodiscard]] std::pair MessageSpace() const { + return std::make_pair(M[0], M[1]); + } + + MSGPACK_DEFINE(k_r1, k_p, k_q, k_r2, k_m, M, randomZeros); +}; +} // namespace heu::lib::algorithms::ashe diff --git a/heu/library/algorithms/ashe/secret_key.h b/heu/library/algorithms/ashe/secret_key.h new file mode 100644 index 00000000..117fe5ce --- /dev/null +++ b/heu/library/algorithms/ashe/secret_key.h @@ -0,0 +1,48 @@ +// Copyright 2022 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. + +#pragma once + +#include + +#include "heu/library/algorithms/util/he_object.h" + +namespace heu::lib::algorithms::ashe { +class SecretKey : public HeObject { + public: + BigInt p_, q_; + + SecretKey(BigInt p, BigInt q) { + this->p_ = std::move(p); + this->q_ = std::move(q); + } + + SecretKey() = default; + + bool operator==(const SecretKey &other) const { + return p_ == other.p_ && q_ == other.q_; + } + + bool operator!=(const SecretKey &other) const { + return !this->operator==(other); + } + + [[nodiscard]] std::string ToString() const override { + return fmt::format("ashe SK, p={}[{}bits], q={}[{}bits]", p_.ToHexString(), + p_.BitCount(), q_.ToHexString(), q_.BitCount()); + } + + MSGPACK_DEFINE(p_, q_); +}; +} // namespace heu::lib::algorithms::ashe diff --git a/heu/library/numpy/numpy.h b/heu/library/numpy/numpy.h index c3cdcae4..d6e112c6 100644 --- a/heu/library/numpy/numpy.h +++ b/heu/library/numpy/numpy.h @@ -57,8 +57,10 @@ class DestinationHeKit : public phe::HeKitPublicBase { public: explicit DestinationHeKit(phe::DestinationHeKit phe_kit) { Setup(phe_kit.GetPublicKey()); - - encryptor_ = std::make_shared(*phe_kit.GetEncryptor()); + auto phe_encryptor = phe_kit.GetEncryptor(); + if (phe_encryptor != nullptr) { + encryptor_ = std::make_shared(*phe_encryptor); + } evaluator_ = std::make_shared(*phe_kit.GetEvaluator()); } diff --git a/heu/library/phe/base/BUILD.bazel b/heu/library/phe/base/BUILD.bazel index 12a5445e..812d4acb 100644 --- a/heu/library/phe/base/BUILD.bazel +++ b/heu/library/phe/base/BUILD.bazel @@ -41,6 +41,7 @@ yacl_cc_library( "//heu/library/algorithms/paillier_ic", "//heu/library/algorithms/paillier_ipcl", "//heu/library/algorithms/paillier_zahlen", + "//heu/library/algorithms/ashe" ], ) diff --git a/heu/library/phe/base/schema.cc b/heu/library/phe/base/schema.cc index fbf8338f..72010c53 100644 --- a/heu/library/phe/base/schema.cc +++ b/heu/library/phe/base/schema.cc @@ -25,8 +25,7 @@ namespace heu::lib::phe { SchemaType::enum_item, { #enum_item, ##__VA_ARGS__ } \ } #define MAP_ITEM_false(...) \ - { \ - } + {} #define MAP_ITEM_HELPER(enable, enum_item, ...) \ MAP_ITEM_##enable(enum_item, ##__VA_ARGS__) #define MAP_ITEM(enable, enum_item, ...) \ @@ -60,6 +59,7 @@ static const std::map, MAP_ITEM(true, DGK, "dgk", "damgard-geisler-kroigaard", "damgard_geisler_kroigaard"), MAP_ITEM(true, DJ, "dj", "damgard-jurik", "damgard_jurik"), + MAP_ITEM(true, AShe, "ashe", "aShe", "aSHE", "ASHE"), // MAP_ITEM(ENABLE, YOUR_ALGO, "one_or_more_name_alias"), }; diff --git a/heu/library/phe/base/schema.h b/heu/library/phe/base/schema.h index 4be56d0b..1b1d5d62 100644 --- a/heu/library/phe/base/schema.h +++ b/heu/library/phe/base/schema.h @@ -16,6 +16,7 @@ #include "msgpack.hpp" +#include "heu/library/algorithms/ashe/ashe.h" #include "heu/library/algorithms/dgk/dgk.h" #include "heu/library/algorithms/dj/dj.h" #include "heu/library/algorithms/elgamal/elgamal.h" @@ -54,6 +55,7 @@ enum class SchemaType : uint8_t { ENUM_ELEMENT(8, true, ElGamal) ENUM_ELEMENT(10, true, DGK) ENUM_ELEMENT(11, true, DJ) + ENUM_ELEMENT(12, true, AShe) // YOUR_ALGO }; // clang-format on @@ -84,7 +86,8 @@ enum class SchemaType : uint8_t { INVOKE(ENABLE_CLUSTAR_FPGA, func_or_macro, ::heu::lib::algorithms::paillier_clustar_fpga, ##__VA_ARGS__) \ INVOKE(true, func_or_macro, ::heu::lib::algorithms::elgamal, ##__VA_ARGS__) \ INVOKE(true, func_or_macro, ::heu::lib::algorithms::dgk, ##__VA_ARGS__) \ - INVOKE(true, func_or_macro, ::heu::lib::algorithms::dj, ##__VA_ARGS__) + INVOKE(true, func_or_macro, ::heu::lib::algorithms::dj, ##__VA_ARGS__) \ + INVOKE(true, func_or_macro, ::heu::lib::algorithms::ashe, ##__VA_ARGS__) // [SPI: Please register your algorithm here] || progress: (4 of 5) // If you add a new schema, change this !! diff --git a/heu/library/phe/phe.cc b/heu/library/phe/phe.cc index 3f4166d9..b80755cb 100644 --- a/heu/library/phe/phe.cc +++ b/heu/library/phe/phe.cc @@ -18,6 +18,55 @@ namespace heu::lib::phe { +template +struct CanConstructWithPkOnly : std::false_type {}; + +template +struct CanConstructWithPkOnly()))>> + : std::true_type {}; + +template ::value> +struct EncryptorCreator; + +template +struct EncryptorCreator { + static Enc create(const PK &pk, const SK &) { return Enc(pk); } +}; + +template +struct EncryptorCreator { + static Enc create(const PK &pk, const SK &sk) { return Enc(pk, sk); } +}; + +template ::value> +struct EncryptorSetupHelper; + +template +struct EncryptorSetupHelper { + static void setupWithPk(std::shared_ptr &encryptor, + SchemaType schema_type, const PK &pk) { + encryptor = std::make_shared(schema_type, Enc(pk)); + } + + static void setupWithSk(std::shared_ptr &, SchemaType, const PK &, + const SK &) {} +}; + +template +struct EncryptorSetupHelper { // symmetric + + static void setupWithPk(std::shared_ptr &, SchemaType, + const PK &) {} + + static void setupWithSk(std::shared_ptr &encryptor, + SchemaType schema_type, const PK &pk, const SK &sk) { + encryptor = std::make_shared(schema_type, Enc(pk, sk)); + } +}; + void HeKitPublicBase::Setup(std::shared_ptr pk) { public_key_ = std::move(pk); @@ -43,17 +92,41 @@ void HeKitSecretBase::Setup(std::shared_ptr pk, schema_type_); } -#define GEN_KEY_AND_INIT(ns) \ - [&](ns::PublicKey &pk) { \ - ns::SecretKey sk; \ - ns::KeyGenerator::Generate(key_size, &sk, &pk); \ - \ - encryptor_ = std::make_shared(schema_type, ns::Encryptor(pk)); \ - decryptor_ = \ - std::make_shared(schema_type, ns::Decryptor(pk, sk)); \ - evaluator_ = std::make_shared(schema_type, ns::Evaluator(pk)); \ - return std::make_shared(std::move(sk)); \ +#define GEN_KEY_AND_INIT(ns) \ + [&](ns::PublicKey &pk) { \ + ns::SecretKey sk; \ + ns::KeyGenerator::Generate(key_size, &sk, &pk); \ + \ + encryptor_ = std::make_shared( \ + schema_type, \ + EncryptorCreator::create( \ + pk, sk)); \ + decryptor_ = \ + std::make_shared(schema_type, ns::Decryptor(pk, sk)); \ + evaluator_ = std::make_shared(schema_type, ns::Evaluator(pk)); \ + return std::make_shared(std::move(sk)); \ + } + +template ::value> +struct DestEncryptorSetupHelper; + +template +struct DestEncryptorSetupHelper { + static void setup(std::shared_ptr &encryptor, + SchemaType schema_type, const PK &pk) { + encryptor = std::make_shared(schema_type, Enc(pk)); + } +}; + +template +struct DestEncryptorSetupHelper { + static void setup(std::shared_ptr &encryptor, + SchemaType schema_type, const PK &) { + (void)schema_type; + encryptor = nullptr; } +}; HeKit::HeKit(SchemaType schema_type, size_t key_size) { auto pk = std::make_shared(schema_type); @@ -62,16 +135,19 @@ HeKit::HeKit(SchemaType schema_type, size_t key_size) { Setup(std::move(pk), std::move(sk)); } -#define GEN_KEY_AND_INIT_DEFAULT(ns) \ - [&](ns::PublicKey &pk) { \ - ns::SecretKey sk; \ - ns::KeyGenerator::Generate(&sk, &pk); \ - \ - encryptor_ = std::make_shared(schema_type, ns::Encryptor(pk)); \ - decryptor_ = \ - std::make_shared(schema_type, ns::Decryptor(pk, sk)); \ - evaluator_ = std::make_shared(schema_type, ns::Evaluator(pk)); \ - return std::make_shared(std::move(sk)); \ +#define GEN_KEY_AND_INIT_DEFAULT(ns) \ + [&](ns::PublicKey &pk) { \ + ns::SecretKey sk; \ + ns::KeyGenerator::Generate(&sk, &pk); \ + \ + encryptor_ = std::make_shared( \ + schema_type, \ + EncryptorCreator::create( \ + pk, sk)); \ + decryptor_ = \ + std::make_shared(schema_type, ns::Decryptor(pk, sk)); \ + evaluator_ = std::make_shared(schema_type, ns::Evaluator(pk)); \ + return std::make_shared(std::move(sk)); \ } HeKit::HeKit(SchemaType schema_type) { @@ -81,18 +157,23 @@ HeKit::HeKit(SchemaType schema_type) { Setup(std::move(pk), std::move(sk)); } -#define HE_SPECIAL_SETUP_BY_PK(ns) \ - [&](const ns::PublicKey &pk1) { \ - evaluator_ = \ - std::make_shared(schema_type_, ns::Evaluator(pk1)); \ - encryptor_ = \ - std::make_shared(schema_type_, ns::Encryptor(pk1)); \ +#define HE_SPECIAL_SETUP_BY_PK(ns) \ + [&](const ns::PublicKey &pk1) { \ + evaluator_ = \ + std::make_shared(schema_type_, ns::Evaluator(pk1)); \ + EncryptorSetupHelper::setupWithPk(encryptor_, schema_type_, \ + pk1); \ } -#define HE_SPECIAL_SETUP_BY_SK(ns) \ - [&](const ns::SecretKey &sk1) { \ - decryptor_ = std::make_shared( \ - schema_type_, ns::Decryptor(public_key_->As(), sk1)); \ +#define HE_SPECIAL_SETUP_BY_SK(ns) \ + [&](const ns::SecretKey &sk1) { \ + const auto &pk1 = public_key_->As(); \ + decryptor_ = \ + std::make_shared(schema_type_, ns::Decryptor(pk1, sk1)); \ + EncryptorSetupHelper::setupWithSk(encryptor_, schema_type_, \ + pk1, sk1); \ } HeKit::HeKit(std::shared_ptr pk, std::shared_ptr sk) { @@ -113,16 +194,24 @@ HeKit::HeKit(yacl::ByteContainerView pk_buffer, secret_key_->Visit(HE_DISPATCH(HE_SPECIAL_SETUP_BY_SK)); } +#define HE_DEST_SETUP_BY_PK(ns) \ + [&](const ns::PublicKey &pk1) { \ + evaluator_ = \ + std::make_shared(schema_type_, ns::Evaluator(pk1)); \ + DestEncryptorSetupHelper::setup( \ + encryptor_, schema_type_, pk1); \ + } + DestinationHeKit::DestinationHeKit(std::shared_ptr pk) { Setup(std::move(pk)); - public_key_->Visit(HE_DISPATCH(HE_SPECIAL_SETUP_BY_PK)); + public_key_->Visit(HE_DISPATCH(HE_DEST_SETUP_BY_PK)); } DestinationHeKit::DestinationHeKit(yacl::ByteContainerView pk_buffer) { auto pk = std::make_shared(); pk->Deserialize(pk_buffer); Setup(std::move(pk)); - public_key_->Visit(HE_DISPATCH(HE_SPECIAL_SETUP_BY_PK)); + public_key_->Visit(HE_DISPATCH(HE_DEST_SETUP_BY_PK)); } } // namespace heu::lib::phe diff --git a/heu/library/phe/phe.h b/heu/library/phe/phe.h index 3b1ca704..2848950b 100644 --- a/heu/library/phe/phe.h +++ b/heu/library/phe/phe.h @@ -94,6 +94,8 @@ class DestinationHeKit : public HeKitPublicBase { return encryptor_; } + [[nodiscard]] bool canEncrypt() const { return encryptor_ != nullptr; } + [[nodiscard]] const std::shared_ptr &GetEvaluator() const { return evaluator_; }