Skip to main content

mc_oblivious_map/
lib.rs

1// Copyright (c) 2018-2023 The MobileCoin Foundation
2
3//! Implementation of a bucketed cuckoo hash table where the arena is an
4//! oblivious RAM
5//! A cuckoo hash table is a data structure that performs access and removal
6//! using a constant number of operations.
7
8//! The trick is that our cuckoo hash table is actually 2 orams with 2 different
9//! hash functions. The key is always hashed twice for access, read, and removal
10//! because we guarantee the element must only reside in the location of its
11//! hash in one of the two orams. This also implies that these operations are
12//! oblivious, because we always do 2 oram queries, once per each oram, and
13//! inherit obliviousness from the oram. In our case, each hash bucket is an
14//! oram bucket that can hold multiple values. See
15//! https://www.ru.is/faculty/ulfar/CuckooHash.pdf for a survey of cuckoo hash
16//! using varying numbers of hash functions and bucket size.
17
18//! Insertion requires variable time and is not usually oblivious. This is ok
19//! for use on things like fog view/ledger where nothing they are inserting is
20//! secret. It is possible to obtain obtain obliviousness for things like
21//! fog-ingest. Please see documentation in [vartime_write_extended()] and
22//! [`ObliviousHashMap::access_and_insert()`]
23
24#![no_std]
25#![deny(unsafe_code)]
26#![deny(missing_docs)]
27
28extern crate alloc;
29
30use aligned_cmov::{
31    subtle::{Choice, ConstantTimeEq, ConstantTimeGreater},
32    typenum, A64Bytes, A8Bytes, ArrayLength, AsAlignedChunks, CMov,
33};
34use alloc::vec::Vec;
35use core::{
36    hash::BuildHasher,
37    marker::PhantomData,
38    ops::{Add, Sub},
39};
40use generic_array::sequence::Split;
41use mc_oblivious_traits::{
42    log2_ceil, OMapCreator, ORAMCreator, ObliviousHashMap, OMAP_FOUND, OMAP_INVALID_KEY,
43    OMAP_NOT_FOUND, OMAP_OVERFLOW, ORAM,
44};
45use rand_core::CryptoRng;
46use typenum::{PartialDiv, Sum, U8};
47
48mod build_hasher;
49use build_hasher::SipBuildHasher;
50
51/// In this implementation, the cuckoo hashing step is permitted to repeat at
52/// most 6 times before we give up. In experiments this lead to about ~75%
53/// memory utilitzation. This will depend on a lot of factors such as how big is
54/// the block size relative to key-size + value-size, and possibly also on the
55/// capacity. In the future we might want to make this configurable.
56const MAX_EVICTION_RETRIES: usize = 6;
57
58/// A bucketed cuckoo hash table built on top of oblivious storage.
59///
60/// The Block stored by ORAM is considered as a bucket in the hashing algorithm.
61/// The bucket gets broken up into aligned chunks of size KeySize + ValueSize,
62/// so the number of items in a bucket is BlockSize / (KeySize + ValueSize)
63pub struct CuckooHashTable<KeySize, ValueSize, BlockSize, RngType, O>
64where
65    KeySize: ArrayLength + Add<ValueSize> + PartialDiv<U8> + 'static,
66    ValueSize: ArrayLength + PartialDiv<U8>,
67    BlockSize: ArrayLength + PartialDiv<U8>,
68    RngType: CryptoRng + Send + Sync + 'static,
69    O: ORAM<BlockSize> + Send + Sync + 'static,
70    Sum<KeySize, ValueSize>: ArrayLength + Sub<KeySize, Output = ValueSize> + PartialDiv<U8>,
71{
72    /// The number of items in the table right now
73    num_items: u64,
74    /// The number of buckets in ONE of the two orams, must be a power of two
75    num_buckets: u64,
76    /// Key for the first hash function
77    hash1: SipBuildHasher,
78    /// Key for the second hash function
79    hash2: SipBuildHasher,
80    /// Oblivious storage for arena corresponding to first hash function
81    oram1: O,
82    /// Oblivious storage for arena corresponding to second hash function
83    oram2: O,
84    /// Rng used to make random eviction decisions
85    rng: RngType,
86    // phantom data
87    _key_size: PhantomData<fn() -> KeySize>,
88    _value_size: PhantomData<fn() -> ValueSize>,
89    _block_size: PhantomData<fn() -> BlockSize>,
90}
91
92impl<KeySize, ValueSize, BlockSize, RngType, O>
93    CuckooHashTable<KeySize, ValueSize, BlockSize, RngType, O>
94where
95    KeySize: ArrayLength + Add<ValueSize> + PartialDiv<U8> + 'static,
96    ValueSize: ArrayLength + PartialDiv<U8>,
97    BlockSize: ArrayLength + PartialDiv<U8>,
98    RngType: CryptoRng + Send + Sync + 'static,
99    O: ORAM<BlockSize> + Send + Sync + 'static,
100    Sum<KeySize, ValueSize>: ArrayLength + Sub<KeySize, Output = ValueSize> + PartialDiv<U8>,
101{
102    /// Create a new hashmap
103    /// The ORAM should be default initialized or bad things will happen
104    pub fn new<OC, M>(desired_capacity: u64, stash_size: usize, mut maker: M) -> Self
105    where
106        OC: ORAMCreator<BlockSize, RngType, Output = O>,
107        M: 'static + FnMut() -> RngType,
108    {
109        assert!(Self::BUCKET_CAPACITY > 0, "Block size is insufficient to store even one item. For good performance it should be enough to store several items.");
110        // We will have two ORAM arenas, and two hash functions. num_buckets is the
111        // number of buckets that one of the arenas should have, to achieve
112        // total desired capacity across both arenas, then rounded up to a power
113        // of two.
114        let num_buckets = 1u64 << log2_ceil((desired_capacity / Self::BUCKET_CAPACITY) / 2);
115        debug_assert!(
116            num_buckets & (num_buckets - 1) == 0,
117            "num_buckets must be a power of two"
118        );
119
120        let oram1 = OC::create(num_buckets, stash_size, &mut maker);
121        debug_assert!(num_buckets <= oram1.len(), "unexpected oram capacity");
122        let oram2 = OC::create(num_buckets, stash_size, &mut maker);
123        debug_assert!(num_buckets <= oram2.len(), "unexpected oram capacity");
124        debug_assert!(
125            oram1.len() == oram2.len(),
126            "Orams didn't have the same length, not expected"
127        );
128
129        let mut rng = maker();
130        let hash1 = SipBuildHasher::from_rng(&mut rng);
131        let hash2 = SipBuildHasher::from_rng(&mut rng);
132
133        Self {
134            num_items: 0,
135            num_buckets,
136            hash1,
137            hash2,
138            oram1,
139            oram2,
140            rng,
141            _key_size: Default::default(),
142            _value_size: Default::default(),
143            _block_size: Default::default(),
144        }
145    }
146
147    fn hash_query(&self, query: &A8Bytes<KeySize>) -> [u64; 2] {
148        let result1 = self.hash1.hash_one(query.as_slice()) & (self.num_buckets - 1);
149        let result2 = self.hash2.hash_one(query.as_slice()) & (self.num_buckets - 1);
150
151        [result1, result2]
152    }
153
154    // Given a block (stored in ORAM), which we think of as a hash-table bucket,
155    // check if any of the entries have key matching query, and count how many are
156    // zeroes This function is constant-time
157    fn count_before_insert(query: &A8Bytes<KeySize>, block: &A64Bytes<BlockSize>) -> (Choice, u32) {
158        let mut found = Choice::from(0);
159        let mut empty_count = 0u32;
160
161        let pairs: &[A8Bytes<Sum<KeySize, ValueSize>>] = block.as_aligned_chunks();
162        for pair in pairs {
163            let (key, _): (&A8Bytes<KeySize>, &A8Bytes<ValueSize>) = pair.split();
164            found |= key.ct_eq(query);
165            empty_count += key.ct_eq(&A8Bytes::<KeySize>::default()).unwrap_u8() as u32;
166        }
167
168        (found, empty_count)
169    }
170
171    // Given a block (stored in ORAM), which we think of as a hash-table bucket,
172    // branchlessly insert key-value pair into it, if condition is true
173    //
174    // - Interpret block as aligned KeySize + ValueSize chunks
175    // - Find the first one if any that has key of all zeroes or matching query, and
176    //   cmov value on top of that.
177    fn insert_to_block(
178        condition: Choice,
179        query: &A8Bytes<KeySize>,
180        new_value: &A8Bytes<ValueSize>,
181        block: &mut A64Bytes<BlockSize>,
182    ) {
183        // key_buf is initially query, and is zeroes for every match there-after
184        let mut key_buf = query.clone();
185        let pairs: &mut [A8Bytes<Sum<KeySize, ValueSize>>] = block.as_mut_aligned_chunks();
186        for pair in pairs {
187            let (key, value): (&mut A8Bytes<KeySize>, &mut A8Bytes<ValueSize>) =
188                <&mut A8Bytes<Sum<KeySize, ValueSize>> as Split<u8, KeySize>>::split(pair);
189            let test = condition & (key.ct_eq(query) | key.ct_eq(&A8Bytes::<KeySize>::default()));
190            key.cmov(test, &key_buf);
191            // This ensures that if we find the key a second time, or, after finding an
192            // empty space, the cell is zeroed to make space for other things
193            key_buf.cmov(test, &A8Bytes::default());
194            value.cmov(test, new_value);
195        }
196    }
197
198    const BUCKET_CAPACITY: u64 = (BlockSize::U64 / (KeySize::U64 + ValueSize::U64));
199}
200
201impl<KeySize, ValueSize, BlockSize, RngType, O> ObliviousHashMap<KeySize, ValueSize>
202    for CuckooHashTable<KeySize, ValueSize, BlockSize, RngType, O>
203where
204    KeySize: ArrayLength + Add<ValueSize> + PartialDiv<U8> + 'static,
205    ValueSize: ArrayLength + PartialDiv<U8>,
206    BlockSize: ArrayLength + PartialDiv<U8>,
207    RngType: CryptoRng + Send + Sync + 'static,
208    O: ORAM<BlockSize> + Send + Sync + 'static,
209    Sum<KeySize, ValueSize>: ArrayLength + Sub<KeySize, Output = ValueSize> + PartialDiv<U8>,
210{
211    fn len(&self) -> u64 {
212        self.num_items
213    }
214    fn capacity(&self) -> u64 {
215        2 * self.num_buckets * Self::BUCKET_CAPACITY
216    }
217
218    /// To read:
219    /// Early return if the query is all zero bytes
220    /// Hash the query (twice)
221    /// Load the corresponding blocks from ORAM (one at a time)
222    /// Interpret block as [(KeySize, ValueSize)]
223    /// Ct-compare the found key with the query
224    /// If successful, cmov OMAP_FOUND onto result_code and cmov value onto
225    /// result. Return result_code and result after scanning both loaded
226    /// blocks
227    fn read(&mut self, query: &A8Bytes<KeySize>, output: &mut A8Bytes<ValueSize>) -> u32 {
228        // Early return for invalid key
229        if bool::from(query.ct_eq(&A8Bytes::<KeySize>::default())) {
230            return OMAP_INVALID_KEY;
231        }
232
233        let mut result_code = OMAP_NOT_FOUND;
234        let hashes = self.hash_query(query);
235        self.oram1.access(hashes[0], |block| {
236            let pairs: &[A8Bytes<Sum<KeySize, ValueSize>>] = block.as_aligned_chunks();
237            for pair in pairs {
238                let (key, value): (&A8Bytes<KeySize>, &A8Bytes<ValueSize>) = pair.split();
239                let test = query.ct_eq(key);
240                result_code.cmov(test, &OMAP_FOUND);
241                output.cmov(test, value);
242            }
243        });
244        self.oram2.access(hashes[1], |block| {
245            let pairs: &[A8Bytes<Sum<KeySize, ValueSize>>] = block.as_aligned_chunks();
246            for pair in pairs {
247                let (key, value): (&A8Bytes<KeySize>, &A8Bytes<ValueSize>) = pair.split();
248                let test = query.ct_eq(key);
249                result_code.cmov(test, &OMAP_FOUND);
250                output.cmov(test, value);
251            }
252        });
253        result_code
254    }
255
256    /// For access:
257    /// Access (without deleting) must be fully oblivious, unlike write
258    /// - Checkout both buckets, scan them for the query, copying onto a stack
259    ///   buffer
260    /// - Run callback at the stack buffer
261    /// - Scan the buckets again and overwrite the old buffer
262    /// - If the callback asked to delete the item, also zero out its key.
263    ///
264    /// We would like to be constant-time about whether deletes happened or not,
265    /// but changing the number of items in the map is a side-effect which we
266    /// leak some information about, so we can't be totally principled about
267    /// that.
268    fn access_and_remove<F: FnOnce(u32, &mut A8Bytes<ValueSize>) -> Choice>(
269        &mut self,
270        query: &A8Bytes<KeySize>,
271        f: F,
272    ) {
273        let mut callback_buffer = A8Bytes::<ValueSize>::default();
274        // All zeroes is an invalid key, because we use that as a sentinel to detect
275        // free spots in the hashmap. In this scenario, we give the callback
276        // OMAP_INVALID_KEY and pass an empty buffer which is later discarded.
277        let query_is_valid = !query.ct_eq(&A8Bytes::<KeySize>::default());
278
279        let hashes = self.hash_query(query);
280        let oram1 = &mut self.oram1;
281        let oram2 = &mut self.oram2;
282
283        oram1.access(hashes[0], |block1| {
284            oram2.access(hashes[1], |block2| {
285                // If we never find the item, we will pass OMAP_NOT_FOUND to callback,
286                // and we will not insert it, no matter what the callback does.
287                let mut result_code = OMAP_INVALID_KEY;
288                result_code.cmov(query_is_valid, &OMAP_NOT_FOUND);
289                for block in &[&block1, &block2] {
290                    let pairs: &[A8Bytes<Sum<KeySize, ValueSize>>] = block.as_aligned_chunks();
291                    for pair in pairs {
292                        let (key, val): (&A8Bytes<KeySize>, &A8Bytes<ValueSize>) = pair.split();
293                        // Test if query = key, but also, always do nothing if query is invalid
294                        let test = query.ct_eq(key) & query_is_valid;
295                        callback_buffer.cmov(test, val);
296                        debug_assert!(
297                            result_code != OMAP_FOUND || bool::from(!test),
298                            "key should not be found twice"
299                        );
300                        result_code.cmov(test, &OMAP_FOUND);
301                    }
302                }
303                let should_delete = f(result_code, &mut callback_buffer);
304
305                for block in &mut [block1, block2] {
306                    let pairs: &mut [A8Bytes<Sum<KeySize, ValueSize>>] =
307                        block.as_mut_aligned_chunks();
308                    for pair in pairs {
309                        let (key, val): (&mut A8Bytes<KeySize>, &mut A8Bytes<ValueSize>) =
310                            pair.split();
311                        // Test if query = key, but also, always do nothing if query is invalid
312                        let test = query.ct_eq(key) & query_is_valid;
313                        val.cmov(test, &callback_buffer);
314                        // Zero the key if the callback wants us to delete this item
315                        key.cmov(test & should_delete, &Default::default());
316                    }
317                }
318            });
319        });
320    }
321
322    fn remove(&mut self, query: &A8Bytes<KeySize>) -> u32 {
323        // Early return for invalid key
324        if bool::from(query.ct_eq(&A8Bytes::<KeySize>::default())) {
325            return OMAP_INVALID_KEY;
326        }
327        let mut result_code = OMAP_NOT_FOUND;
328        let hashes = self.hash_query(query);
329        self.oram1.access(hashes[0], |block| {
330            let pairs: &mut [A8Bytes<Sum<KeySize, ValueSize>>] = block.as_mut_aligned_chunks();
331            for pair in pairs {
332                let (key, _): (&mut A8Bytes<KeySize>, &mut A8Bytes<ValueSize>) = pair.split();
333                let test = query.as_slice().ct_eq(key.as_slice());
334                key.cmov(test, &Default::default());
335                result_code.cmov(test, &OMAP_FOUND);
336            }
337        });
338        self.oram2.access(hashes[1], |block| {
339            let pairs: &mut [A8Bytes<Sum<KeySize, ValueSize>>] = block.as_mut_aligned_chunks();
340            for pair in pairs {
341                let (key, _): (&mut A8Bytes<KeySize>, &mut A8Bytes<ValueSize>) = pair.split();
342                let test = query.as_slice().ct_eq(key.as_slice());
343                key.cmov(test, &Default::default());
344                result_code.cmov(test, &OMAP_FOUND);
345            }
346        });
347        result_code
348    }
349
350    /// For writing: The insertion algorithm is, hash the item twice and load
351    /// its buckets. We always add to the less loaded of the two buckets,
352    /// breaking ties to the right, that is, prefering to write to oram2. If
353    /// BOTH buckets overflow, then we choose an item at random from oram1
354    /// bucket and kick it out, then we hash that item and insert it into the
355    /// other bucket where it can go, repeating the process if necessary. If
356    /// after a few tries it doesn't work, we give up, roll everything back, and
357    /// return OMAP_OVERFLOW.
358    ///
359    /// This is amortized constant time, but not constant time due
360    /// to this relocation. The intuition is that if we treat the hash functons
361    /// as random functions, we can treat the graph connecting buckets that are
362    /// both the target of an element as a random graph and with high
363    /// probability a sparse random graph doesn't have any cycles. See the proof
364    /// here: https://cs.stanford.edu/~rishig/courses/ref/l13a.pdf
365    ///
366    /// The access function is an alternative that allows modifying values in
367    /// the map without taking a variable amount of time.
368    ///
369    /// Note: this function is not generally oblivious with respect to the item
370    /// being inserted. In many cases, inserts don't need to be oblivious. It is
371    /// possible to use this function in such a way as to perform an completely
372    /// oblivious insert. We will explain that now. This function has some
373    /// limited data oblivious guarantees:
374    /// - If we restrict attention to calling vartime_write on items q which are
375    ///   already in the OMAP, then it is oblivious with respect to such q.
376    /// - If we restrict attention to calling vartime_write on items q which are
377    ///   not already in the OMAP, then it is oblivious with respect to these,
378    ///   assuming that the hash function used in the map has a strong PRF
379    ///   property.
380    /// - However, calling the function DOES reveal if q is already in the map
381    ///   or not. To avoid leaking this information, and perform a fully
382    ///   oblivious write-or-insert operation, the caller should use the
383    ///   access_and_insert function from the
384    ///   [`ObliviousHashMap::access_and_insert()`] trait.
385    fn vartime_write_extended(
386        &mut self,
387        query: &A8Bytes<KeySize>,
388        new_value: &A8Bytes<ValueSize>,
389        allow_overwrite: Choice,
390        allow_sideeffects_and_eviction: Choice,
391    ) -> u32 {
392        // Early return for invalid key
393        if bool::from(query.ct_eq(&A8Bytes::<KeySize>::default())) {
394            return OMAP_INVALID_KEY;
395        }
396
397        // The result_code we will return
398        let mut result_code = OMAP_NOT_FOUND;
399
400        // If after access we have to evict something, it is stored temporarily here
401        // Its bucket is pushed to eviction_from, and its index in the bucket to
402        // eviction_indices
403        let mut evicted_key = A8Bytes::<KeySize>::default();
404        let mut evicted_val = A8Bytes::<ValueSize>::default();
405
406        // The number of times we will try evicting before giving up
407        let mut eviction_retries = MAX_EVICTION_RETRIES;
408        // The buckets (hashes) from which we evicted items came, so that we can go back
409        // and restore them if we give up.
410        let mut eviction_from = Vec::<u64>::with_capacity(eviction_retries);
411        // The indices at which we evict items, so that we can go back and restore them
412        // if we give up.
413        let mut eviction_indices = Vec::<usize>::with_capacity(eviction_retries);
414
415        // Compute the hashes for this query
416        let hashes = self.hash_query(query);
417
418        // Get a let binding to self.rng
419        let rng = &mut self.rng;
420        let oram1 = &mut self.oram1;
421        let oram2 = &mut self.oram2;
422        oram1.access(hashes[0], |block1| {
423            oram2.access(hashes[1], |block2| {
424                // Note: These calls don't need to be constant-time to meet the requirement,
425                // but we already had the code that way, and it may serve as "defense-in-depth".
426                // If they show up in profiling, then we can make variable time versions
427                let (block1_found, block1_empty_count) = Self::count_before_insert(query, block1);
428                let (block2_found, block2_empty_count) = Self::count_before_insert(query, block2);
429                debug_assert!(
430                    !bool::from(block1_found & block2_found),
431                    "key should not be found twice"
432                );
433
434                let found = block1_found | block2_found;
435                result_code.cmov(found, &OMAP_FOUND);
436
437                // Scope for "condition" variable
438                {
439                    // condition is false when side-effects are disallowed, OR
440                    // if we found the item and we aren't allowed to overwrite
441                    let condition = allow_sideeffects_and_eviction & (allow_overwrite | !found);
442                    // write_to_block1 is true when we should prefer to write to block1 over block2
443                    // watch the case that hashes[0] == hashes[1] !
444                    // in that case we prefer to modify block2 (in oram2), arbitrarily.
445                    // So, if block2_found, we should be false, even if block1_found.
446                    // And if not found in either place and block1_empty_count ==
447                    // block2_empty_count, prefer block2.
448                    let write_to_block1 = !block2_found
449                        & (block1_found | block1_empty_count.ct_gt(&block2_empty_count));
450
451                    Self::insert_to_block(condition & write_to_block1, query, new_value, block1);
452                    Self::insert_to_block(condition & !write_to_block1, query, new_value, block2);
453                }
454
455                // If it wasn't found and both blocks were full,
456                // then we set result_code to overflow and evict something at random,
457                // from block1. This is done because the item will likely end in block2
458                // after that, which is what we want, per "Unbalanced allocations" papers
459                // that show that this improves performance.
460                // Skip this if eviction is disallowed
461                if bool::from(
462                    !found
463                        & block1_empty_count.ct_eq(&0)
464                        & block2_empty_count.ct_eq(&0)
465                        & allow_sideeffects_and_eviction,
466                ) {
467                    result_code = OMAP_OVERFLOW;
468
469                    let random = rng.next_u32();
470                    // index determines which entry in that bucket we evict
471                    let index = (random % (Self::BUCKET_CAPACITY as u32)) as usize;
472
473                    eviction_from.push(hashes[0]);
474                    eviction_indices.push(index);
475
476                    let pairs: &mut [A8Bytes<Sum<KeySize, ValueSize>>] =
477                        block1.as_mut_aligned_chunks();
478                    let pair = &mut pairs[index];
479                    let (key, val): (&mut A8Bytes<KeySize>, &mut A8Bytes<ValueSize>) = pair.split();
480                    evicted_key = key.clone();
481                    evicted_val = val.clone();
482                    // Note: This is a side-effect but we don't reach this line if allow_sideeffects
483                    // is false
484                    debug_assert!(bool::from(allow_sideeffects_and_eviction));
485                    *key = query.clone();
486                    *val = new_value.clone();
487                }
488            });
489        });
490
491        // Overflow handling loop
492        while result_code == OMAP_OVERFLOW {
493            if eviction_retries > 0 {
494                let last_evicted_from = eviction_from[eviction_from.len() - 1];
495
496                // We always start evicting from bucket 1, and alternate, per cuckoo hashing
497                // algo, so the next_oram (minus one) is eviction_from.len() %
498                // 2.
499                let next_oram = eviction_from.len() % 2;
500
501                let hashes = self.hash_query(&evicted_key);
502                // last_evicted_from should be the other place that evicted_key could go
503                debug_assert!(hashes[1 - next_oram] == last_evicted_from);
504                let dest = hashes[next_oram];
505
506                // Access what is hopefully a block with a vacant spot
507                let rng = &mut self.rng;
508                // Get a reference to the oram we will insert to. This is a branch,
509                // but the access patterns are completely predicatable based on number of passes
510                // through this loop.
511                let oram = if next_oram == 0 {
512                    &mut self.oram1
513                } else {
514                    &mut self.oram2
515                };
516
517                oram.access(dest, |block| {
518                    let pairs: &mut [A8Bytes<Sum<KeySize, ValueSize>>] =
519                        block.as_mut_aligned_chunks();
520                    debug_assert!(pairs.len() == Self::BUCKET_CAPACITY as usize);
521
522                    // If we find a vacant spot in this block, then insert evicted key and val there
523                    let mut found_vacant = Choice::from(0);
524                    for pair in pairs.iter_mut() {
525                        let (key, val): (&mut A8Bytes<KeySize>, &mut A8Bytes<ValueSize>) =
526                            pair.split();
527                        debug_assert!(
528                            key != &evicted_key,
529                            "evicted key should not be present anywhere"
530                        );
531                        let is_vacant = key.ct_eq(&A8Bytes::<KeySize>::default());
532                        // Note: This is a side-effect, but this code is unreachable if
533                        // allow_sideeffects is false.
534                        debug_assert!(bool::from(allow_sideeffects_and_eviction));
535                        let cond = !found_vacant & is_vacant;
536                        key.cmov(cond, &evicted_key);
537                        val.cmov(cond, &evicted_val);
538                        found_vacant |= is_vacant;
539                    }
540
541                    // If we found a vacant spot, then the result is not OMAP_OVERFLOW anymore, we
542                    // are done
543                    if bool::from(found_vacant) {
544                        // This will cause us to exit the while loop
545                        result_code = OMAP_NOT_FOUND;
546                    } else {
547                        // This block was full also, so we repeat the eviction process
548                        let index = (rng.next_u32() % (Self::BUCKET_CAPACITY as u32)) as usize;
549
550                        let pair = &mut pairs[index];
551                        let (key, val): (&mut A8Bytes<KeySize>, &mut A8Bytes<ValueSize>) =
552                            pair.split();
553
554                        // Evict this key and value for the previously evicted key and value.
555                        // Note: This is a side-effect, but this code is unreachable if
556                        // allow_sideeffects is false.
557                        debug_assert!(bool::from(allow_sideeffects_and_eviction));
558                        core::mem::swap(key, &mut evicted_key);
559                        core::mem::swap(val, &mut evicted_val);
560
561                        eviction_from.push(dest);
562                        eviction_indices.push(index);
563                    }
564                });
565
566                eviction_retries = eviction_retries.wrapping_sub(1);
567            } else {
568                // We have given up trying to evict things, we will now roll back
569                debug_assert!(eviction_from.len() == eviction_indices.len());
570                while !eviction_indices.is_empty() {
571                    // We always start evicting from bucket 1, and alternate, per cuckoo hashing
572                    // algo, so the next_oram we WOULD insert to (minus one) is
573                    // eviction_from.len() % 2.
574                    let next_oram = eviction_from.len() % 2;
575
576                    // Note: these unwraps are "okay" because we test is_empty(),
577                    // and we debug_assert that they have the same length.
578                    // If llvm is not eliminating the panicking path then we may
579                    // want to do this differently.
580                    let evicted_index = eviction_indices.pop().unwrap();
581                    let evicted_from = eviction_from.pop().unwrap();
582                    debug_assert!(
583                        self.hash_query(&evicted_key)[1 - next_oram] == evicted_from,
584                        "The evicted key doesn't hash to the spot we thought we evicted it from"
585                    );
586                    // Get a reference to the oram we took this item from. This is a branch,
587                    // but the access patterns are completely predicatable based on number of passes
588                    // through this loop.
589                    let oram = if next_oram == 0 {
590                        &mut self.oram2
591                    } else {
592                        &mut self.oram1
593                    };
594
595                    oram.access(evicted_from, |block| {
596                        let pairs: &mut [A8Bytes<Sum<KeySize, ValueSize>>] =
597                            block.as_mut_aligned_chunks();
598                        let pair = &mut pairs[evicted_index];
599                        let (key, val): (&mut A8Bytes<KeySize>, &mut A8Bytes<ValueSize>) =
600                            pair.split();
601
602                        // Perform a swap, to undo the previous swap.
603                        // Note that this code is unreachable if side-effects are forbidden
604                        debug_assert!(bool::from(allow_sideeffects_and_eviction));
605                        core::mem::swap(key, &mut evicted_key);
606                        core::mem::swap(val, &mut evicted_val);
607                    })
608                }
609
610                debug_assert!(&evicted_key == query, "After rolling back evictions, we didn't end up with the initially inserted item coming back");
611                // At this point the evicted_key should be the item we initially inserted, if
612                // rollback worked. We can now return OMAP_OVERFLOW because the
613                // semantics of the map didn't change, we simply
614                // failed to insert.
615                return OMAP_OVERFLOW;
616            }
617        }
618
619        // Adjust num_items if we inserted a new item successfully
620        self.num_items += (result_code.ct_eq(&OMAP_NOT_FOUND) & allow_sideeffects_and_eviction)
621            .unwrap_u8() as u64;
622        result_code
623    }
624}
625
626/// Factory implementing OMapCreator for this type, based on any ORAM Creator.
627/// Otherwise it is very difficult to invoke CuckooHashTable::new() generically
628pub struct CuckooHashTableCreator<BlockSize, RngType, OC>
629where
630    BlockSize: ArrayLength,
631    RngType: CryptoRng + Send + Sync + 'static,
632    OC: ORAMCreator<BlockSize, RngType>,
633    OC::Output: ORAM<BlockSize> + Send + Sync + 'static,
634{
635    _block_size: PhantomData<fn() -> BlockSize>,
636    _rng_type: PhantomData<fn() -> RngType>,
637    _oc: PhantomData<fn() -> OC>,
638}
639
640impl<KeySize, ValueSize, BlockSize, RngType, OC> OMapCreator<KeySize, ValueSize, RngType>
641    for CuckooHashTableCreator<BlockSize, RngType, OC>
642where
643    KeySize: ArrayLength + Add<ValueSize> + PartialDiv<U8> + 'static,
644    ValueSize: ArrayLength + PartialDiv<U8> + 'static,
645    BlockSize: ArrayLength + PartialDiv<U8> + 'static,
646    RngType: CryptoRng + Send + Sync + 'static,
647    OC: ORAMCreator<BlockSize, RngType>,
648    OC::Output: ORAM<BlockSize> + Send + Sync + 'static,
649    Sum<KeySize, ValueSize>: ArrayLength + Sub<KeySize, Output = ValueSize> + PartialDiv<U8>,
650{
651    type Output = CuckooHashTable<KeySize, ValueSize, BlockSize, RngType, OC::Output>;
652
653    fn create<M: 'static + FnMut() -> RngType>(
654        size: u64,
655        stash_size: usize,
656        rng_maker: M,
657    ) -> Self::Output {
658        Self::Output::new::<OC, M>(size, stash_size, rng_maker)
659    }
660}
661
662#[cfg(test)]
663mod testing {
664    use super::*;
665    use mc_oblivious_ram::PathORAM4096Z4Creator;
666    use mc_oblivious_traits::{
667        rng_maker, testing, HeapORAMStorageCreator, OMapCreator, OMAP_FOUND, OMAP_NOT_FOUND,
668    };
669    use test_helper::{run_with_several_seeds, RngType};
670    use typenum::{U1024, U8};
671
672    extern crate std;
673
674    const STASH_SIZE: usize = 16;
675
676    type ORAMCreatorZ4 = PathORAM4096Z4Creator<RngType, HeapORAMStorageCreator>;
677    type CuckooCreatorZ4 = CuckooHashTableCreator<U1024, RngType, ORAMCreatorZ4>;
678
679    /// Make a8-bytes that are initialized to a particular byte value
680    /// This makes tests shorter to write
681    fn a8_8<N: ArrayLength>(src: u8) -> A8Bytes<N> {
682        let mut result = A8Bytes::<N>::default();
683        for byte in result.as_mut_slice() {
684            *byte = src;
685        }
686        result
687    }
688
689    #[test]
690    fn sanity_check_omap_z4_4() {
691        run_with_several_seeds(|rng| {
692            // This should be ~1 underlying bucket
693            let mut omap = <CuckooCreatorZ4 as OMapCreator<U8, U8, RngType>>::create(
694                4,
695                STASH_SIZE,
696                rng_maker(rng),
697            );
698
699            let mut temp = A8Bytes::<U8>::default();
700
701            assert_eq!(OMAP_NOT_FOUND, omap.read(&a8_8(1), &mut temp));
702            assert_eq!(
703                OMAP_NOT_FOUND,
704                omap.vartime_write(&a8_8(1), &a8_8(2), 0.into())
705            );
706            assert_eq!(OMAP_FOUND, omap.read(&a8_8(1), &mut temp));
707            assert_eq!(&temp, &a8_8(2));
708            assert_eq!(OMAP_FOUND, omap.vartime_write(&a8_8(1), &a8_8(3), 1.into()));
709            assert_eq!(OMAP_FOUND, omap.read(&a8_8(1), &mut temp));
710            assert_eq!(&temp, &a8_8(3));
711
712            assert_eq!(OMAP_NOT_FOUND, omap.read(&a8_8(2), &mut temp));
713            assert_eq!(
714                &temp,
715                &a8_8(3),
716                "omap.read must not modify the output on not_found"
717            );
718            assert_eq!(
719                OMAP_NOT_FOUND,
720                omap.vartime_write(&a8_8(2), &a8_8(20), 1.into())
721            );
722            assert_eq!(OMAP_FOUND, omap.read(&a8_8(2), &mut temp));
723            assert_eq!(&temp, &a8_8(20));
724            assert_eq!(
725                OMAP_FOUND,
726                omap.vartime_write(&a8_8(2), &a8_8(30), 0.into())
727            );
728            assert_eq!(OMAP_FOUND, omap.read(&a8_8(2), &mut temp));
729            assert_eq!(
730                &temp,
731                &a8_8(20),
732                "omap.write must not modify when overwrite is disallowed"
733            );
734            assert_eq!(
735                OMAP_FOUND,
736                omap.vartime_write(&a8_8(2), &a8_8(30), 1.into())
737            );
738            assert_eq!(OMAP_FOUND, omap.read(&a8_8(2), &mut temp));
739            assert_eq!(&temp, &a8_8(30));
740        })
741    }
742
743    #[test]
744    fn sanity_check_omap_z4_256() {
745        run_with_several_seeds(|rng| {
746            // This should be ~ 4 underlying buckets
747            let mut omap = <CuckooCreatorZ4 as OMapCreator<U8, U8, RngType>>::create(
748                256,
749                STASH_SIZE,
750                rng_maker(rng),
751            );
752
753            let mut temp = A8Bytes::<U8>::default();
754
755            assert_eq!(OMAP_NOT_FOUND, omap.read(&a8_8(1), &mut temp));
756            assert_eq!(
757                OMAP_NOT_FOUND,
758                omap.vartime_write(&a8_8(1), &a8_8(2), 0.into())
759            );
760            assert_eq!(OMAP_FOUND, omap.read(&a8_8(1), &mut temp));
761            assert_eq!(&temp, &a8_8(2));
762            assert_eq!(OMAP_FOUND, omap.vartime_write(&a8_8(1), &a8_8(3), 1.into()));
763            assert_eq!(OMAP_FOUND, omap.read(&a8_8(1), &mut temp));
764            assert_eq!(&temp, &a8_8(3));
765
766            assert_eq!(OMAP_NOT_FOUND, omap.read(&a8_8(2), &mut temp));
767            assert_eq!(
768                &temp,
769                &a8_8(3),
770                "omap.read must not modify the output on not_found"
771            );
772            assert_eq!(
773                OMAP_NOT_FOUND,
774                omap.vartime_write(&a8_8(2), &a8_8(20), 1.into())
775            );
776            assert_eq!(OMAP_FOUND, omap.read(&a8_8(2), &mut temp));
777            assert_eq!(&temp, &a8_8(20));
778            assert_eq!(
779                OMAP_FOUND,
780                omap.vartime_write(&a8_8(2), &a8_8(30), 0.into())
781            );
782            assert_eq!(OMAP_FOUND, omap.read(&a8_8(2), &mut temp));
783            assert_eq!(
784                &temp,
785                &a8_8(20),
786                "omap.write must not modify when overwrite is disallowed"
787            );
788            assert_eq!(
789                OMAP_FOUND,
790                omap.vartime_write(&a8_8(2), &a8_8(30), 1.into())
791            );
792            assert_eq!(OMAP_FOUND, omap.read(&a8_8(2), &mut temp));
793            assert_eq!(&temp, &a8_8(30));
794        })
795    }
796
797    #[test]
798    fn sanity_check_omap_access_api_z4_256() {
799        run_with_several_seeds(|rng| {
800            // This should be ~ 4 underlying buckets
801            let mut omap = <CuckooCreatorZ4 as OMapCreator<U8, U8, RngType>>::create(
802                256,
803                STASH_SIZE,
804                rng_maker(rng),
805            );
806
807            omap.access(&a8_8(1), |code, _buffer| {
808                assert_eq!(code, OMAP_NOT_FOUND);
809            });
810
811            assert_eq!(
812                OMAP_NOT_FOUND,
813                omap.vartime_write(&a8_8(1), &a8_8(2), 0.into())
814            );
815
816            omap.access(&a8_8(1), |code, buffer| {
817                assert_eq!(code, OMAP_FOUND);
818                assert_eq!(buffer, &a8_8(2));
819            });
820
821            assert_eq!(OMAP_FOUND, omap.vartime_write(&a8_8(1), &a8_8(3), 1.into()));
822
823            omap.access(&a8_8(1), |code, buffer| {
824                assert_eq!(code, OMAP_FOUND);
825                assert_eq!(buffer, &a8_8(3));
826            });
827
828            omap.access(&a8_8(2), |code, _buffer| {
829                assert_eq!(code, OMAP_NOT_FOUND);
830            });
831
832            assert_eq!(
833                OMAP_NOT_FOUND,
834                omap.vartime_write(&a8_8(2), &a8_8(20), 1.into())
835            );
836
837            omap.access(&a8_8(2), |code, buffer| {
838                assert_eq!(code, OMAP_FOUND);
839                assert_eq!(buffer, &a8_8(20));
840            });
841
842            assert_eq!(
843                OMAP_FOUND,
844                omap.vartime_write(&a8_8(2), &a8_8(30), 0.into())
845            );
846
847            omap.access_and_remove(&a8_8(2), |code, buffer| {
848                assert_eq!(code, OMAP_FOUND);
849                assert_eq!(
850                    buffer,
851                    &a8_8(20),
852                    "omap.write must not modify when overwrite is disallowed"
853                );
854                Choice::from(0)
855            });
856
857            omap.access_and_remove(&a8_8(2), |code, buffer| {
858                assert_eq!(code, OMAP_FOUND);
859                assert_eq!(
860                    buffer,
861                    &a8_8(20),
862                    "omap.access_and_remove should not delete when delete is false"
863                );
864                Choice::from(0)
865            });
866
867            assert_eq!(
868                OMAP_FOUND,
869                omap.vartime_write(&a8_8(2), &a8_8(30), 1.into())
870            );
871
872            omap.access_and_remove(&a8_8(2), |code, buffer| {
873                assert_eq!(code, OMAP_FOUND);
874                assert_eq!(
875                    buffer,
876                    &a8_8(30),
877                    "omap.write must modify when overwrite is allowed"
878                );
879                Choice::from(1)
880            });
881
882            omap.access(&a8_8(2), |code, buffer| {
883                assert_eq!(
884                    code, OMAP_NOT_FOUND,
885                    "omap.access_and_remove should delete when delete is true"
886                );
887                assert_eq!(
888                    buffer,
889                    &a8_8(0),
890                    "when data is not present, access should give us all zeroes buffer"
891                );
892            });
893        })
894    }
895
896    #[test]
897    fn sanity_check_omap_z4_524288() {
898        run_with_several_seeds(|rng| {
899            // This should be ~ 8192 underlying buckets
900            let mut omap = <CuckooCreatorZ4 as OMapCreator<U8, U8, RngType>>::create(
901                524288,
902                STASH_SIZE,
903                rng_maker(rng),
904            );
905
906            let mut temp = A8Bytes::<U8>::default();
907
908            assert_eq!(OMAP_NOT_FOUND, omap.read(&a8_8(1), &mut temp));
909            assert_eq!(
910                OMAP_NOT_FOUND,
911                omap.vartime_write(&a8_8(1), &a8_8(2), 0.into())
912            );
913            assert_eq!(OMAP_FOUND, omap.read(&a8_8(1), &mut temp));
914            assert_eq!(&temp, &a8_8(2));
915            assert_eq!(OMAP_FOUND, omap.vartime_write(&a8_8(1), &a8_8(3), 1.into()));
916            assert_eq!(OMAP_FOUND, omap.read(&a8_8(1), &mut temp));
917            assert_eq!(&temp, &a8_8(3));
918
919            assert_eq!(OMAP_NOT_FOUND, omap.read(&a8_8(2), &mut temp));
920            assert_eq!(
921                &temp,
922                &a8_8(3),
923                "omap.read must not modify the output on not_found"
924            );
925            assert_eq!(
926                OMAP_NOT_FOUND,
927                omap.vartime_write(&a8_8(2), &a8_8(20), 1.into())
928            );
929            assert_eq!(OMAP_FOUND, omap.read(&a8_8(2), &mut temp));
930            assert_eq!(&temp, &a8_8(20));
931            assert_eq!(
932                OMAP_FOUND,
933                omap.vartime_write(&a8_8(2), &a8_8(30), 0.into())
934            );
935            assert_eq!(OMAP_FOUND, omap.read(&a8_8(2), &mut temp));
936            assert_eq!(
937                &temp,
938                &a8_8(20),
939                "omap.write must not modify when overwrite is disallowed"
940            );
941            assert_eq!(
942                OMAP_FOUND,
943                omap.vartime_write(&a8_8(2), &a8_8(30), 1.into())
944            );
945            assert_eq!(OMAP_FOUND, omap.read(&a8_8(2), &mut temp));
946            assert_eq!(&temp, &a8_8(30));
947        })
948    }
949
950    #[test]
951    fn sanity_check_omap_z4_2097152() {
952        run_with_several_seeds(|rng| {
953            // This should be ~32768 underlying buckets
954            let mut omap = <CuckooCreatorZ4 as OMapCreator<U8, U8, RngType>>::create(
955                2097152,
956                STASH_SIZE,
957                rng_maker(rng),
958            );
959
960            let mut temp = A8Bytes::<U8>::default();
961
962            assert_eq!(OMAP_NOT_FOUND, omap.read(&a8_8(1), &mut temp));
963            assert_eq!(
964                OMAP_NOT_FOUND,
965                omap.vartime_write(&a8_8(1), &a8_8(2), 1.into())
966            );
967            assert_eq!(OMAP_FOUND, omap.read(&a8_8(1), &mut temp));
968            assert_eq!(&temp, &a8_8(2));
969            assert_eq!(OMAP_FOUND, omap.vartime_write(&a8_8(1), &a8_8(3), 1.into()));
970            assert_eq!(OMAP_FOUND, omap.read(&a8_8(1), &mut temp));
971            assert_eq!(&temp, &a8_8(3));
972
973            assert_eq!(OMAP_NOT_FOUND, omap.read(&a8_8(2), &mut temp));
974            assert_eq!(
975                &temp,
976                &a8_8(3),
977                "omap.read must not modify the output on not_found"
978            );
979            assert_eq!(
980                OMAP_NOT_FOUND,
981                omap.vartime_write(&a8_8(2), &a8_8(20), 1.into())
982            );
983            assert_eq!(OMAP_FOUND, omap.read(&a8_8(2), &mut temp));
984            assert_eq!(&temp, &a8_8(20));
985            assert_eq!(
986                OMAP_FOUND,
987                omap.vartime_write(&a8_8(2), &a8_8(30), 0.into())
988            );
989            assert_eq!(OMAP_FOUND, omap.read(&a8_8(2), &mut temp));
990            assert_eq!(
991                &temp,
992                &a8_8(20),
993                "omap.write must not modify when overwrite is disallowed"
994            );
995            assert_eq!(
996                OMAP_FOUND,
997                omap.vartime_write(&a8_8(2), &a8_8(30), 1.into())
998            );
999            assert_eq!(OMAP_FOUND, omap.read(&a8_8(2), &mut temp));
1000            assert_eq!(&temp, &a8_8(30));
1001        })
1002    }
1003
1004    #[test]
1005    #[cfg(not(debug_assertions))]
1006    fn sanity_check_omap_z4_262144() {
1007        run_with_several_seeds(|rng| {
1008            use typenum::{U16, U280};
1009            // This should be ~65538 underlying buckets
1010            let mut omap = <CuckooCreatorZ4 as OMapCreator<U16, U280, RngType>>::create(
1011                262144,
1012                STASH_SIZE,
1013                rng_maker(rng),
1014            );
1015
1016            let mut temp = A8Bytes::<U280>::default();
1017
1018            assert_eq!(OMAP_NOT_FOUND, omap.read(&a8_8(1), &mut temp));
1019            assert_eq!(
1020                OMAP_NOT_FOUND,
1021                omap.vartime_write(&a8_8(1), &a8_8(2), 1.into())
1022            );
1023            assert_eq!(OMAP_FOUND, omap.read(&a8_8(1), &mut temp));
1024            assert_eq!(&temp, &a8_8(2));
1025            assert_eq!(OMAP_FOUND, omap.vartime_write(&a8_8(1), &a8_8(3), 1.into()));
1026            assert_eq!(OMAP_FOUND, omap.read(&a8_8(1), &mut temp));
1027            assert_eq!(&temp, &a8_8(3));
1028
1029            assert_eq!(OMAP_NOT_FOUND, omap.read(&a8_8(2), &mut temp));
1030            assert_eq!(
1031                &temp,
1032                &a8_8(3),
1033                "omap.read must not modify the output on not_found"
1034            );
1035            assert_eq!(
1036                OMAP_NOT_FOUND,
1037                omap.vartime_write(&a8_8(2), &a8_8(20), 1.into())
1038            );
1039            assert_eq!(OMAP_FOUND, omap.read(&a8_8(2), &mut temp));
1040            assert_eq!(&temp, &a8_8(20));
1041            assert_eq!(
1042                OMAP_FOUND,
1043                omap.vartime_write(&a8_8(2), &a8_8(30), 0.into())
1044            );
1045            assert_eq!(OMAP_FOUND, omap.read(&a8_8(2), &mut temp));
1046            assert_eq!(
1047                &temp,
1048                &a8_8(20),
1049                "omap.write must not modify when overwrite is disallowed"
1050            );
1051            assert_eq!(
1052                OMAP_FOUND,
1053                omap.vartime_write(&a8_8(2), &a8_8(30), 1.into())
1054            );
1055            assert_eq!(OMAP_FOUND, omap.read(&a8_8(2), &mut temp));
1056            assert_eq!(&temp, &a8_8(30));
1057        })
1058    }
1059
1060    // Run the exercise omap tests for 200 rounds in a map with 256 items
1061    #[test]
1062    fn exercise_omap_two_choice_path_oram_z4_256() {
1063        run_with_several_seeds(|rng| {
1064            let mut maker = rng_maker(rng);
1065            let mut rng = maker();
1066            // This should be ~4 underlying buckets
1067            let mut omap =
1068                <CuckooCreatorZ4 as OMapCreator<U8, U8, RngType>>::create(256, STASH_SIZE, maker);
1069            testing::exercise_omap(200, &mut omap, &mut rng);
1070        });
1071
1072        run_with_several_seeds(|rng| {
1073            let mut maker = rng_maker(rng);
1074            let mut rng = maker();
1075            let mut omap =
1076                <CuckooCreatorZ4 as OMapCreator<U8, U8, RngType>>::create(256, STASH_SIZE, maker);
1077            testing::exercise_omap_counter_table(200, &mut omap, &mut rng);
1078        });
1079    }
1080
1081    // Run the exercise omap tests for 400 rounds in a map with 3072 items
1082    #[test]
1083    fn exercise_omap_two_choice_path_oram_z4_3072() {
1084        run_with_several_seeds(|rng| {
1085            use typenum::{U16, U280};
1086            let mut maker = rng_maker(rng);
1087            let _ = maker();
1088            let mut rng = maker();
1089            // This should be ~1024 underlying buckets
1090            let mut omap = <CuckooCreatorZ4 as OMapCreator<U16, U280, RngType>>::create(
1091                3072, STASH_SIZE, maker,
1092            );
1093
1094            testing::exercise_omap(400, &mut omap, &mut rng);
1095        });
1096
1097        run_with_several_seeds(|rng| {
1098            use typenum::U16;
1099            let mut maker = rng_maker(rng);
1100            let _ = maker();
1101            let mut rng = maker();
1102            let mut omap =
1103                <CuckooCreatorZ4 as OMapCreator<U16, U8, RngType>>::create(3072, STASH_SIZE, maker);
1104
1105            testing::exercise_omap_counter_table(400, &mut omap, &mut rng);
1106        });
1107    }
1108
1109    // Run the exercise omap tests for 2000 rounds in a map with 65536 items
1110    #[test]
1111    #[cfg(not(debug_assertions))]
1112    fn exercise_omap_two_choice_path_oram_z4_65536() {
1113        run_with_several_seeds(|rng| {
1114            let mut maker = rng_maker(rng);
1115            let mut rng = maker();
1116            // This should be ~1024 underlying buckets
1117            let mut omap =
1118                <CuckooCreatorZ4 as OMapCreator<U8, U8, RngType>>::create(65536, STASH_SIZE, maker);
1119
1120            testing::exercise_omap(2000, &mut omap, &mut rng);
1121        });
1122
1123        run_with_several_seeds(|rng| {
1124            let mut maker = rng_maker(rng);
1125            let mut rng = maker();
1126            let mut omap =
1127                <CuckooCreatorZ4 as OMapCreator<U8, U8, RngType>>::create(65536, STASH_SIZE, maker);
1128
1129            testing::exercise_omap_counter_table(2000, &mut omap, &mut rng);
1130        });
1131    }
1132
1133    // Run the exercise omap tests for 16_000 rounds in a map with 524288 items
1134    #[test]
1135    #[cfg(not(debug_assertions))]
1136    fn exercise_omap_two_choice_path_oram_z4_524288() {
1137        run_with_several_seeds(|rng| {
1138            let mut maker = rng_maker(rng);
1139            let mut rng = maker();
1140            // This should be ~8192 underlying buckets
1141            let mut omap = <CuckooCreatorZ4 as OMapCreator<U8, U8, RngType>>::create(
1142                524288, STASH_SIZE, maker,
1143            );
1144
1145            testing::exercise_omap(16_000, &mut omap, &mut rng);
1146        });
1147
1148        run_with_several_seeds(|rng| {
1149            let mut maker = rng_maker(rng);
1150            let mut rng = maker();
1151            let mut omap = <CuckooCreatorZ4 as OMapCreator<U8, U8, RngType>>::create(
1152                524288, STASH_SIZE, maker,
1153            );
1154
1155            testing::exercise_omap_counter_table(16_000, &mut omap, &mut rng);
1156        });
1157    }
1158
1159    // Run the exercise omap tests for 16_000 rounds in a map with 2097152 items
1160    #[test]
1161    #[cfg(not(debug_assertions))]
1162    fn exercise_omap_two_choice_path_oram_z4_2097152() {
1163        run_with_several_seeds(|rng| {
1164            let mut maker = rng_maker(rng);
1165            let mut rng = maker();
1166            // This should be ~32768 underlying buckets
1167            let mut omap = <CuckooCreatorZ4 as OMapCreator<U8, U8, RngType>>::create(
1168                2097152, STASH_SIZE, maker,
1169            );
1170
1171            testing::exercise_omap(16_000, &mut omap, &mut rng);
1172        });
1173
1174        run_with_several_seeds(|rng| {
1175            let mut maker = rng_maker(rng);
1176            let mut rng = maker();
1177            let mut omap = <CuckooCreatorZ4 as OMapCreator<U8, U8, RngType>>::create(
1178                2097152, STASH_SIZE, maker,
1179            );
1180
1181            testing::exercise_omap_counter_table(16_000, &mut omap, &mut rng);
1182        });
1183    }
1184
1185    // Load the omap to 70% capacity in a map with 16384 items
1186    #[test]
1187    #[cfg(not(debug_assertions))]
1188    fn omap_70_capacity_two_choice_path_oram_z4_16384() {
1189        use typenum::U248;
1190        run_with_several_seeds(|rng| {
1191            // This should be ~4096 underlying buckets
1192            let mut omap = <CuckooCreatorZ4 as OMapCreator<U8, U248, RngType>>::create(
1193                16384,
1194                STASH_SIZE,
1195                rng_maker(rng),
1196            );
1197
1198            let mut temp = A8Bytes::<U8>::default();
1199            let val = A8Bytes::<U248>::default();
1200            for idx in 1u64..(omap.capacity() * 7 / 10) {
1201                temp.copy_from_slice(&idx.to_le_bytes());
1202                let result_code = omap.vartime_write(&temp, &val, 0.into());
1203                assert!(OMAP_OVERFLOW != result_code);
1204                assert!(OMAP_NOT_FOUND == result_code);
1205                assert_eq!(omap.len(), idx);
1206            }
1207        });
1208    }
1209
1210    // Test that the omap has correct roll-back semantics around overflow, for 16384
1211    // items To see ratio at which hashmap overflow begins, run
1212    // cargo test --release -p mc-oblivious-map -- overflow --nocapture
1213    #[test]
1214    #[cfg(not(debug_assertions))]
1215    fn omap_overflow_semantics_two_choice_path_oram_z4_16384() {
1216        use std::println;
1217        use typenum::U248;
1218        run_with_several_seeds(|rng| {
1219            // This shoudl be ~4096 underlying buckets
1220            let mut omap = <CuckooCreatorZ4 as OMapCreator<U8, U248, RngType>>::create(
1221                16384,
1222                STASH_SIZE,
1223                rng_maker(rng),
1224            );
1225
1226            let len = testing::test_omap_overflow(&mut omap);
1227            let cap = omap.capacity();
1228            let fraction = (len as f32) * 100f32 / (cap as f32);
1229            println!("Overflowed at {} / {} = {}%", len, cap, fraction);
1230        })
1231    }
1232
1233    // Test that the omap has correct roll-back semantics around overflow, for 32768
1234    // items To see ratio at which hashmap overflow begins, run
1235    // cargo test --release -p mc-oblivious-map -- overflow --nocapture
1236    #[test]
1237    #[cfg(not(debug_assertions))]
1238    fn omap_overflow_semantics_two_choice_path_oram_z4_32768() {
1239        use std::println;
1240        use typenum::U248;
1241        run_with_several_seeds(|rng| {
1242            // This should be ~8192 underlying buckets
1243            let mut omap = <CuckooCreatorZ4 as OMapCreator<U8, U248, RngType>>::create(
1244                32768,
1245                STASH_SIZE,
1246                rng_maker(rng),
1247            );
1248
1249            let len = testing::test_omap_overflow(&mut omap);
1250            let cap = omap.capacity();
1251            let fraction = (len as f32) * 100f32 / (cap as f32);
1252            println!("Overflowed at {} / {} = {}%", len, cap, fraction);
1253        })
1254    }
1255
1256    // Test that the omap has correct roll-back semantics around overflow, for 65536
1257    // items To see ratio at which hashmap overflow begins, run
1258    // cargo test --release -p mc-oblivious-map -- overflow --nocapture
1259    #[test]
1260    #[cfg(not(debug_assertions))]
1261    fn omap_overflow_semantics_two_choice_path_oram_z4_65536() {
1262        use std::println;
1263        use typenum::U248;
1264        let large_stash_size = STASH_SIZE * 2;
1265        run_with_several_seeds(|rng| {
1266            // This should be ~16384 underlying buckets
1267            let mut omap = <CuckooCreatorZ4 as OMapCreator<U8, U248, RngType>>::create(
1268                65536,
1269                large_stash_size,
1270                rng_maker(rng),
1271            );
1272
1273            let len = testing::test_omap_overflow(&mut omap);
1274            let cap = omap.capacity();
1275            let fraction = (len as f32) * 100f32 / (cap as f32);
1276            println!("Overflowed at {} / {} = {}%", len, cap, fraction);
1277        })
1278    }
1279}