Skip to content

Commit 2f18619

Browse files
committed
Replace AHash (DoS-resistant) by FxHash (simpler dependency)
In the datetime module, we hash a set of keys known at compile time and hence are not subject to DoS issues. In the borrow module, the keys include pointer addresses which means that even if a program exposes its usage of NumPy arrays to remote user control, a DoS attack would also need detailed control over the placement of memory allocations in which case outright memory exhaustion seems to be a simpler avenue of attack.
1 parent 347888d commit 2f18619

File tree

3 files changed

+12
-10
lines changed

3 files changed

+12
-10
lines changed

Cargo.toml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,6 @@ keywords = ["python", "numpy", "ffi", "pyo3"]
1515
license = "BSD-2-Clause"
1616

1717
[dependencies]
18-
ahash = "0.7"
1918
half = { version = "2.0", default-features = false, optional = true }
2019
libc = "0.2"
2120
nalgebra = { version = "0.31", default-features = false, optional = true }
@@ -24,6 +23,7 @@ num-integer = "0.1"
2423
num-traits = "0.2"
2524
ndarray = ">= 0.13, < 0.16"
2625
pyo3 = { version = "0.17", default-features = false, features = ["macros"] }
26+
rustc-hash = "1.1"
2727

2828
[dev-dependencies]
2929
pyo3 = { version = "0.17", default-features = false, features = ["auto-initialize"] }

src/borrow.rs

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -167,12 +167,12 @@ use std::fmt;
167167
use std::mem::size_of;
168168
use std::ops::Deref;
169169

170-
use ahash::AHashMap;
171170
use ndarray::{
172171
ArrayView, ArrayViewMut, Dimension, IntoDimension, Ix1, Ix2, Ix3, Ix4, Ix5, Ix6, IxDyn,
173172
};
174173
use num_integer::gcd;
175174
use pyo3::{FromPyObject, PyAny, PyResult, Python};
175+
use rustc_hash::FxHashMap;
176176

177177
use crate::array::PyArray;
178178
use crate::cold;
@@ -237,7 +237,7 @@ impl BorrowKey {
237237
}
238238
}
239239

240-
type BorrowFlagsInner = AHashMap<*mut u8, AHashMap<BorrowKey, isize>>;
240+
type BorrowFlagsInner = FxHashMap<*mut u8, FxHashMap<BorrowKey, isize>>;
241241

242242
struct BorrowFlags(UnsafeCell<Option<BorrowFlagsInner>>);
243243

@@ -250,7 +250,7 @@ impl BorrowFlags {
250250

251251
#[allow(clippy::mut_from_ref)]
252252
unsafe fn get(&self) -> &mut BorrowFlagsInner {
253-
(*self.0.get()).get_or_insert_with(AHashMap::new)
253+
(*self.0.get()).get_or_insert_with(Default::default)
254254
}
255255

256256
fn acquire(&self, _py: Python, address: *mut u8, key: BorrowKey) -> Result<(), BorrowError> {
@@ -287,7 +287,8 @@ impl BorrowFlags {
287287
}
288288
}
289289
Entry::Vacant(entry) => {
290-
let mut same_base_arrays = AHashMap::with_capacity(1);
290+
let mut same_base_arrays =
291+
FxHashMap::with_capacity_and_hasher(1, Default::default());
291292
same_base_arrays.insert(key, 1);
292293
entry.insert(same_base_arrays);
293294
}
@@ -349,7 +350,8 @@ impl BorrowFlags {
349350
}
350351
}
351352
Entry::Vacant(entry) => {
352-
let mut same_base_arrays = AHashMap::with_capacity(1);
353+
let mut same_base_arrays =
354+
FxHashMap::with_capacity_and_hasher(1, Default::default());
353355
same_base_arrays.insert(key, -1);
354356
entry.insert(same_base_arrays);
355357
}

src/datetime.rs

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -63,8 +63,8 @@ use std::fmt;
6363
use std::hash::Hash;
6464
use std::marker::PhantomData;
6565

66-
use ahash::AHashMap;
6766
use pyo3::{Py, Python};
67+
use rustc_hash::FxHashMap;
6868

6969
use crate::dtype::{Element, PyArrayDescr};
7070
use crate::npyffi::{PyArray_DatetimeDTypeMetaData, NPY_DATETIMEUNIT, NPY_TYPES};
@@ -206,7 +206,7 @@ impl<U: Unit> fmt::Debug for Timedelta<U> {
206206

207207
struct TypeDescriptors {
208208
npy_type: NPY_TYPES,
209-
dtypes: UnsafeCell<Option<AHashMap<NPY_DATETIMEUNIT, Py<PyArrayDescr>>>>,
209+
dtypes: UnsafeCell<Option<FxHashMap<NPY_DATETIMEUNIT, Py<PyArrayDescr>>>>,
210210
}
211211

212212
unsafe impl Sync for TypeDescriptors {}
@@ -221,8 +221,8 @@ impl TypeDescriptors {
221221
}
222222

223223
#[allow(clippy::mut_from_ref)]
224-
unsafe fn get(&self) -> &mut AHashMap<NPY_DATETIMEUNIT, Py<PyArrayDescr>> {
225-
(*self.dtypes.get()).get_or_insert_with(AHashMap::new)
224+
unsafe fn get(&self) -> &mut FxHashMap<NPY_DATETIMEUNIT, Py<PyArrayDescr>> {
225+
(*self.dtypes.get()).get_or_insert_with(Default::default)
226226
}
227227

228228
#[allow(clippy::wrong_self_convention)]

0 commit comments

Comments
 (0)