Skip to content

Commit 2c73d2b

Browse files
XuehaiPanXuehaiPan
committed
Update tests
1 parent 6c6db0e commit 2c73d2b

File tree

2 files changed

+210
-82
lines changed

2 files changed

+210
-82
lines changed

‎tests/test_invalid_holder_access.cpp‎

Lines changed: 92 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -4,14 +4,9 @@
44
#include <memory>
55
#include <vector>
66

7-
class VectorOwns4PythonObjects {
7+
class VecOwnsObjs {
88
public:
9-
void append(const py::object &obj) {
10-
if (size() >= 4) {
11-
throw std::out_of_range("Index out of range");
12-
}
13-
vec.emplace_back(obj);
14-
}
9+
void append(const py::object &obj) { vec.emplace_back(obj); }
1510

1611
void set_item(py::ssize_t i, const py::object &obj) {
1712
if (!(i >= 0 && i < size())) {
@@ -31,63 +26,120 @@ class VectorOwns4PythonObjects {
3126

3227
bool is_empty() const { return vec.empty(); }
3328

34-
void sanity_check() const {
35-
auto current_size = size();
36-
if (current_size < 0 || current_size > 4) {
37-
throw std::out_of_range("Invalid size");
38-
}
39-
}
40-
4129
static int tp_traverse(PyObject *self_base, visitproc visit, void *arg) {
4230
#if PY_VERSION_HEX >= 0x03090000 // Python 3.9
4331
Py_VISIT(Py_TYPE(self_base));
4432
#endif
45-
auto *const instance = reinterpret_cast<py::detail::instance *>(self_base);
46-
if (!instance->get_value_and_holder().holder_constructed()) {
47-
// The holder has not been constructed yet. Skip the traversal to avoid segmentation
48-
// faults.
49-
return 0;
33+
if (should_check_holder_initialization) {
34+
auto *const instance = reinterpret_cast<py::detail::instance *>(self_base);
35+
if (!instance->get_value_and_holder().holder_constructed()) {
36+
// The holder has not been constructed yet. Skip the traversal to avoid
37+
// segmentation faults.
38+
return 0;
39+
}
5040
}
51-
auto &self = py::cast<VectorOwns4PythonObjects &>(py::handle{self_base});
41+
auto &self = py::cast<VecOwnsObjs &>(py::handle{self_base});
5242
for (const auto &obj : self.vec) {
5343
Py_VISIT(obj.ptr());
5444
}
5545
return 0;
5646
}
5747

48+
static int tp_clear(PyObject *self_base) {
49+
if (should_check_holder_initialization) {
50+
auto *const instance = reinterpret_cast<py::detail::instance *>(self_base);
51+
if (!instance->get_value_and_holder().holder_constructed()) {
52+
// The holder has not been constructed yet. Skip the traversal to avoid
53+
// segmentation faults.
54+
return 0;
55+
}
56+
}
57+
auto &self = py::cast<VecOwnsObjs &>(py::handle{self_base});
58+
for (auto &obj : self.vec) {
59+
Py_CLEAR(obj.ptr());
60+
}
61+
self.vec.clear();
62+
return 0;
63+
}
64+
65+
py::object get_state() const {
66+
py::list state{};
67+
for (const auto &item : vec) {
68+
state.append(item);
69+
}
70+
return py::tuple(state);
71+
}
72+
73+
static bool get_should_check_holder_initialization() {
74+
return should_check_holder_initialization;
75+
}
76+
77+
static void set_should_check_holder_initialization(bool value) {
78+
should_check_holder_initialization = value;
79+
}
80+
81+
static bool get_should_raise_error_on_set_state() { return should_raise_error_on_set_state; }
82+
83+
static void set_should_raise_error_on_set_state(bool value) {
84+
should_raise_error_on_set_state = value;
85+
}
86+
87+
static bool should_check_holder_initialization;
88+
static bool should_raise_error_on_set_state;
89+
5890
private:
5991
std::vector<py::object> vec{};
6092
};
6193

94+
bool VecOwnsObjs::should_check_holder_initialization = false;
95+
bool VecOwnsObjs::should_raise_error_on_set_state = false;
96+
6297
TEST_SUBMODULE(invalid_holder_access, m) {
6398
m.doc() = "Test invalid holder access";
6499

65100
#if defined(PYBIND11_CPP14)
66-
m.def("create_vector", []() -> std::unique_ptr<VectorOwns4PythonObjects> {
67-
auto vec = std::make_unique<VectorOwns4PythonObjects>();
68-
vec->append(py::none());
69-
vec->append(py::int_(1));
70-
vec->append(py::str("test"));
71-
vec->append(py::tuple());
101+
m.def("create_vector", [](const py::iterable &iterable) -> std::unique_ptr<VecOwnsObjs> {
102+
auto vec = std::make_unique<VecOwnsObjs>();
103+
for (const auto &item : iterable) {
104+
vec->append(py::reinterpret_borrow<py::object>(item));
105+
}
72106
return vec;
73107
});
74108
#endif
75109

76-
py::class_<VectorOwns4PythonObjects>(
77-
m,
78-
"VectorOwns4PythonObjects",
79-
py::custom_type_setup([](PyHeapTypeObject *heap_type) -> void {
110+
py::class_<VecOwnsObjs>(
111+
m, "VecOwnsObjs", py::custom_type_setup([](PyHeapTypeObject *heap_type) -> void {
80112
auto *const type = &heap_type->ht_type;
81113
type->tp_flags |= Py_TPFLAGS_HAVE_GC;
82-
type->tp_traverse = &VectorOwns4PythonObjects::tp_traverse;
114+
type->tp_traverse = &VecOwnsObjs::tp_traverse;
115+
type->tp_clear = &VecOwnsObjs::tp_clear;
83116
}))
84-
.def("append", &VectorOwns4PythonObjects::append, py::arg("obj"))
85-
.def("set_item", &VectorOwns4PythonObjects::set_item, py::arg("i"), py::arg("obj"))
86-
.def("get_item", &VectorOwns4PythonObjects::get_item, py::arg("i"))
87-
.def("size", &VectorOwns4PythonObjects::size)
88-
.def("is_empty", &VectorOwns4PythonObjects::is_empty)
89-
.def("__setitem__", &VectorOwns4PythonObjects::set_item, py::arg("i"), py::arg("obj"))
90-
.def("__getitem__", &VectorOwns4PythonObjects::get_item, py::arg("i"))
91-
.def("__len__", &VectorOwns4PythonObjects::size)
92-
.def("sanity_check", &VectorOwns4PythonObjects::sanity_check);
117+
.def_static("set_should_check_holder_initialization",
118+
&VecOwnsObjs::set_should_check_holder_initialization,
119+
py::arg("value"))
120+
.def_static("set_should_raise_error_on_set_state",
121+
&VecOwnsObjs::set_should_raise_error_on_set_state,
122+
py::arg("value"))
123+
#if defined(PYBIND11_CPP14)
124+
.def(py::pickle([](const VecOwnsObjs &self) -> py::object { return self.get_state(); },
125+
[](const py::object &state) -> std::unique_ptr<VecOwnsObjs> {
126+
if (!py::isinstance<py::tuple>(state)) {
127+
throw std::runtime_error("Invalid state");
128+
}
129+
auto vec = std::make_unique<VecOwnsObjs>();
130+
if (VecOwnsObjs::get_should_raise_error_on_set_state()) {
131+
throw std::runtime_error("raise error on set_state for testing");
132+
}
133+
for (const auto &item : state) {
134+
vec->append(py::reinterpret_borrow<py::object>(item));
135+
}
136+
return vec;
137+
}),
138+
py::arg("state"))
139+
#endif
140+
.def("append", &VecOwnsObjs::append, py::arg("obj"))
141+
.def("is_empty", &VecOwnsObjs::is_empty)
142+
.def("__setitem__", &VecOwnsObjs::set_item, py::arg("i"), py::arg("obj"))
143+
.def("__getitem__", &VecOwnsObjs::get_item, py::arg("i"))
144+
.def("__len__", &VecOwnsObjs::size);
93145
}

‎tests/test_invalid_holder_access.py‎

Lines changed: 118 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import gc
44
import multiprocessing
5+
import pickle
56
import signal
67
import sys
78
import weakref
@@ -23,24 +24,34 @@
2324
spawn_context = None
2425

2526

26-
@pytest.mark.skipif(
27-
pybind11_tests.cpp_std_num < 14,
28-
reason="std::{unique_ptr,make_unique} not available in C++11",
29-
)
30-
def test_create_vector():
31-
vec = m.create_vector()
32-
assert vec.size() == 4
33-
assert not vec.is_empty()
34-
assert vec[0] is None
35-
assert vec[1] == 1
36-
assert vec[2] == "test"
37-
assert vec[3] == ()
27+
def check_segfault(target):
28+
"""Check if the target function causes a segmentation fault.
29+
30+
The check should be done in a separate process to avoid crashing the main
31+
process with the segfault.
32+
"""
33+
process = spawn_context.Process(target=target)
34+
process.start()
35+
process.join()
36+
rc = abs(process.exitcode)
37+
if 128 < rc < 256:
38+
rc -= 128
39+
assert rc in (
40+
0,
41+
signal.SIGSEGV,
42+
signal.SIGABRT,
43+
0xC0000005, # STATUS_ACCESS_VIOLATION on Windows
44+
)
45+
if rc != 0:
46+
raise SystemError(
47+
"Segmentation Fault: The C++ compiler initializes container incorrectly."
48+
)
3849

3950

4051
def test_no_init():
4152
with pytest.raises(TypeError, match=r"No constructor defined"):
42-
m.VectorOwns4PythonObjects()
43-
vec = m.VectorOwns4PythonObjects.__new__(m.VectorOwns4PythonObjects)
53+
m.VecOwnsObjs()
54+
vec = m.VecOwnsObjs.__new__(m.VecOwnsObjs)
4455
with pytest.raises(TypeError, match=r"No constructor defined"):
4556
vec.__init__()
4657

@@ -50,18 +61,14 @@ def manual_new_target():
5061
# detect uninitialized memory bugs.
5162
for _ in range(32):
5263
# The holder is a pointer variable while the C++ ctor is not called.
53-
vec = m.VectorOwns4PythonObjects.__new__(m.VectorOwns4PythonObjects)
64+
vec = m.VecOwnsObjs.__new__(m.VecOwnsObjs)
5465
if vec.is_empty(): # <= this line could cause a segfault
5566
# The C++ compiler initializes container correctly.
56-
assert vec.size() == 0
67+
assert len(vec) == 0
5768
else:
5869
# The program is not supposed to reach here. It will abort with
5970
# SIGSEGV on `vec.is_empty()`.
6071
sys.exit(signal.SIGSEGV) # manually trigger SIGSEGV if not raised
61-
vec.append(1)
62-
vec.append(2)
63-
vec.append(3)
64-
vec.append(4)
6572

6673

6774
# This test might succeed on some platforms with some compilers, but it is not
@@ -73,35 +80,104 @@ def manual_new_target():
7380
reason="Requires multiprocessing",
7481
)
7582
def test_manual_new():
76-
process = spawn_context.Process(
77-
target=manual_new_target,
78-
name="manual_new_target",
79-
)
80-
process.start()
81-
process.join()
82-
rc = abs(process.exitcode)
83-
if 128 < rc < 256:
84-
rc -= 128
85-
assert rc in (
86-
0,
87-
signal.SIGSEGV,
88-
signal.SIGABRT,
89-
0xC0000005, # STATUS_ACCESS_VIOLATION on Windows
90-
)
91-
if rc != 0:
92-
raise SystemError(
93-
"Segmentation Fault: The C++ compiler initializes container incorrectly."
94-
)
83+
"""
84+
Manually calling the constructor (__new__) of a class that has C++ STL
85+
container will allocate memory without initializing it can cause a
86+
segmentation fault.
87+
"""
88+
check_segfault(manual_new_target)
89+
90+
91+
@pytest.mark.skipif(
92+
pybind11_tests.cpp_std_num < 14,
93+
reason="std::{unique_ptr,make_unique} not available in C++11",
94+
)
95+
def test_set_state_with_error_no_segfault_if_gc_checks_holder_has_initialized():
96+
"""
97+
The ``tp_traverse`` and ``tp_clear`` functions should check if the holder
98+
has been initialized before traversing or clearing the C++ STL container.
99+
"""
100+
m.VecOwnsObjs.set_should_check_holder_initialization(True)
101+
102+
vec = m.create_vector((1, 2, 3, 4))
103+
104+
m.VecOwnsObjs.set_should_raise_error_on_set_state(False)
105+
pickle.loads(pickle.dumps(vec))
106+
107+
# During the unpickling process, Python firstly allocates the object with
108+
# the `__new__` method and then calls the `__setstate__`` method to set the
109+
# state of the object.
110+
#
111+
# obj = cls.__new__(cls)
112+
# obj.__setstate__(state)
113+
#
114+
# The `__init__` method is not called during unpickling.
115+
# So, if the `__setstate__` method raises an error, the object will be in
116+
# an uninitialized state.
117+
m.VecOwnsObjs.set_should_raise_error_on_set_state(True)
118+
serialized = pickle.dumps(vec)
119+
with pytest.raises(
120+
RuntimeError,
121+
match="raise error on set_state for testing",
122+
):
123+
pickle.loads(serialized)
124+
125+
126+
def unpicklable_with_error_target():
127+
m.VecOwnsObjs.set_should_check_holder_initialization(False)
128+
m.VecOwnsObjs.set_should_raise_error_on_set_state(True)
129+
130+
vec = m.create_vector((1, 2, 3, 4))
131+
serialized = pickle.dumps(vec)
132+
133+
# Repeatedly trigger allocation without initialization (raw malloc'ed) to
134+
# detect uninitialized memory bugs.
135+
for _ in range(32):
136+
# During the unpickling process, Python firstly allocates the object with
137+
# the `__new__` method and then calls the `__setstate__`` method to set the
138+
# state of the object.
139+
#
140+
# obj = cls.__new__(cls)
141+
# obj.__setstate__(state)
142+
#
143+
# The `__init__` method is not called during unpickling.
144+
# So, if the `__setstate__` method raises an error, the object will be in
145+
# an uninitialized state. The GC will traverse the uninitialized C++ STL
146+
# container and cause a segmentation fault.
147+
try: # noqa: SIM105
148+
pickle.loads(serialized)
149+
except RuntimeError:
150+
pass
151+
152+
153+
# This test might succeed on some platforms with some compilers, but it is not
154+
# guaranteed to work everywhere. It is marked as non-strict xfail.
155+
@pytest.mark.xfail(reason=XFAIL_REASON, raises=SystemError, strict=False)
156+
@pytest.mark.skipif(spawn_context is None, reason="spawn context not available")
157+
@pytest.mark.skipif(
158+
pybind11_tests.cpp_std_num < 14,
159+
reason="std::{unique_ptr,make_unique} not available in C++11",
160+
)
161+
def test_set_state_with_error_will_segfault_if_gc_does_not_check_holder_has_initialized():
162+
m.VecOwnsObjs.set_should_check_holder_initialization(False)
163+
164+
vec = m.create_vector((1, 2, 3, 4))
165+
166+
m.VecOwnsObjs.set_should_raise_error_on_set_state(False)
167+
pickle.loads(pickle.dumps(vec))
168+
169+
# See comments above.
170+
check_segfault(unpicklable_with_error_target)
95171

96172

97173
@pytest.mark.skipif("env.PYPY or env.GRAALPY", reason="Cannot reliably trigger GC")
98174
@pytest.mark.skipif(
99175
pybind11_tests.cpp_std_num < 14,
100176
reason="std::{unique_ptr,make_unique} not available in C++11",
101177
)
102-
def test_gc_traverse():
103-
vec = m.create_vector()
104-
vec[3] = (vec, vec)
178+
def test_gc():
179+
vec = m.create_vector(())
180+
vec.append((vec, vec))
105181

106182
wr = weakref.ref(vec)
107183
assert wr() is vec

0 commit comments

Comments
 (0)