snipe off some unsafe code (#922)

davidhewitt · web-flow · commit 34fbd845c120 · 2023-08-24T22:30:45.000+01:00
diff --git a/src/errors/validation_exception.rs b/src/errors/validation_exception.rs
@@ -3,11 +3,10 @@ use std::fmt::{Display, Write};
 use std::str::from_utf8;
 
 use pyo3::exceptions::{PyKeyError, PyTypeError, PyValueError};
-use pyo3::ffi::Py_ssize_t;
+use pyo3::intern;
 use pyo3::once_cell::GILOnceCell;
 use pyo3::prelude::*;
 use pyo3::types::{PyDict, PyList, PyString};
-use pyo3::{ffi, intern};
 use serde::ser::{Error, SerializeMap, SerializeSeq};
 use serde::{Serialize, Serializer};
 
@@ -173,22 +172,27 @@ impl ValidationError {
     #[pyo3(signature = (*, include_url = true, include_context = true))]
     pub fn errors(&self, py: Python, include_url: bool, include_context: bool) -> PyResult<Py<PyList>> {
         let url_prefix = get_url_prefix(py, include_url);
-        // taken approximately from the pyo3, but modified to return the error during iteration
-        // https://github.com/PyO3/pyo3/blob/a3edbf4fcd595f0e234c87d4705eb600a9779130/src/types/list.rs#L27-L55
-        unsafe {
-            let ptr = ffi::PyList_New(self.line_errors.len() as Py_ssize_t);
-
-            // We create the `Py` pointer here for two reasons:
-            // - panics if the ptr is null
-            // - its Drop cleans up the list if user code or the asserts panic.
-            let list: Py<PyList> = Py::from_owned_ptr(py, ptr);
-
-            for (index, line_error) in (0_isize..).zip(&self.line_errors) {
-                let item = line_error.as_dict(py, url_prefix, include_context, &self.error_mode)?;
-                ffi::PyList_SET_ITEM(ptr, index, item.into_ptr());
-            }
-
-            Ok(list)
+        let mut iteration_error = None;
+        let list = PyList::new(
+            py,
+            // PyList::new takes ExactSizeIterator, so if an error occurs during iteration we
+            // fill the list with None before returning the error; the list will then be thrown
+            // away safely.
+            self.line_errors.iter().map(|e| -> PyObject {
+                if iteration_error.is_some() {
+                    return py.None();
+                }
+                e.as_dict(py, url_prefix, include_context, &self.error_mode)
+                    .unwrap_or_else(|err| {
+                        iteration_error = Some(err);
+                        py.None()
+                    })
+            }),
+        );
+        if let Some(err) = iteration_error {
+            Err(err)
+        } else {
+            Ok(list.into())
         }
     }
 
diff --git a/src/input/input_python.rs b/src/input/input_python.rs
@@ -187,6 +187,8 @@ impl<'a> Input<'a> for PyAny {
             let str = py_str.to_str()?;
             serde_json::from_str(str).map_err(|e| map_json_err(self, e))
         } else if let Ok(py_byte_array) = self.downcast::<PyByteArray>() {
+            // Safety: from_slice does not run arbitrary Python code and the GIL is held so the
+            // bytes array will not be mutated while from_slice is reading it
             serde_json::from_slice(unsafe { py_byte_array.as_bytes() }).map_err(|e| map_json_err(self, e))
         } else {
             Err(ValError::new(ErrorTypeDefaults::JsonType, self))
@@ -235,13 +237,15 @@ impl<'a> Input<'a> for PyAny {
             };
             Ok(str.into())
         } else if let Ok(py_byte_array) = self.downcast::<PyByteArray>() {
-            // see https://docs.rs/pyo3/latest/pyo3/types/struct.PyByteArray.html#method.as_bytes
-            // for why this is marked unsafe
-            let str = match from_utf8(unsafe { py_byte_array.as_bytes() }) {
-                Ok(s) => s,
+            // Safety: the gil is held while from_utf8 is running so py_byte_array is not mutated,
+            // and we immediately copy the bytes into a new Python string
+            let s = match from_utf8(unsafe { py_byte_array.as_bytes() }) {
+                // Why Python not Rust? to avoid an unnecessary allocation on the Rust side, the
+                // final output needs to be Python anyway.
+                Ok(s) => PyString::new(self.py(), s),
                 Err(_) => return Err(ValError::new(ErrorTypeDefaults::StringUnicode, self)),
             };
-            Ok(str.into())
+            Ok(s.into())
         } else {
             Err(ValError::new(ErrorTypeDefaults::StringType, self))
         }
@@ -337,9 +341,8 @@ impl<'a> Input<'a> for PyAny {
         }
     }
     fn strict_float(&'a self) -> ValResult<EitherFloat<'a>> {
-        if PyFloat::is_exact_type_of(self) {
-            // Safety: self is PyFloat
-            Ok(EitherFloat::Py(unsafe { self.downcast_unchecked::<PyFloat>() }))
+        if let Ok(py_float) = self.downcast_exact::<PyFloat>() {
+            Ok(EitherFloat::Py(py_float))
         } else if let Ok(float) = self.extract::<f64>() {
             // bools are cast to floats as either 0.0 or 1.0, so check for bool type in this specific case
             if (float == 0.0 || float == 1.0) && PyBool::is_exact_type_of(self) {
@@ -353,9 +356,8 @@ impl<'a> Input<'a> for PyAny {
     }
 
     fn lax_float(&'a self) -> ValResult<EitherFloat<'a>> {
-        if PyFloat::is_exact_type_of(self) {
-            // Safety: self is PyFloat
-            Ok(EitherFloat::Py(unsafe { self.downcast_unchecked::<PyFloat>() }))
+        if let Ok(py_float) = self.downcast_exact() {
+            Ok(EitherFloat::Py(py_float))
         } else if let Some(cow_str) = maybe_as_string(self, ErrorTypeDefaults::FloatParsing)? {
             str_as_float(self, &cow_str)
         } else if let Ok(float) = self.extract::<f64>() {
diff --git a/src/input/return_enums.rs b/src/input/return_enums.rs
@@ -224,12 +224,11 @@ impl BuildSet for &PySet {
 
 impl BuildSet for &PyFrozenSet {
     fn build_add(&self, item: PyObject) -> PyResult<()> {
-        unsafe {
-            py_error_on_minusone(
-                self.py(),
-                ffi::PySet_Add(self.as_ptr(), item.to_object(self.py()).as_ptr()),
-            )
-        }
+        py_error_on_minusone(self.py(), unsafe {
+            // Safety: self.as_ptr() the _only_ pointer to the `frozenset`, and it's allowed
+            // to mutate this via the C API when nothing else can refer to it.
+            ffi::PySet_Add(self.as_ptr(), item.to_object(self.py()).as_ptr())
+        })
     }
 
     fn build_len(&self) -> usize {
@@ -492,57 +491,32 @@ impl<'py> Iterator for MappingGenericIterator<'py> {
     type Item = ValResult<'py, (&'py PyAny, &'py PyAny)>;
 
     fn next(&mut self) -> Option<Self::Item> {
-        let item = match self.iter.next() {
-            Some(Err(e)) => return Some(Err(mapping_err(e, self.iter.py(), self.input))),
-            Some(Ok(item)) => item,
-            None => return None,
-        };
-        let tuple: &PyTuple = match item.downcast() {
-            Ok(tuple) => tuple,
-            Err(_) => {
-                return Some(Err(ValError::new(
+        Some(match self.iter.next()? {
+            Ok(item) => item.extract().map_err(|_| {
+                ValError::new(
                     ErrorType::MappingType {
                         error: MAPPING_TUPLE_ERROR.into(),
                         context: None,
                     },
                     self.input,
-                )))
-            }
-        };
-        if tuple.len() != 2 {
-            return Some(Err(ValError::new(
-                ErrorType::MappingType {
-                    error: MAPPING_TUPLE_ERROR.into(),
-                    context: None,
-                },
-                self.input,
-            )));
-        };
-        #[cfg(PyPy)]
-        let key = tuple.get_item(0).unwrap();
-        #[cfg(PyPy)]
-        let value = tuple.get_item(1).unwrap();
-        #[cfg(not(PyPy))]
-        let key = unsafe { tuple.get_item_unchecked(0) };
-        #[cfg(not(PyPy))]
-        let value = unsafe { tuple.get_item_unchecked(1) };
-        Some(Ok((key, value)))
+                )
+            }),
+            Err(e) => Err(mapping_err(e, self.iter.py(), self.input)),
+        })
     }
-    // size_hint is omitted as it isn't needed
 }
 
 pub struct AttributesGenericIterator<'py> {
     object: &'py PyAny,
-    attributes: &'py PyList,
-    index: usize,
+    // PyO3 should export this type upstream
+    attributes_iterator: <&'py PyList as IntoIterator>::IntoIter,
 }
 
 impl<'py> AttributesGenericIterator<'py> {
     pub fn new(py_any: &'py PyAny) -> ValResult<'py, Self> {
         Ok(Self {
             object: py_any,
-            attributes: py_any.dir(),
-            index: 0,
+            attributes_iterator: py_any.dir().into_iter(),
         })
     }
 }
@@ -553,37 +527,31 @@ impl<'py> Iterator for AttributesGenericIterator<'py> {
     fn next(&mut self) -> Option<Self::Item> {
         // loop until we find an attribute who's name does not start with underscore,
         // or we get to the end of the list of attributes
-        while self.index < self.attributes.len() {
-            #[cfg(PyPy)]
-            let name: &PyAny = self.attributes.get_item(self.index).unwrap();
-            #[cfg(not(PyPy))]
-            let name: &PyAny = unsafe { self.attributes.get_item_unchecked(self.index) };
-            self.index += 1;
-            // from benchmarks this is 14x faster than using the python `startswith` method
-            let name_cow = match name.downcast::<PyString>() {
-                Ok(name) => name.to_string_lossy(),
-                Err(e) => return Some(Err(e.into())),
-            };
-            if !name_cow.as_ref().starts_with('_') {
-                // getattr is most likely to fail due to an exception in a @property, skip
-                if let Ok(attr) = self.object.getattr(name_cow.as_ref()) {
-                    // we don't want bound methods to be included, is there a better way to check?
-                    // ref https://stackoverflow.com/a/18955425/949890
-                    let is_bound = matches!(attr.hasattr(intern!(attr.py(), "__self__")), Ok(true));
-                    // the PyFunction::is_type_of(attr) catches `staticmethod`, but also any other function,
-                    // I think that's better than including static methods in the yielded attributes,
-                    // if someone really wants fields, they can use an explicit field, or a function to modify input
-                    #[cfg(not(PyPy))]
-                    if !is_bound && !PyFunction::is_type_of(attr) {
-                        return Some(Ok((name, attr)));
-                    }
-                    // MASSIVE HACK! PyFunction doesn't exist for PyPy,
-                    // is_instance_of::<PyFunction> crashes with a null pointer, hence this hack, see
-                    // https://github.com/pydantic/pydantic-core/pull/161#discussion_r917257635
-                    #[cfg(PyPy)]
-                    if !is_bound && attr.get_type().to_string() != "<class 'function'>" {
-                        return Some(Ok((name, attr)));
-                    }
+        let name = self.attributes_iterator.next()?;
+        // from benchmarks this is 14x faster than using the python `startswith` method
+        let name_cow = match name.downcast::<PyString>() {
+            Ok(name) => name.to_string_lossy(),
+            Err(e) => return Some(Err(e.into())),
+        };
+        if !name_cow.as_ref().starts_with('_') {
+            // getattr is most likely to fail due to an exception in a @property, skip
+            if let Ok(attr) = self.object.getattr(name_cow.as_ref()) {
+                // we don't want bound methods to be included, is there a better way to check?
+                // ref https://stackoverflow.com/a/18955425/949890
+                let is_bound = matches!(attr.hasattr(intern!(attr.py(), "__self__")), Ok(true));
+                // the PyFunction::is_type_of(attr) catches `staticmethod`, but also any other function,
+                // I think that's better than including static methods in the yielded attributes,
+                // if someone really wants fields, they can use an explicit field, or a function to modify input
+                #[cfg(not(PyPy))]
+                if !is_bound && !PyFunction::is_type_of(attr) {
+                    return Some(Ok((name, attr)));
+                }
+                // MASSIVE HACK! PyFunction doesn't exist for PyPy,
+                // is_instance_of::<PyFunction> crashes with a null pointer, hence this hack, see
+                // https://github.com/pydantic/pydantic-core/pull/161#discussion_r917257635
+                #[cfg(PyPy)]
+                if !is_bound && attr.get_type().to_string() != "<class 'function'>" {
+                    return Some(Ok((name, attr)));
                 }
             }
         }
@@ -621,12 +589,7 @@ pub enum GenericIterator {
 
 impl From<JsonArray> for GenericIterator {
     fn from(array: JsonArray) -> Self {
-        let length = array.len();
-        let json_iter = GenericJsonIterator {
-            array,
-            length,
-            index: 0,
-        };
+        let json_iter = GenericJsonIterator { array, index: 0 };
         Self::JsonArray(json_iter)
     }
 }
@@ -674,14 +637,15 @@ impl GenericPyIterator {
 #[derive(Debug, Clone)]
 pub struct GenericJsonIterator {
     array: JsonArray,
-    length: usize,
     index: usize,
 }
 
 impl GenericJsonIterator {
     pub fn next(&mut self, _py: Python) -> PyResult<Option<(&JsonInput, usize)>> {
-        if self.index < self.length {
-            let next = unsafe { self.array.get_unchecked(self.index) };
+        if self.index < self.array.len() {
+            // panic here is impossible due to bounds check above; compiler should be
+            // able to optimize it away even
+            let next = &self.array[self.index];
             let a = (next, self.index);
             self.index += 1;
             Ok(Some(a))
@@ -940,13 +904,7 @@ impl<'a> EitherFloat<'a> {
     pub fn as_f64(self) -> f64 {
         match self {
             EitherFloat::F64(f) => f,
-
-            EitherFloat::Py(f) => {
-                {
-                    // Safety: known to be a python float
-                    unsafe { ffi::PyFloat_AS_DOUBLE(f.as_ptr()) }
-                }
-            }
+            EitherFloat::Py(f) => f.value(),
         }
     }
 }
diff --git a/src/serializers/infer.rs b/src/serializers/infer.rs
@@ -133,8 +133,8 @@ pub(crate) fn infer_to_python_known(
                 .map(|s| s.into_py(py))?,
             ObType::Bytearray => {
                 let py_byte_array: &PyByteArray = value.downcast()?;
-                // see https://docs.rs/pyo3/latest/pyo3/types/struct.PyByteArray.html#method.as_bytes
-                // for why this is marked unsafe
+                // Safety: the GIL is held while bytes_to_string is running; it doesn't run
+                // arbitrary Python code, so py_byte_array cannot be mutated.
                 let bytes = unsafe { py_byte_array.as_bytes() };
                 extra
                     .config
@@ -428,8 +428,12 @@ pub(crate) fn infer_serialize_known<S: Serializer>(
         }
         ObType::Bytearray => {
             let py_byte_array: &PyByteArray = value.downcast().map_err(py_err_se_err)?;
-            let bytes = unsafe { py_byte_array.as_bytes() };
-            extra.config.bytes_mode.serialize_bytes(bytes, serializer)
+            // Safety: the GIL is held while serialize_bytes is running; it doesn't run
+            // arbitrary Python code, so py_byte_array cannot be mutated.
+            extra
+                .config
+                .bytes_mode
+                .serialize_bytes(unsafe { py_byte_array.as_bytes() }, serializer)
         }
         ObType::Dict => serialize_dict!(value.downcast::<PyDict>().map_err(py_err_se_err)?),
         ObType::List => serialize_seq_filter!(PyList),
@@ -581,8 +585,15 @@ pub(crate) fn infer_json_key_known<'py>(ob_type: ObType, key: &'py PyAny, extra:
             .bytes_to_string(key.py(), key.downcast::<PyBytes>()?.as_bytes()),
         ObType::Bytearray => {
             let py_byte_array: &PyByteArray = key.downcast()?;
-            let bytes = unsafe { py_byte_array.as_bytes() };
-            extra.config.bytes_mode.bytes_to_string(key.py(), bytes)
+            // Safety: the GIL is held while serialize_bytes is running; it doesn't run
+            // arbitrary Python code, so py_byte_array cannot be mutated during the call.
+            //
+            // We copy the bytes into a new buffer immediately afterwards
+            extra
+                .config
+                .bytes_mode
+                .bytes_to_string(key.py(), unsafe { py_byte_array.as_bytes() })
+                .map(|cow| Cow::Owned(cow.into_owned()))
         }
         ObType::Datetime => {
             let py_dt: &PyDateTime = key.downcast()?;
diff --git a/src/serializers/ob_type.rs b/src/serializers/ob_type.rs
