fuzzytags-sim/src/datasets.rs

use crate::oracle::Oracle;
use crate::server::SimulatedServer;
use fuzzytags::{RootSecret, Tag, TaggingKey};
use rand::Rng;
use serde::Deserialize;
use tracing::event;
use tracing::span;
use tracing::Level;

pub trait TemporalDataset {
    fn register_with_server<R>(&self, server: &mut SimulatedServer, rng: &mut R, min_p: usize, max_p: usize, oracle: &mut Oracle)
    where
        R: Rng;
    fn playthough_traffic<R>(&self, server: &mut SimulatedServer, oracle: &mut Oracle, sample: usize, skip: usize, entangled_prob: f64, rng: &mut R)
    where
        R: Rng;
}

#[derive(Clone, Debug, Deserialize)]
struct TemporalSocialNetworkRecord {
    src_node: usize,
    dst_node: usize,
    timestamp: u64,
}

pub struct CsvDataset {
    root_secrets: Vec<RootSecret<24>>,
    tagging_keys: Vec<TaggingKey<24>>,
    records: Vec<TemporalSocialNetworkRecord>,
}

impl CsvDataset {
    pub fn load_dataset(filename: &str) -> CsvDataset {
        let mut rdr = csv::Reader::from_path(filename).unwrap();

        let mut num_recipients = 0;
        let mut records: Vec<TemporalSocialNetworkRecord> = vec![];
        for result in rdr.deserialize() {
            // Notice that we need to provide a type hint for automatic
            // deserialization.
            let record: Result<TemporalSocialNetworkRecord, csv::Error> = result;
            match record {
                Ok(record) => {
                    if record.dst_node > num_recipients {
                        num_recipients = record.dst_node;
                    }
                    if record.src_node > num_recipients {
                        num_recipients = record.src_node;
                    }
                    records.push(record.clone());
                }
                Err(err) => {
                    panic!("invalid data record found in {}", filename)
                }
            };
        }

        // generate a root secret for each member of the network
        let mut root_secrets = vec![];
        let mut tagging_keys = vec![];
        for i in 0..num_recipients + 1 {
            let secret = RootSecret::<24>::generate();
            let tagging_key = secret.tagging_key();
            root_secrets.push(secret);
            tagging_keys.push(tagging_key)
        }

        CsvDataset { root_secrets, tagging_keys, records }
    }

    pub fn num_parties(&self) -> usize {
        self.root_secrets.len()
    }

    pub fn num_records(&self) -> usize {
        self.records.len()
    }
}

impl TemporalDataset for CsvDataset {
    fn register_with_server<R>(&self, server: &mut SimulatedServer, rng: &mut R, min_p: usize, max_p: usize, oracle: &mut Oracle)
    where
        R: Rng,
    {
        for secret in self.root_secrets.iter() {
            let n = rng.gen_range(min_p..max_p);
            let span = span!(Level::INFO, "register", party = secret.tagging_key().id().as_str());
            let _enter = span.enter();
            let detection_key = secret.extract_detection_key(n);
            event!(Level::TRACE, "create detection key {detection_key}", detection_key = detection_key.id().as_str());
            event!(Level::TRACE, "register with server");
            server.register_key(&detection_key, &secret.tagging_key());
            oracle.register_party(secret.tagging_key().id());
        }
    }

    fn playthough_traffic<R>(&self, server: &mut SimulatedServer, oracle: &mut Oracle, sample: usize, skip: usize, entangled_prob: f64, rng: &mut R)
    where
        R: Rng,
    {
        /// TODO timestamps?
        for (i, record) in self.records.iter().skip(skip).take(sample).enumerate() {
            if i % 100 == 0 {
                let progress = (i - skip) as f64 / (sample as f64);
                let days = (((record.timestamp as f64 / 60.0) / 60.0) / 24.0);
                event!(Level::INFO, "progress..{:.2} ({} days)", progress * 100.0, days);
            }
            // We pretend that the server will always have access to the sender, even though
            // in practical deployments we could mitigate this somewhat using Tor / mixnet.
            let tagging_key_src = &self.tagging_keys[record.src_node];
            let tagging_key_dst = &self.tagging_keys[record.dst_node];

            let entangle = rng.gen_bool(entangled_prob);
            if entangle {
                let entangled_party = rng.gen_range(0..self.num_parties());
                let tagging_key_entangled = &self.tagging_keys[entangled_party];

                event!(Level::TRACE, "entangled send {party}", party = tagging_key_dst.id().as_str());

                let tag = TaggingKey::<24>::generate_entangled_tag(vec![tagging_key_dst.clone(), tagging_key_entangled.clone()], 8);
                event!(Level::TRACE, "message sent server {tag}", tag = tag.to_string());
                server.add_message(tag, tagging_key_src);

                oracle.add_event(tagging_key_src.id(), tagging_key_dst.id(), Some(tagging_key_entangled.id()), 0.0);
            } else {
                event!(Level::TRACE, "regular send {party}", party = tagging_key_dst.id().as_str());
                let tag = tagging_key_dst.generate_tag();
                event!(Level::TRACE, "message sent server {tag}", tag = tag.to_string());
                server.add_message(tag, tagging_key_src);
                oracle.add_event(tagging_key_src.id(), tagging_key_dst.id(), None, 0.0);
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use crate::datasets::CsvDataset;

    #[test]
    fn it_works() {
        let dataset = CsvDataset::load_dataset("datasets/email-Eu-core-temporal.txt");
        assert_eq!(332334, dataset.num_records());
    }
}