123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251 |
- use rand::rngs::ThreadRng;
- 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::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::PreProcSingleRespMsg;
- use crate::VecData;
- enum Command {
- PreProc(usize),
- PreProcFinish(Vec<PreProcSingleRespMsg>),
- Query(usize),
- }
- enum Response {
- PubParams(Vec<u8>),
- PreProcMsg(Vec<u8>),
- PreProcDone,
- QueryMsg(Vec<u8>),
- }
- // 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 {
- 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 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.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_out_state.push(PreProcOutSingleState { 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();
- }
- 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 + nextstate.rand_idx - 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();
- }
- _ => 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(&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"),
- };
- ret
- }
- 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();
- let ret = match self.outgoing_resp.recv() {
- Ok(Response::QueryMsg(x)) => x,
- _ => panic!("Received something unexpected in preproc"),
- };
- 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(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)
- }
|