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    // 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 out without server. After
 43    // receiving others public key shares client independently aggregates the share
 44    // and produces 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 shares 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, given that private inputs are
 75    // encrypted using collective public key, the private inputs are directly
 76    // encrypted under the ideal RLWE secret `s`.
 77    let c0_a = thread_rng().gen::<u8>();
 78    let c0_enc = pk.encrypt(vec![c0_a].as_slice());
 79    let c1_a = thread_rng().gen::<u8>();
 80    let c1_enc = pk.encrypt(vec![c1_a].as_slice());
 81    let c2_a = thread_rng().gen::<u8>();
 82    let c2_enc = pk.encrypt(vec![c2_a].as_slice());
 83    let c3_a = thread_rng().gen::<u8>();
 84    let c3_enc = pk.encrypt(vec![c3_a].as_slice());
 85
 86    // Clients upload their server key along with private encrypted inputs to
 87    // the server
 88
 89    // Server side //
 90
 91    // Server receives server key shares from each client and proceeds to
 92    // aggregated the shares and produce the server key
 93    let server_key = aggregate_server_key_shares(&server_key_shares);
 94    server_key.set_server_key();
 95
 96    // Server proceeds to extract clients private inputs
 97    //
 98    // Clients encrypt their FheUint8s inputs packed in a batched ciphertext.
 99    // The server must extract clients private inputs from the batch ciphertext
100    // either (1) using `extract_at(index)` to extract `index`^{th} FheUint8
101    // ciphertext (2) `extract_all()` to extract all available FheUint8s (3)
102    // `extract_many(many)` to extract first `many` available FheUint8s
103    let c0_a_enc = c0_enc.extract_at(0);
104    let c1_a_enc = c1_enc.extract_at(0);
105    let c2_a_enc = c2_enc.extract_at(0);
106    let c3_a_enc = c3_enc.extract_at(0);
107
108    // Server proceeds to evaluate function1 on clients private inputs
109    let ct_out_f1 = function1_fhe(&c0_a_enc, &c1_a_enc, &c2_a_enc, &c3_a_enc);
110
111    // After server has finished evaluating the circuit on client private
112    // inputs. Clients can proceed to multi-party decryption protocol to
113    // decryption output ciphertext
114
115    // Client Side //
116
117    // In multi-party decryption protocol, client must come online, download the
118    // output ciphertext from the server, product "output ciphertext" dependent
119    // decryption share, and send it to other parties. After receiving
120    // decryption shares of other parties, client independently aggregates the
121    // decrytion shares and decrypts the output ciphertext.
122
123    // Client generate decryption shares
124    let decryption_shares = cks
125        .iter()
126        .map(|k| k.gen_decryption_share(&ct_out_f1))
127        .collect_vec();
128
129    // After receiving decryption shares from other parties, client aggregates the
130    // shares and decryption output ciphertext
131    let out_f1 = cks[0].aggregate_decryption_shares(&ct_out_f1, &decryption_shares);
132
133    // Check correctness of function1 output
134    let want_f1 = function1(c0_a, c1_a, c2_a, c3_a);
135    assert!(out_f1 == want_f1);
136
137    // --------
138
139    // Once server key is produced it can be re-used across different functions
140    // with different private client inputs for the same set of clients.
141    //
142    // Here we run `function2_fhe` for the same of clients but with different
143    // private inputs. Clients do not need to participate in the 2 round
144    // protocol again, instead they only upload their new private inputs to the
145    // server.
146
147    // Clients encrypt their private inputs
148    let c0_a = thread_rng().gen::<u8>();
149    let c0_enc = pk.encrypt(vec![c0_a].as_slice());
150    let c1_a = thread_rng().gen::<u8>();
151    let c1_enc = pk.encrypt(vec![c1_a].as_slice());
152    let c2_a = thread_rng().gen::<u8>();
153    let c2_enc = pk.encrypt(vec![c2_a].as_slice());
154    let c3_a = thread_rng().gen::<u8>();
155    let c3_enc = pk.encrypt(vec![c3_a].as_slice());
156
157    // Clients uploads only their new private inputs to the server
158
159    // Server side //
160
161    // Server receives private inputs from the clients, extract them, and
162    // proceeds to evaluate `function2_fhe`
163    let c0_a_enc = c0_enc.extract_at(0);
164    let c1_a_enc = c1_enc.extract_at(0);
165    let c2_a_enc = c2_enc.extract_at(0);
166    let c3_a_enc = c3_enc.extract_at(0);
167
168    let ct_out_f2 = function2_fhe(&c0_a_enc, &c1_a_enc, &c2_a_enc, &c3_a_enc);
169
170    // Client side //
171
172    // Clients generate decryption shares for `ct_out_f2`
173    let decryption_shares = cks
174        .iter()
175        .map(|k| k.gen_decryption_share(&ct_out_f2))
176        .collect_vec();
177
178    // Clients aggregate decryption shares and decrypt `ct_out_f2`
179    let out_f2 = cks[0].aggregate_decryption_shares(&ct_out_f2, &decryption_shares);
180
181    // We check correctness of function2
182    let want_f2 = function2(c0_a, c1_a, c2_a, c3_a);
183    assert!(want_f2 == out_f2);
184}