smpc/src/garble_classic.rs

use core::fmt::Display;
use rand::rngs::OsRng;
use rand::{CryptoRng, Rng, RngCore};
use std::collections::HashMap;
use std::ops::{BitAnd, BitOr, BitXor};
use tiny_keccak::{Hasher, Shake};

fn encrypt(key_a: [u8; 16], key_b: [u8; 16], out_key: [u8; 16]) -> [u8; 32] {
    let mut shake = Shake::v256();
    shake.update(&key_a);
    shake.update(&key_b);
    let mut enc_out = [0; 32];
    shake.finalize(&mut enc_out);
    for i in 0..16 {
        enc_out[i + 16] ^= out_key[i];
    }
    enc_out
}

fn decrypt(key_a: [u8; 16], key_b: [u8; 16], encrypted: [u8; 32]) -> Option<[u8; 16]> {
    let mut shake = Shake::v256();
    shake.update(&key_a);
    shake.update(&key_b);
    let mut dec_out = [0; 32];
    shake.finalize(&mut dec_out);
    if dec_out[..16] != encrypted[..16] {
        return None;
    }
    for i in 0..32 {
        dec_out[i] ^= encrypted[i];
    }
    let mut key = [0; 16];
    key[..16].copy_from_slice(&dec_out[16..(16 + 16)]);
    Some(key)
}

#[derive(Clone, Debug)]
struct Wire {
    sym: Option<String>,
    keys: [[u8; 16]; 2],
    secret: bool,
}

#[derive(Clone, Debug)]
struct Circuit {
    left: Box<CircOption>,
    right: Box<CircOption>,
    out: Wire,
    table: [[[u8; 32]; 2]; 2],
    sym: String,
}

#[derive(Clone, Debug)]
enum CircOption {
    Circuit(Circuit),
    Wire(Wire),
}

impl Display for Wire {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match &self.sym {
            Some(sym) => write!(f, "{}", sym),
            None => write!(f, ""),
        }
    }
}

impl Display for CircOption {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            CircOption::Wire(w) => write!(f, "{}", w),
            CircOption::Circuit(c) => write!(f, "{}", c),
        }
    }
}

impl Display for Circuit {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{} {} {}", self.left, self.sym, self.right)
    }
}

impl Wire {
    fn new_rng(sym: String, secret: bool, rng: &mut (impl RngCore + CryptoRng)) -> Self {
        let mut keys = [[0; 16]; 2];
        rng.fill_bytes(&mut keys[0]);
        rng.fill_bytes(&mut keys[1]);
        Self {
            sym: Some(sym),
            keys,
            secret,
        }
    }

    fn new(sym: String, secret: bool) -> Self {
        Self::new_rng(sym, secret, &mut OsRng)
    }

    fn to_circ(&self) -> CircOption {
        CircOption::Wire(self.clone())
    }

    fn empty(rng: &mut (impl RngCore + CryptoRng)) -> Self {
        let mut keys = [[0; 16]; 2];
        rng.fill_bytes(&mut keys[0]);
        rng.fill_bytes(&mut keys[1]);
        Self {
            sym: None,
            keys,
            secret: false,
        }
    }
}

impl Circuit {
    fn new(
        left: CircOption,
        right: CircOption,
        value: [[bool; 2]; 2],
        sym: String,
        rng: &mut (impl RngCore + CryptoRng),
    ) -> Self {
        let linp_keys = match left {
            CircOption::Wire(ref w) => w.keys,
            CircOption::Circuit(ref c) => c.out.keys,
        };

        let rinp_keys = match right {
            CircOption::Wire(ref w) => w.keys,
            CircOption::Circuit(ref c) => c.out.keys,
        };

        let out = Wire::empty(rng);

        let mut table = [[[0; 32]; 2]; 2];

        let i = rng.gen::<bool>() as usize;
        let j = rng.gen::<bool>() as usize;
        for left_bit in [i, 1 - i] {
            for right_bit in [j, 1 - j] {
                let enc = encrypt(
                    linp_keys[left_bit],
                    rinp_keys[right_bit],
                    out.keys[value[left_bit][right_bit] as usize],
                );

                table[left_bit][right_bit] = enc;
            }
        }

        Self {
            left: Box::new(left),
            right: Box::new(right),
            table,
            sym,
            out,
        }
    }

    fn to_circ(&self) -> CircOption {
        CircOption::Circuit(self.clone())
    }
}

impl CircOption {
    fn evaluate_re(&self, vals: HashMap<String, bool>) -> [u8; 16] {
        match self {
            CircOption::Circuit(c) => {
                let lv = c.left.evaluate_re(vals.clone());
                let rv = c.right.evaluate_re(vals);
                let mut out = [0; 16];
                for i in [0, 1] {
                    for j in [0, 1] {
                        if let Some(x) = decrypt(lv, rv, c.table[i][j]) {
                            out = x;
                        }
                    }
                }
                out
            }
            CircOption::Wire(w) => match &w.sym {
                // Errr OT? hello?
                Some(c) => w.keys[(vals[c]) as usize],
                None => panic!("This should not happen."),
            },
        }
    }

    pub fn evaluate(&self, vals: HashMap<String, bool>) -> bool {
        let out_label = self.evaluate_re(vals);

        match self {
            CircOption::Circuit(c) => {
                if out_label == c.out.keys[0] {
                    false
                } else if out_label == c.out.keys[1] {
                    true
                } else {
                    panic!("Error!");
                }
            }
            CircOption::Wire(w) => {
                if out_label == w.keys[0] {
                    false
                } else if out_label == w.keys[1] {
                    true
                } else {
                    panic!("Error!");
                }
            }
        }
    }
}

impl BitAnd for CircOption {
    type Output = CircOption;

    fn bitand(self, rhs: Self) -> Self::Output {
        CircOption::Circuit(Circuit::new(
            self,
            rhs,
            [[false, false], [false, true]],
            "&".to_string(),
            &mut OsRng,
        ))
    }
}

impl BitOr for CircOption {
    type Output = CircOption;

    fn bitor(self, rhs: Self) -> Self::Output {
        CircOption::Circuit(Circuit::new(
            self,
            rhs,
            [[false, true], [true, true]],
            "|".to_string(),
            &mut OsRng,
        ))
    }
}

impl BitXor for CircOption {
    type Output = CircOption;

    fn bitxor(self, rhs: Self) -> Self::Output {
        CircOption::Circuit(Circuit::new(
            self,
            rhs,
            [[false, true], [true, false]],
            "^".to_string(),
            &mut OsRng,
        ))
    }
}

#[cfg(test)]
mod test {

    use crate::garble_classic::Wire;
    use crate::garble_classic::{decrypt, encrypt};
    use std::collections::HashMap;

    #[test]
    fn enc() {
        let a = decrypt([2; 16], [4; 16], encrypt([2; 16], [4; 16], [9; 16]));
        assert!(a == Some([9; 16]));
    }

    #[test]
    fn circuit() {
        let a = Wire::new("a".to_string(), true).to_circ();
        let b = Wire::new("b".to_string(), false).to_circ();
        let c = Wire::new("c".to_string(), false).to_circ();
        let circ = (a | b) & c;

        let mut h = HashMap::new();
        h.insert("a".to_string(), false);
        h.insert("b".to_string(), false);
        h.insert("c".to_string(), true);
        let out = circ.evaluate(h);
        assert!(out == false);
    }
}