Add tests to ensure MTE tags are preserved across FFI boundaries

Added run-make tests to verify that, between a Rust-C FFI boundary in both directions,
any MTE tags included in a pointer are preserved for the following pointer types, as
well as any information stored using TBI:
- int
- float
- string
- function

10 files changed, 342 insertions, 0 deletions

diff --git a/tests/run-make/mte-ffi/bar.h b/tests/run-make/mte-ffi/bar.h
new file mode 100644
index 00000000000..a2292ae02a3
--- /dev/null
+++ b/tests/run-make/mte-ffi/bar.h
@@ -0,0 +1,43 @@
+#ifndef __BAR_H
+#define __BAR_H
+
+#include <sys/mman.h>
+#include <sys/auxv.h>
+#include <sys/prctl.h>
+#include <unistd.h>
+#include <stdio.h>
+
+// Set the allocation tag on the destination address using the STG instruction.
+#define set_tag(tagged_addr) do {                                      \
+    asm volatile("stg %0, [%0]" : : "r" (tagged_addr) : "memory"); \
+} while (0)
+
+int mte_enabled() {
+    return (getauxval(AT_HWCAP2)) & HWCAP2_MTE;
+}
+
+void *alloc_page() {
+    // Enable MTE with synchronous checking
+    if (prctl(PR_SET_TAGGED_ADDR_CTRL,
+              PR_TAGGED_ADDR_ENABLE | PR_MTE_TCF_SYNC | (0xfffe << PR_MTE_TAG_SHIFT),
+              0, 0, 0))
+    {
+        perror("prctl() failed");
+    }
+
+    // Using `mmap` allows us to ensure that, on systems which support MTE, the allocated
+    // memory is 16-byte aligned for MTE.
+    // This also allows us to explicitly specify whether the region should be protected by
+    // MTE or not.
+    if (mte_enabled()) {
+        void *ptr = mmap(NULL, sysconf(_SC_PAGESIZE),
+                         PROT_READ | PROT_WRITE | PROT_MTE, MAP_PRIVATE | MAP_ANONYMOUS,
+                         -1, 0);
+    } else {
+        void *ptr = mmap(NULL, sysconf(_SC_PAGESIZE),
+                         PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS,
+                         -1, 0);
+    }
+}
+
+#endif // __BAR_H
diff --git a/tests/run-make/mte-ffi/bar_float.c b/tests/run-make/mte-ffi/bar_float.c
new file mode 100644
index 00000000000..a1590f62765
--- /dev/null
+++ b/tests/run-make/mte-ffi/bar_float.c
@@ -0,0 +1,44 @@
+#include <stdio.h>
+#include <stdlib.h>
+#include <stdint.h>
+#include "bar.h"
+
+extern void foo(char*);
+
+void bar(char *ptr) {
+    if (((uintptr_t)ptr >> 56) != 0x1f) {
+        fprintf(stderr, "Top byte corrupted on Rust -> C FFI boundary!\n");
+        exit(1);
+    }
+}
+
+int main(void)
+{
+    float *ptr = alloc_page();
+    if (ptr == MAP_FAILED)
+    {
+        perror("mmap() failed");
+        return EXIT_FAILURE;
+    }
+
+    // Store an arbitrary tag in bits 56-59 of the pointer (where an MTE tag may be),
+    // and a different value in the ignored top 4 bits.
+    ptr = (float *)((uintptr_t)ptr | 0x1fl << 56);
+
+    if (mte_enabled()) {
+        set_tag(ptr);
+    }
+
+    ptr[0] = 2.0f;
+    ptr[1] = 1.5f;
+
+    foo(ptr); // should change the contents of the page and call `bar`
+
+    if (ptr[0] != 0.5f || ptr[1] != 0.2f) {
+        fprintf(stderr, "invalid data in memory; expected '0.5 0.2', got '%f %f'\n",
+                ptr[0], ptr[1]);
+        return EXIT_FAILURE;
+    }
+
+    return 0;
+}
diff --git a/tests/run-make/mte-ffi/bar_function.c b/tests/run-make/mte-ffi/bar_function.c
new file mode 100644
index 00000000000..1fa48d32a0c
--- /dev/null
+++ b/tests/run-make/mte-ffi/bar_function.c
@@ -0,0 +1,39 @@
+#include <stdio.h>
+#include <stdlib.h>
+#include <stdint.h>
+#include "bar.h"
+
+typedef void (*fp)(int (*)());
+
+extern void foo(fp);
+
+void bar(int (*ptr)()) {
+    if (((uintptr_t)ptr >> 56) != 0x2f) {
+        fprintf(stderr, "Top byte corrupted on Rust -> C FFI boundary!\n");
+        exit(1);
+    }
+
+    int r = (*ptr)();
+    if (r != 32) {
+        fprintf(stderr, "invalid return value; expected 32, got '%d'\n", r);
+        exit(1);
+    }
+}
+
+int main(void)
+{
+    fp ptr = alloc_page();
+    if (ptr == MAP_FAILED)
+    {
+        perror("mmap() failed");
+        return EXIT_FAILURE;
+    }
+
+    // Store an arbitrary tag in bits 56-59 of the pointer (where an MTE tag may be),
+    // and a different value in the ignored top 4 bits.
+    ptr = (fp)((uintptr_t)&bar | 0x1fl << 56);
+
+    foo(ptr);
+
+    return 0;
+}
diff --git a/tests/run-make/mte-ffi/bar_int.c b/tests/run-make/mte-ffi/bar_int.c
new file mode 100644
index 00000000000..d1c79e95dc9
--- /dev/null
+++ b/tests/run-make/mte-ffi/bar_int.c
@@ -0,0 +1,47 @@
+#include <stdio.h>
+#include <stdlib.h>
+#include <stdint.h>
+#include "bar.h"
+
+extern void foo(unsigned int *);
+
+void bar(char *ptr) {
+    if (((uintptr_t)ptr >> 56) != 0x1f) {
+        fprintf(stderr, "Top byte corrupted on Rust -> C FFI boundary!\n");
+        exit(1);
+    }
+}
+
+int main(void)
+{
+    // Construct a pointer with an arbitrary tag in bits 56-59, simulating an MTE tag.
+    // It's only necessary that the tag is preserved across FFI bounds for this test.
+    unsigned int *ptr;
+
+    ptr = alloc_page();
+    if (ptr == MAP_FAILED)
+    {
+        perror("mmap() failed");
+        return EXIT_FAILURE;
+    }
+
+    // Store an arbitrary tag in bits 56-59 of the pointer (where an MTE tag may be),
+    // and a different value in the ignored top 4 bits.
+    ptr = (unsigned int *)((uintptr_t)ptr | 0x1fl << 56);
+
+    if (mte_enabled()) {
+        set_tag(ptr);
+    }
+
+    ptr[0] = 61;
+    ptr[1] = 62;
+
+    foo(ptr); // should change the contents of the page to start with 0x63 0x64 and call `bar`
+
+    if (ptr[0] != 0x63 || ptr[1] != 0x64) {
+        fprintf(stderr, "invalid data in memory; expected '63 64', got '%d %d'\n", ptr[0], ptr[1]);
+        return EXIT_FAILURE;
+    }
+
+    return 0;
+}
diff --git a/tests/run-make/mte-ffi/bar_string.c b/tests/run-make/mte-ffi/bar_string.c
new file mode 100644
index 00000000000..5669ffd6695
--- /dev/null
+++ b/tests/run-make/mte-ffi/bar_string.c
@@ -0,0 +1,48 @@
+#include <stdio.h>
+#include <stdlib.h>
+#include <stdint.h>
+#include "bar.h"
+
+extern void foo(char*);
+
+void bar(char *ptr) {
+    if (((uintptr_t)ptr >> 56) != 0x2f) {
+        fprintf(stderr, "Top byte corrupted on Rust -> C FFI boundary!\n");
+        exit(1);
+    }
+
+    if (strcmp(ptr, "cd")) {
+        fprintf(stderr, "invalid data in memory; expected 'cd', got '%s'\n", ptr);
+        exit(1);
+    }
+}
+
+int main(void)
+{
+    // Construct a pointer with an arbitrary tag in bits 56-59, simulating an MTE tag.
+    // It's only necessary that the tag is preserved across FFI bounds for this test.
+    char *ptr;
+
+    ptr = alloc_page();
+    if (ptr == MAP_FAILED)
+    {
+        perror("mmap() failed");
+        return EXIT_FAILURE;
+    }
+
+    // Store an arbitrary tag in bits 56-59 of the pointer (where an MTE tag may be),
+    // and a different value in the ignored top 4 bits.
+    ptr = (unsigned int *)((uintptr_t)ptr | 0x1fl << 56);
+
+    if (mte_enabled()) {
+        set_tag(ptr);
+    }
+
+    ptr[0] = 'a';
+    ptr[1] = 'b';
+    ptr[2] = '\0';
+
+    foo(ptr);
+
+    return 0;
+}
diff --git a/tests/run-make/mte-ffi/foo_float.rs b/tests/run-make/mte-ffi/foo_float.rs
new file mode 100644
index 00000000000..c1bedd52494
--- /dev/null
+++ b/tests/run-make/mte-ffi/foo_float.rs
@@ -0,0 +1,19 @@
+#![crate_type = "cdylib"]
+#![crate_name = "foo"]
+
+use std::os::raw::c_float;
+
+extern "C" {
+    fn bar(ptr: *const c_float);
+}
+
+#[no_mangle]
+pub extern "C" fn foo(ptr: *mut c_float) {
+    assert_eq!((ptr as usize) >> 56, 0x1f);
+
+    unsafe {
+        *ptr = 0.5;
+        *ptr.wrapping_add(1) = 0.2;
+        bar(ptr);
+    }
+}
diff --git a/tests/run-make/mte-ffi/foo_function.rs b/tests/run-make/mte-ffi/foo_function.rs
new file mode 100644
index 00000000000..2c8e0b26238
--- /dev/null
+++ b/tests/run-make/mte-ffi/foo_function.rs
@@ -0,0 +1,17 @@
+#![crate_type = "cdylib"]
+#![crate_name = "foo"]
+
+extern "C" fn ret32() -> i32 {
+    32
+}
+
+#[no_mangle]
+pub extern "C" fn foo(ptr: extern "C" fn(extern "C" fn() -> i32)) {
+    assert_eq!((ptr as usize) >> 56, 0x1f);
+
+    // Store an arbitrary tag in the tag bits, and convert back to the correct pointer type.
+    let p = ((ret32 as usize) | (0x2f << 56)) as *const ();
+    let p: extern "C" fn() -> i32 = unsafe { std::mem::transmute(p) };
+
+    unsafe { ptr(p) }
+}
diff --git a/tests/run-make/mte-ffi/foo_int.rs b/tests/run-make/mte-ffi/foo_int.rs
new file mode 100644
index 00000000000..106d863cb81
--- /dev/null
+++ b/tests/run-make/mte-ffi/foo_int.rs
@@ -0,0 +1,19 @@
+#![crate_type = "cdylib"]
+#![crate_name = "foo"]
+
+use std::os::raw::c_uint;
+
+extern "C" {
+    fn bar(ptr: *const c_uint);
+}
+
+#[no_mangle]
+pub extern "C" fn foo(ptr: *mut c_uint) {
+    assert_eq!((ptr as usize) >> 56, 0x1f);
+
+    unsafe {
+        *ptr = 0x63;
+        *ptr.wrapping_add(1) = 0x64;
+        bar(ptr);
+    }
+}
diff --git a/tests/run-make/mte-ffi/foo_string.rs b/tests/run-make/mte-ffi/foo_string.rs
new file mode 100644
index 00000000000..54744802448
--- /dev/null
+++ b/tests/run-make/mte-ffi/foo_string.rs
@@ -0,0 +1,27 @@
+#![crate_type = "cdylib"]
+#![crate_name = "foo"]
+
+use std::arch::asm;
+use std::ffi::{CStr, CString};
+use std::os::raw::c_char;
+
+extern "C" {
+    fn bar(ptr: *const c_char);
+}
+
+#[no_mangle]
+pub extern "C" fn foo(ptr: *const c_char) {
+    assert_eq!((ptr as usize) >> 56, 0x1f);
+
+    let s = unsafe { CStr::from_ptr(ptr) };
+    assert_eq!(s.to_str().unwrap(), "ab");
+
+    let s = CString::from_vec_with_nul("cd\0".into()).unwrap();
+    let mut p = ((s.as_ptr() as usize) | (0x2f << 56)) as *const c_char;
+    unsafe {
+        #[cfg(target_feature = "mte")]
+        asm!("stg {p}, [{p}]", p = inout(reg) p);
+
+        bar(p);
+    }
+}
diff --git a/tests/run-make/mte-ffi/rmake.rs b/tests/run-make/mte-ffi/rmake.rs
new file mode 100644
index 00000000000..132c12aa7f0
--- /dev/null
+++ b/tests/run-make/mte-ffi/rmake.rs
@@ -0,0 +1,39 @@
+// Tests that MTE tags and values stored in the top byte of a pointer (TBI) are
+// preserved across FFI boundaries (C <-> Rust).
+// This test does not require MTE: whilst the test will use MTE if available, if it is not,
+// arbitrary tag bits are set using TBI.
+
+//@ only-aarch64
+//@ only-linux
+//@ only-gnu
+//@ run-pass
+
+use run_make_support::{cc, dynamic_lib_name, extra_c_flags, run, rustc, target};
+
+fn main() {
+    run_test("int");
+    run_test("float");
+    run_test("string");
+    run_test("function");
+}
+
+fn run_test(variant: &str) {
+    let flags = {
+        let mut flags = extra_c_flags();
+        flags.push("-march=armv8.5-a+memtag");
+        flags
+    };
+    print!("{variant} test...");
+    rustc()
+        .input(format!("foo_{variant}.rs"))
+        .target(target())
+        .linker("aarch64-linux-gnu-gcc")
+        .run();
+    cc().input(format!("bar_{variant}.c"))
+        .input(dynamic_lib_name("foo"))
+        .out_exe("test")
+        .args(&flags)
+        .run();
+    run("test");
+    println!("\tpassed");
+}