main.cpp 7.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168
  1. #include "pir.hpp"
  2. #include "pir_client.hpp"
  3. #include "pir_server.hpp"
  4. #include <seal/seal.h>
  5. #include <chrono>
  6. #include <memory>
  7. #include <random>
  8. #include <cstdint>
  9. #include <cstddef>
  10. using namespace std::chrono;
  11. using namespace std;
  12. using namespace seal;
  13. int main(int argc, char *argv[]) {
  14. uint64_t number_of_items = 1 << 16;
  15. uint64_t size_per_item = 256; // in bytes
  16. uint32_t N = 4096;
  17. // Recommended values: (logt, d) = (20, 2).
  18. uint32_t logt = 20;
  19. uint32_t d = 2;
  20. bool use_symmetric = true; // use symmetric encryption instead of public key (recommended for smaller query)
  21. bool use_batching = true; // pack as many elements as possible into a BFV plaintext (recommended)
  22. bool use_recursive_mod_switching = true;
  23. EncryptionParameters enc_params(scheme_type::bfv);
  24. PirParams pir_params;
  25. // Generates all parameters
  26. cout << "Main: Generating SEAL parameters" << endl;
  27. gen_encryption_params(N, logt, enc_params);
  28. cout << "Main: Verifying SEAL parameters" << endl;
  29. verify_encryption_params(enc_params);
  30. cout << "Main: SEAL parameters are good" << endl;
  31. cout << "Main: Generating PIR parameters" << endl;
  32. gen_pir_params(number_of_items, size_per_item, d, enc_params, pir_params, use_symmetric, use_batching, use_recursive_mod_switching);
  33. print_seal_params(enc_params);
  34. print_pir_params(pir_params);
  35. // Initialize PIR client....
  36. PIRClient client(enc_params, pir_params);
  37. cout << "Main: Generating galois keys for client" << endl;
  38. GaloisKeys galois_keys = client.generate_galois_keys();
  39. // Initialize PIR Server
  40. cout << "Main: Initializing server" << endl;
  41. PIRServer server(enc_params, pir_params);
  42. // Server maps the galois key to client 0. We only have 1 client,
  43. // which is why we associate it with 0. If there are multiple PIR
  44. // clients, you should have each client generate a galois key,
  45. // and assign each client an index or id, then call the procedure below.
  46. server.set_galois_key(0, galois_keys);
  47. cout << "Main: Creating the database with random data (this may take some time) ..." << endl;
  48. // Create test database
  49. auto db(make_unique<uint8_t[]>(number_of_items * size_per_item));
  50. // Copy of the database. We use this at the end to make sure we retrieved
  51. // the correct element.
  52. auto db_copy(make_unique<uint8_t[]>(number_of_items * size_per_item));
  53. random_device rd;
  54. for (uint64_t i = 0; i < number_of_items; i++) {
  55. for (uint64_t j = 0; j < size_per_item; j++) {
  56. uint8_t val = rd() % 256;
  57. db.get()[(i * size_per_item) + j] = val;
  58. db_copy.get()[(i * size_per_item) + j] = val;
  59. }
  60. }
  61. // Measure database setup
  62. auto time_pre_s = high_resolution_clock::now();
  63. server.set_database(move(db), number_of_items, size_per_item);
  64. server.preprocess_database();
  65. cout << "Main: database pre processed " << endl;
  66. auto time_pre_e = high_resolution_clock::now();
  67. auto time_pre_us = duration_cast<microseconds>(time_pre_e - time_pre_s).count();
  68. // Choose an index of an element in the DB
  69. uint64_t ele_index = rd() % number_of_items; // element in DB at random position
  70. uint64_t index = client.get_fv_index(ele_index); // index of FV plaintext
  71. uint64_t offset = client.get_fv_offset(ele_index); // offset in FV plaintext
  72. cout << "Main: element index = " << ele_index << " from [0, " << number_of_items -1 << "]" << endl;
  73. cout << "Main: FV index = " << index << ", FV offset = " << offset << endl;
  74. // Measure query generation
  75. auto time_query_s = high_resolution_clock::now();
  76. PirQuery query = client.generate_query(index);
  77. auto time_query_e = high_resolution_clock::now();
  78. auto time_query_us = duration_cast<microseconds>(time_query_e - time_query_s).count();
  79. cout << "Main: query generated" << endl;
  80. // Measure serialized query generation (useful for sending over the network)
  81. stringstream client_stream;
  82. stringstream server_stream;
  83. auto time_s_query_s = high_resolution_clock::now();
  84. int query_size = client.generate_serialized_query(index, client_stream);
  85. auto time_s_query_e = high_resolution_clock::now();
  86. auto time_s_query_us = duration_cast<microseconds>(time_s_query_e - time_s_query_s).count();
  87. cout << "Main: serialized query generated" << endl;
  88. // Measure query deserialization (useful for receiving over the network)
  89. auto time_deserial_s = high_resolution_clock::now();
  90. PirQuery query2 = server.deserialize_query(client_stream);
  91. auto time_deserial_e = high_resolution_clock::now();
  92. auto time_deserial_us = duration_cast<microseconds>(time_deserial_e - time_deserial_s).count();
  93. cout << "Main: query deserialized" << endl;
  94. // Measure query processing (including expansion)
  95. auto time_server_s = high_resolution_clock::now();
  96. // Answer PIR query from client 0. If there are multiple clients,
  97. // enter the id of the client (to use the associated galois key).
  98. PirReply reply = server.generate_reply(query2, 0);
  99. auto time_server_e = high_resolution_clock::now();
  100. auto time_server_us = duration_cast<microseconds>(time_server_e - time_server_s).count();
  101. cout << "Main: reply generated" << endl;
  102. // Serialize reply (useful for sending over the network)
  103. int reply_size = server.serialize_reply(reply, server_stream);
  104. // Measure response extraction
  105. auto time_decode_s = chrono::high_resolution_clock::now();
  106. vector<uint8_t> elems = client.decode_reply(reply, offset);
  107. auto time_decode_e = chrono::high_resolution_clock::now();
  108. auto time_decode_us = duration_cast<microseconds>(time_decode_e - time_decode_s).count();
  109. cout << "Main: reply decoded" << endl;
  110. assert(elems.size() == size_per_item);
  111. bool failed = false;
  112. // Check that we retrieved the correct element
  113. for (uint32_t i = 0; i < size_per_item; i++) {
  114. if (elems[i] != db_copy.get()[(ele_index * size_per_item) + i]) {
  115. cout << "Main: elems " << (int)elems[i] << ", db "
  116. << (int) db_copy.get()[(ele_index * size_per_item) + i] << endl;
  117. cout << "Main: PIR result wrong at " << i << endl;
  118. failed = true;
  119. }
  120. }
  121. if(failed){
  122. return -1;
  123. }
  124. // Output results
  125. cout << "Main: PIR result correct!" << endl;
  126. cout << "Main: PIRServer pre-processing time: " << time_pre_us / 1000 << " ms" << endl;
  127. cout << "Main: PIRClient query generation time: " << time_query_us / 1000 << " ms" << endl;
  128. cout << "Main: PIRClient serialized query generation time: " << time_s_query_us / 1000 << " ms" << endl;
  129. cout << "Main: PIRServer query deserialization time: " << time_deserial_us << " us" << endl;
  130. cout << "Main: PIRServer reply generation time: " << time_server_us / 1000 << " ms" << endl;
  131. cout << "Main: PIRClient answer decode time: " << time_decode_us / 1000 << " ms" << endl;
  132. cout << "Main: Query size: " << query_size << " bytes" << endl;
  133. cout << "Main: Reply num ciphertexts: " << reply.size() << endl;
  134. cout << "Main: Reply size: " << reply_size << " bytes" << endl;
  135. return 0;
  136. }