pir_client.cpp 9.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275
  1. #include "pir_client.hpp"
  2. using namespace std;
  3. using namespace seal;
  4. using namespace seal::util;
  5. PIRClient::PIRClient(const EncryptionParameters &enc_params,
  6. const PirParams &pir_params) :
  7. enc_params_(enc_params),
  8. pir_params_(pir_params){
  9. context_ = make_shared<SEALContext>(enc_params, true);
  10. keygen_ = make_unique<KeyGenerator>(*context_);
  11. PublicKey public_key;
  12. keygen_->create_public_key(public_key);
  13. SecretKey secret_key = keygen_->secret_key();
  14. if(pir_params_.enable_symmetric){
  15. encryptor_ = make_unique<Encryptor>(*context_, secret_key);
  16. }
  17. else{
  18. encryptor_ = make_unique<Encryptor>(*context_, public_key);
  19. }
  20. decryptor_ = make_unique<Decryptor>(*context_, secret_key);
  21. evaluator_ = make_unique<Evaluator>(*context_);
  22. encoder_ = make_unique<BatchEncoder>(*context_);
  23. }
  24. int PIRClient::generate_serialized_query(uint64_t desiredIndex, std::stringstream &stream) {
  25. int N = enc_params_.poly_modulus_degree();
  26. int output_size = 0;
  27. indices_ = compute_indices(desiredIndex, pir_params_.nvec);
  28. Plaintext pt(enc_params_.poly_modulus_degree());
  29. for (uint32_t i = 0; i < indices_.size(); i++) {
  30. uint32_t num_ptxts = ceil( (pir_params_.nvec[i] + 0.0) / N);
  31. // initialize result.
  32. cout << "Client: index " << i + 1 << "/ " << indices_.size() << " = " << indices_[i] << endl;
  33. cout << "Client: number of ctxts needed for query = " << num_ptxts << endl;
  34. for (uint32_t j =0; j < num_ptxts; j++){
  35. pt.set_zero();
  36. if (indices_[i] >= N*j && indices_[i] <= N*(j+1)){
  37. uint64_t real_index = indices_[i] - N*j;
  38. uint64_t n_i = pir_params_.nvec[i];
  39. uint64_t total = N;
  40. if (j == num_ptxts - 1){
  41. total = n_i % N;
  42. }
  43. uint64_t log_total = ceil(log2(total));
  44. cout << "Client: Inverting " << pow(2, log_total) << endl;
  45. pt[real_index] = invert_mod(pow(2, log_total), enc_params_.plain_modulus());
  46. }
  47. if(pir_params_.enable_symmetric){
  48. output_size += encryptor_->encrypt_symmetric(pt).save(stream);
  49. }
  50. else{
  51. output_size += encryptor_->encrypt(pt).save(stream);
  52. }
  53. }
  54. }
  55. return output_size;
  56. }
  57. PirQuery PIRClient::generate_query(uint64_t desiredIndex) {
  58. indices_ = compute_indices(desiredIndex, pir_params_.nvec);
  59. PirQuery result(pir_params_.d);
  60. int N = enc_params_.poly_modulus_degree();
  61. Plaintext pt(enc_params_.poly_modulus_degree());
  62. for (uint32_t i = 0; i < indices_.size(); i++) {
  63. uint32_t num_ptxts = ceil( (pir_params_.nvec[i] + 0.0) / N);
  64. // initialize result.
  65. cout << "Client: index " << i + 1 << "/ " << indices_.size() << " = " << indices_[i] << endl;
  66. cout << "Client: number of ctxts needed for query = " << num_ptxts << endl;
  67. for (uint32_t j =0; j < num_ptxts; j++){
  68. pt.set_zero();
  69. if (indices_[i] >= N*j && indices_[i] <= N*(j+1)){
  70. uint64_t real_index = indices_[i] - N*j;
  71. uint64_t n_i = pir_params_.nvec[i];
  72. uint64_t total = N;
  73. if (j == num_ptxts - 1){
  74. total = n_i % N;
  75. }
  76. uint64_t log_total = ceil(log2(total));
  77. cout << "Client: Inverting " << pow(2, log_total) << endl;
  78. pt[real_index] = invert_mod(pow(2, log_total), enc_params_.plain_modulus());
  79. }
  80. Ciphertext dest;
  81. if(pir_params_.enable_symmetric){
  82. encryptor_->encrypt_symmetric(pt, dest);
  83. }
  84. else{
  85. encryptor_->encrypt(pt, dest);
  86. }
  87. result[i].push_back(dest);
  88. }
  89. }
  90. return result;
  91. }
  92. uint64_t PIRClient::get_fv_index(uint64_t element_index) {
  93. return static_cast<uint64_t>(element_index / pir_params_.elements_per_plaintext);
  94. }
  95. uint64_t PIRClient::get_fv_offset(uint64_t element_index) {
  96. return element_index % pir_params_.elements_per_plaintext;
  97. }
  98. Plaintext PIRClient::decrypt(Ciphertext ct){
  99. Plaintext pt;
  100. decryptor_->decrypt(ct, pt);
  101. return pt;
  102. }
  103. vector<uint8_t> PIRClient::decode_reply(PirReply &reply, uint64_t offset){
  104. Plaintext result = decode_reply(reply);
  105. return extract_bytes(result, offset);
  106. }
  107. vector<uint64_t> PIRClient::extract_coeffs(Plaintext pt){
  108. vector<uint64_t> coeffs;
  109. encoder_->decode(pt, coeffs);
  110. return coeffs;
  111. }
  112. std::vector<uint64_t> PIRClient::extract_coeffs(seal::Plaintext pt, uint64_t offset){
  113. vector<uint64_t> coeffs;
  114. encoder_->decode(pt, coeffs);
  115. uint32_t logt = floor(log2(enc_params_.plain_modulus().value()));
  116. uint64_t coeffs_per_element = coefficients_per_element(logt, pir_params_.ele_size);
  117. return std::vector<uint64_t>(coeffs.begin() + offset * coeffs_per_element, coeffs.begin() + (offset + 1) * coeffs_per_element);
  118. }
  119. std::vector<uint8_t> PIRClient::extract_bytes(seal::Plaintext pt, uint64_t offset){
  120. uint32_t N = enc_params_.poly_modulus_degree();
  121. uint32_t logt = floor(log2(enc_params_.plain_modulus().value()));
  122. uint32_t bytes_per_ptxt = pir_params_.elements_per_plaintext * pir_params_.ele_size;
  123. // Convert from FV plaintext (polynomial) to database element at the client
  124. vector<uint8_t> elems(bytes_per_ptxt);
  125. vector<uint64_t> coeffs;
  126. encoder_->decode(pt, coeffs);
  127. coeffs_to_bytes(logt, coeffs, elems.data(), bytes_per_ptxt, pir_params_.ele_size);
  128. return std::vector<uint8_t>(elems.begin() + offset * pir_params_.ele_size, elems.begin() + (offset + 1) * pir_params_.ele_size);
  129. }
  130. Plaintext PIRClient::decode_reply(PirReply &reply) {
  131. EncryptionParameters parms;
  132. parms_id_type parms_id;
  133. if(pir_params_.enable_mswitching){
  134. parms = context_->last_context_data()->parms();
  135. parms_id = context_->last_parms_id();
  136. }
  137. else{
  138. parms = context_->first_context_data()->parms();
  139. parms_id = context_->first_parms_id();
  140. }
  141. uint32_t exp_ratio = compute_expansion_ratio(parms);
  142. uint32_t recursion_level = pir_params_.d;
  143. vector<Ciphertext> temp = reply;
  144. uint32_t ciphertext_size = temp[0].size();
  145. uint64_t t = enc_params_.plain_modulus().value();
  146. for (uint32_t i = 0; i < recursion_level; i++) {
  147. cout << "Client: " << i + 1 << "/ " << recursion_level << "-th decryption layer started." << endl;
  148. vector<Ciphertext> newtemp;
  149. vector<Plaintext> tempplain;
  150. for (uint32_t j = 0; j < temp.size(); j++) {
  151. Plaintext ptxt;
  152. decryptor_->decrypt(temp[j], ptxt);
  153. #ifdef DEBUG
  154. cout << "Client: reply noise budget = " << decryptor_->invariant_noise_budget(temp[j]) << endl;
  155. #endif
  156. //cout << "decoded (and scaled) plaintext = " << ptxt.to_string() << endl;
  157. tempplain.push_back(ptxt);
  158. #ifdef DEBUG
  159. cout << "recursion level : " << i << " noise budget : ";
  160. cout << decryptor_->invariant_noise_budget(temp[j]) << endl;
  161. #endif
  162. if ((j + 1) % (exp_ratio * ciphertext_size) == 0 && j > 0) {
  163. // Combine into one ciphertext.
  164. Ciphertext combined(*context_, parms_id);
  165. compose_to_ciphertext(parms, tempplain, combined);
  166. newtemp.push_back(combined);
  167. tempplain.clear();
  168. // cout << "Client: const term of ciphertext = " << combined[0] << endl;
  169. }
  170. }
  171. cout << "Client: done." << endl;
  172. cout << endl;
  173. if (i == recursion_level - 1) {
  174. assert(temp.size() == 1);
  175. return tempplain[0];
  176. } else {
  177. tempplain.clear();
  178. temp = newtemp;
  179. }
  180. }
  181. // This should never be called
  182. assert(0);
  183. Plaintext fail;
  184. return fail;
  185. }
  186. GaloisKeys PIRClient::generate_galois_keys() {
  187. // Generate the Galois keys needed for coeff_select.
  188. vector<uint32_t> galois_elts;
  189. int N = enc_params_.poly_modulus_degree();
  190. int logN = get_power_of_two(N);
  191. //cout << "printing galois elements...";
  192. for (int i = 0; i < logN; i++) {
  193. galois_elts.push_back((N + exponentiate_uint(2, i)) / exponentiate_uint(2, i));
  194. //#ifdef DEBUG
  195. // cout << galois_elts.back() << ", ";
  196. //#endif
  197. }
  198. GaloisKeys gal_keys;
  199. keygen_->create_galois_keys(galois_elts, gal_keys);
  200. return gal_keys;
  201. }
  202. Plaintext PIRClient::replace_element(Plaintext pt, vector<uint64_t> new_element, uint64_t offset){
  203. vector<uint64_t> coeffs = extract_coeffs(pt);
  204. uint32_t logt = floor(log2(enc_params_.plain_modulus().value()));
  205. uint64_t coeffs_per_element = coefficients_per_element(logt, pir_params_.ele_size);
  206. assert(new_element.size() == coeffs_per_element);
  207. for(uint64_t i = 0; i < coeffs_per_element; i++){
  208. coeffs[i + offset * coeffs_per_element] = new_element[i];
  209. }
  210. Plaintext new_pt;
  211. encoder_->encode(coeffs, new_pt);
  212. return new_pt;
  213. }
  214. Ciphertext PIRClient::get_one(){
  215. Plaintext pt("1");
  216. Ciphertext ct;
  217. if(pir_params_.enable_symmetric){
  218. encryptor_->encrypt_symmetric(pt, ct);
  219. }
  220. else{
  221. encryptor_->encrypt(pt, ct);
  222. }
  223. return ct;
  224. }