Repository: bytecodealliance/wasmtime
Pull request #7590 (2023-11-27 23:07:35)

Winch: cleanup stack in br_if in non-fallthrough case
<!--
Please make sure you include the following information:

- If this work has been discussed elsewhere, please include a link to that
  conversation. If it was discussed in an issue, just mention "issue #...".

- Explain why this change is needed. If the details are in an issue already,
  this can be brief.

Our development process is documented in the Wasmtime book:
https://docs.wasmtime.dev/contributing-development-process.html

Please ensure all communication follows the code of conduct:
https://github.com/bytecodealliance/wasmtime/blob/main/CODE_OF_CONDUCT.md
-->
I discovered that executing the following module compiled with Winch causes a segmentation fault:

```wat
(module
  (func (export "") 
    call 1
    call 1
    br_if 0
    drop
  )
  (func (;1;) (result i32)
    i32.const 1
  )
)
```

The cause appears to be that the code that would increment `rsp` to balance the machine stack is jumped over when the `br_if` path is taken.

What appears to work is, if the frame's `sp_offset` is different than the current `sp_offset`, we invert the existing jump and execute instructions to balance the stack before jumping to the frame's label, and reset the `sp_offset` before continuing with the fallthrough code (falling through in terms of the `br_if`, not the machine code jump). In the case where there is no difference between `sp_offset`s, we can use the existing `jne` machine code emission since we don't need to run any stack balancing instructions.

In terms of machine instructions, we'd be changing

```
sub rsp, 4
test <reg> <reg>
jne end
add rsp, 4
end:
add rsp, 8
pop rbp
ret
```

to

```
sub rsp, 4
test <reg> <reg>
jeq fallthrough
add rsp, 4
jmp end
fallthrough:
add rsp, 4
end:
add rsp, 8
pop rbp
ret
```

github-actions[bot]: #### Subscribe to Label Action

cc @saulecabrera

<details>
This issue or pull request has been labeled: "winch"

Thus the following users have been cc'd because of the following labels:

* saulecabrera: winch

To subscribe or unsubscribe from this label, edit the <code>.github/subscribe-to-label.json</code> configuration file.

[Learn more.](https://github.com/bytecodealliance/subscribe-to-label-action)
</details>

Patch:
diff --git a/crates/c-api/include/wasmtime/config.h b/crates/c-api/include/wasmtime/config.h
index b299d048e..eb17b5e25 100644
--- a/crates/c-api/include/wasmtime/config.h
+++ b/crates/c-api/include/wasmtime/config.h
@@ -367,6 +367,17 @@ WASM_API_EXTERN void wasmtime_config_cranelift_flag_enable(wasm_config_t*, const
  */
 WASM_API_EXTERN void wasmtime_config_cranelift_flag_set(wasm_config_t*, const char *key, const char *value);
 
+/**
+ * \brief Configures whether, when on macOS, Mach ports are used for exception handling
+ * instead of traditional Unix-based signal handling.
+ *
+ * This option defaults to true, using Mach ports by default.
+ *
+ * For more information see the Rust documentation at
+ * https://docs.wasmtime.dev/api/wasmtime/struct.Config.html#method.macos_use_mach_ports
+ */
+WASMTIME_CONFIG_PROP(void, macos_use_mach_ports, bool)
+
 
 /**
  * Return the data from a LinearMemory instance.
diff --git a/crates/c-api/src/config.rs b/crates/c-api/src/config.rs
index 201e05196..a2c404810 100644
--- a/crates/c-api/src/config.rs
+++ b/crates/c-api/src/config.rs
@@ -240,6 +240,11 @@ pub unsafe extern "C" fn wasmtime_config_target_set(
     handle_result(c.config.target(target), |_cfg| {})
 }
 
+#[no_mangle]
+pub extern "C" fn wasmtime_config_macos_use_mach_ports_set(c: &mut wasm_config_t, enabled: bool) {
+    c.config.macos_use_mach_ports(enabled);
+}
+
 #[no_mangle]
 pub unsafe extern "C" fn wasmtime_config_cranelift_flag_enable(
     c: &mut wasm_config_t,
diff --git a/crates/runtime/src/mpk/enabled.rs b/crates/runtime/src/mpk/enabled.rs
index f8a045d3a..2ce2de8ea 100644
--- a/crates/runtime/src/mpk/enabled.rs
+++ b/crates/runtime/src/mpk/enabled.rs
@@ -152,6 +152,9 @@ mod tests {
     #[test]
     fn check_is_supported() {
         println!("is pku supported = {}", is_supported());
+        if std::env::var("WASMTIME_TEST_FORCE_MPK").is_ok() {
+            assert!(is_supported());
+        }
     }
 
     #[test]
diff --git a/winch/codegen/src/masm.rs b/winch/codegen/src/masm.rs
index c33f701c5..b4613635a 100644
--- a/winch/codegen/src/masm.rs
+++ b/winch/codegen/src/masm.rs
@@ -24,7 +24,7 @@ pub(crate) enum RemKind {
 }
 
 /// Representation of the stack pointer offset.
-#[derive(Copy, Clone, Eq, PartialEq, Debug)]
+#[derive(Copy, Clone, Eq, PartialEq, Debug, PartialOrd, Ord)]
 pub struct SPOffset(u32);
 
 impl SPOffset {
diff --git a/winch/codegen/src/visitor.rs b/winch/codegen/src/visitor.rs
index 8eebf54c9..8a26cda94 100644
--- a/winch/codegen/src/visitor.rs
+++ b/winch/codegen/src/visitor.rs
@@ -1189,14 +1189,31 @@ where
             self.context.pop_to_reg(self.masm, None)
         };
 
-        self.masm.branch(
-            IntCmpKind::Ne,
-            top.reg.into(),
-            top.reg.into(),
-            *frame.label(),
-            OperandSize::S32,
-        );
+        // Emit instructions to balance the machine stack if the frame has
+        // a different offset.
+        let current_sp_offset = self.masm.sp_offset();
+        let (_, frame_sp_offset) = frame.base_stack_len_and_sp();
+        let (label, cmp, needs_cleanup) = if current_sp_offset > frame_sp_offset {
+            (self.masm.get_label(), IntCmpKind::Eq, true)
+        } else {
+            (*frame.label(), IntCmpKind::Ne, false)
+        };
+
+        self.masm
+            .branch(cmp, top.reg.into(), top.reg.into(), label, OperandSize::S32);
         self.context.free_reg(top);
+
+        if needs_cleanup {
+            // Emit instructions to balance the stack and jump if not falling
+            // through.
+            self.masm.ensure_sp_for_jump(frame_sp_offset);
+            self.masm.jmp(*frame.label());
+
+            // Restore sp_offset to what it was for falling through and emit
+            // fallthrough label.
+            self.masm.reset_stack_pointer(current_sp_offset);
+            self.masm.bind(label);
+        }
     }
 
     fn visit_br_table(&mut self, targets: BrTable<'a>) {


Test patch:
diff --git a/tests/all/main.rs b/tests/all/main.rs
index cffe15555..e308b4b43 100644
--- a/tests/all/main.rs
+++ b/tests/all/main.rs
@@ -92,6 +92,12 @@ pub(crate) fn small_pool_config() -> wasmtime::PoolingAllocationConfig {
     config.total_tables(1);
     config.table_elements(10);
 
+    // When testing, we may choose to start with MPK force-enabled to ensure
+    // we use that functionality.
+    if std::env::var("WASMTIME_TEST_FORCE_MPK").is_ok() {
+        config.memory_protection_keys(wasmtime_runtime::MpkEnabled::Enable);
+    }
+
     #[cfg(feature = "async")]
     config.total_stacks(1);
 
diff --git a/tests/all/wast.rs b/tests/all/wast.rs
index 5b93799ab..01f184cba 100644
--- a/tests/all/wast.rs
+++ b/tests/all/wast.rs
@@ -7,6 +7,7 @@ use wasmtime::{
     Config, Engine, InstanceAllocationStrategy, PoolingAllocationConfig, Store, Strategy,
 };
 use wasmtime_environ::WASM_PAGE_SIZE;
+use wasmtime_runtime::MpkEnabled;
 use wasmtime_wast::WastContext;
 
 include!(concat!(env!("OUT_DIR"), "/wast_testsuite_tests.rs"));
@@ -133,6 +134,13 @@ fn run_wast(wast: &str, strategy: Strategy, pooling: bool) -> anyhow::Result<()>
             .memory_pages(805)
             .max_memories_per_module(if multi_memory { 9 } else { 1 })
             .max_tables_per_module(4);
+
+        // When testing, we may choose to start with MPK force-enabled to ensure
+        // we use that functionality.
+        if std::env::var("WASMTIME_TEST_FORCE_MPK").is_ok() {
+            pool.memory_protection_keys(MpkEnabled::Enable);
+        }
+
         cfg.allocation_strategy(InstanceAllocationStrategy::Pooling(pool));
         Some(lock_pooling())
     } else {
diff --git a/winch/filetests/filetests/x64/br_if/with_machine_stack_entry.wat b/winch/filetests/filetests/x64/br_if/with_machine_stack_entry.wat
new file mode 100644
index 000000000..0ecaac9b0
--- /dev/null
+++ b/winch/filetests/filetests/x64/br_if/with_machine_stack_entry.wat
@@ -0,0 +1,42 @@
+;;! target = "x86_64"
+
+(module
+  (func (export "") 
+    call 1
+    call 1
+    br_if 0
+    drop
+  )
+  (func (;1;) (result i32)
+    i32.const 1
+  )
+)
+;;    0:	 55                   	push	rbp
+;;    1:	 4889e5               	mov	rbp, rsp
+;;    4:	 4883ec08             	sub	rsp, 8
+;;    8:	 4c893424             	mov	qword ptr [rsp], r14
+;;    c:	 4883ec08             	sub	rsp, 8
+;;   10:	 e800000000           	call	0x15
+;;   15:	 4883c408             	add	rsp, 8
+;;   19:	 4883ec04             	sub	rsp, 4
+;;   1d:	 890424               	mov	dword ptr [rsp], eax
+;;   20:	 4883ec04             	sub	rsp, 4
+;;   24:	 e800000000           	call	0x29
+;;   29:	 4883c404             	add	rsp, 4
+;;   2d:	 85c0                 	test	eax, eax
+;;   2f:	 0f8409000000         	je	0x3e
+;;   35:	 4883c404             	add	rsp, 4
+;;   39:	 e904000000           	jmp	0x42
+;;   3e:	 4883c404             	add	rsp, 4
+;;   42:	 4883c408             	add	rsp, 8
+;;   46:	 5d                   	pop	rbp
+;;   47:	 c3                   	ret	
+;;
+;;    0:	 55                   	push	rbp
+;;    1:	 4889e5               	mov	rbp, rsp
+;;    4:	 4883ec08             	sub	rsp, 8
+;;    8:	 4c893424             	mov	qword ptr [rsp], r14
+;;    c:	 b801000000           	mov	eax, 1
+;;   11:	 4883c408             	add	rsp, 8
+;;   15:	 5d                   	pop	rbp
+;;   16:	 c3                   	ret	
diff --git a/winch/filetests/filetests/x64/call/reg_on_stack.wat b/winch/filetests/filetests/x64/call/reg_on_stack.wat
index eb73009df..05f1530c8 100644
--- a/winch/filetests/filetests/x64/call/reg_on_stack.wat
+++ b/winch/filetests/filetests/x64/call/reg_on_stack.wat
@@ -35,8 +35,10 @@
 ;;   56:	 8b0424               	mov	eax, dword ptr [rsp]
 ;;   59:	 4883c404             	add	rsp, 4
 ;;   5d:	 85c9                 	test	ecx, ecx
-;;   5f:	 0f8502000000         	jne	0x67
-;;   65:	 0f0b                 	ud2	
-;;   67:	 4883c410             	add	rsp, 0x10
-;;   6b:	 5d                   	pop	rbp
-;;   6c:	 c3                   	ret	
+;;   5f:	 0f8409000000         	je	0x6e
+;;   65:	 4883c404             	add	rsp, 4
+;;   69:	 e902000000           	jmp	0x70
+;;   6e:	 0f0b                 	ud2	
+;;   70:	 4883c410             	add	rsp, 0x10
+;;   74:	 5d                   	pop	rbp
+;;   75:	 c3                   	ret