mc_oblivious_traits/
lib.rs

1// Copyright (c) 2018-2023 The MobileCoin Foundation
2
3//! Traits for different pieces of ORAM, from the level of block storage up to
4//! an oblivious map.
5//! These are all defined in terms of fixed-length chunks of bytes and the
6//! A8Bytes object from the aligned-cmov crate.
7//!
8//! There is also a naive implementation of the ORAM storage object for tests.
9
10#![no_std]
11#![deny(unsafe_code)]
12
13use core::fmt::{Debug, Display};
14
15extern crate alloc;
16use alloc::vec::Vec;
17
18// Re-export some traits we depend on, so that downstream can ensure that they
19// have the same version as us.
20pub use aligned_cmov::{
21    cswap, subtle, typenum, A64Bytes, A8Bytes, ArrayLength, CMov, GenericArray,
22};
23pub use rand_core::{CryptoRng, RngCore};
24use subtle::{Choice, ConditionallySelectable, ConstantTimeEq};
25
26mod naive_storage;
27pub use naive_storage::{HeapORAMStorage, HeapORAMStorageCreator};
28
29mod linear_scanning;
30pub use linear_scanning::LinearScanningORAM;
31
32mod creators;
33pub use creators::*;
34
35pub mod testing;
36
37/// Represents trusted block storage holding aligned blocks of memory of a
38/// certain size. This is a building block for ORAM.
39///
40/// This object is required to encrypt / mac the memory if it pushes things out
41/// to untrusted, but it is not required to keep the indices a secret when
42/// accessed. This object is not itself an oblivious data structure.
43///
44/// In tests this can simply be Vec.
45/// In production it is planned to be an object that makes OCalls to untrusted,
46/// and which encrypts and macs the memory blocks that it sends to and from
47/// untrusted. This is analogous to the "Intel memory engine" in SGX.
48///
49/// It is anticipated that "tree-top caching" occurs at this layer, so the
50/// initial portion of the storage is in the enclave and the rest is in
51/// untrusted
52///
53/// TODO: Create an API that allows checking out from two branches
54/// simultaneously.
55#[allow(clippy::len_without_is_empty)]
56pub trait ORAMStorage<BlockSize: ArrayLength<u8>, MetaSize: ArrayLength<u8>> {
57    /// Get the number of blocks represented by this block storage
58    /// This is also the bound of the largest valid index
59    fn len(&self) -> u64;
60
61    /// Checkout all blocks on the branch leading to a particular index in the
62    /// tree, copying them and their metadata into two scratch buffers.
63    ///
64    /// Arguments:
65    /// * index: The index of the leaf, a u64 TreeIndex value.
66    /// * dest: The destination data buffer
67    /// * dest_meta: The destination metadata buffer
68    ///
69    /// Requirements:
70    /// * 0 < index <= len
71    /// * index.height() + 1 == dest.len() == dest_meta.len()
72    /// * It is illegal to checkout while there is an existing checkout.
73    fn checkout(
74        &mut self,
75        index: u64,
76        dest: &mut [A64Bytes<BlockSize>],
77        dest_meta: &mut [A8Bytes<MetaSize>],
78    );
79
80    /// Checkin a number of blocks, copying them and their metadata
81    /// from two scratch buffers.
82    ///
83    /// It is illegal to checkin when there is not an existing checkout.
84    /// It is illegal to checkin different blocks than what was checked out.
85    ///
86    /// Arguments:
87    /// * index: The index of the leaf, a u64 TreeIndex.
88    /// * src: The source data buffer
89    /// * src_meta: The source metadata buffer
90    ///
91    /// Note: src and src_meta are mutable, because it is more efficient to
92    /// encrypt them in place than to copy them and then encrypt.
93    /// These buffers are left in an unspecified but valid state.
94    fn checkin(
95        &mut self,
96        index: u64,
97        src: &mut [A64Bytes<BlockSize>],
98        src_meta: &mut [A8Bytes<MetaSize>],
99    );
100}
101
102/// An Oblivious RAM -- that is, an array like [A8Bytes<ValueSize>; N]
103/// which supports access queries *without memory access patterns revealing
104/// what indices were queried*. (Here, N is a runtime parameter set at
105/// construction time.)
106///
107/// The ValueSize parameter indicates the number of bytes in a stored value.
108///
109/// The key-type here is always u64 even if it "could" be smaller.
110/// We think that if keys are actually stored as u32 or u16 in some of the
111/// recursive position maps, that conversion can happen at a different layer of
112/// the system.
113///
114/// TODO: Should there be, perhaps, a separate trait for "resizable" ORAMs?
115/// We don't have a good way for the OMAP to take advantage of that right now.
116#[allow(clippy::len_without_is_empty)]
117#[allow(clippy::upper_case_acronyms)]
118pub trait ORAM<ValueSize: ArrayLength<u8>> {
119    /// Get the number of values logically in the ORAM.
120    /// This is also one more than the largest index that can be legally
121    /// accessed.
122    fn len(&self) -> u64;
123
124    /// Get the number of values in the ORAM's stash for diagnostics. In prod,
125    /// this number should be viewed as secret and not revealed.
126    fn stash_size(&self) -> usize;
127
128    /// Access the ORAM at a position, calling a lambda with the recovered
129    /// value, and returning the result of the lambda.
130    /// This cannot fail, but will panic if index is out of bounds.
131    ///
132    /// This is the lowest-level API that we offer for getting data from the
133    /// ORAM.
134    fn access<T, F: FnOnce(&mut A64Bytes<ValueSize>) -> T>(&mut self, index: u64, func: F) -> T;
135
136    /// High-level helper -- when you only need to read and don't need to write
137    /// a new value, this is simpler than using `access`.
138    /// In most ORAM there will not be a significantly faster implementation of
139    /// this.
140    #[inline]
141    fn read(&mut self, index: u64) -> A64Bytes<ValueSize> {
142        self.access(index, |val| val.clone())
143    }
144
145    /// High-level helper -- when you need to write a value and want the
146    /// previous value, but you don't need to see the previous value when
147    /// deciding what to write, this is simpler than using `access`.
148    /// In most ORAM there will not be a significantly faster implementation of
149    /// this.
150    #[inline]
151    fn write(&mut self, index: u64, new_val: &A64Bytes<ValueSize>) -> A64Bytes<ValueSize> {
152        self.access(index, |val| {
153            let retval = val.clone();
154            *val = new_val.clone();
155            retval
156        })
157    }
158}
159
160/// Trait that helps to debug ORAM.
161/// This should only be used in tests.
162///
163/// This should never be called in production. IMO the best practice is that
164/// implementations of this trait should be gated by `#[cfg(test)]`, or perhaps
165/// `#[cfg(debug_assertions)]`.
166pub trait ORAMDebug<ValueSize: ArrayLength<u8>> {
167    /// Systematically check the data structure invariants, asserting that they
168    /// hold. Also, produce an array representation of the logical state of
169    /// the ORAM.
170    ///
171    /// This should not change the ORAM.
172    ///
173    /// This is returned so that recursive path ORAM can implement
174    /// check_invariants by first asking recursive children to check their
175    /// invariants.
176    fn check_invariants(&self) -> Vec<A64Bytes<ValueSize>>;
177}
178
179/// PositionMap trait conceptually is an array of TreeIndex.
180/// Each value in the map corresponds to a leaf in the complete binary tree,
181/// at some common height.
182///
183/// PositionMap trait must be object-safe so that dyn PositionMap works.
184/// It also only needs to work with integer types, and padding up to u64 is
185/// fine. Therefore we make a new trait which is reduced and only exposes the
186/// things that PathORAM needs from the position map.
187///
188/// TODO: API for resizing it? Changing height?
189#[allow(clippy::len_without_is_empty)]
190pub trait PositionMap {
191    /// The number of keys in the map. The valid keys are in the range 0..len.
192    fn len(&self) -> u64;
193    /// Write a new value to a particular key.
194    /// The new value should be a random nonce from a CSPRNG.
195    /// Returns the old value.
196    /// It is illegal to write to a key that is out of bounds.
197    fn write(&mut self, key: &u64, new_val: &u64) -> u64;
198}
199
200/// Trait for an oblivious hash map, where READING and ACCESSING EXISTING
201/// ENTRIES have a strong oblivious property.
202///
203/// This is different from an ORAM in that it is like a hashmap rather than an
204/// array, and the keys can be byte chunks of any length, and need not be
205/// consecutive.
206///
207/// Oblivious here means: Timings, code and data access patterns are independent
208/// of the value of inputs when calling a function with the oblivious property.
209/// Generally this means the query (and the value).
210///
211/// - Read, access, and remove are strongly oblivious and have constant
212///   execution time. With the exception that, the all zeroes key may be invalid
213///   and they may early return if you use it.
214/// - Write is not strongly oblivious and may take a different amount of time
215///   for different inputs. It may also fail, returning OMAP_OVERFLOW if an item
216///   could not be added.
217/// - Write is the only way to add new values to the map.
218///
219/// In many use-cases, the writes are taking place at keys that are known to the
220/// SGX adversary (node operator).
221/// It would be okay to use e.g. a cuckoo-hashing type strategy, where different
222/// keys may take different amounts of time to insert, as long as the reads are
223/// strongly oblivious, because those are what what correspond to user queries.
224/// As long as the read timing and access patterns are independent of
225/// the write timings and access patterns, it meets the requirement. This is
226/// similar to the situation with "Read-Only ORAM" in the literature.
227///
228/// The access_and_insert function supports additional use cases where new
229/// inserts at secret locations, which might be new or old values, must be
230/// performed obliviously.
231///
232/// The API is designed to make it as easy as possible to write constant-time
233/// code that uses the map. This means, not using Option or Result, using
234/// status_code that implements CMov trait, conditionally writing to output
235/// parameters which can be on the stack, which can avoid copies, and avoiding a
236/// "checkout / checkin" API which while more powerful, is more complicated than
237/// a callback-based API.
238///
239/// We are not trying to mimic the rust hashmap API or the entry API because
240/// those are based on enums, option, etc. and can't be used when constant-time,
241/// branchless code is a requirement. For planned use-cases, read, write and
242/// access are sufficient.
243///
244/// TODO: Should there be an explicit resize API? How should that work?
245pub trait ObliviousHashMap<KeySize: ArrayLength<u8>, ValueSize: ArrayLength<u8>> {
246    /// Get the number of items in the map.
247    fn len(&self) -> u64;
248
249    /// Get the capacity of the map.
250    /// TODO: What should this be for the hashmap?
251    /// At the moment this is number of buckets * bucket size.
252    /// But this is not an achievable value of len in practice.
253    fn capacity(&self) -> u64;
254
255    /// Is the map empty
256    #[inline]
257    fn is_empty(&self) -> bool {
258        self.len() == 0
259    }
260
261    /// Read from the map at some position, without logically modifying it.
262    ///
263    /// Note: This is strongly oblivious, regardless of whether the value was
264    /// found, with the exception that we may early return if the all zeroes
265    /// key is encountered.
266    ///
267    /// Arguments:
268    /// - key: The data to query from the map
269    /// - output: A mutable location to write the found value.
270    ///
271    /// Returns a status code:
272    /// - OMAP_FOUND: The value was found in the map, and output was
273    ///   overwritten.
274    /// - OMAP_NOT_FOUND: The value was not found in the map, and the output
275    ///   buffer was not modified.
276    /// - OMAP_INVALID_KEY: The key was rejected. The map is permitted to reject
277    ///   an all-zeroes key.
278    fn read(&mut self, key: &A8Bytes<KeySize>, output: &mut A8Bytes<ValueSize>) -> u32;
279
280    /// Access from the map at some position, and forward the value to a
281    /// callback, which may modify it.
282    ///
283    /// Note: This is strongly oblivious regardless of whether the value was
284    /// found, with the exception that we may early return if the all zeroes
285    /// key is encountered. The callback must also be oblivious for this
286    /// property to hold.
287    ///
288    /// Arguments:
289    /// - key: The data to query from the map
290    /// - callback: A function to call after the entry has been retrieved on the
291    ///   stack.
292    ///
293    /// Callback:
294    /// The callback is passed a status code and a mutable value.
295    /// The status code indicates if this value already existed in the map
296    /// (OMAP_FOUND) or if there was no match (OMAP_NOT_FOUND), or if the
297    /// query was illegal (OMAP_INVALID_KEY).
298    ///
299    /// This call cannot change the number of items in the map.
300    /// If there was no match then it doesn't matter what the callback does,
301    /// nothing will be added. If there is a match then after the callback
302    /// runs, the contents of the mutable buffer will be
303    /// the value associated to this key.
304    #[inline]
305    fn access<F: FnOnce(u32, &mut A8Bytes<ValueSize>)>(
306        &mut self,
307        key: &A8Bytes<KeySize>,
308        callback: F,
309    ) {
310        // By default we call to access-and-remove and always elect not to remove.
311        // The perf consequence of this is usually very small, so this is a reasonable
312        // default implementation.
313        self.access_and_remove(key, |code, val| -> Choice {
314            callback(code, val);
315            Choice::from(0)
316        })
317    }
318
319    /// Remove an entry from the map, by its key.
320    ///
321    /// Arguments:
322    /// - key: The key value to remove from the map
323    ///
324    /// Returns a status code:
325    /// - OMAP_FOUND: Found an existing value which was removed.
326    /// - OMAP_NOT_FOUND: Did not find an existing value and nothing was
327    ///   removed.
328    /// - OMAP_INVALID_KEY: The key was rejected. The map is permitted to reject
329    ///   an all-zeroes key.
330    fn remove(&mut self, key: &A8Bytes<KeySize>) -> u32;
331
332    /// Access from the map at some position, and forward the value to a
333    /// callback, which may modify it. The callback returns a boolean, which,
334    /// if true, removes the entry.
335    ///
336    /// Note: This is strongly oblivious regardless of whether the value was
337    /// found, with the exception that we may early return if the all zeroes
338    /// key is encountered. The callback must also be oblivious for this
339    /// property to hold.
340    /// We are somewhat oblivious about whether a delete occurred, see note
341    /// below.
342    ///
343    /// Arguments:
344    /// - key: The data to query from the map
345    /// - callback: A function to call after the entry has been retrieved on the
346    ///   stack.
347    ///
348    /// Callback:
349    /// The callback is passed a status code and a mutable value.
350    /// The status code indicates if this value already existed in the map
351    /// (OMAP_FOUND) or if there was no match (OMAP_NOT_FOUND), or if the
352    /// query was illegal (OMAP_INVALID_KEY).
353    /// The return value indicates whether we should also remove the item from
354    /// the map -- if true, it is removed.
355    ///
356    /// If there was no match then it doesn't matter what the callback does,
357    /// nothing will be added or removed. If there is a match then after the
358    /// callback runs, the contents of the mutable buffer will be
359    /// the value associated to this key.
360    ///
361    /// Note: Over time, there will be some information leakage if many things
362    /// are removed or not removed by this call, because the number of items in
363    /// the map is changing, and so if there are subsequent calls to insert
364    /// things to the map, they will perform cuckoo steps with greater or
365    /// lesser probability, depending roughly on how full the map is.
366    ///
367    /// If this is unacceptable, another strategy might involve adding junk
368    /// values or not (obliviously) whenever items are deleted. (This would
369    /// be another possible call, because this has significant perf
370    /// consequences as well.)
371    fn access_and_remove<F: FnOnce(u32, &mut A8Bytes<ValueSize>) -> Choice>(
372        &mut self,
373        key: &A8Bytes<KeySize>,
374        callback: F,
375    );
376
377    /// Write to the map at a position.
378    ///
379    /// Note: This call IS NOT strongly constant-time, it may take different
380    /// amounts of time depending on if the item is or isn't in the map already,
381    /// and how full the map is.
382    ///
383    /// This allows to implement the map using e.g. a cuckoo hashing strategy.
384    ///
385    /// Arguments:
386    /// - key:             The key to store data at
387    /// - value:           The data to store
388    /// - allow_overwrite: Whether to allow overwriting an existing entry. If
389    ///   false, then when OMAP_FOUND occurs, the map will not be modified
390    ///
391    /// Returns a status code indicating if the key/value pair was
392    /// - not found (OMAP_NOT_FOUND), and so was added successfully,
393    /// - found (OMAP_FOUND), and so was overwritten,
394    /// - a table overflow occurred (OMAP_OVERFLOW) and so the operation failed,
395    /// - the key was rejected (OMAP_INVALID_KEY) and so the operation failed.
396    ///   the map implementation is allowed to reject the all zeroes key.
397    ///
398    /// Security propreties:
399    /// - This call is completely oblivious with respect to value and
400    ///   allow_overwrite flag
401    /// - The call is not completely oblivious with respect to the key
402    /// - When the keys are both in the map, the call is constant-time.
403    /// - When both keys are not in the map, the call takes a variable amount of
404    ///   time, but the distribution of access patterns is indistinguishable for
405    ///   both keys. This requires the assumption that the hash function used to
406    ///   construct the table models a PRF.
407    ///
408    /// This operation does not return the old value.
409    #[inline]
410    fn vartime_write(
411        &mut self,
412        key: &A8Bytes<KeySize>,
413        value: &A8Bytes<ValueSize>,
414        allow_overwrite: Choice,
415    ) -> u32 {
416        self.vartime_write_extended(key, value, allow_overwrite, Choice::from(1))
417    }
418
419    /// Write to the map at a position, OR, perform the same access patterns of
420    /// this without actually writing to the map.
421    ///
422    /// This has the same semantics as `vartime_write`, except, it takes an
423    /// extra flag `allow_sideeffects_and_eviction` which, if false,
424    /// prevents all side-effects to the map, and prevents the eviction
425    /// procedure from happening, while still performing the same memory
426    /// access patterns as a "simple" write to an existing entry of the map.
427    ///
428    /// Arguments:
429    /// - key:             The key to store data at
430    /// - value:           The data to store
431    /// - allow_overwrite: Whether to allow overwriting an existing entry. If
432    ///   false, then when OMAP_FOUND occurs, the map will not be modified
433    /// - allow_sideffects_and_eviction: If false, then no side-effects are
434    ///   performed for the map, and access patterns are the same as if we
435    ///   performed a write at a key that is already in the map. The return code
436    ///   is the same as if we had performed a read operation at this key.
437    ///
438    /// Returns a status code indicating if the key/value pair was
439    /// - not found (OMAP_NOT_FOUND), and so was added successfully (if
440    ///   condition was true)
441    /// - found (OMAP_FOUND), and so was overwritten (if condition and
442    ///   allow_overwrite were true)
443    /// - a table overflow occurred (OMAP_OVERFLOW) and so the operation failed,
444    /// - the key was rejected (OMAP_INVALID_KEY) and so the operation failed.
445    ///   the map implementation is allowed to reject the all zeroes key.
446    fn vartime_write_extended(
447        &mut self,
448        key: &A8Bytes<KeySize>,
449        value: &A8Bytes<ValueSize>,
450        allow_overwrite: Choice,
451        allow_sideeffects_and_eviction: Choice,
452    ) -> u32;
453
454    /// Access the map at a position, inserting the item if it doesn't exist,
455    /// AND obliviously inserting a new item with a default value and random key
456    /// if the targetted item DOES exist.
457    ///
458    /// This complex call allows to both modify the map, and insert
459    /// new items into the map, on the fly, obliviously.
460    ///
461    /// This call is rather niche and you should only use it if you need it
462    /// because it will be slower than a call to access or read, or
463    /// vartime_write. It is implemented here as a trait function in terms
464    /// of `read` and `vartime_write_extended`, but a specific implementation
465    /// may be able to do it faster.
466    ///
467    /// This call ALWAYS inserts a new item into the map, and so if used
468    /// repeatedly, WILL eventually cause the map to overflow.
469    ///
470    /// This call cannot remove an item from the map after accessing it.
471    ///
472    /// The oblivious property here requires that the chance that a random key
473    /// already exists in the map in negligibly small. This assumption is
474    /// justified if the key size is e.g. 32 bytes. If the key size is only
475    /// 8 bytes, then you will have less than a 64-bit security level for
476    /// the oblivious property, as the map gets more and more full. It is
477    /// questionable to use it for maps with a key size of less than 16
478    /// bytes.
479    ///
480    /// Arguments:
481    /// - key: The key to store data at
482    /// - default_value: The default value to insert into the map.
483    /// - rng: An rng for the operation
484    /// - callback:  A function to call after the value has been retrieved on
485    ///   the stack.
486    ///
487    /// Callback:
488    /// The callback is passed a status code indicating if the key/value pair
489    /// was
490    /// - not found (OMAP_NOT_FOUND), and so was added successfully,
491    /// - found (OMAP_FOUND), and so was overwritten,
492    /// And, it is passed the mutable buffer on the stack.
493    /// This buffer either contains the value associated with this key in the
494    /// map, or contains default_value. This buffer will be stored to the
495    /// map after the callback returns, unless the key was invalid or
496    /// overflow occurred.
497    ///
498    /// Returns:
499    /// A status code for the operation as whole:
500    /// - not found (OMAP_NOT_FOUND), the key was not found in the map, and so
501    ///   was added successfully
502    /// - found (OMAP_FOUND), the key was found in the map, and was overwritten
503    ///   successfully
504    /// - a table overflow occurred (OMAP_OVERFLOW), either when inserting the
505    ///   new item OR the random item. In this case the table is left in a
506    ///   consistent state but the write specified by the callback may or may
507    ///   not have been rolled back, and the random write may or may not have
508    ///   been rolled back.
509    /// - the key passed was invalid (OMAP_INVALID_KEY)
510    ///
511    /// Security propreties:
512    /// - This call fails fast, returning early if OMAP_INVALID_KEY or
513    ///   OMAP_OVERFLOW occurs.
514    /// - OMAP_OVERFLOW is equally likely to occur no matter whether the key was
515    ///   already in the map or not. (This requires the assumption that the hash
516    ///   function used is a PRF, like SipHash.)
517    /// - This call is completely oblivious with respect to default_value
518    ///   parameter.
519    /// - This call is completely data-oblivious when considering two keys that
520    ///   are both in the map, or are both not in the map.
521    /// - This call is data-oblivious when considering two keys, one of which is in the map and
522    ///   one of which is not, with an error parameter of self.len() / 2^{KeySize}
523    ///   That is, the adversary can get at most this much advantage in distinguishing the two
524    ///   scenarios based on access patterns. This analysis requires the assumption that
525    ///   the hash function used in the map is a PRF. (https://en.wikipedia.org/wiki/SipHash).
526    fn access_and_insert<F: FnOnce(u32, &mut A8Bytes<ValueSize>), R: RngCore + CryptoRng>(
527        &mut self,
528        key: &A8Bytes<KeySize>,
529        default_value: &A8Bytes<ValueSize>,
530        rng: &mut R,
531        callback: F,
532    ) -> u32 {
533        let mut buffer = default_value.clone();
534        let read_code = self.read(key, &mut buffer);
535        if read_code == OMAP_INVALID_KEY {
536            return read_code;
537        }
538        debug_assert!(
539            read_code == OMAP_FOUND || read_code == OMAP_NOT_FOUND,
540            "unexpected status code value: {}",
541            read_code
542        );
543
544        callback(read_code, &mut buffer);
545
546        // Prepare two writes which will be performed in a random order
547        // The first write writes back the buffer evaluated by callback.
548        let mut first_write_key = key.clone();
549        let mut first_write_val = buffer;
550        let mut first_write_allow_overwrite = Choice::from(1);
551        let mut first_write_allow_sideeffects_and_eviction = Choice::from(1);
552        // The second write writes default_value to a random place.
553        // Initialize second_write_key to random, and resample if we get invalid key.
554        let mut second_write_key = A8Bytes::<KeySize>::default();
555        while bool::from(second_write_key.ct_eq(&A8Bytes::<KeySize>::default())) {
556            rng.fill_bytes(second_write_key.as_mut());
557        }
558        let mut second_write_val = default_value.clone();
559        // Overwrite is not allowed to avoid corrupting the map by writing over
560        // previously existing entries, or the first write.
561        let mut second_write_allow_overwrite = Choice::from(0);
562        // This only actually occurs if the read had OMAP_FOUND.
563        // If the read had NOT_FOUND, then the first write has a chance to do eviction
564        // procedure, and this write should look like a write to a FOUND
565        // location. If the first had FOUND, then this write is with high
566        // probability a write to a new location, and so has the same
567        // characteristics around eviction procedure.
568        let mut second_write_allow_sideeffects_and_eviction = read_code.ct_eq(&OMAP_FOUND);
569
570        // Swap the order of the two writes with equal probability, in constant-time
571        // conditional_swap comes from subtle crate, cswap comes from aligned-cmov
572        let swap = (rng.next_u32() & 1).ct_eq(&0);
573        cswap(swap, &mut first_write_key, &mut second_write_key);
574        cswap(swap, &mut first_write_val, &mut second_write_val);
575        ConditionallySelectable::conditional_swap(
576            &mut first_write_allow_overwrite,
577            &mut second_write_allow_overwrite,
578            swap,
579        );
580        ConditionallySelectable::conditional_swap(
581            &mut first_write_allow_sideeffects_and_eviction,
582            &mut second_write_allow_sideeffects_and_eviction,
583            swap,
584        );
585
586        // Perform both writes
587        let first_write_code = self.vartime_write_extended(
588            &first_write_key,
589            &first_write_val,
590            first_write_allow_overwrite,
591            first_write_allow_sideeffects_and_eviction,
592        );
593        debug_assert!(
594            first_write_code != OMAP_INVALID_KEY,
595            "unexpected status code value: {}",
596            first_write_code
597        );
598        let second_write_code = self.vartime_write_extended(
599            &second_write_key,
600            &second_write_val,
601            second_write_allow_overwrite,
602            second_write_allow_sideeffects_and_eviction,
603        );
604        debug_assert!(
605            second_write_code != OMAP_INVALID_KEY,
606            "unexpected status code value: {}",
607            second_write_code
608        );
609
610        // We return read_code unless one of the two operations overflowed
611        let mut result = read_code;
612        result.cmov(first_write_code.ct_eq(&OMAP_OVERFLOW), &OMAP_OVERFLOW);
613        result.cmov(second_write_code.ct_eq(&OMAP_OVERFLOW), &OMAP_OVERFLOW);
614        result
615    }
616}
617
618// Status codes associated to ObliviousHashMap trait.
619//
620// These are represented as u32 and not a rust enum, because there is little
621// point to using a rust enum here.
622//
623// The main ergonomic benefit of an enum is that you can use match expressions.
624// However, when you have constant-time requirements, you cannot use match
625// expressions because they leak. Since u32 implements our CMov trait it is
626// simpler just to use that directly.
627//
628// Most of the time you will not want to use match with these codes anyways, it
629// looks more like foo.cmov(status == OMAP_FOUND, &bar);
630
631/// Status code for the case that an OMAP operation found and returned a value.
632pub const OMAP_FOUND: u32 = 0;
633/// Status code for the case that an OMAP operation did not find a value that
634/// was searched for.
635pub const OMAP_NOT_FOUND: u32 = 1;
636/// Status code for the case that an OMAP wanted to add a new value but could
637/// not because the hash table overflowed, and so the operation failed.
638pub const OMAP_OVERFLOW: u32 = 2;
639/// Status code for the case that the OMAP algorithm rejected the key. The all
640/// zeroes key may be invalid for instance.
641pub const OMAP_INVALID_KEY: u32 = 3;
642
643/// Utility function for logs base 2 rounded up, implemented as const fn
644#[inline]
645pub const fn log2_ceil(arg: u64) -> u32 {
646    if arg == 0 {
647        return 0;
648    }
649    (!0u64).count_ones() - (arg - 1).leading_zeros()
650}
651
652#[cfg(test)]
653mod test {
654    use super::*;
655    // Sanity check the log2_ceil function
656    #[test]
657    fn test_log2_ceil() {
658        assert_eq!(0, log2_ceil(0));
659        assert_eq!(0, log2_ceil(1));
660        assert_eq!(1, log2_ceil(2));
661        assert_eq!(2, log2_ceil(3));
662        assert_eq!(2, log2_ceil(4));
663        assert_eq!(3, log2_ceil(5));
664        assert_eq!(3, log2_ceil(8));
665        assert_eq!(4, log2_ceil(9));
666        assert_eq!(4, log2_ceil(16));
667        assert_eq!(5, log2_ceil(17));
668    }
669}