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