spiral-spir/src/main.rs

use std::env;
use std::time::Instant;

use rand::RngCore;

use spiral_spir::client::Client;
use spiral_spir::init;
use spiral_spir::server::Server;
use spiral_spir::DbEntry;

fn main() {
    let args: Vec<String> = env::args().collect();
    if args.len() < 2 || args.len() > 5 {
        println!(
            "Usage: {} r [num_threads [num_preproc [num_pirs]]]\nr = log_2(num_records)",
            args[0]
        );
        return;
    }
    let r: usize = args[1].parse().unwrap();
    let mut num_threads = 1usize;
    let mut num_preproc = 1usize;
    let num_pirs: usize;
    if args.len() > 2 {
        num_threads = args[2].parse().unwrap();
    }
    if args.len() > 3 {
        num_preproc = args[3].parse().unwrap();
    }
    if args.len() > 4 {
        num_pirs = args[4].parse().unwrap();
    } else {
        num_pirs = num_preproc;
    }
    let num_records = 1 << r;
    let num_records_mask = num_records - 1;

    println!("===== ONE-TIME SETUP =====\n");

    let otsetup_start = Instant::now();

    init(num_threads);
    let (client, pub_params) = Client::new(r);
    let pub_params_len = pub_params.len();
    let server = Server::new(r, pub_params);

    let otsetup_us = otsetup_start.elapsed().as_micros();
    println!("One-time setup: {} µs", otsetup_us);
    println!("pub_params len = {}", pub_params_len);

    println!("\n===== PREPROCESSING =====\n");

    println!("num_preproc = {}", num_preproc);

    let preproc_client_start = Instant::now();
    let preproc_msg = client.preproc(num_preproc);
    let preproc_client_us = preproc_client_start.elapsed().as_micros();

    println!("Preprocessing client: {} µs", preproc_client_us);
    println!("preproc_msg len = {}", preproc_msg.len());

    let preproc_server_start = Instant::now();
    let preproc_resp = server.preproc_process(&preproc_msg);
    let preproc_server_us = preproc_server_start.elapsed().as_micros();

    println!("Preprocessing server: {} µs", preproc_server_us);
    println!("preproc_resp len = {}", preproc_resp.len());

    let preproc_finish_start = Instant::now();
    client.preproc_finish(&preproc_resp);
    let preproc_finish_us = preproc_finish_start.elapsed().as_micros();

    println!("Preprocessing client finish: {} µs", preproc_finish_us);

    // Create a database with recognizable contents
    let db: Vec<DbEntry> = ((0 as DbEntry)..(num_records as DbEntry))
        .map(|x| 10000001 * x)
        .collect();
    let dbptr = db.as_ptr();

    let mut rng = rand::thread_rng();
    for i in 1..num_pirs + 1 {
        println!("\n===== SPIR QUERY {} =====\n", i);

        let idx = (rng.next_u64() as usize) & num_records_mask;
        let query_client_start = Instant::now();
        let query_msg = client.query(idx);
        let query_client_us = query_client_start.elapsed().as_micros();

        println!("Query client: {} µs", query_client_us);
        println!("query_msg len = {}", query_msg.len());

        let query_server_start = Instant::now();
        let query_resp = server.query_process(&query_msg, dbptr, 100, 20);
        let query_server_us = query_server_start.elapsed().as_micros();

        println!("Query server: {} µs", query_server_us);
        println!("query_resp len = {}", query_resp.len());

        let query_finish_start = Instant::now();
        let entry = client.query_finish(&query_resp);
        let query_finish_us = query_finish_start.elapsed().as_micros();

        println!("Query client finish: {} µs", query_finish_us);
        println!("idx = {}; entry = {}", idx, entry);
    }

    /*
        let spiral_params = params::get_spiral_params(r);
        let mut rng = rand::thread_rng();
        print_params_summary(&spiral_params);
        println!("OT one-time setup: {} µs", otsetup_us);

        // One-time setup for the Spiral client
        let spc_otsetup_start = Instant::now();
        let mut clientrng = rand::thread_rng();
        let mut client = Client::init(&spiral_params, &mut clientrng);
        let pub_params = client.generate_keys();
        let pub_params_buf = pub_params.serialize();
        let spc_otsetup_us = spc_otsetup_start.elapsed().as_micros();
        let spiral_blocking_factor = spiral_params.db_item_size / mem::size_of::<DbEntry>();
        println!(
            "Spiral client one-time setup: {} µs, {} bytes",
            spc_otsetup_us,
            pub_params_buf.len()
        );

        println!("\n===== PREPROCESSING =====\n");

        // Spiral preprocessing: create a PIR lookup for an element at a
        // random location
        let spc_query_start = Instant::now();
        let rand_idx = (rng.next_u64() as usize) % num_records;
        let rand_pir_idx = rand_idx / spiral_blocking_factor;
        println!("rand_idx = {} rand_pir_idx = {}", rand_idx, rand_pir_idx);
        let spc_query = client.generate_query(rand_pir_idx);
        let spc_query_buf = spc_query.serialize();
        let spc_query_us = spc_query_start.elapsed().as_micros();
        println!(
            "Spiral query: {} µs, {} bytes",
            spc_query_us,
            spc_query_buf.len()
        );

        // Create the database encryption keys and do the OT to fetch the
        // right one, but don't actually encrypt the database yet
        let dbkeys = gen_db_enc_keys(r);
        let otkeyreq_start = Instant::now();
        let (keystate, keyquery) = otkey_request(rand_idx, r);
        let keyquerysize = keyquery.len() * keyquery[0].len();
        let otkeyreq_us = otkeyreq_start.elapsed().as_micros();
        let otkeysrv_start = Instant::now();
        let keyresponse = otkey_serve(keyquery, &dbkeys);
        let keyrespsize = keyresponse.len() * keyresponse[0].len();
        let otkeysrv_us = otkeysrv_start.elapsed().as_micros();
        let otkeyrcv_start = Instant::now();
        let otkey = otkey_receive(keystate, &keyresponse);
        let otkeyrcv_us = otkeyrcv_start.elapsed().as_micros();
        println!("key OT query in {} µs, {} bytes", otkeyreq_us, keyquerysize);
        println!("key OT serve in {} µs, {} bytes", otkeysrv_us, keyrespsize);
        println!("key OT receive in {} µs", otkeyrcv_us);

        // Create a database with recognizable contents
        let db: Vec<DbEntry> = ((0 as DbEntry)..(num_records as DbEntry))
            .map(|x| 10000001 * x)
            .collect();

        println!("\n===== RUNTIME =====\n");

        // Pick the record we actually want to query
        let q = (rng.next_u64() as usize) % num_records;

        // Compute the offset from the record index we're actually looking
        // for to the random one we picked earlier.  Tell it to the server,
        // who will rotate right the database by that amount before
        // encrypting it.
        let idx_offset = (num_records + rand_idx - q) % num_records;

        println!("Send to server {} bytes", 8 /* sizeof(idx_offset) */);

        // The server rotates, blinds, and encrypts the database
        let blind: DbEntry = 20;
        let encdb_start = Instant::now();
        let encdb = encdb_xor_keys(&db, &dbkeys, r, idx_offset, blind, num_threads);
        let encdb_us = encdb_start.elapsed().as_micros();
        println!("Server encrypt database {} µs", encdb_us);

        // Load the encrypted database into Spiral
        let sps_loaddb_start = Instant::now();
        let sps_db = load_db_from_slice_mt(&spiral_params, &encdb, num_threads);
        let sps_loaddb_us = sps_loaddb_start.elapsed().as_micros();
        println!("Server load database {} µs", sps_loaddb_us);

        // Do the PIR query
        let sps_query_start = Instant::now();
        let sps_query = Query::deserialize(&spiral_params, &spc_query_buf);
        let sps_response = process_query(&spiral_params, &pub_params, &sps_query, sps_db.as_slice());
        let sps_query_us = sps_query_start.elapsed().as_micros();
        println!(
            "Server compute response {} µs, {} bytes (*including* the above expansion time)",
            sps_query_us,
            sps_response.len()
        );

        // Decode the response to yield the whole Spiral block
        let spc_recv_start = Instant::now();
        let encdbblock = client.decode_response(sps_response.as_slice());
        // Extract the one encrypted DbEntry we were looking for (and the
        // only one we are able to decrypt)
        let entry_in_block = rand_idx % spiral_blocking_factor;
        let loc_in_block = entry_in_block * mem::size_of::<DbEntry>();
        let loc_in_block_end = (entry_in_block + 1) * mem::size_of::<DbEntry>();
        let encdbentry = DbEntry::from_le_bytes(
            encdbblock[loc_in_block..loc_in_block_end]
                .try_into()
                .unwrap(),
        );
        let decdbentry = otkey_decrypt(&otkey, rand_idx, encdbentry);
        let spc_recv_us = spc_recv_start.elapsed().as_micros();
        println!("Client decode response {} µs", spc_recv_us);
        println!("index = {}, Response = {}", q, decdbentry);
    */
}