spiral-spir/src/client.rs

use rand::RngCore;

use std::collections::VecDeque;
use std::mem;
use std::os::raw::c_uchar;
use std::sync::mpsc::*;
use std::thread::*;

use aes::Block;

use rayon::prelude::*;

use subtle::Choice;

use curve25519_dalek::scalar::Scalar;

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

enum Command {
    PreProc(usize),
    PreProcFinish(Vec<PreProcSingleRespMsg>),
    Query(usize),
    QueryFinish(Vec<u8>),
}

enum Response {
    PubParams(Vec<u8>),
    PreProcMsg(Vec<u8>),
    PreProcDone,
    QueryMsg(Vec<u8>),
    QueryDone(DbEntry),
}

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

// The internal client state for a single preprocess ready to be used
struct PreProcSingleState {
    rand_idx: usize,
    ot_key: Block,
}

pub struct Client {
    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);
        spawn(move || {
            let spiral_params = params::get_spiral_params(r);
            let mut spiral_client = spiral_rs::client::Client::init(&spiral_params);
            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 outstanding preprocessing queries
            let mut preproc_out_state: Vec<PreProcOutSingleState> = Vec::new();

            // State for preprocessing queries ready to be used
            let mut preproc_state: VecDeque<PreProcSingleState> = VecDeque::new();

            // State for outstanding active queries
            let mut query_state: VecDeque<PreProcSingleState> = VecDeque::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_out_state.is_empty());
                        let mut preproc_msg: Vec<PreProcSingleMsg> = Vec::new();
                        (0..num_preproc)
                        .into_par_iter()
                        .map(|_| {
                            let mut rng = rand::thread_rng();
                            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);
                            (PreProcOutSingleState { rand_idx, ot_state },
                            PreProcSingleMsg {
                                ot_query,
                                spc_query,
                            })
                        })
                        .unzip_into_vecs(&mut preproc_out_state, &mut preproc_msg);
                        let ret: Vec<u8> = bincode::serialize(&preproc_msg).unwrap();
                        outgoing_resp_send.send(Response::PreProcMsg(ret)).unwrap();
                    }
                    Ok(Command::PreProcFinish(srvresp)) => {
                        let num_preproc = srvresp.len();
                        assert!(preproc_out_state.len() == num_preproc);
                        let mut newstate: VecDeque<PreProcSingleState> = preproc_out_state
                            .into_par_iter()
                            .zip(srvresp)
                            .map(|(c, s)| {
                                let ot_key = otkey_receive(c.ot_state, &s.ot_resp);
                                PreProcSingleState {
                                    rand_idx: c.rand_idx,
                                    ot_key,
                                }
                            })
                            .collect();
                        preproc_state.append(&mut newstate);
                        preproc_out_state = Vec::new();
                        outgoing_resp_send.send(Response::PreProcDone).unwrap();
                    }
                    Ok(Command::Query(idx)) => {
                        // panic if there are no preproc states
                        // available
                        let nextstate = preproc_state.pop_front().unwrap();
                        let offset = (num_records + idx - nextstate.rand_idx) & num_records_mask;
                        let mut querymsg: Vec<u8> = Vec::new();
                        querymsg.extend(offset.to_le_bytes());
                        query_state.push_back(nextstate);
                        outgoing_resp_send
                            .send(Response::QueryMsg(querymsg))
                            .unwrap();
                    }
                    Ok(Command::QueryFinish(msg)) => {
                        // panic if there is no outstanding state
                        let nextstate = query_state.pop_front().unwrap();
                        let encdbblock = spiral_client.decode_response(msg.as_slice());
                        // Extract the one encrypted DbEntry we were
                        // looking for (and the only one we are able to
                        // decrypt)
                        let entry_in_block = nextstate.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 =
                            dbentry_decrypt(&nextstate.ot_key, nextstate.rand_idx, encdbentry);
                        outgoing_resp_send
                            .send(Response::QueryDone(decdbentry))
                            .unwrap();
                    }
                    // When adding new messages, the following line is
                    // useful during development
                    // _ => 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 {
                incoming_cmd,
                outgoing_resp,
            },
            pub_params,
        )
    }

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

    pub fn preproc_finish(&self, msg: &[u8]) {
        self.incoming_cmd
            .send(Command::PreProcFinish(bincode::deserialize(msg).unwrap()))
            .unwrap();
        match self.outgoing_resp.recv() {
            Ok(Response::PreProcDone) => (),
            _ => panic!("Received something unexpected in preproc_finish"),
        }
    }

    pub fn query(&self, idx: usize) -> Vec<u8> {
        self.incoming_cmd.send(Command::Query(idx)).unwrap();
        match self.outgoing_resp.recv() {
            Ok(Response::QueryMsg(x)) => x,
            _ => panic!("Received something unexpected in preproc"),
        }
    }

    pub fn query_finish(&self, msg: &[u8]) -> DbEntry {
        self.incoming_cmd
            .send(Command::QueryFinish(msg.to_vec()))
            .unwrap();
        match self.outgoing_resp.recv() {
            Ok(Response::QueryDone(entry)) => entry,
            _ => panic!("Received something unexpected in preproc_finish"),
        }
    }
}

#[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(clientptr: *mut Client, num_preproc: u32) -> VecData {
    let client = unsafe {
        assert!(!clientptr.is_null());
        &mut *clientptr
    };
    let retvec = client.preproc(num_preproc as usize);
    to_vecdata(retvec)
}

#[no_mangle]
pub extern "C" fn spir_client_preproc_finish(
    clientptr: *mut Client,
    msgdata: *const c_uchar,
    msglen: usize,
) {
    let client = unsafe {
        assert!(!clientptr.is_null());
        &mut *clientptr
    };
    let msg_slice = unsafe {
        assert!(!msgdata.is_null());
        std::slice::from_raw_parts(msgdata, msglen)
    };
    client.preproc_finish(msg_slice);
}

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

#[no_mangle]
pub extern "C" fn spir_client_query_finish(
    clientptr: *mut Client,
    msgdata: *const c_uchar,
    msglen: usize,
) -> DbEntry {
    let client = unsafe {
        assert!(!clientptr.is_null());
        &mut *clientptr
    };
    let msg_slice = unsafe {
        assert!(!msgdata.is_null());
        std::slice::from_raw_parts(msgdata, msglen)
    };
    client.query_finish(msg_slice)
}