spiral-spir/src/client.rs

use rand::rngs::ThreadRng;
use rand::RngCore;

use std::mem;
use std::sync::mpsc::*;
use std::thread::*;

use subtle::Choice;

use curve25519_dalek::ristretto::RistrettoPoint;
use curve25519_dalek::scalar::Scalar;

use crate::ot::*;
use crate::params;
use crate::to_vecdata;
use crate::DbEntry;
use crate::PreProcSingleMsg;
use crate::VecData;

enum Command {
    PreProc(usize),
    PreProcResp(Vec<u8>),
}

enum Response {
    PubParams(Vec<u8>),
    PreProcMsg(Vec<u8>),
}

// The internal client state for a single preprocess query
struct PreProcSingleState {
    rand_idx: usize,
    ot_state: Vec<(Choice, Scalar)>,
}

pub struct Client {
    r: usize,
    thread_handle: JoinHandle<()>,
    incoming_cmd: SyncSender<Command>,
    outgoing_resp: Receiver<Response>,
}

impl Client {
    pub fn new(r: usize) -> (Self, Vec<u8>) {
        let (incoming_cmd, incoming_cmd_recv) = sync_channel(0);
        let (outgoing_resp_send, outgoing_resp) = sync_channel(0);
        let thread_handle = spawn(move || {
            let spiral_params = params::get_spiral_params(r);
            let mut clientrng = rand::thread_rng();
            let mut rng = rand::thread_rng();
            let mut spiral_client = spiral_rs::client::Client::init(&spiral_params, &mut clientrng);
            let num_records = 1 << r;
            let num_records_mask = num_records - 1;
            let spiral_blocking_factor = spiral_params.db_item_size / mem::size_of::<DbEntry>();

            // The first communication is the pub_params
            let pub_params = spiral_client.generate_keys().serialize();
            outgoing_resp_send
                .send(Response::PubParams(pub_params))
                .unwrap();

            // State for preprocessing queries
            let mut preproc_state: Vec<PreProcSingleState> = Vec::new();

            // Wait for commands
            loop {
                match incoming_cmd_recv.recv() {
                    Err(_) => break,
                    Ok(Command::PreProc(num_preproc)) => {
                        // Ensure we don't already have outstanding
                        // preprocessing state
                        assert!(preproc_state.len() == 0);
                        let mut preproc_msg: Vec<PreProcSingleMsg> = Vec::new();
                        for _ in 0..num_preproc {
                            let rand_idx = (rng.next_u64() as usize) & num_records_mask;
                            let rand_pir_idx = rand_idx / spiral_blocking_factor;
                            let spc_query = spiral_client.generate_query(rand_pir_idx).serialize();
                            let (ot_state, ot_query) = otkey_request(rand_idx, r);
                            preproc_state.push(PreProcSingleState { rand_idx, ot_state });
                            preproc_msg.push(PreProcSingleMsg {
                                ot_query,
                                spc_query,
                            });
                        }
                        let ret: Vec<u8> = bincode::serialize(&preproc_msg).unwrap();
                        outgoing_resp_send.send(Response::PreProcMsg(ret)).unwrap();
                    }
                    _ => panic!("Received something unexpected in client loop"),
                }
            }
        });
        let pub_params = match outgoing_resp.recv() {
            Ok(Response::PubParams(x)) => x,
            _ => panic!("Received something unexpected in client new"),
        };

        (
            Client {
                r,
                thread_handle,
                incoming_cmd,
                outgoing_resp,
            },
            pub_params,
        )
    }

    pub fn preproc_PIRs(&self, num_preproc: usize) -> Vec<u8> {
        self.incoming_cmd
            .send(Command::PreProc(num_preproc))
            .unwrap();
        let ret = match self.outgoing_resp.recv() {
            Ok(Response::PreProcMsg(x)) => x,
            _ => panic!("Received something unexpected in preproc_PIRs"),
        };
        ret
    }
}

#[repr(C)]
pub struct ClientNewRet {
    client: *mut Client,
    pub_params: VecData,
}

#[no_mangle]
pub extern "C" fn spir_client_new(r: u8) -> ClientNewRet {
    let (client, pub_params) = Client::new(r as usize);
    ClientNewRet {
        client: Box::into_raw(Box::new(client)),
        pub_params: to_vecdata(pub_params),
    }
}

#[no_mangle]
pub extern "C" fn spir_client_free(client: *mut Client) {
    if client.is_null() {
        return;
    }
    unsafe {
        Box::from_raw(client);
    }
}

#[no_mangle]
pub extern "C" fn spir_client_preproc_PIRs(clientptr: *mut Client, num_preproc: u32) -> VecData {
    let client = unsafe {
        assert!(!clientptr.is_null());
        &mut *clientptr
    };
    let retvec = client.preproc_PIRs(num_preproc as usize);
    to_vecdata(retvec)
}