runtime/libia2: enable MTE

immunant · Jun 4, 2024 · 9897cd9 · 9897cd9
1 parent 9ad8945
commit 9897cd9
Show file tree

Hide file tree

Showing 4 changed files with 28 additions and 9 deletions.
diff --git a/runtime/libia2/ia2.c b/runtime/libia2/ia2.c
@@ -127,6 +127,7 @@ int ia2_mprotect_with_tag(void *addr, size_t len, int prot, int tag) {
     int res = mprotect(addr, len, prot | PROT_MTE);
     if (res != 0) {
         /* Skip memory tagging if mprotect returned an error */
+        printf("mprotect failed with %d\n", res);
         return res;
     }
     assert((len % PAGE_SIZE) == 0);
@@ -231,9 +232,13 @@ static bool in_extra_libraries(struct dl_phdr_info *info, const char *extra_libr
 
 /// Map ELF segment flags to mprotect access flags
 static int segment_flags_to_access_flags(Elf64_Word flags) {
-  return ((flags & PF_X) != 0 ? PROT_EXEC : 0) |
-         ((flags & PF_W) != 0 ? PROT_WRITE : 0) |
-         ((flags & PF_R) != 0 ? PROT_READ : 0);
+  return
+#if LIBIA2_AARCH64
+      PROT_MTE |
+#endif
+      ((flags & PF_X) != 0 ? PROT_EXEC : 0) |
+      ((flags & PF_W) != 0 ? PROT_WRITE : 0) |
+      ((flags & PF_R) != 0 ? PROT_READ : 0);
 }
 
 int protect_tls_pages(struct dl_phdr_info *info, size_t size, void *data) {

diff --git a/runtime/libia2/include/ia2_compartment_init.inc b/runtime/libia2/include/ia2_compartment_init.inc
@@ -50,9 +50,9 @@ extern int ia2_n_pkeys_to_alloc;
 extern char __start_ia2_shared_data __attribute__((visibility("hidden"))),
     __stop_ia2_shared_data __attribute__((visibility("hidden")));
 
-void ensure_pkeys_allocated(int *n_to_alloc);
+void ia2_set_up_tags(int *n_to_alloc);
 __attribute__((constructor)) static void COMPARTMENT_IDENT(init_pkey)() {
-  ensure_pkeys_allocated(&ia2_n_pkeys_to_alloc);
+  ia2_set_up_tags(&ia2_n_pkeys_to_alloc);
   struct IA2SharedSection shared_sections[2] = {{
                                                     &__start_ia2_shared_data,
                                                     &__stop_ia2_shared_data,

diff --git a/runtime/libia2/include/ia2_internal.h b/runtime/libia2/include/ia2_internal.h
@@ -297,7 +297,7 @@ int ia2_mprotect_with_tag(void *addr, size_t len, int prot, int tag);
 #endif
 char *allocate_stack(int i);
 void verify_tls_padding(void);
-void ensure_pkeys_allocated(int *n_to_alloc);
+void ia2_set_up_tags(int *n_to_alloc);
 _Noreturn void ia2_reinit_stack_err(int i);
 
 #define _IA2_INIT_RUNTIME(n)                                                   \
@@ -323,7 +323,7 @@ _Noreturn void ia2_reinit_stack_err(int i);
                                                                                \
   __attribute__((constructor)) static void ia2_init(void) {                    \
     /* Set up global resources. */                                             \
-    ensure_pkeys_allocated(&ia2_n_pkeys_to_alloc);                             \
+    ia2_set_up_tags(&ia2_n_pkeys_to_alloc);                                    \
     /* Initialize stacks for the main thread/ */                               \
     init_stacks_and_setup_tls();                                               \
     REPEATB##n(setup_destructors_for_compartment, nop_macro);                  \

diff --git a/runtime/libia2/init.c b/runtime/libia2/init.c
@@ -1,4 +1,6 @@
 #include "ia2_internal.h"
+#include <sys/auxv.h>
+#include <sys/prctl.h>
 
 /* The 0th compartment is unprivileged and does not protect its memory, */
 /* so declare its stack pointer in the shared object that sets up the */
@@ -43,8 +45,8 @@ void verify_tls_padding(void) {
   }
 }
 
-/* Ensure that all required pkeys are allocated or no-op on aarch64. */
-void ensure_pkeys_allocated(int *n_to_alloc) {
+/* Allocates the required pkeys on x86 or enables MTE on aarch64 */
+void ia2_set_up_tags(int *n_to_alloc) {
 #if LIBIA2_X86_64
   if (*n_to_alloc != 0) {
     for (int pkey = 1; pkey <= *n_to_alloc; pkey++) {
@@ -62,6 +64,18 @@ void ensure_pkeys_allocated(int *n_to_alloc) {
     }
     *n_to_alloc = 0;
   }
+#elif LIBIA2_AARCH64
+  if (!(getauxval(AT_HWCAP2) & HWCAP2_MTE)) {
+    printf("MTE is not supported\n");
+    exit(-1);
+  }
+  int res = prctl(PR_SET_TAGGED_ADDR_CTRL,
+                  PR_TAGGED_ADDR_ENABLE | PR_MTE_TCF_SYNC | (0xFFFE << PR_MTE_TAG_SHIFT),
+                  0, 0, 0);
+  if (res) {
+    printf("prctl(PR_SET_TAGGED_ADDR_CTRL) failed to enable MTE: (%s)\n", errno_s);
+    exit(-1);
+  }
 #endif
 }