diff --git a/Lib/test/test_marshal.py b/Lib/test/test_marshal.py index c595e8cf14f1e15..e92ae729ae473c2 100644 --- a/Lib/test/test_marshal.py +++ b/Lib/test/test_marshal.py @@ -317,6 +317,40 @@ def test_recursion_limit(self): last.append([0]) self.assertRaises(ValueError, marshal.dumps, head) + # gh-155901: the limit also applies to nesting through set elements. + head = 0 + for i in range(MAX_MARSHAL_STACK_DEPTH - 1): + head = frozenset({head}) + new_head = marshal.loads(marshal.dumps(head)) + depth = 0 + while isinstance(new_head, frozenset): + new_head, = new_head + depth += 1 + self.assertEqual(depth, MAX_MARSHAL_STACK_DEPTH - 1) + self.assertEqual(new_head, 0) + + head = frozenset({head}) + self.assertRaises(ValueError, marshal.dumps, head) + + def test_nested_sets(self): + # gh-155901: marshalling nested sets took time exponential in the + # nesting depth, because every set element was marshalled twice + # (once to compute a sort key and once to write it out). + obj = frozenset() + for i in range(100): + obj = frozenset({obj, i}) + for version in range(marshal.version + 1): + with self.subTest(version=version): + self.assertEqual(marshal.loads(marshal.dumps(obj, version)), obj) + + # Nested sets shared between elements of the same set. + obj = frozenset() + for i in range(10): + obj = frozenset({obj, (obj, i), i}) + for version in range(marshal.version + 1): + with self.subTest(version=version): + self.assertEqual(marshal.loads(marshal.dumps(obj, version)), obj) + def test_reference_loop_list(self): a = [] a.append(a) @@ -535,6 +569,9 @@ def test_deterministic_sets(self): "float('nan'), b'a', b'b', b'c', 'x', 'y', 'z'", # Also test for bad interactions with backreferencing: "('Spam', 0), ('Spam', 1), ('Spam', 2), ('Spam', 3), ('Spam', 4), ('Spam', 5)", + # Nested sets are sorted only once (gh-155901); make sure + # that the cached order does not depend on the hash seed: + "frozenset({'a', 'b', 'c'}), frozenset({'d', 'e', 'f'}), 'g', 'h', 'i', 'j'", ): s = f"{kind}([{elements}])" with self.subTest(s): diff --git a/Misc/NEWS.d/next/Core_and_Builtins/2026-09-08-00-55-00.gh-issue-155901.IFlnAB.rst b/Misc/NEWS.d/next/Core_and_Builtins/2026-09-08-00-55-00.gh-issue-155901.IFlnAB.rst new file mode 100644 index 000000000000000..b5479e5e6a4c32c --- /dev/null +++ b/Misc/NEWS.d/next/Core_and_Builtins/2026-09-08-00-55-00.gh-issue-155901.IFlnAB.rst @@ -0,0 +1,2 @@ +Fix :func:`marshal.dumps` taking time exponential in the nesting depth when +marshalling nested sets. diff --git a/Python/marshal.c b/Python/marshal.c index 1897d700c055bd3..6dbf662fb2743ff 100644 --- a/Python/marshal.c +++ b/Python/marshal.c @@ -117,6 +117,9 @@ typedef struct { const char *end; char *buf; _Py_hashtable_t *hashtable; + /* Maps sets to tuples of their elements in marshalling order. + See w_set_order(). */ + _Py_hashtable_t *set_orders; int version; int allow_code; } WFILE; @@ -275,7 +278,7 @@ w_short_pstring(const void *s, Py_ssize_t n, WFILE *p) } while(0) static PyObject * -_PyMarshal_WriteObjectToString(PyObject *x, int version, int allow_code); +w_set_order(PyObject *v, WFILE *p); #define _r_digits(bitsize) \ static void \ @@ -659,59 +662,24 @@ w_complex_object(PyObject *v, char flag, WFILE *p) } } else if (PyAnySet_CheckExact(v)) { - PyObject *value; - Py_ssize_t pos = 0; - Py_hash_t hash; - - if (PyFrozenSet_CheckExact(v)) - W_TYPE(TYPE_FROZENSET, p); - else - W_TYPE(TYPE_SET, p); - n = PySet_GET_SIZE(v); - W_SIZE(n, p); // bpo-37596: To support reproducible builds, sets and frozensets need // to have their elements serialized in a consistent order (even when // they have been scrambled by hash randomization). To ensure this, we // use an order equivalent to sorted(v, key=marshal.dumps): - PyObject *pairs = PyList_New(n); - if (pairs == NULL) { - p->error = WFERR_NOMEMORY; + PyObject *order = w_set_order(v, p); + if (order == NULL) { + assert(p->error != WFERR_OK); return; } - Py_ssize_t i = 0; - Py_BEGIN_CRITICAL_SECTION(v); - while (_PySet_NextEntryRef(v, &pos, &value, &hash)) { - PyObject *dump = _PyMarshal_WriteObjectToString(value, - p->version, p->allow_code); - if (dump == NULL) { - p->error = WFERR_UNMARSHALLABLE; - Py_DECREF(value); - break; - } - PyObject *pair = _PyTuple_FromPairSteal(dump, value); - if (pair == NULL) { - p->error = WFERR_NOMEMORY; - break; - } - PyList_SET_ITEM(pairs, i++, pair); - } - Py_END_CRITICAL_SECTION(); - if (p->error == WFERR_UNMARSHALLABLE || p->error == WFERR_NOMEMORY) { - Py_DECREF(pairs); - return; - } - assert(i == n); - if (PyList_Sort(pairs)) { - p->error = WFERR_NOMEMORY; - Py_DECREF(pairs); - return; - } - for (Py_ssize_t i = 0; i < n; i++) { - PyObject *pair = PyList_GET_ITEM(pairs, i); - value = PyTuple_GET_ITEM(pair, 1); - w_object(value, p); + if (PyFrozenSet_CheckExact(v)) + W_TYPE(TYPE_FROZENSET, p); + else + W_TYPE(TYPE_SET, p); + n = PyTuple_GET_SIZE(order); + W_SIZE(n, p); + for (i = 0; i < n; i++) { + w_object(PyTuple_GET_ITEM(order, i), p); } - Py_DECREF(pairs); } else if (PyCode_Check(v)) { if (!p->allow_code) { @@ -806,6 +774,149 @@ w_clear_refs(WFILE *wf) } } +/* Marshal a set element on its own, to be used as a sort key by + w_set_order(). The nested writer has its own reference table, so the + key does not depend on the position of the element in the stream, but + it shares the nesting depth and the set order cache with p. + Return a new reference, or NULL with p->error set. */ +static PyObject * +w_set_key(PyObject *value, WFILE *p) +{ + WFILE wf; + + assert(p->set_orders != NULL); + memset(&wf, 0, sizeof(wf)); + wf.str = PyBytes_FromStringAndSize((char *)NULL, 50); + if (wf.str == NULL) { + p->error = WFERR_NOMEMORY; + return NULL; + } + wf.ptr = wf.buf = PyBytes_AS_STRING(wf.str); + wf.end = wf.ptr + PyBytes_GET_SIZE(wf.str); + wf.error = WFERR_OK; + wf.depth = p->depth; + wf.version = p->version; + wf.allow_code = p->allow_code; + wf.set_orders = p->set_orders; + if (w_init_refs(&wf, wf.version) < 0) { + Py_DECREF(wf.str); + p->error = WFERR_NOMEMORY; + return NULL; + } + w_object(value, &wf); + w_clear_refs(&wf); + if (wf.error != WFERR_OK) { + Py_XDECREF(wf.str); + p->error = wf.error; + return NULL; + } + if (wf.str == NULL) { + /* w_reserve() failed to grow the buffer. */ + p->error = WFERR_NOMEMORY; + return NULL; + } + const char *base = PyBytes_AS_STRING(wf.str); + if (_PyBytes_Resize(&wf.str, (Py_ssize_t)(wf.ptr - base)) < 0) { + p->error = WFERR_NOMEMORY; + return NULL; + } + return wf.str; +} + +/* Return the elements of the set v in the order in which they are + marshalled: sorted by their marshalled form, so that the output does + not depend on hash randomization (see bpo-37596). + + The order is cached in p->set_orders. The sort key of an element is + its own marshalled form, so without the cache a nested set would be + sorted again, and its elements marshalled again, for every enclosing + set, taking time exponential in the nesting depth (see gh-155901). + + The cache does not own a reference to the set: it is kept alive by the + object being marshalled, and an extra reference would make w_ref() + treat a uniquely referenced set as shared, changing the output. + + Return a borrowed reference to a tuple, or NULL with p->error set. */ +static PyObject * +w_set_order(PyObject *v, WFILE *p) +{ + if (p->set_orders == NULL) { + p->set_orders = _Py_hashtable_new_full(_Py_hashtable_hash_ptr, + _Py_hashtable_compare_direct, + NULL, w_decref_entry, NULL); + if (p->set_orders == NULL) { + p->error = WFERR_NOMEMORY; + return NULL; + } + } + PyObject *order = _Py_hashtable_get(p->set_orders, v); + if (order != NULL) { + return order; + } + + Py_ssize_t n = PySet_GET_SIZE(v); + PyObject *pairs = PyList_New(n); + if (pairs == NULL) { + p->error = WFERR_NOMEMORY; + return NULL; + } + Py_ssize_t i = 0; + Py_ssize_t pos = 0; + PyObject *value; + Py_hash_t hash; + Py_BEGIN_CRITICAL_SECTION(v); + while (_PySet_NextEntryRef(v, &pos, &value, &hash)) { + PyObject *key = w_set_key(value, p); + if (key == NULL) { + Py_DECREF(value); + break; + } + PyObject *pair = _PyTuple_FromPairSteal(key, value); + if (pair == NULL) { + p->error = WFERR_NOMEMORY; + break; + } + PyList_SET_ITEM(pairs, i++, pair); + } + Py_END_CRITICAL_SECTION(); + if (p->error != WFERR_OK) { + Py_DECREF(pairs); + return NULL; + } + assert(i == n); + if (PyList_Sort(pairs) < 0) { + p->error = WFERR_NOMEMORY; + Py_DECREF(pairs); + return NULL; + } + order = PyTuple_New(n); + if (order == NULL) { + p->error = WFERR_NOMEMORY; + Py_DECREF(pairs); + return NULL; + } + for (i = 0; i < n; i++) { + PyObject *pair = PyList_GET_ITEM(pairs, i); + PyTuple_SET_ITEM(order, i, Py_NewRef(PyTuple_GET_ITEM(pair, 1))); + } + Py_DECREF(pairs); + if (_Py_hashtable_set(p->set_orders, v, order) < 0) { + Py_DECREF(order); + p->error = WFERR_NOMEMORY; + return NULL; + } + return order; +} + +static void +w_clear_set_orders(WFILE *wf) +{ + if (wf->set_orders != NULL) { + _Py_hashtable_destroy(wf->set_orders); + wf->set_orders = NULL; + } +} + /* Set the exception indicator according to the recorded error. */ static void w_set_exception(WFILE *p) @@ -874,6 +985,7 @@ PyMarshal_WriteObjectToFile(PyObject *x, FILE *fp, int version) } w_object(x, &wf); w_clear_refs(&wf); + w_clear_set_orders(&wf); w_flush(&wf); if (wf.error != WFERR_OK) { w_set_exception(&wf); @@ -2018,6 +2130,7 @@ _PyMarshal_WriteObjectToString(PyObject *x, int version, int allow_code) } w_object(x, &wf); w_clear_refs(&wf); + w_clear_set_orders(&wf); if (wf.str != NULL) { const char *base = PyBytes_AS_STRING(wf.str); if (_PyBytes_Resize(&wf.str, (Py_ssize_t)(wf.ptr - base)) < 0)