From b5ec9da926f8f8ea1064922a98f85a8e3a90a2be Mon Sep 17 00:00:00 2001 From: Sayan Nandan Date: Wed, 18 Jan 2023 04:37:58 -0800 Subject: [PATCH] Add ord index impls and fix UB Added all basic impls for the ord index, also fixed dealloc on a nullptr in drop impl --- server/src/engine/core/def.rs | 7 + server/src/engine/core/idx.rs | 383 ++++++++++++++++++++++++++++++++-- 2 files changed, 374 insertions(+), 16 deletions(-) diff --git a/server/src/engine/core/def.rs b/server/src/engine/core/def.rs index 8b6ca346..f57639f0 100644 --- a/server/src/engine/core/def.rs +++ b/server/src/engine/core/def.rs @@ -45,6 +45,9 @@ impl AsKey for T { } } +pub trait AsKeyRef: Hash + Eq {} +impl AsKeyRef for T {} + /// Any type implementing this trait can be used as a value inside memory engine structures pub trait AsValue: Clone { /// Read the value @@ -93,6 +96,8 @@ where fn mt_init() -> Self; /// Initialize a pre-loaded instance of the MTIndex fn mt_init_with(s: Self) -> Self; + /// Clears all the entries in the MTIndex + fn mt_clear(&self); // write /// Returns true if the entry was inserted successfully; returns false if the uniqueness constraint is /// violated @@ -164,6 +169,8 @@ pub trait STIndex { fn st_init() -> Self; /// Initialize a pre-loaded instance of the MTIndex fn st_init_with(s: Self) -> Self; + /// Clears all the entries in the STIndex + fn st_clear(&mut self); // write /// Returns true if the entry was inserted successfully; returns false if the uniqueness constraint is /// violated diff --git a/server/src/engine/core/idx.rs b/server/src/engine/core/idx.rs index 5ca05a27..7605060a 100644 --- a/server/src/engine/core/idx.rs +++ b/server/src/engine/core/idx.rs @@ -24,12 +24,17 @@ * */ -use super::def::{AsKey, AsValue}; +use super::def::{AsKey, AsKeyRef, AsValue}; use std::{ alloc::{alloc as std_alloc, dealloc as std_dealloc, Layout}, borrow::Borrow, - collections::{hash_map::RandomState, HashMap as StdMap}, + collections::{ + hash_map::{Iter, Keys as StdMapIterKey, RandomState, Values as StdMapIterVal}, + HashMap as StdMap, + }, + fmt::{self, Debug}, hash::{BuildHasher, Hash, Hasher}, + iter::FusedIterator, mem, ptr::{self, NonNull}, }; @@ -47,7 +52,6 @@ pub type IndexSTSeq = IndexSTSeqDll; -- Sayan (@ohsayan) // Jan. 16 '23 */ -#[derive(Debug)] #[repr(transparent)] /// # WARNING: Segfault/UAF alert /// @@ -55,7 +59,14 @@ pub type IndexSTSeq = IndexSTSeqDll; /// you're unsure about it's validity. For example, if you simply `==` this or attempt to use it an a hashmap, /// you can segfault. IFF, the ptr is valid will it not segfault struct IndexSTSeqDllKeyptr { - p: *mut K, + p: *const K, +} + +impl IndexSTSeqDllKeyptr { + #[inline(always)] + fn new(r: &K) -> Self { + Self { p: r as *const _ } + } } impl Hash for IndexSTSeqDllKeyptr { @@ -90,7 +101,7 @@ impl PartialEq for IndexSTSeqDllKeyptr { impl Eq for IndexSTSeqDllKeyptr {} // stupid type for trait impl conflict riddance -#[derive(Debug, Hash, PartialEq)] +#[derive(Debug, Hash, PartialEq, Eq)] #[repr(transparent)] struct IndexSTSeqDllQref(Q); @@ -140,7 +151,7 @@ impl IndexSTSeqDllNode { unsafe { // UNSAFE(@ohsayan): aight shut up, it's a malloc let ptr = std_alloc(Self::LAYOUT) as *mut Self; - assert!(ptr.is_null(), "damn the allocator failed"); + assert!(!ptr.is_null(), "damn the allocator failed"); ptr } } @@ -169,7 +180,7 @@ impl IndexSTSeqDllNode { Self::_alloc::(Self::new(k, v, p, n)) } #[inline(always)] - unsafe fn dealloc(slf: *mut Self) { + unsafe fn _drop(slf: *mut Self) { let _ = Box::from_raw(slf); } #[inline(always)] @@ -189,6 +200,13 @@ impl IndexSTSeqDllNode { (*from).n = to; (*(*to).n).p = to; } + #[inline(always)] + fn alloc_box(node: IndexSTSeqDllNode) -> NonNull> { + unsafe { + // UNSAFE(@ohsayan): Safe because of box alloc + NonNull::new_unchecked(Box::into_raw(Box::new(node))) + } + } } type IndexSTSeqDllNodePtr = NonNull>; @@ -238,11 +256,14 @@ impl IndexSTSeqDll { impl IndexSTSeqDll { #[inline(always)] fn ensure_sentinel(&mut self) { - let ptr = IndexSTSeqDllNode::_alloc_with_garbage(); - unsafe { - self.h = ptr; - (*ptr).p = ptr; - (*ptr).n = ptr; + if self.h.is_null() { + let ptr = IndexSTSeqDllNode::_alloc_with_garbage(); + unsafe { + // UNSAFE(@ohsayan): Fresh alloc + self.h = ptr; + (*ptr).p = ptr; + (*ptr).n = ptr; + } } } #[inline(always)] @@ -250,25 +271,48 @@ impl IndexSTSeqDll { /// /// Head must not be null unsafe fn drop_nodes_full(&mut self) { - let mut c = self.h; - while !c.is_null() { + // don't drop sentinenl + let mut c = (*self.h).n; + while c != self.h { let nx = (*c).n; - IndexSTSeqDllNode::dealloc(c); + IndexSTSeqDllNode::_drop(c); c = nx; } } #[inline(always)] fn vacuum_free(&mut self) { unsafe { + // UNSAFE(@ohsayan): All nullck let mut c = self.f; while !c.is_null() { let nx = (*c).n; - IndexSTSeqDllNode::dealloc_headless(nx); + IndexSTSeqDllNode::dealloc_headless(c); c = nx; } } self.f = ptr::null_mut(); } + #[inline(always)] + fn recycle_or_alloc(&mut self, node: IndexSTSeqDllNode) -> IndexSTSeqDllNodePtr { + if self.f.is_null() { + IndexSTSeqDllNode::alloc_box(node) + } else { + unsafe { + // UNSAFE(@ohsayan): Safe because we already did a nullptr check + let f = self.f; + self.f = (*self.f).n; + ptr::write(f, node); + IndexSTSeqDllNodePtr::new_unchecked(f) + } + } + } + #[inline(always)] + /// NOTE: `&mut Self` for aliasing + /// ## Safety + /// Ensure head is non null + unsafe fn link(&mut self, node: IndexSTSeqDllNodePtr) { + IndexSTSeqDllNode::link(self.h, node.as_ptr()) + } } impl IndexSTSeqDll { @@ -279,3 +323,310 @@ impl IndexSTSeqDll { self.vacuum_free(); } } + +impl IndexSTSeqDll { + const GET_REFRESH: bool = true; + const GET_BYPASS: bool = false; + #[inline(always)] + fn _insert(&mut self, k: K, v: V) -> bool { + if self.m.contains_key(&IndexSTSeqDllKeyptr::new(&k)) { + return false; + } + self.__insert(k, v) + } + fn _get(&self, k: &Q) -> Option<&V> + where + K: Borrow, + Q: AsKeyRef, + { + self.m + .get(unsafe { + // UNSAFE(@ohsayan): Ref with correct bounds + IndexSTSeqDllQref::from_ref(k) + }) + .map(|e| unsafe { + // UNSAFE(@ohsayan): ref is non-null and ensures aliasing reqs + &(e.as_ref()).read_value().v + }) + } + #[inline(always)] + fn _update(&mut self, k: &Q, v: V) -> Option + where + K: Borrow, + Q: AsKeyRef, + { + match self.m.get(unsafe { + // UNSAFE(@ohsayan): Just got a ref with the right bounds + IndexSTSeqDllQref::from_ref(k) + }) { + Some(e) => unsafe { + // UNSAFE(@ohsayan): Impl guarantees that entry presence == nullck head + self.__update(*e, v) + }, + None => return None, + } + } + #[inline(always)] + fn _upsert(&mut self, k: K, v: V) -> Option { + match self.m.get(&IndexSTSeqDllKeyptr::new(&k)) { + Some(e) => unsafe { + // UNSAFE(@ohsayan): Impl guarantees that entry presence == nullck head + self.__update(*e, v) + }, + None => { + let _ = self.__insert(k, v); + None + } + } + } + #[inline(always)] + fn _remove(&mut self, k: &Q) -> Option + where + K: Borrow, + Q: AsKeyRef + ?Sized, + { + self.m + .remove(unsafe { + // UNSAFE(@ohsayan): good trait bounds and type + IndexSTSeqDllQref::from_ref(k) + }) + .map(|n| unsafe { + let n = n.as_ptr(); + // UNSAFE(@ohsayan): Correct init and aligned to K + drop(ptr::read(&(*n).k)); + // UNSAFE(@ohsayan): Correct init and aligned to V + let v = ptr::read(&(*n).v); + // UNSAFE(@ohsayan): non-null guaranteed by as_ptr + IndexSTSeqDllNode::unlink(n); + (*n).n = self.f; + self.f = n; + v + }) + } + #[inline(always)] + fn __insert(&mut self, k: K, v: V) -> bool { + self.ensure_sentinel(); + let node = self.recycle_or_alloc(IndexSTSeqDllNode::new_null(k, v)); + let kptr = unsafe { + // UNSAFE(@ohsayan): All g, we allocated it rn + IndexSTSeqDllKeyptr::new(&node.as_ref().k) + }; + let _ = self.m.insert(kptr, node); + unsafe { + // UNSAFE(@ohsayan): sentinel check done + self.link(node); + } + true + } + #[inline(always)] + /// ## Safety + /// + /// Has sentinel + unsafe fn __update(&mut self, e: NonNull>, v: V) -> Option { + let old = unsafe { + // UNSAFE(@ohsayan): Same type layout, alignments and non-null + ptr::replace(&mut (*e.as_ptr()).v, v) + }; + self._refresh(e); + Some(old) + } + #[inline(always)] + /// ## Safety + /// + /// Has sentinel + unsafe fn _refresh(&mut self, e: NonNull>) { + // UNSAFE(@ohsayan): Since it's in the collection, it is a valid ptr + IndexSTSeqDllNode::unlink(e.as_ptr()); + // UNSAFE(@ohsayan): As we found a node, our impl guarantees that the head is not-null + self.link(e); + } + #[inline(always)] + fn _clear(&mut self) { + self.m.clear(); + if !self.h.is_null() { + unsafe { + // UNSAFE(@ohsayan): nullck + self.drop_nodes_full(); + // UNSAFE(@ohsayan): Drop won't kill sentinel; link back to self + (*self.h).p = self.h; + (*self.h).n = self.h; + } + } + } + #[inline(always)] + fn _iter_unord_kv<'a>(&'a self) -> IndexSTSeqDllIterUnordKV<'a, K, V> { + IndexSTSeqDllIterUnordKV::new(&self.m) + } + #[inline(always)] + fn _iter_unord_k<'a>(&'a self) -> IndexSTSeqDllIterUnordK<'a, K, V> { + IndexSTSeqDllIterUnordK::new(&self.m) + } + #[inline(always)] + fn _iter_unord_v<'a>(&'a self) -> IndexSTSeqDllIterUnordV<'a, K, V> { + IndexSTSeqDllIterUnordV::new(&self.m) + } +} + +impl Drop for IndexSTSeqDll { + fn drop(&mut self) { + if !self.h.is_null() { + unsafe { + // UNSAFE(@ohsayan): nullck + self.drop_nodes_full(); + // UNSAFE(@ohsayan): nullck: drop doesn't clear sentinel + IndexSTSeqDllNode::dealloc_headless(self.h); + } + } + self.vacuum_free(); + } +} + +unsafe impl Send for IndexSTSeqDll {} +unsafe impl Sync for IndexSTSeqDll {} + +macro_rules! unsafe_marker_impl { + ($ty:ty) => { + unsafe impl<'a, K: Send, V: Send> Send for $ty {} + unsafe impl<'a, K: Sync, V: Sync> Sync for $ty {} + }; +} + +pub struct IndexSTSeqDllIterUnordKV<'a, K: 'a, V: 'a> { + i: Iter<'a, IndexSTSeqDllKeyptr, IndexSTSeqDllNodePtr>, +} + +// UNSAFE(@ohsayan): aliasing guarantees correctness +unsafe_marker_impl!(IndexSTSeqDllIterUnordKV<'a, K, V>); + +impl<'a, K: 'a, V: 'a> IndexSTSeqDllIterUnordKV<'a, K, V> { + #[inline(always)] + fn new(m: &'a StdMap, NonNull>, S>) -> Self { + Self { i: m.iter() } + } +} + +impl Clone for IndexSTSeqDllIterUnordKV<'_, K, V> { + fn clone(&self) -> Self { + Self { i: self.i.clone() } + } +} + +impl<'a, K, V> Iterator for IndexSTSeqDllIterUnordKV<'a, K, V> { + type Item = (&'a K, &'a V); + + fn next(&mut self) -> Option { + self.i.next().map(|(_, n)| { + let n = n.as_ptr(); + unsafe { + // UNSAFE(@ohsayan): nullck + (&(*n).k, &(*n).v) + } + }) + } +} + +impl<'a, K, V> ExactSizeIterator for IndexSTSeqDllIterUnordKV<'a, K, V> { + fn len(&self) -> usize { + self.i.len() + } +} + +impl<'a, K, V> FusedIterator for IndexSTSeqDllIterUnordKV<'a, K, V> {} + +impl Debug for IndexSTSeqDllIterUnordKV<'_, K, V> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_list().entries(self.clone()).finish() + } +} + +pub struct IndexSTSeqDllIterUnordK<'a, K: 'a, V: 'a> { + k: StdMapIterKey<'a, IndexSTSeqDllKeyptr, IndexSTSeqDllNodePtr>, +} + +// UNSAFE(@ohsayan): aliasing guarantees correctness +unsafe_marker_impl!(IndexSTSeqDllIterUnordK<'a, K, V>); + +impl<'a, K: 'a, V: 'a> IndexSTSeqDllIterUnordK<'a, K, V> { + #[inline(always)] + fn new(m: &'a StdMap, NonNull>, S>) -> Self { + Self { k: m.keys() } + } +} + +impl Clone for IndexSTSeqDllIterUnordK<'_, K, V> { + fn clone(&self) -> Self { + Self { k: self.k.clone() } + } +} + +impl<'a, K, V> Iterator for IndexSTSeqDllIterUnordK<'a, K, V> { + type Item = &'a K; + fn next(&mut self) -> Option { + self.k.next().map(|k| { + unsafe { + // UNSAFE(@ohsayan): nullck + &*(*k).p + } + }) + } +} + +impl<'a, K, V> ExactSizeIterator for IndexSTSeqDllIterUnordK<'a, K, V> { + fn len(&self) -> usize { + self.k.len() + } +} + +impl<'a, K, V> FusedIterator for IndexSTSeqDllIterUnordK<'a, K, V> {} + +impl<'a, K: Debug, V> Debug for IndexSTSeqDllIterUnordK<'_, K, V> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_list().entries(self.clone()).finish() + } +} + +pub struct IndexSTSeqDllIterUnordV<'a, K: 'a, V: 'a> { + v: StdMapIterVal<'a, IndexSTSeqDllKeyptr, IndexSTSeqDllNodePtr>, +} + +// UNSAFE(@ohsayan): aliasing guarantees correctness +unsafe_marker_impl!(IndexSTSeqDllIterUnordV<'a, K, V>); + +impl<'a, K: 'a, V: 'a> IndexSTSeqDllIterUnordV<'a, K, V> { + #[inline(always)] + fn new(m: &'a StdMap, NonNull>, S>) -> Self { + Self { v: m.values() } + } +} + +impl Clone for IndexSTSeqDllIterUnordV<'_, K, V> { + fn clone(&self) -> Self { + Self { v: self.v.clone() } + } +} + +impl<'a, K, V> Iterator for IndexSTSeqDllIterUnordV<'a, K, V> { + type Item = &'a V; + fn next(&mut self) -> Option { + self.v.next().map(|k| { + unsafe { + // UNSAFE(@ohsayan): nullck + &(*k.as_ptr()).v + } + }) + } +} + +impl<'a, K, V> ExactSizeIterator for IndexSTSeqDllIterUnordV<'a, K, V> { + fn len(&self) -> usize { + self.v.len() + } +} + +impl<'a, K, V> FusedIterator for IndexSTSeqDllIterUnordV<'a, K, V> {} + +impl<'a, K, V: Debug> Debug for IndexSTSeqDllIterUnordV<'_, K, V> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_list().entries(self.clone()).finish() + } +}