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    // Select parameter set
 23    set_parameter_set(ParameterSelector::InteractiveLTE4Party);
 24
 25    // set application's common reference seed
 26    let mut seed = [0u8; 32];
 27    thread_rng().fill_bytes(&mut seed);
 28    set_common_reference_seed(seed);
 29
 30    let no_of_parties = 4;
 31
 32    // Client side //
 33
 34    // Clients generate their private keys
 35    let cks = (0..no_of_parties)
 36        .into_iter()
 37        .map(|_| gen_client_key())
 38        .collect_vec();
 39
 40    // -- Round 1 -- //
 41    // In round 1 each client generates their share for the collective public key.
 42    // They send public key shares to each other with or out without the server.
 43    // After receiving others public key shares clients independently aggregate
 44    // the shares and produce the collective public key `pk`
 45
 46    let pk_shares = cks
 47        .iter()
 48        .map(|k| interactive_multi_party_round1_share(k))
 49        .collect_vec();
 50
 51    // Clients aggregate public key shares to produce collective public key `pk`
 52    let pk = aggregate_public_key_shares(&pk_shares);
 53
 54    // -- Round 2 -- //
 55    // In round 2 each client generates server key share using the public key `pk`.
 56    // Clients may also encrypt their private inputs using collective public key
 57    // `pk`. Each client then uploads their server key share and private input
 58    // ciphertexts to the server.
 59
 60    // Clients generate server key shares
 61    //
 62    // We assign user_id 0 to client 0, user_id 1 to client 1, user_id 2 to client
 63    // 2, and user_id 4 to client 4.
 64    //
 65    // Note that `user_id`'s must be unique among the clients and must be less than
 66    // total number of clients.
 67    let server_key_shares = cks
 68        .iter()
 69        .enumerate()
 70        .map(|(user_id, k)| gen_mp_keys_phase2(k, user_id, no_of_parties, &pk))
 71        .collect_vec();
 72
 73    // Each client encrypts their private inputs using the collective public key
 74    // `pk`. Unlike non-inteactive MPC protocol, private inputs are
 75    // encrypted using collective public key.
 76    let c0_a = thread_rng().gen::<u8>();
 77    let c0_enc = pk.encrypt(vec![c0_a].as_slice());
 78    let c1_a = thread_rng().gen::<u8>();
 79    let c1_enc = pk.encrypt(vec![c1_a].as_slice());
 80    let c2_a = thread_rng().gen::<u8>();
 81    let c2_enc = pk.encrypt(vec![c2_a].as_slice());
 82    let c3_a = thread_rng().gen::<u8>();
 83    let c3_enc = pk.encrypt(vec![c3_a].as_slice());
 84
 85    // Clients upload their server key along with private encrypted inputs to
 86    // the server
 87
 88    // Server side //
 89
 90    // Server receives server key shares from each client and proceeds to
 91    // aggregate the shares and produce the server key
 92    let server_key = aggregate_server_key_shares(&server_key_shares);
 93    server_key.set_server_key();
 94
 95    // Server proceeds to extract clients private inputs
 96    //
 97    // Clients encrypt their FheUint8s inputs packed in a batched ciphertext.
 98    // The server must extract clients private inputs from the batch ciphertext
 99    // either (1) using `extract_at(index)` to extract `index`^{th} FheUint8
100    // ciphertext (2) or using `extract_all()` to extract all available FheUint8s
101    // (3) or using `extract_many(many)` to extract first `many` available FheUint8s
102    let c0_a_enc = c0_enc.extract_at(0);
103    let c1_a_enc = c1_enc.extract_at(0);
104    let c2_a_enc = c2_enc.extract_at(0);
105    let c3_a_enc = c3_enc.extract_at(0);
106
107    // Server proceeds to evaluate function1 on clients private inputs
108    let ct_out_f1 = function1_fhe(&c0_a_enc, &c1_a_enc, &c2_a_enc, &c3_a_enc);
109
110    // After server has finished evaluating the circuit on client private
111    // inputs, clients can proceed to multi-party decryption protocol to
112    // decrypt output ciphertext
113
114    // Client Side //
115
116    // In multi-party decryption protocol, client must come online, download the
117    // output ciphertext from the server, product "output ciphertext" dependent
118    // decryption share, and send it to other parties. After receiving
119    // decryption shares of other parties, clients independently aggregate the
120    // decrytion shares and decrypt the output ciphertext.
121
122    // Clients generate decryption shares
123    let decryption_shares = cks
124        .iter()
125        .map(|k| k.gen_decryption_share(&ct_out_f1))
126        .collect_vec();
127
128    // After receiving decryption shares from other parties, clients aggregate the
129    // shares and decrypt output ciphertext
130    let out_f1 = cks[0].aggregate_decryption_shares(&ct_out_f1, &decryption_shares);
131
132    // Check correctness of function1 output
133    let want_f1 = function1(c0_a, c1_a, c2_a, c3_a);
134    assert!(out_f1 == want_f1);
135
136    // --------
137
138    // Once server key is produced it can be re-used across different functions
139    // with different private client inputs for the same set of clients.
140    //
141    // Here we run `function2_fhe` for the same of clients but with different
142    // private inputs. Clients do not need to participate in the 2 round
143    // protocol again, instead they only upload their new private inputs to the
144    // server.
145
146    // Clients encrypt their private inputs
147    let c0_a = thread_rng().gen::<u8>();
148    let c0_enc = pk.encrypt(vec![c0_a].as_slice());
149    let c1_a = thread_rng().gen::<u8>();
150    let c1_enc = pk.encrypt(vec![c1_a].as_slice());
151    let c2_a = thread_rng().gen::<u8>();
152    let c2_enc = pk.encrypt(vec![c2_a].as_slice());
153    let c3_a = thread_rng().gen::<u8>();
154    let c3_enc = pk.encrypt(vec![c3_a].as_slice());
155
156    // Clients uploads only their new private inputs to the server
157
158    // Server side //
159
160    // Server receives private inputs from the clients, extracts them, and
161    // proceeds to evaluate `function2_fhe`
162    let c0_a_enc = c0_enc.extract_at(0);
163    let c1_a_enc = c1_enc.extract_at(0);
164    let c2_a_enc = c2_enc.extract_at(0);
165    let c3_a_enc = c3_enc.extract_at(0);
166
167    let ct_out_f2 = function2_fhe(&c0_a_enc, &c1_a_enc, &c2_a_enc, &c3_a_enc);
168
169    // Client side //
170
171    // Clients generate decryption shares for `ct_out_f2`
172    let decryption_shares = cks
173        .iter()
174        .map(|k| k.gen_decryption_share(&ct_out_f2))
175        .collect_vec();
176
177    // Clients aggregate decryption shares and decrypt `ct_out_f2`
178    let out_f2 = cks[0].aggregate_decryption_shares(&ct_out_f2, &decryption_shares);
179
180    // We check correctness of function2
181    let want_f2 = function2(c0_a, c1_a, c2_a, c3_a);
182    assert!(want_f2 == out_f2);
183}