spiral-spir/src/server.rs

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

use rayon::prelude::*;

use spiral_rs::client::PublicParameters;
use spiral_rs::client::Query;
use spiral_rs::params::Params;

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

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

enum Response {
    PreProcResp(Vec<u8>),
}

// The internal client state for a single preprocess query
struct PreProcSingleState<'a> {
    db_keys: Vec<[u8; 16]>,
    query: Query<'a>,
}

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

impl Server {
    pub fn new(r: usize, pub_params: Vec<u8>) -> Self {
        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 pub_params = PublicParameters::deserialize(&spiral_params, &pub_params);

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

            // Wait for commands
            loop {
                match incoming_cmd_recv.recv() {
                    Err(_) => break,
                    Ok(Command::PreProcMsg(cliquery)) => {
                        let num_preproc = cliquery.len();
                        let mut resp_state: Vec<PreProcSingleState> =
                            Vec::with_capacity(num_preproc);
                        let mut resp_msg: Vec<PreProcSingleRespMsg> =
                            Vec::with_capacity(num_preproc);
                        cliquery
                            .into_par_iter()
                            .map(|q| {
                                let db_keys = gen_db_enc_keys(r);
                                let query = Query::deserialize(&spiral_params, &q.spc_query);
                                let ot_resp = otkey_serve(q.ot_query, &db_keys);
                                (
                                    PreProcSingleState { db_keys, query },
                                    PreProcSingleRespMsg { ot_resp },
                                )
                            })
                            .unzip_into_vecs(&mut resp_state, &mut resp_msg);
                        preproc_state.append(&mut VecDeque::from(resp_state));
                        let ret: Vec<u8> = bincode::serialize(&resp_msg).unwrap();
                        outgoing_resp_send.send(Response::PreProcResp(ret)).unwrap();
                    }
                    _ => panic!("Received something unexpected"),
                }
            }
        });
        Server {
            r,
            thread_handle,
            incoming_cmd,
            outgoing_resp,
        }
    }

    pub fn preproc_process(&self, msg: &[u8]) -> Vec<u8> {
        self.incoming_cmd
            .send(Command::PreProcMsg(bincode::deserialize(msg).unwrap()))
            .unwrap();
        let ret = match self.outgoing_resp.recv() {
            Ok(Response::PreProcResp(x)) => x,
            _ => panic!("Received something unexpected in preproc_process"),
        };
        ret
    }
}

#[no_mangle]
pub extern "C" fn spir_server_new(
    r: u8,
    pub_params: *const c_uchar,
    pub_params_len: usize,
) -> *mut Server {
    let pub_params_slice = unsafe {
        assert!(!pub_params.is_null());
        std::slice::from_raw_parts(pub_params, pub_params_len)
    };
    let mut pub_params_vec: Vec<u8> = Vec::new();
    pub_params_vec.extend_from_slice(pub_params_slice);
    let server = Server::new(r as usize, pub_params_vec);
    Box::into_raw(Box::new(server))
}

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

#[no_mangle]
pub extern "C" fn spir_server_preproc_process(
    serverptr: *mut Server,
    msgdata: *const c_uchar,
    msglen: usize,
) -> VecData {
    let server = unsafe {
        assert!(!serverptr.is_null());
        &mut *serverptr
    };
    let msg_slice = unsafe {
        assert!(!msgdata.is_null());
        std::slice::from_raw_parts(msgdata, msglen)
    };
    let retvec = server.preproc_process(&msg_slice);
    to_vecdata(retvec)
}