utils: add Fp with FromPrf, remove prf module

utils/Cargo.toml

@@ -8,7 +8,7 @@ edition = "2021"
 [dependencies]
 aes = "0.8.1"
 blake3 = "1.3.3"
-ff = "0.13.0"
+ff = { version = "0.13.0", features = ["derive"] }
 num = "0.4.0"
 rand = "0.8.5"
 rand_chacha = "0.3.1"

utils/src/field.rs

@@ -0,0 +1,72 @@
+use blake3;
+use ff::PrimeField;
+use rand::{thread_rng, Rng};
+
+/// Prime field with modulus
+/// p = 340282366920938462946865773367900766209.
+#[derive(PrimeField)]
+#[PrimeFieldModulus = "340282366920938462946865773367900766209"]
+#[PrimeFieldGenerator = "7"]
+#[PrimeFieldReprEndianness = "little"]
+pub struct Fp([u64; 3]);
+
+pub trait FromPrf {
+    type PrfKey: Copy;
+    /// PRF key generation
+    fn prf_key_gen() -> Self::PrfKey;
+    /// PRF into Fp
+    fn prf(key: &Self::PrfKey, input: u64) -> Self;
+}
+
+pub trait FromHash {
+    /// Hash into Fp
+    fn hash(input: u64) -> Self;
+}
+
+impl Fp {
+    fn from_xof(mut xof: blake3::OutputReader) -> Self {
+        assert_eq!(Self::NUM_BITS, 128);
+        loop {
+            let tmp = {
+                let mut repr = [0u64; 3];
+                for i in 0..2 {
+                    let mut bytes = [0u8; 8];
+                    xof.fill(&mut bytes);
+                    repr[i] = u64::from_le_bytes(bytes);
+                }
+                Self(repr)
+            };
+
+            if tmp.is_valid() {
+                return tmp;
+            }
+        }
+    }
+}
+
+impl FromPrf for Fp {
+    type PrfKey = [u8; blake3::KEY_LEN];
+
+    /// PRF key generation
+    fn prf_key_gen() -> Self::PrfKey {
+        thread_rng().gen()
+    }
+
+    /// PRF into Fp
+    fn prf(key: &Self::PrfKey, input: u64) -> Self {
+        let mut hasher = blake3::Hasher::new_keyed(&key);
+        hasher.update(&input.to_be_bytes());
+        let xof = hasher.finalize_xof();
+        Self::from_xof(xof)
+    }
+}
+
+impl FromHash for Fp {
+    /// Hash into Fp
+    fn hash(input: u64) -> Self {
+        let mut hasher = blake3::Hasher::new();
+        hasher.update(&input.to_be_bytes());
+        let xof = hasher.finalize_xof();
+        Self::from_xof(xof)
+    }
+}

utils/src/lib.rs

@@ -1,6 +1,6 @@
 pub mod bit_decompose;
+pub mod field;
 pub mod fixed_key_aes;
 pub mod p_ot;
 pub mod permutation;
-pub mod prf;
 pub mod pseudorandom_conversion;

utils/src/p_ot.rs

@@ -1,23 +1,23 @@
+use crate::field::FromPrf;
 use crate::permutation::Permutation;
-use crate::prf::{Prf, PrfKey};
 use core::marker::PhantomData;
-use ff::{Field, FromUniformBytes};
+use ff::Field;
 
-pub struct POTKeyParty<F, Perm> {
+pub struct POTKeyParty<F: FromPrf, Perm> {
     /// log of the database size
     log_domain_size: u32,
     /// if init was run
     is_initialized: bool,
     /// PRF key of the Index Party
-    prf_key_i: Option<PrfKey>,
+    prf_key_i: Option<<F as FromPrf>::PrfKey>,
     /// PRF key of the Receiver Party
-    prf_key_r: Option<PrfKey>,
+    prf_key_r: Option<<F as FromPrf>::PrfKey>,
     /// Permutation
     permutation: Option<Perm>,
     _phantom: PhantomData<F>,
 }
 
-impl<F: Field + FromUniformBytes<{ Prf::OUT_LEN }>, Perm: Permutation> POTKeyParty<F, Perm> {
+impl<F: Field + FromPrf, Perm: Permutation> POTKeyParty<F, Perm> {
     pub fn new(log_domain_size: u32) -> Self {
         Self {
             log_domain_size,
@@ -29,10 +29,10 @@ impl<F: Field + FromUniformBytes<{ Prf::OUT_LEN }>, Perm: Permutation> POTKeyPar
         }
     }
 
-    pub fn init(&mut self) -> ((PrfKey, Perm::Key), PrfKey) {
+    pub fn init(&mut self) -> ((F::PrfKey, Perm::Key), F::PrfKey) {
         assert!(!self.is_initialized);
-        self.prf_key_i = Some(Prf::key_gen());
-        self.prf_key_r = Some(Prf::key_gen());
+        self.prf_key_i = Some(F::prf_key_gen());
+        self.prf_key_r = Some(F::prf_key_gen());
         let permutation_key = Perm::sample(self.log_domain_size);
         self.permutation = Some(Perm::from_key(permutation_key));
         self.is_initialized = true;
@@ -48,26 +48,25 @@ impl<F: Field + FromUniformBytes<{ Prf::OUT_LEN }>, Perm: Permutation> POTKeyPar
         (0..n)
             .map(|x| {
                 let pi_x = self.permutation.as_ref().unwrap().permute(x) as u64;
-                Prf::eval::<F>(&self.prf_key_i.unwrap(), pi_x)
-                    + Prf::eval::<F>(&self.prf_key_r.unwrap(), pi_x)
+                F::prf(&self.prf_key_i.unwrap(), pi_x) + F::prf(&self.prf_key_r.unwrap(), pi_x)
             })
             .collect()
     }
 }
 
-pub struct POTIndexParty<F, Perm> {
+pub struct POTIndexParty<F: FromPrf, Perm> {
     /// log of the database size
     log_domain_size: u32,
     /// if init was run
     is_initialized: bool,
     /// PRF key of the Index Party
-    prf_key_i: Option<PrfKey>,
+    prf_key_i: Option<<F as FromPrf>::PrfKey>,
     /// Permutation
     permutation: Option<Perm>,
     _phantom: PhantomData<F>,
 }
 
-impl<F: Field + FromUniformBytes<{ Prf::OUT_LEN }>, Perm: Permutation> POTIndexParty<F, Perm> {
+impl<F: Field + FromPrf, Perm: Permutation> POTIndexParty<F, Perm> {
     pub fn new(log_domain_size: u32) -> Self {
         Self {
             log_domain_size,
@@ -78,30 +77,31 @@ impl<F: Field + FromUniformBytes<{ Prf::OUT_LEN }>, Perm: Permutation> POTIndexP
         }
     }
 
-    pub fn init(&mut self, prf_key_i: PrfKey, permutation_key: Perm::Key) {
+    pub fn init(&mut self, prf_key_i: F::PrfKey, permutation_key: Perm::Key) {
         assert!(!self.is_initialized);
         self.prf_key_i = Some(prf_key_i);
         self.permutation = Some(Perm::from_key(permutation_key));
         self.is_initialized = true;
     }
 
-    pub fn access(&self, index: u64) -> (F, u64) {
+    pub fn access(&self, index: u64) -> (u64, F) {
+        assert!(index < (1 << self.log_domain_size));
         let pi_x = self.permutation.as_ref().unwrap().permute(index as usize) as u64;
-        (Prf::eval(&self.prf_key_i.unwrap(), pi_x), pi_x)
+        (pi_x, F::prf(&self.prf_key_i.unwrap(), pi_x))
     }
 }
 
-pub struct POTReceiverParty<F> {
+pub struct POTReceiverParty<F: FromPrf> {
     /// log of the database size
     log_domain_size: u32,
     /// if init was run
     is_initialized: bool,
     /// PRF key of the Receiver Party
-    prf_key_r: Option<PrfKey>,
+    prf_key_r: Option<<F as FromPrf>::PrfKey>,
     _phantom: PhantomData<F>,
 }
 
-impl<F: Field + FromUniformBytes<{ Prf::OUT_LEN }>> POTReceiverParty<F> {
+impl<F: Field + FromPrf> POTReceiverParty<F> {
     pub fn new(log_domain_size: u32) -> Self {
         Self {
             log_domain_size,
@@ -111,13 +111,14 @@ impl<F: Field + FromUniformBytes<{ Prf::OUT_LEN }>> POTReceiverParty<F> {
         }
     }
 
-    pub fn init(&mut self, prf_key_r: PrfKey) {
+    pub fn init(&mut self, prf_key_r: F::PrfKey) {
         assert!(!self.is_initialized);
         self.prf_key_r = Some(prf_key_r);
         self.is_initialized = true;
     }
 
     pub fn access(&self, permuted_index: u64, output_share: F) -> F {
-        Prf::eval::<F>(&self.prf_key_r.unwrap(), permuted_index) + output_share
+        assert!(permuted_index < (1 << self.log_domain_size));
+        F::prf(&self.prf_key_r.unwrap(), permuted_index) + output_share
     }
 }

utils/src/prf.rs

@@ -1,21 +0,0 @@
-use blake3;
-use ff::FromUniformBytes;
-use rand::{thread_rng, Rng};
-
-#[derive(Clone, Copy, Debug)]
-pub struct PrfKey([u8; blake3::KEY_LEN]);
-
-pub struct Prf {}
-
-impl Prf {
-    pub const OUT_LEN: usize = blake3::OUT_LEN;
-
-    pub fn key_gen() -> PrfKey {
-        PrfKey(thread_rng().gen())
-    }
-
-    pub fn eval<F: FromUniformBytes<{ blake3::OUT_LEN }>>(key: &PrfKey, index: u64) -> F {
-        let hash = blake3::keyed_hash(&key.0, &index.to_be_bytes());
-        F::from_uniform_bytes(hash.as_bytes())
-    }
-}