Add support for POP_MARK

README.md

@@ -13,7 +13,7 @@ for “fickle,” … or something else. Divining its meaning is a personal jour
 in discretion and is left as an exercise to the reader.
 
 Learn more about it in our [blog post](https://blog.trailofbits.com/2021/03/15/never-a-dill-moment-exploiting-machine-learning-pickle-files/)
-and [DEF CON 2021 talk](https://www.youtube.com/watch?v=bZ0m_H_dEJI).
+and [DEF CON AI Village 2021 talk](https://www.youtube.com/watch?v=bZ0m_H_dEJI).
 
 ## Installation
 

fickling/fickle.py

@@ -1032,6 +1032,21 @@ def run(self, interpreter: Interpreter):
         interpreter.stack.pop()
 
 
+class PopMark(Opcode):
+    name = "POP_MARK"
+
+    def run(self, interpreter: Interpreter):
+        objs = []
+        while interpreter.stack:
+            obj = interpreter.stack.pop()
+            if isinstance(obj, MarkObject):
+                break
+            objs.append(obj)
+        else:
+            raise ValueError("Exhausted the stack while searching for a MarkObject!")
+        return objs
+
+
 class ShortBinUnicode(DynamicLength, ConstantOpcode):
     name = "SHORT_BINUNICODE"
     priority = 5000

test/test_crashes.py

@@ -2,9 +2,9 @@
 This test module checks against inputs that previously caused crashes
 """
 
+import io
 from base64 import b64decode
 from functools import wraps
-from pickle import dumps, loads
 from sys import version_info
 from unittest import TestCase
 
@@ -13,7 +13,6 @@
 
 from astunparse import unparse
 
-from fickling import fickle as fpickle
 from fickling.fickle import Pickled
 
 
@@ -53,3 +52,13 @@ class TestCrashes(TestCase):
     def test_stable_diffusion(self):
         """Reproduces https://github.com/trailofbits/fickling/issues/22"""
         pass
+
+    @unparse_test(
+        io.BytesIO(
+            b"\x80\x04\x95\x82\x00\x00\x00\x00\x00\x00\x00(\x8c\x05numpy\x8c\x06poly1d\x93\x94\x8c\x05numpy\x8c\x04size\x93\x94\x8c\x05numpy\x8c\x0c__builtins__\x93\x94h\x00N\x85R\x94h\x03\x94h\x02\x94h\x04(\x8c\x05shapeh\x05dbh\x01h\x04\x8c\x04eval\x86R\x8c\x1d__import__('os').system('id')\x85R1N."
+        )
+    )
+    def test_pop_mark(self):
+        """Tests the correctness of the POP_MARK opcode by using the bytecode from https://github.com/mindspore-ai/mindspore/issues/183
+        This can be simplified to allow for the correctness of additional opcodes to be tested"""
+        pass

test/test_pickle.py

@@ -12,9 +12,9 @@
     from astunparse import unparse
 
 from fickling import fickle as fpickle
-from fickling.cli import main
-from fickling.fickle import Pickled, Interpreter, StackedPickle
 from fickling.analysis import check_safety
+from fickling.cli import main
+from fickling.fickle import Interpreter, Pickled, StackedPickle
 
 
 def get_result(pickled: Pickled):
@@ -177,7 +177,7 @@ def test_insert_stacked(self):
             tmpfile.write(dumps(1234567))
             tmpfile.close()
 
-            # Make sure that it fails if we try and inject into the forth stacked pickle (there are only 3)
+            # Ensure it fails if we try and inject into the fourth stacked pickle (there are only 3)
             self.assertNotEqual(
                 main(["", tmpfile.name, "--inject", 'print("foo")', "--inject-target", "3"]), 0
             )