guard against arbitrary rust recursion depth via coroutine abuse

Author: kyren

src/fuel.rs

@@ -1,3 +1,9 @@
+use std::ops::{Deref, DerefMut};
+
+use thiserror::Error;
+
+pub const RECURSION_LIMIT: u8 = u8::MAX;
+
 /// A counter for tracking the amount of time spent in `Thread::step` and in callbacks.
 ///
 /// The fuel unit is *approximately* one VM instruction, but this is just a rough estimate
@@ -9,6 +15,7 @@
 pub struct Fuel {
     fuel: i32,
     interrupted: bool,
+    recursion_level: u8,
 }
 
 impl Fuel {
@@ -20,6 +27,7 @@ impl Fuel {
         Self {
             fuel,
             interrupted: false,
+            recursion_level: 0,
         }
     }
 
@@ -79,4 +87,61 @@ impl Fuel {
     pub fn should_continue(&self) -> bool {
         self.fuel > 0 && !self.interrupted
     }
+
+    /// Mark that we are about to run a Rust callback that is potentially controlled by untrusted
+    /// code.
+    ///
+    /// Increments the current recursion level if it is below `RECURSION_LIMIT`. If the recursion
+    /// level would rise above the limit, this returns a recursion error, otherwise returns a
+    /// guard that restores the previous recursion level on drop. This prevents untrusted code from
+    /// consuming arbitrary Rust stack depth and execution time by tricking Rust code into recursing
+    /// endlessly.
+    ///
+    /// By default, the recursion level is automatically incremented whenever `Thread::step` is
+    /// about to trigger a callback, so this is should almost never be necessary to call explicitly.
+    ///
+    /// With the normal stdlib, arbitrary recursion is only possible by (ab)using coroutines. Normal
+    /// Lua recursion and Rust code "calling" Lua code via a `Sequence` poll does not actually use
+    /// the real Rust call stack, and cannot lead to using unbounded time or unbounded Rust stack
+    /// space. Coroutines create their own inner `Thread`s and step them inside a `Sequence`, so
+    /// they can eventually trigger a recursion limit in pathological cases.
+    pub fn recurse(&mut self) -> Result<Recurse<'_>, RecursionLimit> {
+        if self.recursion_level == RECURSION_LIMIT {
+            return Err(RecursionLimit);
+        }
+
+        self.recursion_level += 1;
+        Ok(Recurse(self))
+    }
+
+    /// Returns the current Rust callback recursion level.
+    pub fn recursion_level(&self) -> u8 {
+        self.recursion_level
+    }
+}
+
+#[derive(Debug, Copy, Clone, Error)]
+#[error("callback recursion limit reached")]
+pub struct RecursionLimit;
+
+pub struct Recurse<'a>(&'a mut Fuel);
+
+impl<'a> Drop for Recurse<'a> {
+    fn drop(&mut self) {
+        self.recursion_level -= 1;
+    }
+}
+
+impl<'a> Deref for Recurse<'a> {
+    type Target = Fuel;
+
+    fn deref(&self) -> &Fuel {
+        self.0
+    }
+}
+
+impl<'a> DerefMut for Recurse<'a> {
+    fn deref_mut(&mut self) -> &mut Fuel {
+        self.0
+    }
 }

src/lib.rs

@@ -39,7 +39,7 @@ pub use self::{
     stack::Stack,
     string::{String, StringError},
     table::{InvalidTableKey, Table},
-    thread::{BadThreadMode, Thread, ThreadError, ThreadMode},
+    thread::{BadThreadMode, Thread, ThreadMode, VMError},
     userdata::{AnyUserData, BadUserDataType},
     value::Value,
 };

src/thread/error.rs

@@ -50,13 +50,15 @@ pub struct BadThreadMode {
 }
 
 #[derive(Debug, Copy, Clone, Error)]
-pub enum ThreadError {
+pub enum VMError {
     #[error("{}", if *.0 {
         "operation expects variable stack"
     } else {
         "unexpected variable stack during operation"
     })]
-    ExpectedVariable(bool),
+    ExpectedVariableStack(bool),
     #[error(transparent)]
     BadType(#[from] TypeError),
+    #[error("_ENV upvalue is only allowed on top-level closure")]
+    BadEnvUpValue,
 }

src/thread/mod.rs

@@ -3,7 +3,7 @@ mod thread;
 mod vm;
 
 pub use self::{
-    error::{BadThreadMode, BinaryOperatorError, ThreadError},
+    error::{BadThreadMode, BinaryOperatorError, VMError},
     thread::{Thread, ThreadMode},
 };
 

src/thread/thread.rs

@@ -16,8 +16,7 @@ use crate::{
     meta_ops,
     types::{RegisterIndex, VarCount},
     AnyCallback, AnySequence, BadThreadMode, CallbackReturn, Closure, Context, Error,
-    FromMultiValue, Fuel, Function, IntoMultiValue, SequencePoll, Stack, ThreadError, TypeError,
-    Value,
+    FromMultiValue, Fuel, Function, IntoMultiValue, SequencePoll, Stack, TypeError, VMError, Value,
 };
 
 use super::run_vm;
@@ -195,13 +194,22 @@ impl<'gc> Thread<'gc> {
         while state.mode() == ThreadMode::Normal {
             match state.frames.pop().expect("no frame to step") {
                 Frame::Callback(callback) => {
-                    fuel.consume_fuel(FUEL_PER_CALLBACK);
+                    let mut rfuel = match fuel.recurse() {
+                        Ok(r) => r,
+                        Err(err) => {
+                            state.unwind(&ctx, err.into());
+                            continue;
+                        }
+                    };
+
+                    rfuel.consume_fuel(FUEL_PER_CALLBACK);
                     state.frames.push(Frame::Calling);
 
                     assert!(state.error.is_none());
                     let mut stack = mem::replace(&mut state.external_stack, Stack::new(&ctx));
                     drop(state);
-                    let seq = callback.call(ctx, fuel, &mut stack);
+                    let seq = callback.call(ctx, &mut rfuel, &mut stack);
+                    drop(rfuel);
                     state = self.0.borrow_mut(&ctx);
                     state.external_stack = stack;
 
@@ -216,18 +224,27 @@ impl<'gc> Thread<'gc> {
                     }
                 }
                 Frame::Sequence(mut sequence) => {
-                    fuel.consume_fuel(FUEL_PER_SEQ_STEP);
+                    let mut rfuel = match fuel.recurse() {
+                        Ok(r) => r,
+                        Err(err) => {
+                            state.unwind(&ctx, err.into());
+                            continue;
+                        }
+                    };
+
+                    rfuel.consume_fuel(FUEL_PER_SEQ_STEP);
                     state.frames.push(Frame::Calling);
 
                     let mut stack = mem::replace(&mut state.external_stack, Stack::new(&ctx));
                     let error = state.error.take();
                     drop(state);
                     let fin = if let Some(error) = error {
                         assert!(stack.is_empty());
-                        sequence.error(ctx, fuel, error, &mut stack)
+                        sequence.error(ctx, &mut rfuel, error, &mut stack)
                     } else {
-                        sequence.poll(ctx, fuel, &mut stack)
+                        sequence.poll(ctx, &mut rfuel, &mut stack)
                     };
+                    drop(rfuel);
                     state = self.0.borrow_mut(&ctx);
                     state.external_stack = stack;
 
@@ -414,11 +431,7 @@ impl<'gc, 'a> LuaFrame<'gc, 'a> {
     }
 
     // Place the current frame's varargs at the given register, expecting the given count
-    pub(crate) fn varargs(
-        &mut self,
-        dest: RegisterIndex,
-        count: VarCount,
-    ) -> Result<(), ThreadError> {
+    pub(crate) fn varargs(&mut self, dest: RegisterIndex, count: VarCount) -> Result<(), VMError> {
         match self.state.frames.last_mut() {
             Some(Frame::Lua {
                 bottom,
@@ -427,7 +440,7 @@ impl<'gc, 'a> LuaFrame<'gc, 'a> {
                 ..
             }) => {
                 if *is_variable {
-                    return Err(ThreadError::ExpectedVariable(false));
+                    return Err(VMError::ExpectedVariableStack(false));
                 }
 
                 let varargs_start = *bottom + 1;
@@ -459,7 +472,7 @@ impl<'gc, 'a> LuaFrame<'gc, 'a> {
         mc: &Mutation<'gc>,
         table_base: RegisterIndex,
         count: VarCount,
-    ) -> Result<(), ThreadError> {
+    ) -> Result<(), VMError> {
         let Some(&mut Frame::Lua {
             base,
             ref mut is_variable,
@@ -471,7 +484,7 @@ impl<'gc, 'a> LuaFrame<'gc, 'a> {
         };
 
         if count.is_variable() != *is_variable {
-            return Err(ThreadError::ExpectedVariable(count.is_variable()));
+            return Err(VMError::ExpectedVariableStack(count.is_variable()));
         }
 
         let table_ind = base + table_base.0 as usize;
@@ -527,7 +540,7 @@ impl<'gc, 'a> LuaFrame<'gc, 'a> {
         func: RegisterIndex,
         args: VarCount,
         returns: VarCount,
-    ) -> Result<(), ThreadError> {
+    ) -> Result<(), VMError> {
         match self.state.frames.last_mut() {
             Some(Frame::Lua {
                 expected_return,
@@ -536,7 +549,7 @@ impl<'gc, 'a> LuaFrame<'gc, 'a> {
                 ..
             }) => {
                 if *is_variable != args.is_variable() {
-                    return Err(ThreadError::ExpectedVariable(args.is_variable()));
+                    return Err(VMError::ExpectedVariableStack(args.is_variable()));
                 }
 
                 *expected_return = Some(LuaReturn::Normal(returns));
@@ -598,7 +611,7 @@ impl<'gc, 'a> LuaFrame<'gc, 'a> {
         func: RegisterIndex,
         arg_count: u8,
         returns: VarCount,
-    ) -> Result<(), ThreadError> {
+    ) -> Result<(), VMError> {
         match self.state.frames.last_mut() {
             Some(Frame::Lua {
                 expected_return,
@@ -607,7 +620,7 @@ impl<'gc, 'a> LuaFrame<'gc, 'a> {
                 ..
             }) => {
                 if *is_variable {
-                    return Err(ThreadError::ExpectedVariable(false));
+                    return Err(VMError::ExpectedVariableStack(false));
                 }
 
                 consume_call_fuel(self.fuel, arg_count as usize);
@@ -672,7 +685,7 @@ impl<'gc, 'a> LuaFrame<'gc, 'a> {
         func: Function<'gc>,
         args: &[Value<'gc>],
         ret_index: Option<RegisterIndex>,
-    ) -> Result<(), ThreadError> {
+    ) -> Result<(), VMError> {
         match self.state.frames.last_mut() {
             Some(Frame::Lua {
                 expected_return,
@@ -682,7 +695,7 @@ impl<'gc, 'a> LuaFrame<'gc, 'a> {
                 ..
             }) => {
                 if *is_variable {
-                    return Err(ThreadError::ExpectedVariable(false));
+                    return Err(VMError::ExpectedVariableStack(false));
                 }
 
                 consume_call_fuel(self.fuel, args.len());
@@ -738,7 +751,7 @@ impl<'gc, 'a> LuaFrame<'gc, 'a> {
         ctx: Context<'gc>,
         func: RegisterIndex,
         args: VarCount,
-    ) -> Result<(), ThreadError> {
+    ) -> Result<(), VMError> {
         match self.state.frames.last() {
             Some(&Frame::Lua {
                 bottom,
@@ -747,7 +760,7 @@ impl<'gc, 'a> LuaFrame<'gc, 'a> {
                 ..
             }) => {
                 if is_variable != args.is_variable() {
-                    return Err(ThreadError::ExpectedVariable(args.is_variable()));
+                    return Err(VMError::ExpectedVariableStack(args.is_variable()));
                 }
 
                 self.state.close_upvalues(&ctx, bottom);
@@ -814,7 +827,7 @@ impl<'gc, 'a> LuaFrame<'gc, 'a> {
         mc: &Mutation<'gc>,
         start: RegisterIndex,
         count: VarCount,
-    ) -> Result<(), ThreadError> {
+    ) -> Result<(), VMError> {
         match self.state.frames.pop() {
             Some(Frame::Lua {
                 bottom,
@@ -823,7 +836,7 @@ impl<'gc, 'a> LuaFrame<'gc, 'a> {
                 ..
             }) => {
                 if is_variable != count.is_variable() {
-                    return Err(ThreadError::ExpectedVariable(count.is_variable()));
+                    return Err(VMError::ExpectedVariableStack(count.is_variable()));
                 }
                 self.state.close_upvalues(mc, bottom);
 

src/thread/vm.rs

@@ -11,7 +11,7 @@ use crate::{
     Closure, Constant, Context, Function, RuntimeError, String, Table, Value,
 };
 
-use super::{BinaryOperatorError, LuaFrame};
+use super::{BinaryOperatorError, LuaFrame, VMError};
 
 // Runs the VM for the given number of instructions or until the current LuaFrame may have been
 // changed.
@@ -221,7 +221,7 @@ pub(crate) fn run_vm<'gc>(
                 for &desc in proto.upvalues.iter() {
                     match desc {
                         UpValueDescriptor::Environment => {
-                            panic!("_ENV upvalue is only allowed on top-level closure");
+                            return Err(VMError::BadEnvUpValue.into());
                         }
                         UpValueDescriptor::ParentLocal(reg) => {
                             upvalues.push(registers.open_upvalue(&ctx, reg));

tests/callback.rs

@@ -9,7 +9,8 @@ fn callback() -> Result<(), StaticError> {
     let mut lua = Lua::core();
 
     lua.try_run(|ctx| {
-        let callback = AnyCallback::from_fn(&ctx, |_, _, stack| {
+        let callback = AnyCallback::from_fn(&ctx, |_, fuel, stack| {
+            assert_eq!(fuel.recursion_level(), 1);
             stack.push_back(Value::Integer(42));
             Ok(CallbackReturn::Return)
         });
@@ -22,7 +23,9 @@ fn callback() -> Result<(), StaticError> {
             ctx,
             &br#"
                 local a, b, c = callback(1, 2)
-                return a == 1 and b == 2 and c == 42
+                assert(a == 1 and b == 2 and c == 42)
+                local d, e, f = callback(3, 4)
+                assert(d == 3 and e == 4 and f == 42)
             "#[..],
         )?;
 
@@ -31,7 +34,7 @@ fn callback() -> Result<(), StaticError> {
         Ok(ctx.state.registry.stash(&ctx, thread))
     })?;
 
-    assert!(lua.run_thread::<bool>(&thread)?);
+    lua.run_thread::<()>(&thread)?;
     Ok(())
 }