1use itertools::Itertools;
  2use phantom_zone::*;
  3use rand::{thread_rng, Rng, RngCore};
  4
  5struct Location<T>(T, T);
  6
  7impl<T> Location<T> {
  8    fn new(x: T, y: T) -> Self {
  9        Location(x, y)
 10    }
 11
 12    fn x(&self) -> &T {
 13        &self.0
 14    }
 15    fn y(&self) -> &T {
 16        &self.1
 17    }
 18}
 19
 20fn should_meet(a: &Location<u8>, b: &Location<u8>, b_threshold: &u8) -> bool {
 21    let diff_x = a.x() - b.x();
 22    let diff_y = a.y() - b.y();
 23    let d_sq = &(&diff_x * &diff_x) + &(&diff_y * &diff_y);
 24
 25    d_sq.le(b_threshold)
 26}
 27
 28/// Calculates distance square between a's and b's location. Returns a boolean
 29/// indicating whether diatance sqaure is <= `b_threshold`.
 30fn should_meet_fhe(
 31    a: &Location<FheUint8>,
 32    b: &Location<FheUint8>,
 33    b_threshold: &FheUint8,
 34) -> FheBool {
 35    let diff_x = a.x() - b.x();
 36    let diff_y = a.y() - b.y();
 37    let d_sq = &(&diff_x * &diff_x) + &(&diff_y * &diff_y);
 38
 39    d_sq.le(b_threshold)
 40}
 41
 42// Ever wondered who are the long distance friends (friends of friends or
 43// friends of friends of friends...) that live nearby ? But how do you find
 44// them? Surely no-one will simply reveal their exact location just because
 45// there's a slight chance that a long distance friend lives nearby.
 46//
 47// Here we write a simple application with two users `a` and `b`. User `a` wants
 48// to find (long distance) friends that live in their neighbourhood. User `b` is
 49// open to meeting new friends within some distance of their location. Both user
 50// `a` and `b` encrypt their locations and upload their encrypted locations to
 51// the server. User `b` also encrypts the distance square threshold within which
 52// they are interested in meeting new friends and sends encrypted distance
 53// square threshold to the server.
 54// The server calculates the square of the distance between user a's location
 55// and user b's location and produces encrypted boolean output indicating
 56// whether square of distance is <= user b's supplied distance square threshold.
 57// User `a` then comes online, downloads output ciphertext, produces their
 58// decryption share for user `b`, and uploads the decryption share to the
 59// server. User `b` comes online, downloads output ciphertext and user a's
 60// decryption share, produces their own decryption share, and then decrypts the
 61// encrypted boolean output. If the output is `True`, it indicates user `a` is
 62// within the distance square threshold defined by user `b`.
 63fn main() {
 64    set_parameter_set(ParameterSelector::NonInteractiveLTE2Party);
 65
 66    // set application's common reference seed
 67    let mut seed = [0u8; 32];
 68    thread_rng().fill_bytes(&mut seed);
 69    set_common_reference_seed(seed);
 70
 71    let no_of_parties = 2;
 72
 73    // Client Side //
 74
 75    // Generate client keys
 76    let cks = (0..no_of_parties).map(|_| gen_client_key()).collect_vec();
 77
 78    // We assign user_id 0 to user `a` and user_id 1 user `b`
 79    let a_id = 0;
 80    let b_id = 1;
 81    let user_a_secret = &cks[0];
 82    let user_b_secret = &cks[1];
 83
 84    // User `a` and `b` generate server key shares
 85    let a_server_key_share = gen_server_key_share(a_id, no_of_parties, user_a_secret);
 86    let b_server_key_share = gen_server_key_share(b_id, no_of_parties, user_b_secret);
 87
 88    // User `a` and `b` encrypt their locations
 89    let user_a_secret = &cks[0];
 90    let user_a_location = Location::new(thread_rng().gen::<u8>(), thread_rng().gen::<u8>());
 91    let user_a_enc =
 92        user_a_secret.encrypt(vec![*user_a_location.x(), *user_a_location.y()].as_slice());
 93
 94    let user_b_location = Location::new(thread_rng().gen::<u8>(), thread_rng().gen::<u8>());
 95    // User `b` also encrypts the distance square threshold
 96    let user_b_threshold = 40;
 97    let user_b_enc = user_b_secret
 98        .encrypt(vec![*user_b_location.x(), *user_b_location.y(), user_b_threshold].as_slice());
 99
100    // Server Side //
101
102    // Both user `a` and `b` upload their private inputs and server key shares to
103    // the server in single shot message
104    let server_key = aggregate_server_key_shares(&vec![a_server_key_share, b_server_key_share]);
105    server_key.set_server_key();
106
107    // Server parses private inputs from user `a` and `b`
108    let user_a_location_enc = {
109        let c = user_a_enc.unseed::<Vec<Vec<u64>>>().key_switch(a_id);
110        Location::new(c.extract_at(0), c.extract_at(1))
111    };
112    let (user_b_location_enc, user_b_threshold_enc) = {
113        let c = user_b_enc.unseed::<Vec<Vec<u64>>>().key_switch(b_id);
114        (
115            Location::new(c.extract_at(0), c.extract_at(1)),
116            c.extract_at(2),
117        )
118    };
119
120    // run the circuit
121    let out_c = should_meet_fhe(
122        &user_a_location_enc,
123        &user_b_location_enc,
124        &user_b_threshold_enc,
125    );
126
127    // Client Side //
128
129    // user `a` comes online, downloads `out_c`, produces a decryption share, and
130    // uploads the decryption share to the server.
131    let a_dec_share = user_a_secret.gen_decryption_share(&out_c);
132
133    // user `b` comes online downloads user `a`'s decryption share, generates their
134    // own decryption share, decrypts the output ciphertext. If the output is
135    // True, user `b` contacts user `a` to meet.
136    let b_dec_share = user_b_secret.gen_decryption_share(&out_c);
137    let out_bool =
138        user_b_secret.aggregate_decryption_shares(&out_c, &vec![b_dec_share, a_dec_share]);
139
140    assert_eq!(
141        out_bool,
142        should_meet(&user_a_location, &user_b_location, &user_b_threshold)
143    );
144
145    if out_bool {
146        println!("A lives nearby. B should meet A.");
147    } else {
148        println!("A lives too far away!")
149    }
150}