Fix miscompilation for maliciously-crafted use of locals

src/backend.rs

@@ -263,15 +263,32 @@ impl Value {
     }
 }
 
+/// A value on the logical stack. The logical stack is the value stack as it
+/// is visible to the WebAssembly, whereas the physical stack is the stack as
+/// it exists on the machine (i.e. as offsets in memory relative to `rsp`).
 #[derive(Debug, Copy, Clone, PartialEq, Eq)]
 enum StackValue {
+    /// A local (arguments are also locals). Note that when setting a local
+    /// any values of this local stored on the stack are rendered invalid.
+    /// See `manifest_local` for how we manage this.
     Local(u32),
+    /// A temporary, stored in a register
     Temp(GPR),
+    /// An immediate value, created with `i32.const`/`i64.const` etc.
     Immediate(i64),
+    /// This value is on the physical stack and so should be accessed
+    /// with the `pop` instruction.
     Pop,
 }
 
 impl StackValue {
+    /// Returns either the location that this value can be accessed at
+    /// if possible. If this value is `Pop`, you can only access it by
+    /// popping the physical stack and so this function returns `None`.
+    ///
+    /// Of course, we could calculate the location of the value on the
+    /// physical stack, but that would be unncessary computation for
+    /// our usecases.
     fn location(&self, locals: &Locals) -> Option<ValueLocation> {
         match *self {
             StackValue::Local(loc) => Some(locals.get(loc)),
@@ -282,13 +299,20 @@ impl StackValue {
     }
 }
 
+/// A store for the local values for our function, including arguments.
 #[derive(Debug, Default, Clone)]
-struct Locals {
+pub struct Locals {
     // TODO: Store all places that the value can be read, so we can optimise
     //       passing (register) arguments along into a noop after saving their
     //       values.
+    /// All arguments in registers, check `ARGS_IN_GPRS` for the list of
+    /// registers that this can contain. If we need to move the argument
+    /// out of a register (for example, because we're calling a function)
+    /// we note that down here, so we don't have to move it back afterwards.
     register_arguments: ArrayVec<[ValueLocation; ARGS_IN_GPRS.len()]>,
+    /// The number of arguments stored on the stack.
     num_stack_args: u32,
+    /// The number of local stack slots, i.e. the amount of stack space reserved for locals.
     num_local_stack_slots: u32,
 }
 
@@ -340,10 +364,10 @@ pub struct BlockState {
     return_register: Option<GPR>,
     regs: Registers,
     /// This is the _current_ locals, since we can shuffle them about during function calls.
-    /// We will restore this to be the same state as the `Locals` in `Context<T>` at the end
+    pub locals: Locals,
-    /// of a block.
+    /// In non-linear control flow (ifs and loops) we have to set the locals to the state that
-    locals: Locals,
+    /// the next block that we enter will expect them in.
-    end_locals: Option<Locals>,
+    pub end_locals: Option<Locals>,
 }
 
 type Stack = Vec<StackValue>;
@@ -366,7 +390,7 @@ pub struct Context<'a> {
     asm: &'a mut Assembler,
     func_starts: &'a Vec<(Option<AssemblyOffset>, DynamicLabel)>,
     /// Each push and pop on the value stack increments or decrements this value by 1 respectively.
-    block_state: BlockState,
+    pub block_state: BlockState,
     has_memory: bool,
 }
 
@@ -1064,13 +1088,7 @@ impl Context<'_> {
 
     // We use `put` instead of `pop` since with `BrIf` it's possible
     // that the block will continue after returning.
-    pub fn return_from_block(&mut self, arity: u32, is_function_end: bool) {
+    pub fn return_from_block(&mut self, arity: u32) {
-        // This should just be an optimisation, passing `false` should always result
-        // in correct code.
-        if !is_function_end {
-            self.restore_locals();
-        }
-
         if arity == 0 {
             return;
         }
@@ -1092,7 +1110,7 @@ impl Context<'_> {
         }
     }
 
-    pub fn start_block(&mut self, is_loop: bool) -> BlockState {
+    pub fn start_block(&mut self) -> BlockState {
         use std::mem;
 
         // OPTIMISATION: We cannot use the parent's stack values (it is disallowed by the spec)
@@ -1105,10 +1123,6 @@ impl Context<'_> {
         let mut current_state = self.block_state.clone();
         current_state.stack = out_stack;
 
-        if is_loop {
-            self.block_state.end_locals = Some(current_state.locals.clone());
-        }
-
         self.block_state.return_register = None;
         current_state
     }
@@ -1140,6 +1154,8 @@ impl Context<'_> {
             );
         }
 
+        self.restore_locals();
+
         let return_reg = self.block_state.return_register;
         let locals = mem::replace(&mut self.block_state.locals, Default::default());
         self.block_state = parent_block_state;
@@ -1153,33 +1169,32 @@ impl Context<'_> {
         }
     }
 
-    fn restore_locals(&mut self) {
+    pub fn restore_locals(&mut self) {
-        if let Some(end_registers) = self
+        if let Some(end_locals) = self.block_state.end_locals.clone() {
-            .block_state
+            self.restore_locals_to(&end_locals);
-            .end_locals
+        }
-            .as_ref()
+    }
-            .map(|l| l.register_arguments.clone())
-        {
-            for (src, dst) in self
-                .block_state
-                .locals
-                .register_arguments
-                .clone()
-                .iter()
-                .zip(&end_registers)
-            {
-                self.copy_value(*src, *dst);
-            }
 
-            for (src, dst) in self
+    pub fn restore_locals_to(&mut self, locals: &Locals) {
-                .block_state
+        for (src, dst) in self
-                .locals
+            .block_state
-                .register_arguments
+            .locals
-                .iter_mut()
+            .register_arguments
-                .zip(&end_registers)
+            .clone()
-            {
+            .iter()
-                *src = *dst;
+            .zip(&locals.register_arguments)
-            }
+        {
+            self.copy_value(*src, *dst);
+        }
+
+        for (src, dst) in self
+            .block_state
+            .locals
+            .register_arguments
+            .iter_mut()
+            .zip(&locals.register_arguments)
+        {
+            *src = *dst;
         }
     }
 

src/function_body.rs

@@ -40,6 +40,13 @@ impl ControlFrameKind {
         }
     }
 
+    fn is_loop(&self) -> bool {
+        match *self {
+            ControlFrameKind::Loop { .. } => true,
+            _ => false,
+        }
+    }
+
     fn branch_target(&self) -> Label {
         match *self {
             ControlFrameKind::Block { end_label } => end_label,
@@ -101,13 +108,30 @@ pub fn translate(
         control_frames: &mut Vec<ControlFrame>,
         idx: usize,
     ) {
-        control_frames
+        let control_frame = control_frames.get_mut(idx).expect("wrong depth");
-            .last_mut()
-            .expect("Control stack is empty!")
-            .mark_unreachable();
 
-        let control_frame = control_frames.get(idx).expect("wrong depth");
+        if control_frame.kind.is_loop() {
-        ctx.return_from_block(control_frame.arity(), idx == 0);
+            ctx.restore_locals_to(&control_frame.block_state.locals);
+        } else {
+            // We can't do any execution after the function end so we just skip this logic
+            // if we're breaking out of the whole function.
+            if idx != 0 {
+                // Workaround for borrowck limitations
+                let should_set = if let Some(locals) = control_frame.block_state.end_locals.as_ref()
+                {
+                    ctx.restore_locals_to(locals);
+                    false
+                } else {
+                    true
+                };
+
+                if should_set {
+                    control_frame.block_state.end_locals = Some(ctx.block_state.locals.clone());
+                }
+            }
+
+            ctx.return_from_block(control_frame.arity());
+        }
 
         ctx.br(control_frame.kind.branch_target());
     }
@@ -141,7 +165,7 @@ pub fn translate(
     // Upon entering the function implicit frame for function body is pushed. It has the same
     // result type as the function itself. Branching to it is equivalent to returning from the function.
     let epilogue_label = ctx.create_label();
-    let function_block_state = ctx.start_block(false);
+    let function_block_state = ctx.start_block();
     control_frames.push(ControlFrame::new(
         ControlFrameKind::Block {
             end_label: epilogue_label,
@@ -180,7 +204,7 @@ pub fn translate(
             }
             Operator::Block { ty } => {
                 let label = ctx.create_label();
-                let state = ctx.start_block(false);
+                let state = ctx.start_block();
                 control_frames.push(ControlFrame::new(
                     ControlFrameKind::Block { end_label: label },
                     state,
@@ -188,23 +212,31 @@ pub fn translate(
                 ));
             }
             Operator::Return => {
+                control_frames
+                    .last_mut()
+                    .expect("Control stack is empty!")
+                    .mark_unreachable();
+
                 break_from_control_frame_with_id(ctx, &mut control_frames, 0);
             }
             Operator::Br { relative_depth } => {
+                control_frames
+                    .last_mut()
+                    .expect("Control stack is empty!")
+                    .mark_unreachable();
+
                 let idx = control_frames.len() - 1 - relative_depth as usize;
 
                 break_from_control_frame_with_id(ctx, &mut control_frames, idx);
             }
             Operator::BrIf { relative_depth } => {
                 let idx = control_frames.len() - 1 - relative_depth as usize;
-                let control_frame = control_frames.get(idx).expect("wrong depth");
 
                 let if_not = ctx.create_label();
 
                 ctx.jump_if_false(if_not);
 
-                ctx.return_from_block(control_frame.arity(), idx == 0);
+                break_from_control_frame_with_id(ctx, &mut control_frames, idx);
-                ctx.br(control_frame.kind.branch_target());
 
                 ctx.define_label(if_not);
             }
@@ -213,7 +245,7 @@ pub fn translate(
                 let if_not = ctx.create_label();
 
                 ctx.jump_if_false(if_not);
-                let state = ctx.start_block(false);
+                let state = ctx.start_block();
 
                 control_frames.push(ControlFrame::new(
                     ControlFrameKind::IfTrue { end_label, if_not },
@@ -225,7 +257,7 @@ pub fn translate(
                 let header = ctx.create_label();
 
                 ctx.define_label(header);
-                let state = ctx.start_block(true);
+                let state = ctx.start_block();
 
                 control_frames.push(ControlFrame::new(
                     ControlFrameKind::Loop { header },
@@ -241,7 +273,7 @@ pub fn translate(
                         block_state,
                         ..
                     }) => {
-                        ctx.return_from_block(arity(ty), false);
+                        ctx.return_from_block(arity(ty));
                         ctx.reset_block(block_state.clone());
 
                         // Finalize `then` block by jumping to the `end_label`.
@@ -280,7 +312,7 @@ pub fn translate(
 
                 // Don't bother generating this code if we're in unreachable code
                 if !control_frame.unreachable {
-                    ctx.return_from_block(arity, control_frames.is_empty());
+                    ctx.return_from_block(arity);
                 }
 
                 let block_end = control_frame.kind.block_end();