Skip to content

Commit 850935f

Browse files
ovrovr
committed
feat(native): Improve safety for Python into CLRepr conversion
1 parent 7e6037c commit 850935f

File tree

3 files changed

+122
-23
lines changed

3 files changed

+122
-23
lines changed

‎packages/cubejs-backend-native/src/cross/clrepr_python.rs‎

Lines changed: 38 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,10 @@
11
use crate::cross::clrepr::CLReprObject;
22
use crate::cross::{CLRepr, CLReprObjectKind, StringType};
3-
use pyo3::exceptions::{PyNotImplementedError, PyTypeError};
3+
use crate::python::utils::PyAnyHelpers;
4+
use pyo3::exceptions::{PyException, PyNotImplementedError, PyTypeError};
45
use pyo3::types::{
5-
PyBool, PyComplex, PyDate, PyDict, PyFloat, PyFrame, PyFunction, PyInt, PyList, PySequence,
6-
PySet, PyString, PyTraceback, PyTuple,
6+
PyBool, PyComplex, PyDate, PyDateTime, PyDelta, PyDict, PyFloat, PyFrame, PyFrozenSet,
7+
PyFunction, PyInt, PyList, PySequence, PySet, PyString, PyTraceback, PyTuple,
78
};
89
use pyo3::{Py, PyAny, PyErr, PyObject, Python, ToPyObject};
910

@@ -12,7 +13,7 @@ pub enum PythonRef {
1213
PyObject(PyObject),
1314
PyFunction(Py<PyFunction>),
1415
/// Special type to transfer functions through JavaScript
15-
/// In JS it's an external object. It's not the same as Function.
16+
/// In JS it's an external object.
1617
PyExternalFunction(Py<PyFunction>),
1718
}
1819

@@ -87,10 +88,28 @@ impl CLRepr {
8788
return Err(PyErr::new::<PyTypeError, _>(
8889
"Unable to represent PyComplex type as CLR from Python".to_string(),
8990
));
91+
} else if v.get_type().is_subclass_of::<PyDateTime>()? {
92+
return Err(PyErr::new::<PyTypeError, _>(
93+
"Unable to represent PyDateTime type as CLR from Python".to_string(),
94+
));
9095
} else if v.get_type().is_subclass_of::<PyDate>()? {
9196
return Err(PyErr::new::<PyTypeError, _>(
9297
"Unable to represent PyDate type as CLR from Python".to_string(),
9398
));
99+
} else if v.get_type().is_subclass_of::<PyFrozenSet>()? {
100+
let set = v.downcast::<PyFrozenSet>()?;
101+
102+
return Err(PyErr::new::<PyTypeError, _>(format!(
103+
"Unable to represent PyFrozenSet type as CLR from Python, value: {:?}",
104+
set
105+
)));
106+
} else if v.get_type().is_subclass_of::<PyException>()? {
107+
let exception = v.downcast::<PyException>()?;
108+
109+
return Err(PyErr::new::<PyTypeError, _>(format!(
110+
"Unable to represent PyException type as CLR from Python, value: {:?}",
111+
exception
112+
)));
94113
} else if v.get_type().is_subclass_of::<PyFrame>()? {
95114
let frame = v.downcast::<PyFrame>()?;
96115

@@ -105,17 +124,22 @@ impl CLRepr {
105124
"Unable to represent PyTraceback type as CLR from Python, value: {:?}",
106125
trb
107126
)));
108-
} else {
109-
let is_sequence = unsafe { pyo3::ffi::PySequence_Check(v.as_ptr()) == 1 };
110-
if is_sequence {
111-
let seq = v.downcast::<PySequence>()?;
112-
113-
return Err(PyErr::new::<PyTypeError, _>(format!(
114-
"Unable to represent PySequence type as CLR from Python, value: {:?}",
115-
seq
116-
)));
117-
}
127+
} else if v.get_type().is_subclass_of::<PyDelta>()? {
128+
let delta = v.downcast::<PyDelta>()?;
129+
130+
return Err(PyErr::new::<PyTypeError, _>(format!(
131+
"Unable to represent PyDelta type as CLR from Python, value: {:?}",
132+
delta
133+
)));
134+
} else if v.is_sequence()? {
135+
let seq = v.downcast::<PySequence>()?;
118136

137+
return Err(PyErr::new::<PyTypeError, _>(format!(
138+
"Unable to represent PySequence type as CLR from Python, value: {:?}",
139+
seq
140+
)));
141+
} else {
142+
// Fallback to PyObject, it will lead to throw error in the JS side
119143
Self::PythonRef(PythonRef::PyObject(v.into()))
120144
})
121145
}

‎packages/cubejs-backend-native/src/python/runtime.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
use crate::cross::CLRepr;
22
use crate::python::neon_py::*;
3+
use crate::python::utils::PyAnyHelpers;
34
use crate::tokio_runtime_node;
45
use cubesql::CubeError;
56
use log::{error, trace};
@@ -143,8 +144,7 @@ impl PyRuntime {
143144
let py_args = PyTuple::new(py, prep_tuple);
144145
let call_res = fun.call(py, py_args, py_kwargs)?;
145146

146-
let is_coroutine = unsafe { pyo3::ffi::PyCoro_CheckExact(call_res.as_ptr()) == 1 };
147-
if is_coroutine {
147+
if call_res.is_coroutine()? {
148148
let fut = pyo3_asyncio::tokio::into_future(call_res.as_ref(py))?;
149149
Ok(PyScheduledFunResult::Poll(Box::pin(fut)))
150150
} else {

‎packages/cubejs-backend-native/src/python/utils.rs‎

Lines changed: 82 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,9 @@
11
use crate::cross::*;
2-
use pyo3::exceptions::PyNotImplementedError;
2+
use pyo3::exceptions::{PyNotImplementedError, PySystemError};
3+
use pyo3::ffi;
34
use pyo3::prelude::*;
45
use pyo3::types::{PyFunction, PyString, PyTuple};
6+
use std::ffi::c_int;
57

68
pub fn python_fn_call_sync(py_fun: &Py<PyFunction>, arguments: Vec<CLRepr>) -> PyResult<CLRepr> {
79
Python::with_gil(|py| {
@@ -15,8 +17,7 @@ pub fn python_fn_call_sync(py_fun: &Py<PyFunction>, arguments: Vec<CLRepr>) -> P
1517

1618
let call_res = py_fun.call1(py, tuple)?;
1719

18-
let is_coroutine = unsafe { pyo3::ffi::PyCoro_CheckExact(call_res.as_ptr()) == 1 };
19-
if is_coroutine {
20+
if call_res.is_coroutine()? {
2021
Err(PyErr::new::<PyNotImplementedError, _>(
2122
"Calling function with async response is not supported",
2223
))
@@ -38,8 +39,7 @@ pub fn python_obj_call_sync(py_fun: &PyObject, arguments: Vec<CLRepr>) -> PyResu
3839

3940
let call_res = py_fun.call1(py, tuple)?;
4041

41-
let is_coroutine = unsafe { pyo3::ffi::PyCoro_CheckExact(call_res.as_ptr()) == 1 };
42-
if is_coroutine {
42+
if call_res.is_coroutine()? {
4343
Err(PyErr::new::<PyNotImplementedError, _>(
4444
"Calling object with async response is not supported",
4545
))
@@ -68,8 +68,7 @@ where
6868

6969
let call_res = py_fun.call_method1(py, name, tuple)?;
7070

71-
let is_coroutine = unsafe { pyo3::ffi::PyCoro_CheckExact(call_res.as_ptr()) == 1 };
72-
if is_coroutine {
71+
if call_res.is_coroutine()? {
7372
Err(PyErr::new::<PyNotImplementedError, _>(
7473
"Calling object method with async response is not supported",
7574
))
@@ -78,3 +77,79 @@ where
7877
}
7978
})
8079
}
80+
81+
pub trait PyAnyHelpers {
82+
fn is_sequence(&self) -> PyResult<bool>;
83+
84+
fn is_coroutine(&self) -> PyResult<bool>;
85+
}
86+
87+
pub(crate) fn internal_error_on_minusone(result: c_int) -> PyResult<()> {
88+
if result != -1 {
89+
Ok(())
90+
} else {
91+
Err(PyErr::new::<PySystemError, _>(
92+
"Error on call via ffi, result is -1",
93+
))
94+
}
95+
}
96+
97+
impl PyAnyHelpers for PyAny {
98+
fn is_sequence(&self) -> PyResult<bool> {
99+
let ptr = self.as_ptr();
100+
if ptr.is_null() {
101+
return Err(PyErr::new::<PySystemError, _>(
102+
"Unable to verify that object is sequence, because ptr to PyAny is null",
103+
));
104+
}
105+
106+
let v = unsafe { ffi::PySequence_Check(ptr) };
107+
internal_error_on_minusone(v)?;
108+
109+
Ok(v != 0)
110+
}
111+
112+
fn is_coroutine(&self) -> PyResult<bool> {
113+
let ptr = self.as_ptr();
114+
if ptr.is_null() {
115+
return Err(PyErr::new::<PySystemError, _>(
116+
"Unable to verify that object is sequence, because ptr to PyAny is null",
117+
));
118+
}
119+
120+
let v = unsafe { ffi::PyCoro_CheckExact(ptr) };
121+
internal_error_on_minusone(v)?;
122+
123+
Ok(v != 0)
124+
}
125+
}
126+
127+
impl PyAnyHelpers for PyObject {
128+
fn is_sequence(&self) -> PyResult<bool> {
129+
let ptr = self.as_ptr();
130+
if ptr.is_null() {
131+
return Err(PyErr::new::<PySystemError, _>(
132+
"Unable to verify that object is sequence, because ptr to PyAny is null",
133+
));
134+
}
135+
136+
let v = unsafe { ffi::PySequence_Check(ptr) };
137+
internal_error_on_minusone(v)?;
138+
139+
Ok(v != 0)
140+
}
141+
142+
fn is_coroutine(&self) -> PyResult<bool> {
143+
let ptr = self.as_ptr();
144+
if ptr.is_null() {
145+
return Err(PyErr::new::<PySystemError, _>(
146+
"Unable to verify that object is sequence, because ptr to PyAny is null",
147+
));
148+
}
149+
150+
let v = unsafe { ffi::PyCoro_CheckExact(ptr) };
151+
internal_error_on_minusone(v)?;
152+
153+
Ok(v != 0)
154+
}
155+
}

0 commit comments

Comments
 (0)