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