Skip to content

Commit

Permalink
fix benches and limit cache to "interesting" values
Browse files Browse the repository at this point in the history 
Only values that are searched once are then cached. This limits the size
of the cache to the size of the Entry Table (worst case).
  • Loading branch information
@tbrezot
tbrezot committed Jul 1, 2024
1 parent 2f31f1e commit 6301d8f
Show file tree
Hide file tree
Showing 5 changed files with 109 additions and 78 deletions.
69 changes: 51 additions & 18 deletions benches/benches.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,19 +12,19 @@ use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};
use findex_bis::{dummy_decode, dummy_encode, Findex, Index, KvStore};
use futures::executor::block_on;

fn bench_search(c: &mut Criterion) {
fn build_benchmarking_index(rng: &mut impl CryptoRngCore) -> Vec<([u8; 8], HashSet<[u8; 8]>)> {
(0..6)
.map(|i| {
let kw = rng.next_u64().to_be_bytes();
let vals = (1..10_i64.pow(i) as usize)
.map(|_| rng.next_u64().to_be_bytes())
.collect::<HashSet<_>>();
(kw, vals)
})
.collect()
}
fn build_benchmarking_index(rng: &mut impl CryptoRngCore) -> Vec<([u8; 8], HashSet<[u8; 8]>)> {
(0..4)
.map(|i| {
let kw = rng.next_u64().to_be_bytes();
let vals = (0..10_i64.pow(i) as usize)
.map(|_| rng.next_u64().to_be_bytes())
.collect::<HashSet<_>>();
(kw, vals)
})
.collect()
}

fn bench_search(c: &mut Criterion) {
let mut rng = CsRng::from_entropy();
let index = build_benchmarking_index(&mut rng);
let seed = Secret::random(&mut rng);
Expand All @@ -36,23 +36,23 @@ fn bench_search(c: &mut Criterion) {
dummy_encode,
dummy_decode,
);
findex.insert(index.clone().into_iter());
block_on(findex.insert(index.clone().into_iter())).unwrap();

{
let mut group = c.benchmark_group("Multiple bindings search");
for (i, (kw, _)) in index.clone().into_iter().enumerate() {
for (i, (kw, vals)) in index.clone().into_iter().enumerate() {
let n = 10i32.pow(i as u32) as usize;
group.bench_function(BenchmarkId::from_parameter(n), |b| {
b.iter(|| {
block_on(findex.search([kw].into_iter())).expect("search failed");
let res = block_on(findex.search([kw].into_iter())).expect("search failed");
});
});
}
}

{
let mut group = c.benchmark_group("Multiple keywords search (1 binding)");
for i in 0..4 {
for i in 0..=4 {
let n = 10i32.pow(i) as usize;
let kws = vec![index[0].0; n];
group.bench_function(BenchmarkId::from_parameter(n), |b| {
Expand All @@ -69,7 +69,7 @@ fn bench_search(c: &mut Criterion) {

{
let mut group = c.benchmark_group("Multiple keywords search (1000 bindings)");
for i in 0..4 {
for i in 0..=4 {
let n = 10i32.pow(i) as usize;
let kws = vec![index[3].0; n];
group.bench_function(BenchmarkId::from_parameter(n), |b| {
Expand All @@ -85,10 +85,43 @@ fn bench_search(c: &mut Criterion) {
}
}

fn bench_insert(c: &mut Criterion) {
let mut rng = CsRng::from_entropy();
let index = build_benchmarking_index(&mut rng);
let seed = Secret::random(&mut rng);
{
let mut group = c.benchmark_group("Multiple bindings insert (same keyword)");
for (i, (kw, vals)) in index.clone().into_iter().enumerate() {
let n = 10i32.pow(i as u32) as usize;
group.bench_function(BenchmarkId::from_parameter(n), |b| {
b.iter_batched(
|| {
let rng = CsRng::from_entropy();
let seed = seed.clone();
let vals = vals.clone();
(seed, rng, kw, vals)
},
|(seed, rng, kw, vals)| {
let findex = Findex::new(
seed,
Arc::new(Mutex::new(rng)),
KvStore::default(),
dummy_encode,
dummy_decode,
);
block_on(findex.insert([(kw, vals)].into_iter())).expect("search failed");
},
criterion::BatchSize::SmallInput,
);
});
}
}
}

criterion_group!(
name = benches;
config = Criterion::default().sample_size(5000);
targets = bench_search,
targets = bench_search, bench_insert,
);

criterion_main!(benches);
1 change: 0 additions & 1 deletion src/address.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -53,7 +53,6 @@ impl<const LENGTH: usize> Add<u64> for Address<LENGTH> {
pos += 1;

if (pos % LENGTH) == 0 {
println!("circling");
carry -= 1;
}
}
Expand Down
1 change: 1 addition & 0 deletions src/error.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ pub enum Error<Address: Debug, MemoryError: std::error::Error> {
Memory(MemoryError),
Conversion(String),
MissingValue(Address),
EncryptionLayer(String),
}

impl<Address: Debug, MemoryError: std::error::Error> Display for Error<Address, MemoryError> {
Expand Down
111 changes: 56 additions & 55 deletions src/obf.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ use std::{

use crate::{address::Address, error::Error, Stm, ADDRESS_LENGTH, KEY_LENGTH};
use aes::{
cipher::{generic_array::GenericArray, BlockDecrypt, BlockEncrypt, KeyInit},
cipher::{generic_array::GenericArray, BlockEncrypt, KeyInit},
Aes256,
};
use cosmian_crypto_core::{
Expand All @@ -18,7 +18,7 @@ use cosmian_crypto_core::{
pub struct MemoryEncryptionLayer<Memory: Stm<Address = Address<ADDRESS_LENGTH>, Word = Vec<u8>>> {
permutation_key: SymmetricKey<KEY_LENGTH>,
encryption_key: SymmetricKey<32>,
cache: Arc<Mutex<HashMap<(Address<ADDRESS_LENGTH>, Vec<u8>), Vec<u8>>>>,
cache: Arc<Mutex<HashMap<Address<ADDRESS_LENGTH>, HashMap<Vec<u8>, Vec<u8>>>>>,
rng: Arc<Mutex<CsRng>>,
stm: Memory,
}
Expand All @@ -45,79 +45,89 @@ impl<Memory: Stm<Address = Address<ADDRESS_LENGTH>, Word = Vec<u8>>> MemoryEncry
}

/// Retains values cached for the given keys only.
pub fn retain_cached_keys(&self, keys: &HashSet<(Memory::Address, Vec<u8>)>) {
pub fn retain_cached_keys(&self, keys: &HashSet<Memory::Address>) {
self.cache
.lock()
.expect("poisoned mutex")
.deref_mut()
.retain(|k, _| keys.contains(k));
}

/// Gets the encrypted value from the cache and computes it upon cache miss.
pub fn encrypt(
/// Decrypts the given value and caches the ciphertext.
fn find_or_encrypt(
&self,
ptx: &[u8],
tok: &Memory::Address,
) -> Result<Vec<u8>, Error<Memory::Address, Memory::Error>> {
let k = (tok.clone(), ptx.to_vec());
if let Some(ctx) = self.find(&k) {
if let Some(ctx) = self.find(tok, ptx) {
Ok(ctx)
} else {
let nonce = Nonce::<{ Aes256Gcm::NONCE_LENGTH }>::new(&mut *self.rng());
let ctx = Aes256Gcm::new(&self.encryption_key)
.encrypt(&nonce, ptx, Some(tok))
.map_err(Error::Encryption)?;
let ctx = [nonce.as_bytes(), &ctx].concat();
self.bind(k, ctx.clone());
let ctx = self.encrypt(ptx, tok)?;
self.bind(tok.clone(), ptx.to_vec(), ctx.clone());
Ok(ctx)
}
}

/// Decrypts the given value and caches the ciphertext.
pub fn decrypt(
fn bind_and_decrypt(
&self,
ctx: Vec<u8>,
tok: Memory::Address,
) -> Result<Vec<u8>, Error<Memory::Address, Memory::Error>> {
let ptx = self.decrypt(&ctx, &tok)?;
self.bind(tok, ptx.clone(), ctx);
Ok(ptx)
}

fn encrypt(
&self,
ptx: &[u8],
tok: &Memory::Address,
) -> Result<Vec<u8>, Error<Memory::Address, Memory::Error>> {
let nonce = Nonce::<{ Aes256Gcm::NONCE_LENGTH }>::new(&mut *self.rng());
let ctx = Aes256Gcm::new(&self.encryption_key)
.encrypt(&nonce, ptx, Some(tok))
.map_err(Error::Encryption)?;
Ok([nonce.as_bytes(), &ctx].concat())
}

#[inline]
fn decrypt(
&self,
ctx: &[u8],
tok: &Memory::Address,
) -> Result<Vec<u8>, Error<Memory::Address, Memory::Error>> {
if ctx.len() < Aes256Gcm::NONCE_LENGTH {
return Err(Error::Parsing("ciphertext too small".to_string()));
}
let nonce = Nonce::try_from_slice(&ctx[..Aes256Gcm::NONCE_LENGTH])
.map_err(|e| Error::Parsing(e.to_string()))?;
let ptx = Aes256Gcm::new(&self.encryption_key)
.decrypt(&nonce, &ctx[Aes256Gcm::NONCE_LENGTH..], Some(&tok))
.map_err(Error::Encryption)?;
self.bind((tok, ptx.clone()), ctx);
Ok(ptx)
Aes256Gcm::new(&self.encryption_key)
.decrypt(&nonce, &ctx[Aes256Gcm::NONCE_LENGTH..], Some(tok))
.map_err(Error::Encryption)
}

pub fn reorder(&self, mut a: Memory::Address) -> Memory::Address {
Aes256::new(GenericArray::from_slice(&self.permutation_key))
.decrypt_block(GenericArray::from_mut_slice(&mut a));
a
}

pub fn permute(&self, mut a: Memory::Address) -> Memory::Address {
fn permute(&self, mut a: Memory::Address) -> Memory::Address {
Aes256::new(GenericArray::from_slice(&self.permutation_key))
.encrypt_block(GenericArray::from_mut_slice(&mut a));
a
}

pub fn bind(&self, k: (Memory::Address, Vec<u8>), v: Vec<u8>) {
self.cache
.lock()
.expect("poisoned lock")
.deref_mut()
.insert(k, v);
fn bind(&self, tok: Memory::Address, ptx: Vec<u8>, ctx: Vec<u8>) {
let mut cache = self.cache.lock().expect("poisoned lock");
if let Some(bindings) = cache.get_mut(&tok) {
bindings.insert(ptx, ctx);
}
}

pub fn find(&self, k: &(Address<ADDRESS_LENGTH>, Vec<u8>)) -> Option<Vec<u8>> {
self.cache
.lock()
.expect("poisoned lock")
.deref_mut()
.get(k)
.cloned()
fn find(&self, tok: &Address<ADDRESS_LENGTH>, ref_ptx: &[u8]) -> Option<Vec<u8>> {
let mut cache = self.cache.lock().expect("poisoned lock");
if let Some(bindings) = cache.get(tok) {
bindings.get(ref_ptx).cloned()
} else {
cache.insert(tok.clone(), HashMap::new());
None
}
}
}

Expand All @@ -139,7 +149,7 @@ impl<Memory: Stm<Address = Address<ADDRESS_LENGTH>, Word = Vec<u8>>> Stm
bindings
.into_iter()
.zip(tokens)
.map(|(ctx, tok)| ctx.map(|ctx| self.decrypt(ctx, tok.clone())).transpose())
.map(|(ctx, tok)| ctx.map(|ctx| self.decrypt(&ctx, &tok)).transpose())
.collect()
}

Expand All @@ -148,26 +158,17 @@ impl<Memory: Stm<Address = Address<ADDRESS_LENGTH>, Word = Vec<u8>>> Stm
guard: (Self::Address, Option<Self::Word>),
bindings: Vec<(Self::Address, Self::Word)>,
) -> Result<Option<Self::Word>, Self::Error> {
let (a, v) = guard;
let tok = self.permute(a);
let old = v.and_then(|v| self.find(&(tok.clone(), v)));

let tok = self.permute(guard.0.clone());
let old = guard.1.and_then(|ptx| self.find(&tok, &ptx));
let bindings = bindings
.into_iter()
.map(|(a, v)| {
let tok = self.permute(a);
self.encrypt(&v, &tok).map(|ctx| (tok, ctx))
self.find_or_encrypt(&v, &tok).map(|ctx| (tok, ctx))
})
.collect::<Result<Vec<_>, _>>()?;

let ctx = self.stm.guarded_write((tok.clone(), old), bindings).await?;

if let Some(ctx) = ctx {
let ptx = self.decrypt(ctx.clone(), tok)?;
Ok(Some(ptx))
} else {
Ok(None)
}
let cur = self.stm.guarded_write((tok.clone(), old), bindings).await?;
cur.map(|ctx| self.bind_and_decrypt(ctx, tok)).transpose()
}
}

Expand All @@ -194,7 +195,7 @@ mod tests {
let tok = Address::<ADDRESS_LENGTH>::random(&mut rng);
let ptx = vec![1];
let ctx = obf.encrypt(&ptx, &tok).unwrap();
let res = obf.decrypt(ctx, tok).unwrap();
let res = obf.decrypt(&ctx, &tok).unwrap();
assert_eq!(ptx.len(), res.len());
assert_eq!(ptx, res);
}
Expand Down
5 changes: 1 addition & 4 deletions src/ovec.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -142,10 +142,7 @@ impl<

let cur_header = res[0]
.clone()
.map(|v| {
println!("{v:?}");
Header::try_from(v.as_slice())
})
.map(|v| Header::try_from(v.as_slice()))
.transpose()
.map_err(Error::Conversion)?
.unwrap_or_default();
Expand Down

0 comments on commit 6301d8f

Please sign in to comment.