py3/rust: cpython-ext: support memoryview as PySimpleBuf

Summary: Otherwise we got RustPanic when clindex or dagindex reads mmapped changelog.i.

Reviewed By: sfilipco

Differential Revision: D19581189

fbshipit-source-id: 3ee74a1bd000d58272551ae404dcfe7f957bb2c0
This commit is contained in:
Xavier Deguillard 2020-01-27 16:48:13 -08:00 committed by Facebook Github Bot
parent ad58839ca1
commit f16bb04977

View File

@ -51,20 +51,30 @@ pub struct SimplePyBuf<T>(cpy::Py_buffer, PhantomData<T>);
unsafe impl<T> Send for SimplePyBuf<T> {}
unsafe impl<T> Sync for SimplePyBuf<T> {}
unsafe fn is_buffer(obj: &PyObject) -> bool {
if cpy::PyByteArray_Check(obj.as_ptr()) == 0 && cpy::PyBytes_Check(obj.as_ptr()) == 0 {
#[cfg(feature = "python2")]
return cpy::PyBuffer_Check(obj.as_ptr()) == 0;
#[cfg(feature = "python3")]
unsafe fn is_safe_type(obj: &PyObject) -> bool {
if cpy::PyByteArray_Check(obj.as_ptr()) != 0 {
return true;
}
if cpy::PyBytes_Check(obj.as_ptr()) != 0 {
return true;
}
#[cfg(feature = "python2")]
{
if cpy::PyBuffer_Check(obj.as_ptr()) != 0 {
return true;
}
}
#[cfg(feature = "python3")]
{
if cpy::PyMemoryView_Check(obj.as_ptr()) != 0 {
return true;
}
}
return false;
}
impl<T: Copy> SimplePyBuf<T> {
pub fn new(_py: Python<'_>, obj: &PyObject) -> Self {
pub fn new(py: Python<'_>, obj: &PyObject) -> Self {
// Note about GC on obj:
//
// Practically, obj here is some low-level, non-container ones like
@ -87,8 +97,9 @@ impl<T: Copy> SimplePyBuf<T> {
// whitelist those two types. Beware that `PyBuffer_Check` won't guarnatee
// its inner object is also immutable.
unsafe {
if is_buffer(obj) {
panic!("potentially unsafe type");
if !is_safe_type(obj) {
let ty = obj.get_type(py);
panic!("potentially unsafe type for SimplePyBuf: {}", ty.name(py));
}
let mut buf = mem::zeroed::<SimplePyBuf<T>>();