RustPython · youknowone · May 27, 2022 · Feb 22, 2023
diff --git a/compiler/codegen/src/compile.rs b/compiler/codegen/src/compile.rs
@@ -17,7 +17,7 @@ use num_traits::ToPrimitive;
 use rustpython_ast as ast;
 use rustpython_compiler_core::{
     self as bytecode, Arg as OpArgMarker, CodeObject, ConstantData, Instruction, Location, NameIdx,
-    OpArg, OpArgType,
+    OpArg, OpArgType, UnsafeCodeObject,
 };
 use std::borrow::Cow;
 

diff --git a/compiler/codegen/src/ir.rs b/compiler/codegen/src/ir.rs
@@ -151,6 +151,7 @@ impl CodeInfo {
             locations.clear()
         }
 
+        // SAFETY: we assume that the compiler produces correct output
         CodeObject {
             flags,
             posonlyarg_count,

diff --git a/compiler/core/src/bytecode.rs b/compiler/core/src/bytecode.rs
@@ -62,7 +62,11 @@ impl ConstantBag for BasicBag {
 
 /// Primary container of a single code object. Each python function has
 /// a codeobject. Also a module has a codeobject.
-#[derive(Clone, Serialize, Deserialize)]
+/// Once UnsafeCodeObject is turned into CodeObject, it guarantees correctness of all
+/// aspects of the contained bytecode, in order for the VM to make optimizations
+/// that skip safety checks.
+#[non_exhaustive] // to prevent manual construction by-passing UnsafeCodeObject
+#[derive(Clone, serde::Serialize, serde::Deserialize)]
 pub struct CodeObject<C: Constant = ConstantData> {
     pub instructions: Box<[CodeUnit]>,
     pub locations: Box<[Location]>,
@@ -864,6 +868,17 @@ impl<N: AsRef<str>> fmt::Debug for Arguments<'_, N> {
 }
 
 impl<C: Constant> CodeObject<C> {
+    /// Create a new CodeObject
+    ///
+    /// # Safety
+    ///
+    /// Caller is responsible for ensuring that `code` is correct python bytecode.
+    /// The easiest way to do so is get a CodeObject from the
+    /// `rustpython-compiler` crate.
+    pub unsafe fn new(code: UnsafeCodeObject<C>) -> Self {
+        code.0
+    }
+
     /// Get all arguments of the code object
     /// like inspect.getargs
     pub fn arg_names(&self) -> Arguments<C::Name> {
@@ -1031,6 +1046,16 @@ impl<C: Constant> CodeObject<C> {
             cell2arg: self.cell2arg.clone(),
         }
     }
+
+    /// Serialize this bytecode to bytes.
+    pub fn to_bytes(&self) -> Vec<u8>
+    where
+        C: serde::Serialize,
+        C::Name: serde::Serialize,
+    {
+        let data = bincode::serialize(&self).expect("CodeObject is not serializable");
+        lz4_flex::compress_prepend_size(&data)
+    }
 }
 
 /// Error that occurs during code deserialization
@@ -1054,7 +1079,7 @@ impl fmt::Display for CodeDeserializeError {
 
 impl std::error::Error for CodeDeserializeError {}
 
-impl CodeObject<ConstantData> {
+impl UnsafeCodeObject<ConstantData> {
     /// Load a code object from bytes
     pub fn from_bytes(data: &[u8]) -> Result<Self, CodeDeserializeError> {
         use lz4_flex::block::DecompressError;
@@ -1072,12 +1097,6 @@ impl CodeObject<ConstantData> {
         })?;
         Ok(data)
     }
-
-    /// Serialize this bytecode to bytes.
-    pub fn to_bytes(&self) -> Vec<u8> {
-        let data = bincode::serialize(&self).expect("CodeObject is not serializable");
-        lz4_flex::compress_prepend_size(&data)
-    }
 }
 
 impl<C: Constant> fmt::Display for CodeObject<C> {
@@ -1446,6 +1465,21 @@ impl<C: Constant> InstrDisplayContext for CodeObject<C> {
             .as_ref()
     }
 }
+impl<C: Constant> InstrDisplayContext for UnsafeCodeObject<C> {
+    type Constant = C;
+    fn get_constant(&self, i: usize) -> &C {
+        (**self).get_constant(i)
+    }
+    fn get_name(&self, i: usize) -> &str {
+        (**self).get_name(i)
+    }
+    fn get_varname(&self, i: usize) -> &str {
+        (**self).get_varname(i)
+    }
+    fn get_cellname(&self, i: usize) -> &str {
+        (**self).get_cellname(i)
+    }
+}
 
 impl fmt::Display for ConstantData {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
@@ -1478,7 +1512,10 @@ pub mod frozen_lib {
     use std::io;
 
     /// Decode a library to a iterable of frozen modules
-    pub fn decode_lib(bytes: &[u8]) -> FrozenModulesIter {
+    ///
+    /// # Safety
+    /// `bytes` must be the output from [`encode_lib`].
+    pub unsafe fn decode_lib(bytes: &[u8]) -> FrozenModulesIter {
         let data = lz4_flex::decompress_size_prepended(bytes).unwrap();
         let r = VecReader { data, pos: 0 };
         let mut de = bincode::Deserializer::with_bincode_read(r, options());
@@ -1601,3 +1638,45 @@ pub mod frozen_lib {
         }
     }
 }
+
+/// A wrapper around [`CodeObject`] that doesn't guarantee correctness of all
+/// aspects of the contained bytecode.
+#[derive(serde::Serialize, serde::Deserialize)]
+#[repr(transparent)]
+pub struct UnsafeCodeObject<C: Constant = ConstantData>(
+    #[serde(bound(
+        deserialize = "CodeObject<C>: serde::Deserialize<'de>",
+        serialize = "CodeObject<C>: serde::Serialize"
+    ))]
+    CodeObject<C>,
+);
+
+impl<C: Constant> Clone for UnsafeCodeObject<C>
+where
+    CodeObject<C>: Clone,
+{
+    fn clone(&self) -> Self {
+        Self(self.0.clone())
+    }
+}
+
+impl<C: Constant> fmt::Display for UnsafeCodeObject<C> {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        fmt::Display::fmt(&**self, f)
+    }
+}
+
+impl<C: Constant> fmt::Debug for UnsafeCodeObject<C> {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        fmt::Debug::fmt(&**self, f)
+    }
+}
+
+impl<C: Constant> std::ops::Deref for UnsafeCodeObject<C> {
+    type Target = CodeObject<C>;
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+impl<C: Constant> UnsafeCodeObject<C> {}
diff --git a/derive-impl/src/compile_bytecode.rs b/derive-impl/src/compile_bytecode.rs
@@ -373,8 +373,13 @@ pub fn impl_py_compile(
     let bytes = LitByteStr::new(&bytes, Span::call_site());
 
     let output = quote! {
-        #crate_name::CodeObject::from_bytes(#bytes)
-            .expect("Deserializing CodeObject failed")
+        unsafe {
+            // SAFETY: we just deserialized from a serialized CodeObject
+            #crate_name::CodeObject::new(
+                #crate_name::UnsafeCodeObject::from_bytes(#bytes)
+                    .expect("Deserializing UnsafeCodeObject failed")
+            )
+        }
     };
 
     Ok(output)
@@ -395,7 +400,8 @@ pub fn impl_py_freeze(
     let bytes = LitByteStr::new(&data, Span::call_site());
 
     let output = quote! {
-        #crate_name::frozen_lib::decode_lib(#bytes)
+        // SAFETY: we just deserialized from a serialized CodeObject
+        unsafe { #crate_name::frozen_lib::decode_lib(#bytes) }
     };
 
     Ok(output)

diff --git a/vm/src/builtins/code.rs b/vm/src/builtins/code.rs
@@ -344,33 +344,33 @@ impl PyRef<PyCode> {
             OptionalArg::Missing => self.code.varnames.iter().map(|s| s.to_object()).collect(),
         };
 
-        Ok(PyCode {
-            code: CodeObject {
-                flags: CodeFlags::from_bits_truncate(flags),
-                posonlyarg_count,
-                arg_count,
-                kwonlyarg_count,
-                source_path: source_path.as_object().as_interned_str(vm).unwrap(),
-                first_line_number,
-                obj_name: obj_name.as_object().as_interned_str(vm).unwrap(),
-
-                max_stackdepth: self.code.max_stackdepth,
-                instructions: self.code.instructions.clone(),
-                locations: self.code.locations.clone(),
-                constants: constants.into_iter().map(Literal).collect(),
-                names: names
-                    .into_iter()
-                    .map(|o| o.as_interned_str(vm).unwrap())
-                    .collect(),
-                varnames: varnames
-                    .into_iter()
-                    .map(|o| o.as_interned_str(vm).unwrap())
-                    .collect(),
-                cellvars: self.code.cellvars.clone(),
-                freevars: self.code.freevars.clone(),
-                cell2arg: self.code.cell2arg.clone(),
-            },
-        })
+        // SAFETY: none, really, but this is something cpython lets people do, so ¯\_(ツ)_/¯
+        let code = CodeObject {
+            flags: CodeFlags::from_bits_truncate(flags),
+            posonlyarg_count,
+            arg_count,
+            kwonlyarg_count,
+            source_path: source_path.as_object().as_interned_str(vm).unwrap(),
+            first_line_number,
+            obj_name: obj_name.as_object().as_interned_str(vm).unwrap(),
+
+            max_stackdepth: self.code.max_stackdepth,
+            instructions: self.code.instructions.clone(),
+            locations: self.code.locations.clone(),
+            constants: constants.into_iter().map(Literal).collect(),
+            names: names
+                .into_iter()
+                .map(|o| o.as_interned_str(vm).unwrap())
+                .collect(),
+            varnames: varnames
+                .into_iter()
+                .map(|o| o.as_interned_str(vm).unwrap())
+                .collect(),
+            cellvars: self.code.cellvars.clone(),
+            freevars: self.code.freevars.clone(),
+            cell2arg: self.code.cell2arg.clone(),
+        };
+        Ok(PyCode { code })
     }
 }
 

diff --git a/vm/src/frame.rs b/vm/src/frame.rs
@@ -739,7 +739,7 @@ impl ExecutingFrame<'_> {
                 let item = self.pop_value();
                 let obj = self.nth_value(i.get(arg));
                 let list: &Py<PyList> = unsafe {
-                    // SAFETY: trust compiler
+                    // SAFETY: invariants of CodeObject
                     obj.downcast_unchecked_ref()
                 };
                 list.append(item);
@@ -749,7 +749,7 @@ impl ExecutingFrame<'_> {
                 let item = self.pop_value();
                 let obj = self.nth_value(i.get(arg));
                 let set: &Py<PySet> = unsafe {
-                    // SAFETY: trust compiler
+                    // SAFETY: invariants of CodeObject
                     obj.downcast_unchecked_ref()
                 };
                 set.add(item, vm)?;
@@ -760,7 +760,7 @@ impl ExecutingFrame<'_> {
                 let key = self.pop_value();
                 let obj = self.nth_value(i.get(arg));
                 let dict: &Py<PyDict> = unsafe {
-                    // SAFETY: trust compiler
+                    // SAFETY: invariants of CodeObject
                     obj.downcast_unchecked_ref()
                 };
                 dict.set_item(&*key, value, vm)?;

diff --git a/vm/src/stdlib/marshal.rs b/vm/src/stdlib/marshal.rs
@@ -378,13 +378,16 @@ mod decl {
                     return Err(too_short_error(vm));
                 }
                 let (bytes, buf) = buf.split_at(len);
-                let code = bytecode::CodeObject::from_bytes(bytes).map_err(|e| match e {
+                let code = bytecode::UnsafeCodeObject::from_bytes(bytes).map_err(|e| match e {
                     bytecode::CodeDeserializeError::Eof => vm.new_exception_msg(
                         vm.ctx.exceptions.eof_error.to_owned(),
                         "End of file while deserializing bytecode".to_owned(),
                     ),
                     _ => vm.new_value_error("Couldn't deserialize python bytecode".to_owned()),
                 })?;
+                // SAFETY: none, really 😬 but CPython trusts Python to give it invalid bytecode and do all
+                // kinds of other unsafe actions, so we'll just sorta follow suit
+                let code = unsafe { bytecode::CodeObject::new(code) };
                 (vm.ctx.new_code(code).into(), buf)
             }
         };