1use itertools::Itertools;
  2use phantom_zone::*;
  3use rand::{thread_rng, Rng, RngCore};
  4
  5fn function1(a: u8, b: u8, c: u8, d: u8) -> u8 {
  6    ((a + b) * c) * d
  7}
  8
  9fn function1_fhe(a: &FheUint8, b: &FheUint8, c: &FheUint8, d: &FheUint8) -> FheUint8 {
 10    &(&(a + b) * c) * d
 11}
 12
 13fn function2(a: u8, b: u8, c: u8, d: u8) -> u8 {
 14    (a * b) + (c * d)
 15}
 16
 17fn function2_fhe(a: &FheUint8, b: &FheUint8, c: &FheUint8, d: &FheUint8) -> FheUint8 {
 18    &(a * b) + &(c * d)
 19}
 20
 21fn main() {
 22    set_parameter_set(ParameterSelector::NonInteractiveLTE4Party);
 23
 24    // set application's common reference seed
 25    let mut seed = [0u8; 32];
 26    thread_rng().fill_bytes(&mut seed);
 27    set_common_reference_seed(seed);
 28
 29    let no_of_parties = 4;
 30
 31    // Clide side //
 32
 33    // Generate client keys
 34    let cks = (0..no_of_parties).map(|_| gen_client_key()).collect_vec();
 35
 36    // client 0 encrypts its private inputs
 37    let c0_a = thread_rng().gen::<u8>();
 38    // Clients encrypt their private inputs in a seeded batched ciphertext using
 39    // their private RLWE secret `u_j`.
 40    let c0_enc = cks[0].encrypt(vec![c0_a].as_slice());
 41
 42    // client 1 encrypts its private inputs
 43    let c1_a = thread_rng().gen::<u8>();
 44    let c1_enc = cks[1].encrypt(vec![c1_a].as_slice());
 45
 46    // client 2 encrypts its private inputs
 47    let c2_a = thread_rng().gen::<u8>();
 48    let c2_enc = cks[2].encrypt(vec![c2_a].as_slice());
 49
 50    // client 3 encrypts its private inputs
 51    let c3_a = thread_rng().gen::<u8>();
 52    let c3_enc = cks[3].encrypt(vec![c3_a].as_slice());
 53
 54    // Clients independently generate their server key shares
 55    //
 56    // We assign user_id 0 to client 0, user_id 1 to client 1, user_id 2 to client
 57    // 2, user_id 3 to client 3.
 58    //
 59    // Note that `user_id`s must be unique among the clients and must be less than
 60    // total number of clients.
 61    let server_key_shares = cks
 62        .iter()
 63        .enumerate()
 64        .map(|(id, k)| gen_server_key_share(id, no_of_parties, k))
 65        .collect_vec();
 66
 67    // Each client uploads their server key shares and encrypted private inputs to
 68    // the server in a single shot message.
 69
 70    // Server side //
 71
 72    // Server receives server key shares from each client and proceeds to aggregate
 73    // them to produce the server key. After this point, server can use the server
 74    // key to evaluate any arbitrary function on encrypted private inputs from
 75    // the fixed set of clients
 76
 77    // aggregate server shares and generate the server key
 78    let server_key = aggregate_server_key_shares(&server_key_shares);
 79    server_key.set_server_key();
 80
 81    // Server proceeds to extract private inputs sent by clients
 82    //
 83    // To extract client 0's (with user_id=0) private inputs we first key switch
 84    // client 0's private inputs from theit secret `u_j` to ideal secret of the mpc
 85    // protocol. To indicate we're key switching client 0's private input we
 86    // supply client 0's `user_id` i.e. we call `key_switch(0)`. Then we extract
 87    // the first ciphertext by calling `extract_at(0)`.
 88    //
 89    // Since client 0 only encrypts 1 input in batched ciphertext, calling
 90    // extract_at(index) for `index` > 0 will panic. If client 0 had more private
 91    // inputs then we can either extract them all at once with `extract_all` or
 92    // first `many` of them with `extract_many(many)`
 93    let ct_c0_a = c0_enc.unseed::<Vec<Vec<u64>>>().key_switch(0).extract_at(0);
 94
 95    let ct_c1_a = c1_enc.unseed::<Vec<Vec<u64>>>().key_switch(1).extract_at(0);
 96    let ct_c2_a = c2_enc.unseed::<Vec<Vec<u64>>>().key_switch(2).extract_at(0);
 97    let ct_c3_a = c3_enc.unseed::<Vec<Vec<u64>>>().key_switch(3).extract_at(0);
 98
 99    // After extracting each client's private inputs, server proceeds to evaluate
100    // function1
101    let now = std::time::Instant::now();
102    let ct_out_f1 = function1_fhe(&ct_c0_a, &ct_c1_a, &ct_c2_a, &ct_c3_a);
103    println!("Function1 FHE evaluation time: {:?}", now.elapsed());
104
105    // Server has finished running compute. Clients can proceed to decrypt the
106    // output ciphertext using multi-party decryption.
107
108    // Client side //
109
110    // In multi-party decryption, each client needs to come online, download output
111    // ciphertext from the server, produce "output ciphertext" dependent decryption
112    // share, and send it to other parties (either via p2p or via server). After
113    // receving decryption shares from other parties, clients can independently
114    // decrypt output ciphertext.
115
116    // each client produces decryption share
117    let decryption_shares = cks
118        .iter()
119        .map(|k| k.gen_decryption_share(&ct_out_f1))
120        .collect_vec();
121
122    // With all decryption shares, clients can aggregate the shares and decrypt the
123    // ciphertext
124    let out_f1 = cks[0].aggregate_decryption_shares(&ct_out_f1, &decryption_shares);
125
126    // we check correctness of function1
127    let want_out_f1 = function1(c0_a, c1_a, c2_a, c3_a);
128    assert_eq!(out_f1, want_out_f1);
129
130    // -----------
131
132    // Server key can be re-used for different functions with different private
133    // client inputs for the same set of clients.
134    //
135    // Here we run `function2_fhe` for the same set of client but with new inputs.
136    // Clients only have to upload their private inputs to the server this time.
137
138    // Each client encrypts their private input
139    let c0_a = thread_rng().gen::<u8>();
140    let c0_enc = cks[0].encrypt(vec![c0_a].as_slice());
141    let c1_a = thread_rng().gen::<u8>();
142    let c1_enc = cks[1].encrypt(vec![c1_a].as_slice());
143    let c2_a = thread_rng().gen::<u8>();
144    let c2_enc = cks[2].encrypt(vec![c2_a].as_slice());
145    let c3_a = thread_rng().gen::<u8>();
146    let c3_enc = cks[3].encrypt(vec![c3_a].as_slice());
147
148    // Clients upload only their new private inputs to the server
149
150    // Server side //
151
152    // Server receives clients private inputs and extracts them
153    let ct_c0_a = c0_enc.unseed::<Vec<Vec<u64>>>().key_switch(0).extract_at(0);
154    let ct_c1_a = c1_enc.unseed::<Vec<Vec<u64>>>().key_switch(1).extract_at(0);
155    let ct_c2_a = c2_enc.unseed::<Vec<Vec<u64>>>().key_switch(2).extract_at(0);
156    let ct_c3_a = c3_enc.unseed::<Vec<Vec<u64>>>().key_switch(3).extract_at(0);
157
158    // Server proceeds to evaluate `function2_fhe`
159    let now = std::time::Instant::now();
160    let ct_out_f2 = function2_fhe(&ct_c0_a, &ct_c1_a, &ct_c2_a, &ct_c3_a);
161    println!("Function2 FHE evaluation time: {:?}", now.elapsed());
162
163    // Client side //
164
165    // Each client generates decrytion share for `ct_out_f2`
166    let decryption_shares = cks
167        .iter()
168        .map(|k| k.gen_decryption_share(&ct_out_f2))
169        .collect_vec();
170
171    // Clients independently aggregate the shares and decrypt
172    let out_f2 = cks[0].aggregate_decryption_shares(&ct_out_f2, &decryption_shares);
173
174    // We check correctness of function2
175    let want_out_f2 = function2(c0_a, c1_a, c2_a, c3_a);
176    assert_eq!(out_f2, want_out_f2);
177}