[Python-checkins] bpo-33720: Improve tests for the stack overflow in marshal.loads(). (GH-7336)

[Miss Islington (bot)]

https://github.com/python/cpython/commit/878c4fe40e8954372e9bf80ae555045ecae68e73
commit: 878c4fe40e8954372e9bf80ae555045ecae68e73
branch: 3.6
author: Miss Islington (bot) <31488909+miss-islington at users.noreply.github.com>
committer: GitHub <noreply at github.com>
date: 2018-07-05T01:48:47-07:00
summary:

bpo-33720: Improve tests for the stack overflow in marshal.loads(). (GH-7336)

(cherry picked from commit fc05e68d8fac70349b7ea17ec14e7e0cfa956121)

Co-authored-by: Serhiy Storchaka <storchaka at gmail.com>

files:
M Lib/test/test_marshal.py

diff --git a/Lib/test/test_marshal.py b/Lib/test/test_marshal.py
index 29dda987d0bb..c40cdbfa5bee 100644
--- a/Lib/test/test_marshal.py
+++ b/Lib/test/test_marshal.py
@@ -210,13 +210,24 @@ def test_fuzz(self):
             except Exception:
                 pass
 
-    def test_loads_2x_code(self):
-        s = b'c' + (b'X' * 4*4) + b'{' * 2**20
-        self.assertRaises(ValueError, marshal.loads, s)
-
     def test_loads_recursion(self):
-        s = b'c' + (b'X' * 4*5) + b'{' * 2**20
-        self.assertRaises(ValueError, marshal.loads, s)
+        def run_tests(N, check):
+            # (((...None...),),)
+            check(b')\x01' * N + b'N')
+            check(b'(\x01\x00\x00\x00' * N + b'N')
+            # [[[...None...]]]
+            check(b'[\x01\x00\x00\x00' * N + b'N')
+            # {None: {None: {None: ...None...}}}
+            check(b'{N' * N + b'N' + b'0' * N)
+            # frozenset([frozenset([frozenset([...None...])])])
+            check(b'>\x01\x00\x00\x00' * N + b'N')
+        # Check that the generated marshal data is valid and marshal.loads()
+        # works for moderately deep nesting
+        run_tests(100, marshal.loads)
+        # Very deeply nested structure shouldn't blow the stack
+        def check(s):
+            self.assertRaises(ValueError, marshal.loads, s)
+        run_tests(2**20, check)
 
     def test_recursion_limit(self):
         # Create a deeply nested structure.