From 8b9e2204e41a9297fba68a7e45c2fd573b75a2f8 Mon Sep 17 00:00:00 2001
From: Ayrton Munoz <adm@immunant.com>
Date: Tue, 12 Mar 2024 23:45:37 -0400
Subject: [PATCH 01/31] runtime: Move most of _IA2_INIT_RUNTIME to init.c

---
 runtime/libia2/CMakeLists.txt         |  2 +-
 runtime/libia2/include/ia2_internal.h | 68 ++-----------------------
 runtime/libia2/init.c                 | 71 +++++++++++++++++++++++++++
 3 files changed, 77 insertions(+), 64 deletions(-)
 create mode 100644 runtime/libia2/init.c

diff --git a/runtime/libia2/CMakeLists.txt b/runtime/libia2/CMakeLists.txt
index d1de67c44e..c07b193097 100644
--- a/runtime/libia2/CMakeLists.txt
+++ b/runtime/libia2/CMakeLists.txt
@@ -1,7 +1,7 @@
 cmake_minimum_required(VERSION 3.12)
 project(libia2)
 
-add_library(libia2 ia2.c threads.c main.c exit.c)
+add_library(libia2 ia2.c init.c threads.c main.c exit.c)
 target_compile_options(libia2 PRIVATE "-fPIC")
 
 if(LIBIA2_DEBUG)
diff --git a/runtime/libia2/include/ia2_internal.h b/runtime/libia2/include/ia2_internal.h
index 3f7c86a283..61138e6b73 100644
--- a/runtime/libia2/include/ia2_internal.h
+++ b/runtime/libia2/include/ia2_internal.h
@@ -228,75 +228,17 @@ works as a reasonable signpost no-op. */
   void ia2_setup_destructors_##n(void);                                        \
   ia2_setup_destructors_##n();
 
+char *allocate_stack(int i);
+void verify_tls_padding(void);
+void ensure_pkeys_allocated(int *n_to_alloc);
+_Noreturn void ia2_reinit_stack_err(int i);
+
 #define _IA2_INIT_RUNTIME(n)                                                   \
   int ia2_n_pkeys_to_alloc = n;                                                \
-  /* Allocate a fixed-size stack and protect it with the ith pkey. */          \
-  /* Returns the top of the stack, not the base address of the allocation. */  \
-  char *allocate_stack(int i) {                                                \
-    char *stack = (char *)mmap(NULL, STACK_SIZE, PROT_READ | PROT_WRITE,       \
-                               MAP_PRIVATE | MAP_ANON, -1, 0);                 \
-    if (stack == MAP_FAILED) {                                                 \
-      printf("Failed to allocate stack %d (%s)\n", i, errno_s);                \
-      exit(-1);                                                                \
-    }                                                                          \
-    if (i != 0) {                                                              \
-      int res = pkey_mprotect(stack, STACK_SIZE, PROT_READ | PROT_WRITE, i);   \
-      if (res == -1) {                                                         \
-        printf("Failed to mprotect stack %d (%s)\n", i, errno_s);              \
-        exit(-1);                                                              \
-      }                                                                        \
-    }                                                                          \
-    /* Each stack frame start + 8 is initially 16-byte aligned. */             \
-    return stack + STACK_SIZE - 8;                                             \
-  }                                                                            \
-  /* The 0th compartment is unprivileged and does not protect its memory, */   \
-  /* so declare its stack pointer in the shared object that sets up the */     \
-  /* runtime. */                                                               \
-  /* Ensure that ia2_stackptr_0 is at least a page long to ensure that the */  \
-  /* last page of the TLS segment of compartment 0 does not contain any */     \
-  /* variables that will be used, because the last page-1 bytes may be */      \
-  /* pkey_mprotected by the next compartment depending on sizes/alignment. */  \
   __thread void *ia2_stackptr_0[PAGE_SIZE / sizeof(void *)]                    \
       __attribute__((aligned(4096)));                                          \
                                                                                \
   REPEATB(n, declare_init_tls_fn, nop_macro);                                  \
-  /* Confirm that stack pointers for compartments 0 and 1 are on separate */   \
-  /* pages. */                                                                 \
-  void verify_tls_padding(void) {                                              \
-    /* It's safe to depend on ia2_stackptr_1 existing because all users of */  \
-    /* IA2 will have at least one compartment other than the untrusted one. */ \
-    extern __thread void *ia2_stackptr_1;                                      \
-    if (IA2_ROUND_DOWN((uintptr_t)&ia2_stackptr_1, PAGE_SIZE) ==               \
-        IA2_ROUND_DOWN((uintptr_t)ia2_stackptr_0, PAGE_SIZE)) {                \
-      printf("ia2_stackptr_1 is too close to ia2_stackptr_0\n");               \
-      exit(1);                                                                 \
-    }                                                                          \
-  }                                                                            \
-  /* Ensure that all required pkeys are allocated. */                          \
-  void ensure_pkeys_allocated(int *n_to_alloc) {                               \
-    if (*n_to_alloc != 0) {                                                    \
-      for (int pkey = 1; pkey <= *n_to_alloc; pkey++) {                        \
-        int allocated = pkey_alloc(0, 0);                                      \
-        if (allocated < 0) {                                                   \
-          printf("Failed to allocate protection key %d (%s)\n", pkey,          \
-                 errno_s);                                                     \
-          exit(-1);                                                            \
-        }                                                                      \
-        if (allocated != pkey) {                                               \
-          printf(                                                              \
-              "Failed to allocate protection keys in the expected order\n");   \
-          exit(-1);                                                            \
-        }                                                                      \
-      }                                                                        \
-      *n_to_alloc = 0;                                                         \
-    }                                                                          \
-  }                                                                            \
-  /* Forbid overwriting an existing stack. */                                  \
-  _Noreturn void ia2_reinit_stack_err(int i) {                                 \
-    printf("compartment %d in thread %d tried to allocate existing stack\n",   \
-           i, gettid());                                                       \
-    exit(1);                                                                   \
-  }                                                                            \
                                                                                \
   /* Returns `&ia2_stackptr_N` given a pkru value for the Nth compartment. */  \
   void **ia2_stackptr_for_pkru(uint32_t pkru) {                                \
diff --git a/runtime/libia2/init.c b/runtime/libia2/init.c
new file mode 100644
index 0000000000..b18010d16e
--- /dev/null
+++ b/runtime/libia2/init.c
@@ -0,0 +1,71 @@
+#include "ia2_internal.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 */
+/* runtime. */
+/* Ensure that ia2_stackptr_0 is at least a page long to ensure that the */
+/* last page of the TLS segment of compartment 0 does not contain any */
+/* variables that will be used, because the last page-1 bytes may be */
+/* ia2_memkey_mprotected by the next compartment depending on sizes/alignment. */
+extern __thread void *ia2_stackptr_0[PAGE_SIZE / sizeof(void *)]
+    __attribute__((aligned(4096)));
+
+/* Allocate a fixed-size stack and protect it with the ith pkey. */
+/* Returns the top of the stack, not the base address of the allocation. */
+char *allocate_stack(int i) {
+  char *stack = (char *)mmap(NULL, STACK_SIZE, PROT_READ | PROT_WRITE,
+                             MAP_PRIVATE | MAP_ANON, -1, 0);
+  if (stack == MAP_FAILED) {
+    printf("Failed to allocate stack %d (%s)\n", i, errno_s);
+    exit(-1);
+  }
+  if (i != 0) {
+    int res = pkey_mprotect(stack, STACK_SIZE, PROT_READ | PROT_WRITE, i);
+    if (res == -1) {
+      printf("Failed to mprotect stack %d (%s)\n", i, errno_s);
+      exit(-1);
+    }
+  }
+  /* Each stack frame start + 8 is initially 16-byte aligned. */
+  return stack + STACK_SIZE - 8;
+}
+
+/* Confirm that stack pointers for compartments 0 and 1 are on separate */
+/* pages. */
+void verify_tls_padding(void) {
+  /* It's safe to depend on ia2_stackptr_1 existing because all users of */
+  /* IA2 will have at least one compartment other than the untrusted one. */
+  extern __thread void *ia2_stackptr_1;
+  if (IA2_ROUND_DOWN((uintptr_t)&ia2_stackptr_1, PAGE_SIZE) ==
+      IA2_ROUND_DOWN((uintptr_t)ia2_stackptr_0, PAGE_SIZE)) {
+    printf("ia2_stackptr_1 is too close to ia2_stackptr_0\n");
+    exit(1);
+  }
+}
+
+/* Ensure that all required pkeys are allocated. */
+void ensure_pkeys_allocated(int *n_to_alloc) {
+  if (*n_to_alloc != 0) {
+    for (int pkey = 1; pkey <= *n_to_alloc; pkey++) {
+      int allocated = pkey_alloc(0, 0);
+      if (allocated < 0) {
+        printf("Failed to allocate protection key %d (%s)\n", pkey,
+               errno_s);
+        exit(-1);
+      }
+      if (allocated != pkey) {
+        printf(
+            "Failed to allocate protection keys in the expected order\n");
+        exit(-1);
+      }
+    }
+    *n_to_alloc = 0;
+  }
+}
+
+/* Forbid overwriting an existing stack. */
+_Noreturn void ia2_reinit_stack_err(int i) {
+  printf("compartment %d in thread %d tried to allocate existing stack\n",
+         i, gettid());
+  exit(1);
+}

From 72e8f14c8a09cae7c76596785378ad64f3634701 Mon Sep 17 00:00:00 2001
From: Ayrton Munoz <adm@immunant.com>
Date: Wed, 13 Mar 2024 00:09:45 -0400
Subject: [PATCH 02/31] runtime: Add LIBIA2_AARCH64 flag for libia2

---
 runtime/libia2/CMakeLists.txt | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/runtime/libia2/CMakeLists.txt b/runtime/libia2/CMakeLists.txt
index c07b193097..e1cece4cd5 100644
--- a/runtime/libia2/CMakeLists.txt
+++ b/runtime/libia2/CMakeLists.txt
@@ -4,6 +4,13 @@ project(libia2)
 add_library(libia2 ia2.c init.c threads.c main.c exit.c)
 target_compile_options(libia2 PRIVATE "-fPIC")
 
+if (LIBIA2_AARCH64)
+    target_compile_definitions(libia2 PUBLIC LIBIA2_AARCH64=1)
+    target_compile_options(libia2 PUBLIC "-march=armv8.5-a+memtag" "-ffixed-x18")
+else()
+    target_compile_definitions(libia2 PUBLIC LIBIA2_X86_64=1)
+endif()
+
 if(LIBIA2_DEBUG)
     target_compile_definitions(libia2 PUBLIC LIBIA2_DEBUG=1)
 endif()

From d1b8a517ea6f3df99e646394bd766ec7a3375a10 Mon Sep 17 00:00:00 2001
From: Ayrton Munoz <adm@immunant.com>
Date: Wed, 13 Mar 2024 00:10:26 -0400
Subject: [PATCH 03/31] runtime: Make exit arch-dependent

---
 runtime/libia2/exit.c | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/runtime/libia2/exit.c b/runtime/libia2/exit.c
index 7aa005167c..772b84b590 100644
--- a/runtime/libia2/exit.c
+++ b/runtime/libia2/exit.c
@@ -13,6 +13,7 @@ static void call_libc_exit(int status) {
 
 __attribute__((naked)) void exit(int status) {
   __asm__(
+#if LIBIA2_X86_64
       /* clang-format off */
       "pushq %%rbp\n"
       "movq %%rsp, %%rbp\n"
@@ -29,5 +30,9 @@ __attribute__((naked)) void exit(int status) {
       // Call the real exit function.
       "call call_libc_exit\n"
       /* clang-format on */
+#elif LIBIA2_AARCH64
+#warning "libia2 does not properly wrap `exit` yet"
+      "udf #0\n"
+#endif
       ::);
 }

From 30cd58127ba1bf0bb456efebf27a07729f493e4c Mon Sep 17 00:00:00 2001
From: Ayrton Munoz <adm@immunant.com>
Date: Wed, 13 Mar 2024 00:12:02 -0400
Subject: [PATCH 04/31] runtime: Make libia2 main arch-dependent

---
 runtime/libia2/main.c | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/runtime/libia2/main.c b/runtime/libia2/main.c
index 500fc58a74..17d9866675 100644
--- a/runtime/libia2/main.c
+++ b/runtime/libia2/main.c
@@ -8,6 +8,7 @@ static void *main_sp __attribute__((used)) = 0;
 /* XXX: Assumes main compartment has pkey 1. */
 __attribute__((naked)) int __wrap_main(int argc, char **argv) {
   __asm__(
+#if LIBIA2_X86_64
       /* clang-format off */
       "pushq %%rbp\n"
       "movq %%rsp, %%rbp\n"
@@ -39,5 +40,9 @@ __attribute__((naked)) int __wrap_main(int argc, char **argv) {
       "popq %%rbp\n"
       "ret\n"
       /* clang-format on */
+#elif LIBIA2_AARCH64
+#warning "libia2 does not properly wrap `main` yet"
+      "b __real_main\n"
+#endif
       ::);
 }

From fd1a42e4aefc2b62b7aaef9f53005d926c943c7b Mon Sep 17 00:00:00 2001
From: Ayrton Munoz <adm@immunant.com>
Date: Wed, 13 Mar 2024 00:13:15 -0400
Subject: [PATCH 05/31] runtime: Make ia2_thread_begin arch-dependent

---
 runtime/libia2/threads.c | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/runtime/libia2/threads.c b/runtime/libia2/threads.c
index c341c412dc..05d02935a2 100644
--- a/runtime/libia2/threads.c
+++ b/runtime/libia2/threads.c
@@ -29,6 +29,7 @@ void *ia2_thread_begin(void *arg) {
    * data. */
   /*  sigaltstack(&alt_stack, NULL); */
 
+#if LIBIA2_X86_64
   /* Determine the current compartment so know which stack to use. */
   uint32_t pkru = 0;
   __asm__ volatile(
@@ -68,6 +69,10 @@ void *ia2_thread_begin(void *arg) {
       : "rdi");
   /* clang-format on */
   return result;
+#elif LIBIA2_AARCH64
+#warning "libia2 does not implement ia2_thread_begin yet"
+  __builtin_trap();
+#endif
 }
 
 int __real_pthread_create(pthread_t *restrict thread,

From b262db386e4ec19695554da80e0c276dc819e17b Mon Sep 17 00:00:00 2001
From: Ayrton Munoz <adm@immunant.com>
Date: Wed, 13 Mar 2024 00:26:41 -0400
Subject: [PATCH 06/31] runtime: Make memory tagging arch-dependent

---
 runtime/libia2/ia2.c                  | 91 +++++++++++++++++++++++----
 runtime/libia2/include/ia2_internal.h | 16 +++++
 runtime/libia2/init.c                 |  4 +-
 3 files changed, 98 insertions(+), 13 deletions(-)

diff --git a/runtime/libia2/ia2.c b/runtime/libia2/ia2.c
index a59c95b1dc..2345ef5739 100644
--- a/runtime/libia2/ia2.c
+++ b/runtime/libia2/ia2.c
@@ -8,6 +8,8 @@
 #include "ia2_internal.h"
 #include "ia2.h"
 
+#if LIBIA2_X86_64
+
 __attribute__((__used__)) uint32_t ia2_get_pkru() {
   uint32_t pkru = 0;
   __asm__ volatile("rdpkru" : "=a"(pkru) : "a"(0), "d"(0), "c"(0));
@@ -77,6 +79,69 @@ size_t ia2_get_pkey() {
   }
   }
 }
+size_t ia2_get_tag(void) __attribute__((alias("ia2_get_pkey")));
+
+#elif LIBIA2_AARCH64
+
+size_t ia2_get_x18(void) {
+    size_t x18;
+    asm("mov %0, x18" : "=r"(x18));
+    return x18;
+}
+size_t ia2_get_tag(void) __attribute__((alias("ia2_get_x18")));
+
+// TODO: insert_tag could probably be cleaned up a bit, but I'm not sure if the
+// generated code could be simplified since addg encodes the tag as an imm field
+#define _addg(out_ptr, in_ptr, tag) \
+    asm("addg %0, %1, #0, %2" : "=r"(out_ptr) : "r"(in_ptr), "i"(tag)); \
+
+#define insert_tag(ptr, tag) \
+    ({ \
+        uint64_t _res; \
+        switch (tag) { \
+            case 0: { _addg(_res, ptr, 0); break; } \
+            case 1: { _addg(_res, ptr, 1); break; } \
+            case 2: { _addg(_res, ptr, 2); break; } \
+            case 3: { _addg(_res, ptr, 3); break; } \
+            case 4: { _addg(_res, ptr, 4); break; } \
+            case 5: { _addg(_res, ptr, 5); break; } \
+            case 6: { _addg(_res, ptr, 6); break; } \
+            case 7: { _addg(_res, ptr, 7); break; } \
+            case 8: { _addg(_res, ptr, 8); break; } \
+            case 9: { _addg(_res, ptr, 9); break; } \
+            case 10: { _addg(_res, ptr, 10); break; } \
+            case 11: { _addg(_res, ptr, 11); break; } \
+            case 12: { _addg(_res, ptr, 12); break; } \
+            case 13: { _addg(_res, ptr, 13); break; } \
+            case 14: { _addg(_res, ptr, 14); break; } \
+            case 15: { _addg(_res, ptr, 15); break; } \
+        } \
+        _res; \
+    })
+
+
+#define set_tag(tagged_ptr) \
+    asm volatile("st2g %0, [%0]" :: "r"(tagged_ptr) : "memory");
+
+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 */
+        return res;
+    }
+    assert((len % PAGE_SIZE) == 0);
+    /* Assuming we're using st2g. stgm is undefined at EL0 so it's not an option */
+    const int granule_sz = 32;
+    const int granules_per_page = PAGE_SIZE / 32;
+    size_t tag = ia2_get_tag();
+    for (int i = 0; i < granules_per_page; i++) {
+        // TODO: It may be possible to simplify this to be more efficient using the addg imm offset
+        uint64_t tagged_ptr = insert_tag((uint64_t)addr + (i * granule_sz), tag);
+        set_tag(tagged_ptr);
+    }
+    return 0;
+}
+#endif
 
 // Reserve one extra shared range for the RELRO segment that we are
 // also sharing, in addition to the special shared sections in PhdrSearchArgs.
@@ -229,31 +294,31 @@ int protect_tls_pages(struct dl_phdr_info *info, size_t size, void *data) {
     if (untrusted_stackptr_addr >= start && untrusted_stackptr_addr < end) {
       // Protect TLS region start to the beginning of the untrusted region.
       if (untrusted_stackptr_addr > start_round_down) {
-        int mprotect_err = pkey_mprotect(
+        int mprotect_err = ia2_mprotect_with_tag(
             (void *)start_round_down, untrusted_stackptr_addr - start_round_down,
             PROT_READ | PROT_WRITE, search_args->pkey);
         if (mprotect_err != 0) {
-          printf("pkey_mprotect failed: %s\n", strerror(errno));
+          printf("ia2_mprotect_with_tag failed: %s\n", strerror(errno));
           exit(-1);
         }
       }
       uint64_t after_untrusted_region_start = untrusted_stackptr_addr + 0x1000;
       uint64_t after_untrusted_region_len = end - after_untrusted_region_start;
       if (after_untrusted_region_len > 0) {
-        int mprotect_err = pkey_mprotect((void *)after_untrusted_region_start,
+        int mprotect_err = ia2_mprotect_with_tag((void *)after_untrusted_region_start,
                                          after_untrusted_region_len,
                                          PROT_READ | PROT_WRITE, search_args->pkey);
         if (mprotect_err != 0) {
-          printf("pkey_mprotect failed: %s\n", strerror(errno));
+          printf("ia2_mprotect_with_tag failed: %s\n", strerror(errno));
           exit(-1);
         }
       }
     } else {
       int mprotect_err =
-          pkey_mprotect((void *)start_round_down, len_round_up,
+          ia2_mprotect_with_tag((void *)start_round_down, len_round_up,
                         PROT_READ | PROT_WRITE, search_args->pkey);
       if (mprotect_err != 0) {
-        printf("pkey_mprotect failed: %s\n", strerror(errno));
+        printf("ia2_mprotect_with_tag failed: %s\n", strerror(errno));
         exit(-1);
       }
     }
@@ -270,12 +335,16 @@ int protect_pages(struct dl_phdr_info *info, size_t size, void *data) {
 
   struct PhdrSearchArgs *search_args = (struct PhdrSearchArgs *)data;
 
-  size_t cur_pkey = ia2_get_pkey();
+  size_t cur_pkey = ia2_get_tag();
+#if LIBIA2_X86_64
   if (cur_pkey != search_args->pkey) {
     fprintf(stderr, "Invalid pkey, expected %" PRId32 ", found %zu\n",
             search_args->pkey, cur_pkey);
     abort();
   }
+#elif LIBIA2_AARCH64
+#warning "libia2 missing tag validation in protect_pages"
+#endif
   Elf64_Addr address = (Elf64_Addr)search_args->address;
   bool extra = in_extra_libraries(info, search_args->extra_libraries);
   if (!in_loaded_segment(info, address) && !extra) {
@@ -389,7 +458,7 @@ int protect_pages(struct dl_phdr_info *info, size_t size, void *data) {
 
       if (cur_end > start) {
         // Probe each page to ensure that we can read from it. This ensures that
-        // we at least have read permission to the page before we pkey_mprotect
+        // we at least have read permission to the page before we ia2_mprotect_with_tag
         // it to exclude all other compartments from accessing the page. We only
         // use pkeys to grant both read and write permission together and rely
         // on normal page permissions to create read-only regions, so reading is
@@ -400,12 +469,12 @@ int protect_pages(struct dl_phdr_info *info, size_t size, void *data) {
           volatile char *cur = (volatile char *)start + i;
           (void)*cur;
         }
-        // TODO: Inline pkey_mprotect call and make sure the pkey is in a
+        // TODO: Inline ia2_mprotect_with_tag call and make sure the pkey is in a
         // register here so we can disallow calls to the libc function
-        int mprotect_err = pkey_mprotect((void *)start, cur_end - start,
+        int mprotect_err = ia2_mprotect_with_tag((void *)start, cur_end - start,
                                          access_flags, (int)cur_pkey);
         if (mprotect_err != 0) {
-          printf("pkey_mprotect failed: %s\n", strerror(errno));
+          printf("ia2_mprotect_with_tag failed: %s\n", strerror(errno));
           exit(-1);
         }
 
diff --git a/runtime/libia2/include/ia2_internal.h b/runtime/libia2/include/ia2_internal.h
index 61138e6b73..3f0022f9e3 100644
--- a/runtime/libia2/include/ia2_internal.h
+++ b/runtime/libia2/include/ia2_internal.h
@@ -145,6 +145,7 @@ asm(".macro mov_pkru_eax pkey\n"
 
 #define IA2_ROUND_DOWN(x, y) ((x) & ~((y)-1))
 
+#if LIBIA2_X86_64
 /* clang-format can't handle inline asm in macros */
 /* clang-format off */
 /* Allocate and protect the stack for this thread's i'th compartment. */
@@ -203,7 +204,12 @@ asm(".macro mov_pkru_eax pkey\n"
         : "rdi", "rcx", "rdx", "r10", "r11", "r12");                           \
   }
 /* clang-format on */
+#elif LIBIA2_AARCH64
+#warning "libia2 does not implement ALLOCATE_COMPARTMENT_STACK_AND_SETUP_TLS yet"
+#define ALLOCATE_COMPARTMENT_STACK_AND_SETUP_TLS(i)
+#endif
 
+#if LIBIA2_X86_64
 #define return_stackptr_if_compartment(compartment)                            \
   if (pkru == PKRU(compartment)) {                                             \
     register void *out asm("rax");                                             \
@@ -215,6 +221,10 @@ asm(".macro mov_pkru_eax pkey\n"
         :);                                                                    \
     return out;                                                                \
   }
+#elif LIBIA2_AARCH64
+#warning "libia2 does not implement return_stackptr_if_compartment yet"
+#define return_stackptr_if_compartment(compartment)
+#endif
 
 /* Pass to mmap to signal end of program init */
 #define IA2_FINISH_INIT_MAGIC 0x1a21face1a21faceULL
@@ -228,6 +238,12 @@ works as a reasonable signpost no-op. */
   void ia2_setup_destructors_##n(void);                                        \
   ia2_setup_destructors_##n();
 
+#if LIBIA2_AARCH64
+int ia2_mprotect_with_tag(void *addr, size_t len, int prot, int tag);
+#elif LIBIA2_X86_64
+/* We can't use an alias attribute since this points to a function outside the translation unit */
+#define ia2_mprotect_with_tag pkey_mprotect
+#endif
 char *allocate_stack(int i);
 void verify_tls_padding(void);
 void ensure_pkeys_allocated(int *n_to_alloc);
diff --git a/runtime/libia2/init.c b/runtime/libia2/init.c
index b18010d16e..5dcf25354a 100644
--- a/runtime/libia2/init.c
+++ b/runtime/libia2/init.c
@@ -6,7 +6,7 @@
 /* Ensure that ia2_stackptr_0 is at least a page long to ensure that the */
 /* last page of the TLS segment of compartment 0 does not contain any */
 /* variables that will be used, because the last page-1 bytes may be */
-/* ia2_memkey_mprotected by the next compartment depending on sizes/alignment. */
+/* ia2_mprotect_with_taged by the next compartment depending on sizes/alignment. */
 extern __thread void *ia2_stackptr_0[PAGE_SIZE / sizeof(void *)]
     __attribute__((aligned(4096)));
 
@@ -20,7 +20,7 @@ char *allocate_stack(int i) {
     exit(-1);
   }
   if (i != 0) {
-    int res = pkey_mprotect(stack, STACK_SIZE, PROT_READ | PROT_WRITE, i);
+    int res = ia2_mprotect_with_tag(stack, STACK_SIZE, PROT_READ | PROT_WRITE, i);
     if (res == -1) {
       printf("Failed to mprotect stack %d (%s)\n", i, errno_s);
       exit(-1);

From 05d6d04aa3d1c5e876a11e828609bb95a00ac8e1 Mon Sep 17 00:00:00 2001
From: Ayrton Munoz <adm@immunant.com>
Date: Wed, 13 Mar 2024 01:09:49 -0400
Subject: [PATCH 07/31] runtime: Make scrub registers arch-dependent

---
 runtime/libia2/include/scrub_registers.h | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/runtime/libia2/include/scrub_registers.h b/runtime/libia2/include/scrub_registers.h
index ec2558aea6..e8c6f79272 100644
--- a/runtime/libia2/include/scrub_registers.h
+++ b/runtime/libia2/include/scrub_registers.h
@@ -1,3 +1,4 @@
+#if LIBIA2_X86_64
 // This file defines the feature specific scrub routines.
 //
 // This is implemented via the fairly standard fallthrough approach, and
@@ -70,3 +71,6 @@ asm(".text\n"
     "jmp __libia2_scrub_registers_sse\n"
     "int3\n"
     ".previous\n");
+#elif LIBIA2_AARCH64
+#warning "__libia2_scrub_registers is not implemented for aarch64 yet"
+#endif

From 008bab6bfb88ca3ea25075fd5b943c0e3c892963 Mon Sep 17 00:00:00 2001
From: Ayrton Munoz <adm@immunant.com>
Date: Wed, 13 Mar 2024 01:12:02 -0400
Subject: [PATCH 08/31] runtime: Make compartment init arch-dependent

---
 runtime/libia2/include/ia2_compartment_init.inc | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/runtime/libia2/include/ia2_compartment_init.inc b/runtime/libia2/include/ia2_compartment_init.inc
index c720b43804..91e5506edf 100644
--- a/runtime/libia2/include/ia2_compartment_init.inc
+++ b/runtime/libia2/include/ia2_compartment_init.inc
@@ -66,6 +66,7 @@ __attribute__((constructor)) static void COMPARTMENT_IDENT(init_pkey)() {
       .shared_sections = shared_sections,
   };
 
+#if LIBIA2_X86_64
   __asm__ volatile(
       /* Set PKRU to the compartment's value */
       "xorl %%ecx, %%ecx\n"
@@ -75,7 +76,13 @@ __attribute__((constructor)) static void COMPARTMENT_IDENT(init_pkey)() {
       :
       :
       : "rax", "rcx", "rdx");
+#elif LIBIA2_AARCH64
+#warning "libia2 may use incorrect memory permissions in init_pkey"
+#endif
+
   dl_iterate_phdr(protect_pages, &args);
+
+#if LIBIA2_X86_64
   __asm__ volatile(
       /* Set PKRU to fully untrusted (no access) */
       "xorl %%ecx, %%ecx\n"
@@ -85,6 +92,9 @@ __attribute__((constructor)) static void COMPARTMENT_IDENT(init_pkey)() {
       :
       :
       : "rax", "rcx", "rdx");
+#elif LIBIA2_AARCH64
+#warning "libia2 may use incorrect memory permissions in init_pkey"
+#endif
 
   /* Check that we found all extra libraries */
   const char *cur_pos = args.extra_libraries;

From 057c1b29d9c80553f5359647280cc996d08d99fe Mon Sep 17 00:00:00 2001
From: Ayrton Munoz <adm@immunant.com>
Date: Wed, 13 Mar 2024 01:19:05 -0400
Subject: [PATCH 09/31] runtime: Remove calls to pkey_alloc for aarch64

---
 runtime/libia2/init.c | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/runtime/libia2/init.c b/runtime/libia2/init.c
index 5dcf25354a..65b5334895 100644
--- a/runtime/libia2/init.c
+++ b/runtime/libia2/init.c
@@ -43,8 +43,9 @@ void verify_tls_padding(void) {
   }
 }
 
-/* Ensure that all required pkeys are allocated. */
+/* Ensure that all required pkeys are allocated or no-op on aarch64. */
 void ensure_pkeys_allocated(int *n_to_alloc) {
+#if LIBIA2_X86_64
   if (*n_to_alloc != 0) {
     for (int pkey = 1; pkey <= *n_to_alloc; pkey++) {
       int allocated = pkey_alloc(0, 0);
@@ -61,6 +62,7 @@ void ensure_pkeys_allocated(int *n_to_alloc) {
     }
     *n_to_alloc = 0;
   }
+#endif
 }
 
 /* Forbid overwriting an existing stack. */

From e6aa11a5385e06c09517b26168430badd06f6d43 Mon Sep 17 00:00:00 2001
From: Ayrton Munoz <adm@immunant.com>
Date: Wed, 13 Mar 2024 00:56:53 -0400
Subject: [PATCH 10/31] tools: Add arch flag to rewriter

---
 tools/rewriter/GenCallAsm.cpp     |  2 +-
 tools/rewriter/GenCallAsm.h       |  6 +++++-
 tools/rewriter/SourceRewriter.cpp | 24 +++++++++++++++++++-----
 3 files changed, 25 insertions(+), 7 deletions(-)

diff --git a/tools/rewriter/GenCallAsm.cpp b/tools/rewriter/GenCallAsm.cpp
index 9461d9d607..6e8a1c131e 100644
--- a/tools/rewriter/GenCallAsm.cpp
+++ b/tools/rewriter/GenCallAsm.cpp
@@ -284,7 +284,7 @@ std::string emit_asm_wrapper(const CAbiSignature &sig,
                              const std::string &wrapper_name,
                              const std::optional<std::string> target_name,
                              WrapperKind kind, int caller_pkey, int target_pkey,
-                             bool as_macro) {
+                             Arch arch, bool as_macro) {
 
   // Small sanity check
   assert(caller_pkey != target_pkey);
diff --git a/tools/rewriter/GenCallAsm.h b/tools/rewriter/GenCallAsm.h
index ee21b69d9c..49557c2818 100644
--- a/tools/rewriter/GenCallAsm.h
+++ b/tools/rewriter/GenCallAsm.h
@@ -13,6 +13,10 @@ enum class WrapperKind {
   IndirectCallsite,
 };
 
+enum class Arch {
+    Aarch64, X86
+};
+
 // Generates a wrapper for a function named \p name with the signature \p sig.
 // The WrapperKind parameter \p kind determines the type of call which may
 // affect the order of operations and the layout of the wrapper stack frame.
@@ -28,4 +32,4 @@ std::string emit_asm_wrapper(const CAbiSignature &sig,
                              const std::string &wrapper_name,
                              const std::optional<std::string> target_name,
                              WrapperKind kind, int caller_pkey, int target_pkey,
-                             bool as_macro = false);
+                             Arch arch, bool as_macro = false);
diff --git a/tools/rewriter/SourceRewriter.cpp b/tools/rewriter/SourceRewriter.cpp
index 999548a1c1..8546bbaf59 100644
--- a/tools/rewriter/SourceRewriter.cpp
+++ b/tools/rewriter/SourceRewriter.cpp
@@ -48,6 +48,17 @@ typedef std::string OpaqueStruct;
 static llvm::cl::OptionCategory
     SourceRewriterCategory("Source rewriter options");
 
+static llvm::cl::opt<Arch> Target("arch",
+    llvm::cl::init(Arch::X86),
+    llvm::cl::Optional,
+    llvm::cl::cat(SourceRewriterCategory),
+    llvm::cl::desc("<aarch64 or x86>"),
+    llvm::cl::values(
+        clEnumValN(Arch::X86, "x86", "Generate code for compartmentalization on x86 using MPK"),
+        clEnumValN(Arch::Aarch64, "aarch64", "Generate code for compartmentalization on Aarch64 using MTE")
+    ));
+
+
 static llvm::cl::opt<std::string>
     RootDirectory("root-directory", llvm::cl::Required,
                   llvm::cl::cat(SourceRewriterCategory),
@@ -895,6 +906,9 @@ int main(int argc, const char **argv) {
       new_args.erase(std::remove_if(new_args.begin(), new_args.end(),
                                     [](std::string &x) { return x.starts_with("-DIA2_ENABLE="); }),
                      new_args.end());
+      if (Target == Arch::Aarch64) {
+        new_args.push_back("--target=arm64"s);
+      }
       new_args.push_back("-DIA2_ENABLE=0"s);
       return new_args;
   });
@@ -1089,7 +1103,7 @@ int main(int argc, const char **argv) {
                                  std::to_string(caller_pkey);
       std::string asm_wrapper =
           emit_asm_wrapper(c_abi_sig, wrapper_name, std::nullopt,
-                           WrapperKind::IndirectCallsite, caller_pkey, 0);
+                           WrapperKind::IndirectCallsite, caller_pkey, 0, Target);
       wrapper_out << "asm(\n";
       wrapper_out << asm_wrapper;
       wrapper_out << ");\n";
@@ -1152,7 +1166,7 @@ int main(int argc, const char **argv) {
       }
       std::string asm_wrapper =
           emit_asm_wrapper(c_abi_sig, wrapper_name, fn_name,
-                           WrapperKind::Direct, caller_pkey, target_pkey);
+                           WrapperKind::Direct, caller_pkey, target_pkey, Target);
       wrapper_out << "asm(\n";
       wrapper_out << asm_wrapper;
       wrapper_out << ");\n";
@@ -1175,7 +1189,7 @@ int main(int argc, const char **argv) {
     std::string wrapper_name = "__wrap_"s + fn_name;
     std::string asm_wrapper =
         emit_asm_wrapper(c_abi_sig, wrapper_name, fn_name, WrapperKind::Direct,
-                         0, compartment_pkey);
+                         0, compartment_pkey, Target);
     wrapper_out << "asm(\n";
     wrapper_out << asm_wrapper;
     wrapper_out << ");\n";
@@ -1204,7 +1218,7 @@ int main(int argc, const char **argv) {
       CAbiSignature c_abi_sig = fn_decl_pass.abi_signatures[fn_name];
       std::string asm_wrapper =
           emit_asm_wrapper(c_abi_sig, wrapper_name, fn_name,
-                           WrapperKind::Pointer, 0, target_pkey);
+                           WrapperKind::Pointer, 0, target_pkey, Target);
       wrapper_out << "asm(\n";
       wrapper_out << asm_wrapper;
       wrapper_out << ");\n";
@@ -1233,7 +1247,7 @@ int main(int argc, const char **argv) {
 
         std::string asm_wrapper = emit_asm_wrapper(
             c_abi_sig, wrapper_name, fn_name, WrapperKind::Pointer, 0,
-            target_pkey, true /* as_macro */);
+            target_pkey, Target, true /* as_macro */);
         static_wrappers += "#define IA2_DEFINE_WRAPPER_"s + fn_name + " \\\n";
         static_wrappers += "asm(\\\n";
         static_wrappers += asm_wrapper;

From 03d2ae42b9b34431aec1788aa34b2a4fced6b755 Mon Sep 17 00:00:00 2001
From: Ayrton Munoz <adm@immunant.com>
Date: Wed, 13 Mar 2024 01:15:47 -0400
Subject: [PATCH 11/31] tools: Use arch-dependent illegal instruction in
 rewriter

---
 tools/rewriter/SourceRewriter.cpp | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/tools/rewriter/SourceRewriter.cpp b/tools/rewriter/SourceRewriter.cpp
index 8546bbaf59..4f460aea0a 100644
--- a/tools/rewriter/SourceRewriter.cpp
+++ b/tools/rewriter/SourceRewriter.cpp
@@ -1259,8 +1259,16 @@ int main(int argc, const char **argv) {
       }
     }
   }
+  const char *undef_insn_x86 = "ud2";
+  const char *undef_insn_arm = "udf #0";
+  const char *undef_insn;
+  if (Target == Arch::Aarch64) {
+    undef_insn = undef_insn_arm;
+  } else {
+    undef_insn = undef_insn_x86;
+  }
   header_out << "asm(\"__libia2_abort:\\n\"\n"
-             << "    \"ud2\");\n";
+             << "    \"" << undef_insn << "\");\n";
   header_out << static_wrappers.c_str();
 
   for (int i = 0; i < num_pkeys; i++) {

From 7277e171fa4246fb337f6a031a0140ebde6ec6f6 Mon Sep 17 00:00:00 2001
From: Ayrton Munoz <adm@immunant.com>
Date: Wed, 13 Mar 2024 01:02:49 -0400
Subject: [PATCH 12/31] tools: make rewriter call gate no-op transition on
 aarch64

---
 tools/rewriter/GenCallAsm.cpp | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/tools/rewriter/GenCallAsm.cpp b/tools/rewriter/GenCallAsm.cpp
index 6e8a1c131e..37b10f7bbe 100644
--- a/tools/rewriter/GenCallAsm.cpp
+++ b/tools/rewriter/GenCallAsm.cpp
@@ -387,6 +387,11 @@ std::string emit_asm_wrapper(const CAbiSignature &sig,
   add_asm_line(aw, ".type "s + wrapper_name + ", @function");
   add_asm_line(aw, wrapper_name + ":");
 
+  if (arch == Arch::Aarch64) {
+      if (target_name) {
+        add_asm_line(aw, "b "s + target_name.value());
+      }
+  } else {
   // Save the old frame pointer and set the frame pointer for the call gate
   add_asm_line(aw, "pushq %rbp");
   add_asm_line(aw, "movq %rsp, %rbp");
@@ -618,6 +623,7 @@ std::string emit_asm_wrapper(const CAbiSignature &sig,
   // Return to the caller
   add_comment_line(aw, "Return to the caller");
   add_asm_line(aw, "ret");
+  }
   // Set the symbol size
   add_asm_line(aw, ".size "s + wrapper_name + ", .-" + wrapper_name);
 

From 429cea844203cb7a5126aafa988ce1766d0c5e4c Mon Sep 17 00:00:00 2001
From: Ayrton Munoz <adm@immunant.com>
Date: Thu, 14 Mar 2024 14:31:38 -0400
Subject: [PATCH 13/31] tools: Fix target passed to rewriter

---
 tools/rewriter/SourceRewriter.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tools/rewriter/SourceRewriter.cpp b/tools/rewriter/SourceRewriter.cpp
index 4f460aea0a..f49d57cec0 100644
--- a/tools/rewriter/SourceRewriter.cpp
+++ b/tools/rewriter/SourceRewriter.cpp
@@ -907,7 +907,7 @@ int main(int argc, const char **argv) {
                                     [](std::string &x) { return x.starts_with("-DIA2_ENABLE="); }),
                      new_args.end());
       if (Target == Arch::Aarch64) {
-        new_args.push_back("--target=arm64"s);
+        new_args.push_back("--target=aarch64-linux-gnu"s);
       }
       new_args.push_back("-DIA2_ENABLE=0"s);
       return new_args;

From 8373dcf8e90a37b22b017b24e2b7d7924445264f Mon Sep 17 00:00:00 2001
From: Ayrton Munoz <adm@immunant.com>
Date: Wed, 13 Mar 2024 01:00:24 -0400
Subject: [PATCH 14/31] tests: Add criterion-free test for aarch64

---
 tests/CMakeLists.txt                          |  3 +++
 tests/minimal_no_criterion/CMakeLists.txt     | 11 +++++++++
 tests/minimal_no_criterion/Output/minimal.out |  1 +
 tests/minimal_no_criterion/include/minimal.h  | 10 ++++++++
 tests/minimal_no_criterion/main.c             | 18 +++++++++++++++
 tests/minimal_no_criterion/minimal.c          | 23 +++++++++++++++++++
 6 files changed, 66 insertions(+)
 create mode 100644 tests/minimal_no_criterion/CMakeLists.txt
 create mode 100644 tests/minimal_no_criterion/Output/minimal.out
 create mode 100644 tests/minimal_no_criterion/include/minimal.h
 create mode 100644 tests/minimal_no_criterion/main.c
 create mode 100644 tests/minimal_no_criterion/minimal.c

diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt
index b4a6d0aa6b..c75d292125 100644
--- a/tests/CMakeLists.txt
+++ b/tests/CMakeLists.txt
@@ -43,6 +43,9 @@ add_subdirectory(heap_two_keys)
 # add_subdirectory(libusb)
 add_subdirectory(macro_attr)
 add_subdirectory(minimal)
+# FIXME: This test is only used for aarch64 since we need to build criterion
+# from source to support Aarch64
+add_subdirectory(minimal_no_criterion)
 add_subdirectory(mmap_loop)
 # TODO: support C++ namespaces
 #add_subdirectory(namespaces)
diff --git a/tests/minimal_no_criterion/CMakeLists.txt b/tests/minimal_no_criterion/CMakeLists.txt
new file mode 100644
index 0000000000..e0dfe47b5b
--- /dev/null
+++ b/tests/minimal_no_criterion/CMakeLists.txt
@@ -0,0 +1,11 @@
+# Build the wrapped lib
+define_shared_lib(SRCS minimal.c)
+
+# Build the test
+define_test(
+    SRCS main.c
+    NEEDS_LD_WRAP
+)
+
+# Build the wrapper lib
+define_ia2_wrapper()
diff --git a/tests/minimal_no_criterion/Output/minimal.out b/tests/minimal_no_criterion/Output/minimal.out
new file mode 100644
index 0000000000..257cc5642c
--- /dev/null
+++ b/tests/minimal_no_criterion/Output/minimal.out
@@ -0,0 +1 @@
+foo
diff --git a/tests/minimal_no_criterion/include/minimal.h b/tests/minimal_no_criterion/include/minimal.h
new file mode 100644
index 0000000000..b55f1ff740
--- /dev/null
+++ b/tests/minimal_no_criterion/include/minimal.h
@@ -0,0 +1,10 @@
+#pragma once
+
+// This function does nothing
+void foo();
+
+// This returns an integer
+int return_val();
+
+// This takes an integer
+void arg1(int x);
diff --git a/tests/minimal_no_criterion/main.c b/tests/minimal_no_criterion/main.c
new file mode 100644
index 0000000000..65702cd06e
--- /dev/null
+++ b/tests/minimal_no_criterion/main.c
@@ -0,0 +1,18 @@
+/*
+RUN: sh -c 'if [ ! -s "minimal_call_gates_0.ld" ]; then echo "No link args as expected"; exit 0; fi; echo "Unexpected link args"; exit 1;'
+*/
+
+// Check that readelf shows exactly one executable segment
+
+#include <stdio.h>
+#include "minimal.h"
+#include <ia2.h>
+
+INIT_RUNTIME(1);
+#define IA2_COMPARTMENT 1
+#include <ia2_compartment_init.inc>
+
+int main() {
+    printf("Calling foo");
+    foo();
+}
diff --git a/tests/minimal_no_criterion/minimal.c b/tests/minimal_no_criterion/minimal.c
new file mode 100644
index 0000000000..f8db4f9b0f
--- /dev/null
+++ b/tests/minimal_no_criterion/minimal.c
@@ -0,0 +1,23 @@
+/*
+RUN: cat minimal_call_gates_1.ld | FileCheck --check-prefix=LINKARGS %s
+*/
+
+#include "minimal.h"
+#include <stdio.h>
+
+// LINKARGS: --wrap=arg1
+void arg1(int x) {
+    printf("arg1");
+}
+
+// LINKARGS: --wrap=foo
+void foo() {
+    printf("foo");
+}
+
+// LINKARGS: --wrap=return_val
+int return_val() {
+    printf("return_val");
+    return 1;
+}
+

From 4d268a68b53ab3c8a43707a34be3f38784c1a6ec Mon Sep 17 00:00:00 2001
From: Ayrton Munoz <adm@immunant.com>
Date: Thu, 14 Mar 2024 12:29:34 -0400
Subject: [PATCH 15/31] external/nginx: Use tools target in nginx build

---
 external/nginx/reconfigure | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/external/nginx/reconfigure b/external/nginx/reconfigure
index 2c585dd62f..7f86736280 100755
--- a/external/nginx/reconfigure
+++ b/external/nginx/reconfigure
@@ -29,7 +29,7 @@ fi
 mkdir -p $IA2_BUILD_DIR
 pushd $IA2_BUILD_DIR
 cmake -GNinja $REPO_ROOT ${IA2_CMAKE_FLAGS:-}
-ninja ia2-rewriter libia2 pad-tls ${PA_LIBRARY}
+ninja tools libia2 ${PA_LIBRARY}
 popd
 
 mkdir -p $NGINX_BUILD_DIR

From 74b999e879c245e33c4eb44a3cf46f55c4132b76 Mon Sep 17 00:00:00 2001
From: Ayrton Munoz <adm@immunant.com>
Date: Thu, 14 Mar 2024 16:04:51 -0400
Subject: [PATCH 16/31] external/nginx: Set arch to x86-64 for nginx build

---
 external/nginx/auto/make | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/external/nginx/auto/make b/external/nginx/auto/make
index f98cac815b..b95f80ebb3 100644
--- a/external/nginx/auto/make
+++ b/external/nginx/auto/make
@@ -53,6 +53,7 @@ IA2_CFLAGS = \\
 	-fPIC \\
 	-DPKEY=1 \\
 	-DIA2_ENABLE=\$(IA2_ENABLE) \\
+	-DLIBIA2_X86_64=1 \\
 	-DLIBIA2_DEBUG=1 \\
 	-I \$(C_SYSTEM_INCLUDE) \\
 	-I \$(C_SYSTEM_INCLUDE_FIXED) \\
@@ -91,6 +92,7 @@ IA2_MODULE_CFLAGS = \\
 	-fPIC \\
 	-DPKEY=2 \\
 	-DIA2_ENABLE=\$(IA2_ENABLE) \\
+	-DLIBIA2_X86_64=1 \\
 	-DLIBIA2_DEBUG=1 \\
 	-I \$(C_SYSTEM_INCLUDE) \\
 	-I \$(C_SYSTEM_INCLUDE_FIXED) \\

From ff64490ddbb0ff5851bc7243d8878bdbd3fd3e57 Mon Sep 17 00:00:00 2001
From: Ayrton Munoz <adm@immunant.com>
Date: Thu, 14 Mar 2024 16:23:57 -0400
Subject: [PATCH 17/31] external/nginx: Set arch for nginx call gate

Note that the rewriter emits arch-specific assembly and this just sets the arch
for the scrub registers header.
---
 external/nginx/auto/make | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/external/nginx/auto/make b/external/nginx/auto/make
index b95f80ebb3..d80a85fa35 100644
--- a/external/nginx/auto/make
+++ b/external/nginx/auto/make
@@ -316,7 +316,7 @@ binary:	$NGX_OBJS${ngx_dirsep}nginx$ngx_binext
 
 \$(IA2_CALLGATES_TARGET):
 	\$(CC) -shared -fPIC -Wl,-z,now $NGX_OBJS/../wrapper.c \
-		-I \$(LIBIA2_DIR)/include -o \$(IA2_CALLGATES_TARGET)
+		-I \$(LIBIA2_DIR)/include -DLIBIA2_X86_64=1 -o \$(IA2_CALLGATES_TARGET)
 
 $NGX_OBJS/ngx_rtmp_module_tls_padded.so: $NGX_OBJS/ngx_rtmp_module.so
 	cp $NGX_OBJS/ngx_rtmp_module.so $NGX_OBJS/ngx_rtmp_module_tls_padded.so

From d722ae72e6e0eea612cf445cbdd55f1a26918a9c Mon Sep 17 00:00:00 2001
From: Ayrton Munoz <adm@immunant.com>
Date: Thu, 14 Mar 2024 16:10:38 -0400
Subject: [PATCH 18/31] docs: Update CMake target list

---
 docs/build_instructions.md | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/docs/build_instructions.md b/docs/build_instructions.md
index 4a9b3e49d2..f022369d37 100644
--- a/docs/build_instructions.md
+++ b/docs/build_instructions.md
@@ -46,8 +46,7 @@ cmake ..                                                 \
 ## CMake targets
 
 - `check` - builds and runs the test suite. Pass `-v` to ninja to see build commands and output from failing tests.
-- `ia2-rewriter` - builds the source-code rewriter. Depends on libclang-dev and llvm-dev.
-- `pad-tls` - builds a script for padding ELF headers for TLS segments. Only required for compartmentalized DSOs that use thread-local storage.
+- `tools` - builds the source-code rewriter and pad-tls script. The former depends on libclang-dev and llvm-dev. The latter is only required for compartmentalized DSOs that use thread-local storage.
 - `libia2` - builds the runtime as a static library. This does not include call gate transitions as those are program-specific and generated by the rewriter. 
 - `partition-alloc-padding` - builds the compartment-aware shim for Chromium's PartitionAlloc allocator.
 - `ia2-sandbox` - builds the syscall tracer.

From 9f20a09539d37defdcfdba6523b61ed74d775eec Mon Sep 17 00:00:00 2001
From: Ayrton Munoz <adm@immunant.com>
Date: Wed, 13 Mar 2024 01:45:51 -0400
Subject: [PATCH 19/31] CI: Add basic CI for aarch64

---
 .github/workflows/ci.yml | 31 +++++++++++++++++++++++++++++++
 1 file changed, 31 insertions(+)

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index d70f204663..0b8f44fda9 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -65,3 +65,34 @@ jobs:
           make
           popd
           popd
+  aarch64-wip-test:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Check out code
+        uses: actions/checkout@v3
+      - name: Test ARM build
+        run: |
+          sudo apt-get update
+          sudo apt-get install -y \
+                gcc-aarch64-linux-gnu \
+                g++-aarch64-linux-gnu \
+                clang-15 pkg-config \
+                cmake ninja-build \
+                llvm-15-dev libclang-15-dev \
+                zlib1g-dev \
+                qemu-user-static
+          LLVM_DIR=`llvm-config-15 --cmakedir`
+          Clang_DIR=`realpath $LLVM_DIR/../clang`
+          mkdir build
+          pushd build
+          cmake .. \
+            -DCMAKE_TOOLCHAIN_FILE=../cmake/aarch64-toolchain.cmake \
+            -DClang_DIR=$Clang_DIR                                  \
+            -DCLANG_EXE=`which clang-15`                            \
+            -DLLVM_DIR=$LLVM_DIR                                    \
+            -G Ninja
+          ninja -v minimal_no_criterion_main_wrapped
+          qemu-aarch64 -L /usr/aarch64-linux-gnu/ \
+            tests/minimal_no_criterion/minimal_no_criterion_main_wrapped
+          popd
+

From d9357f47eab8c7b73fb9d549ee3246a3a6643b5e Mon Sep 17 00:00:00 2001
From: Ayrton Munoz <adm@immunant.com>
Date: Thu, 14 Mar 2024 14:39:14 -0400
Subject: [PATCH 20/31] CI: Fix qemu invocation for aarch64

---
 .github/workflows/ci.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 0b8f44fda9..f6d6a16c57 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -92,7 +92,7 @@ jobs:
             -DLLVM_DIR=$LLVM_DIR                                    \
             -G Ninja
           ninja -v minimal_no_criterion_main_wrapped
-          qemu-aarch64 -L /usr/aarch64-linux-gnu/ \
+          qemu-aarch64-static -L /usr/aarch64-linux-gnu/ \
             tests/minimal_no_criterion/minimal_no_criterion_main_wrapped
           popd
 

From 04da869c8c12f7f1a7419c6b386485927abe3a73 Mon Sep 17 00:00:00 2001
From: Ayrton Munoz <adm@immunant.com>
Date: Wed, 13 Mar 2024 00:05:32 -0400
Subject: [PATCH 21/31] cmake: Add aarch64 toolchain file

---
 cmake/aarch64-toolchain.cmake | 14 ++++++++++++++
 1 file changed, 14 insertions(+)
 create mode 100644 cmake/aarch64-toolchain.cmake

diff --git a/cmake/aarch64-toolchain.cmake b/cmake/aarch64-toolchain.cmake
new file mode 100644
index 0000000000..e5b2beb600
--- /dev/null
+++ b/cmake/aarch64-toolchain.cmake
@@ -0,0 +1,14 @@
+set(LIBIA2_AARCH64 TRUE)
+
+set(CMAKE_SYSTEM_NAME Linux)
+
+set(CMAKE_C_COMPILER aarch64-linux-gnu-gcc)
+set(CMAKE_CXX_COMPILER aarch64-linux-gnu-g++)
+
+set(CMAKE_C_HOST_COMPILER gcc)
+set(CMAKE_CXX_HOST_COMPILER g++)
+
+set(CMAKE_FIND_ROOT_PATH_MODE_PROGRAM NEVER)
+
+set(CMAKE_FIND_ROOT_PATH_MODE_LIBRARY ONLY)
+set(CMAKE_FIND_ROOT_PATH_MODE_INCLUDE ONLY)

From 9bec906aadf8279b68c39cb67052e14909c37eef Mon Sep 17 00:00:00 2001
From: Ayrton Munoz <adm@immunant.com>
Date: Thu, 14 Mar 2024 16:51:55 -0400
Subject: [PATCH 22/31] cmake: Add support for building host tools when
 cross-compiling

---
 CMakeLists.txt                | 14 ++++++++++++--
 cmake/ia2.cmake               |  8 ++++----
 tests/CMakeLists.txt          |  2 +-
 tools/rewriter/CMakeLists.txt | 13 +++++++++++--
 4 files changed, 28 insertions(+), 9 deletions(-)

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 1910c81e67..177add11b1 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -1,5 +1,6 @@
 cmake_minimum_required(VERSION 3.13)
 project(IA2Phase2)
+include(ExternalProject)
 
 find_package(PkgConfig REQUIRED)
 
@@ -10,6 +11,15 @@ set(EXTERNAL_DIR ${PROJECT_SOURCE_DIR}/external)
 # runtime needs to be first so it defines libia2_BINARY_DIR
 add_subdirectory(runtime)
 
-add_subdirectory(examples)
+if (NOT LIBIA2_AARCH64)
+    add_subdirectory(examples)
+endif()
+
 add_subdirectory(tests)
-add_subdirectory(tools)
+ExternalProject_Add(tools
+    SOURCE_DIR ${CMAKE_SOURCE_DIR}/tools
+    BINARY_DIR ${CMAKE_BINARY_DIR}/tools
+    CMAKE_ARGS
+        -DClang_DIR=${Clang_DIR}
+        -DLLVM_DIR=${LLVM_DIR}
+    INSTALL_COMMAND "")
diff --git a/cmake/ia2.cmake b/cmake/ia2.cmake
index be8a3486a6..ececcbd462 100644
--- a/cmake/ia2.cmake
+++ b/cmake/ia2.cmake
@@ -94,8 +94,8 @@ function(pad_tls_library INPUT OUTPUT)
   add_custom_command(
     OUTPUT ${CMAKE_CURRENT_BINARY_DIR}/lib${OUTPUT}.so
     COMMAND ${CMAKE_COMMAND} -E copy $<TARGET_FILE:${INPUT}> ${CMAKE_CURRENT_BINARY_DIR}/lib${OUTPUT}.so
-    COMMAND pad-tls --allow-no-tls ${CMAKE_CURRENT_BINARY_DIR}/lib${OUTPUT}.so
-    DEPENDS pad-tls $<TARGET_FILE:${INPUT}>
+    COMMAND ${CMAKE_BINARY_DIR}/tools/pad-tls/pad-tls --allow-no-tls ${CMAKE_CURRENT_BINARY_DIR}/lib${OUTPUT}.so
+    DEPENDS tools $<TARGET_FILE:${INPUT}>
     COMMENT "Padding TLS segment of wrapped library"
   )
   add_custom_target(${OUTPUT}-padding DEPENDS "${CMAKE_CURRENT_BINARY_DIR}/lib${OUTPUT}.so")
@@ -246,7 +246,7 @@ function(add_ia2_call_gates NAME)
 
   add_custom_command(
     OUTPUT ${CALL_GATE_SRC} ${CALL_GATE_HDR} ${LD_ARGS_FILES} ${REWRITTEN_SOURCES}
-    COMMAND ia2-rewriter
+    COMMAND ${CMAKE_BINARY_DIR}/tools/rewriter/ia2-rewriter
         --output-prefix=${REWRITER_OUTPUT_PREFIX}
         --root-directory=${CMAKE_CURRENT_SOURCE_DIR}
         --output-directory=${CMAKE_CURRENT_BINARY_DIR}
@@ -257,7 +257,7 @@ function(add_ia2_call_gates NAME)
         ${ARG_EXTRA_REWRITER_ARGS}
         ${SOURCES}
     WORKING_DIRECTORY ${CMAKE_BINARY_DIR}
-    DEPENDS ia2-rewriter ${SOURCES}
+    DEPENDS tools ${SOURCES}
     VERBATIM
   )
 
diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt
index c75d292125..988d0ba534 100644
--- a/tests/CMakeLists.txt
+++ b/tests/CMakeLists.txt
@@ -20,7 +20,7 @@ configure_lit_site_cfg(
   ${CMAKE_CURRENT_SOURCE_DIR}/lit.cfg.py
   )
 
-set(IA2_TEST_DEPENDS ia2-rewriter)
+set(IA2_TEST_DEPENDS tools)
 
 add_lit_testsuite(check-ia2 "Running the IA2 tests"
   ${CMAKE_CURRENT_BINARY_DIR}
diff --git a/tools/rewriter/CMakeLists.txt b/tools/rewriter/CMakeLists.txt
index 2fc0891fa2..1d5a1d4543 100644
--- a/tools/rewriter/CMakeLists.txt
+++ b/tools/rewriter/CMakeLists.txt
@@ -6,8 +6,9 @@ message(STATUS "Found LLVM: ${LLVM_DIR} (found version \"${LLVM_PACKAGE_VERSION}
 find_package(Clang REQUIRED CONFIG)
 message(STATUS "Found Clang: ${Clang_DIR}")
 
-add_definitions(${CLANG_DEFINITIONS} ${LLVM_DEFINITIONS})
-include_directories(${CLANG_INCLUDE_DIRS} ${LLVM_INCLUDE_DIRS})
+# We use the lit CMake functions from LLVM
+list(APPEND CMAKE_MODULE_PATH "${LLVM_CMAKE_DIR}")
+include(AddLLVM)
 
 add_executable(ia2-rewriter
     DetermineAbi.cpp
@@ -20,6 +21,14 @@ set_target_properties(ia2-rewriter PROPERTIES
     CXX_STANDARD 20
     CXX_STANDARD_REQUIRED ON)
 
+target_compile_definitions(ia2-rewriter PRIVATE
+    ${CLANG_DEFINITIONS}
+    ${LLVM_DEFINITIONS})
+
+target_include_directories(ia2-rewriter PRIVATE
+    ${CLANG_INCLUDE_DIRS}
+    ${LLVM_INCLUDE_DIRS})
+
 target_compile_options(ia2-rewriter PRIVATE "-Wunused")
 target_link_libraries(ia2-rewriter PRIVATE
     clang-cpp

From dd3e078c13747fffe33d888af63322ec3df4c38f Mon Sep 17 00:00:00 2001
From: Ayrton Munoz <adm@immunant.com>
Date: Thu, 14 Mar 2024 16:55:19 -0400
Subject: [PATCH 23/31] cmake: Temporarily remove dependency on ubsan and
 partition-alloc on aarch64

---
 cmake/ia2.cmake | 13 ++++++++++---
 1 file changed, 10 insertions(+), 3 deletions(-)

diff --git a/cmake/ia2.cmake b/cmake/ia2.cmake
index ececcbd462..dd9805eec4 100644
--- a/cmake/ia2.cmake
+++ b/cmake/ia2.cmake
@@ -1,3 +1,10 @@
+if(LIBIA2_AARCH64)
+    set(UBSAN_FLAG "")
+    set(PARTITION_ALLOC "")
+else()
+    set(UBSAN_FLAG "-fsanitize=undefined")
+    set(PARTITION_ALLOC "partition-alloc")
+endif()
 # Creates a compartmentalized IA2 target
 #
 # add_ia2_compartment(<name> <EXECUTABLE|LIBRARY> PKEY <n> SOURCES <src>...
@@ -53,11 +60,11 @@ function(add_ia2_compartment NAME TYPE)
 
   if (ARG_ENABLE_UBSAN)
     # UBSAN requires passing this as both a compiler and linker flag
-    target_compile_options(${NAME} PRIVATE "-fsanitize=undefined")
-    target_link_options(${NAME} PRIVATE "-fsanitize=undefined")
+    target_compile_options(${NAME} PRIVATE ${UBSAN_FLAG})
+    target_link_options(${NAME} PRIVATE ${UBSAN_FLAG})
   endif()
 
-  target_link_libraries(${NAME} PRIVATE dl libia2 partition-alloc)
+  target_link_libraries(${NAME} PRIVATE dl libia2 ${PARTITION_ALLOC})
   target_link_options(${NAME} PRIVATE "-Wl,--export-dynamic")
 
   target_link_libraries(${NAME} PRIVATE ${ARG_LIBRARIES})

From 85f21c543879885790640142e6e55f0afdebd438 Mon Sep 17 00:00:00 2001
From: Ayrton Munoz <adm@immunant.com>
Date: Thu, 14 Mar 2024 16:55:35 -0400
Subject: [PATCH 24/31] cmake: Disable tracer for aarch64

---
 cmake/define-test.cmake | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/cmake/define-test.cmake b/cmake/define-test.cmake
index 4479a66469..bad55b6447 100644
--- a/cmake/define-test.cmake
+++ b/cmake/define-test.cmake
@@ -108,6 +108,10 @@ function(define_test)
         set(DEFINE_TEST_LIBS ${TEST_NAME}_lib)
     endif()
 
+    if (LIBIA2_AARCH64)
+        set(DEFINE_TEST_WITHOUT_SANDBOX TRUE)
+    endif()
+
 #    if(DEFINED DEFINE_TEST_INCLUDE_DIR)
 #        set(RELATIVE_INCLUDE_DIR ${DEFINE_TEST_INCLUDE_DIR})
 #    else()

From 148b251fa2b7592cf5afa3ca2c8b398a7a6a0961 Mon Sep 17 00:00:00 2001
From: Ayrton Munoz <adm@immunant.com>
Date: Thu, 14 Mar 2024 17:00:18 -0400
Subject: [PATCH 25/31] cmake: Add workaround for aarch64 ld.so lookup error

---
 cmake/ia2.cmake | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/cmake/ia2.cmake b/cmake/ia2.cmake
index dd9805eec4..4105055310 100644
--- a/cmake/ia2.cmake
+++ b/cmake/ia2.cmake
@@ -241,6 +241,12 @@ function(add_ia2_call_gates NAME)
         endif()
       endif()
 
+      # FIXME: This shouldn't be necessary but it seems aarch64-gcc < v13 might
+      # default to --as-needed so this is needed to fix some runtime ld.so lookup
+      # error
+      if (LIBIA2_AARCH64)
+          target_link_options(${target} PRIVATE "-Wl,--no-as-needed")
+      endif()
       target_link_libraries(${target} PRIVATE ${CALL_GATE_TARGET})
 
       if("${target_pkey}" GREATER "0")

From 5c59a6e014e48d1dec8ddf08bd9dfa60d2dfbb6e Mon Sep 17 00:00:00 2001
From: Ayrton Munoz <adm@immunant.com>
Date: Thu, 14 Mar 2024 17:01:10 -0400
Subject: [PATCH 26/31] cmake: Add arch flag to rewriter invocation

---
 cmake/ia2.cmake | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/cmake/ia2.cmake b/cmake/ia2.cmake
index 4105055310..eb2f8bca52 100644
--- a/cmake/ia2.cmake
+++ b/cmake/ia2.cmake
@@ -257,12 +257,16 @@ function(add_ia2_call_gates NAME)
     endif()
   endforeach()
 
+  if (LIBIA2_AARCH64)
+      set(ARCH_FLAG "--arch=aarch64")
+  endif()
   add_custom_command(
     OUTPUT ${CALL_GATE_SRC} ${CALL_GATE_HDR} ${LD_ARGS_FILES} ${REWRITTEN_SOURCES}
     COMMAND ${CMAKE_BINARY_DIR}/tools/rewriter/ia2-rewriter
         --output-prefix=${REWRITER_OUTPUT_PREFIX}
         --root-directory=${CMAKE_CURRENT_SOURCE_DIR}
         --output-directory=${CMAKE_CURRENT_BINARY_DIR}
+        ${ARCH_FLAG}
         # Set the build path so the rewriter can find the compile_commands JSON
         -p=${CMAKE_BINARY_DIR}
         --extra-arg=-isystem "--extra-arg=${CLANG_HEADERS_INCLUDE}"

From a8c66694264e83860116431ecd09ff18da8939ac Mon Sep 17 00:00:00 2001
From: Ayrton Munoz <adm@immunant.com>
Date: Thu, 14 Mar 2024 17:12:41 -0400
Subject: [PATCH 27/31] cmake: Add missing include for tests directory

---
 tests/CMakeLists.txt | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt
index 988d0ba534..feafdced90 100644
--- a/tests/CMakeLists.txt
+++ b/tests/CMakeLists.txt
@@ -2,6 +2,7 @@ include(CTest)
 
 enable_testing()
 
+include("../cmake/define-ia2-wrapper.cmake")
 include("../cmake/define-test.cmake")
 
 find_package(LLVM REQUIRED CONFIG)

From 2fb59b809280f90b5be61b743ebed08c50048362 Mon Sep 17 00:00:00 2001
From: Ayrton Munoz <adm@immunant.com>
Date: Thu, 14 Mar 2024 17:17:26 -0400
Subject: [PATCH 28/31] runtime/libia2: Fix CMake dependency on pad-tls

---
 runtime/libia2/CMakeLists.txt | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/runtime/libia2/CMakeLists.txt b/runtime/libia2/CMakeLists.txt
index e1cece4cd5..24837da1ac 100644
--- a/runtime/libia2/CMakeLists.txt
+++ b/runtime/libia2/CMakeLists.txt
@@ -60,8 +60,8 @@ endif()
 add_custom_command(
     OUTPUT ${CMAKE_CURRENT_BINARY_DIR}/libc.so.6
     COMMAND ${CMAKE_COMMAND} -E copy ${LIBC_PATH} ${CMAKE_CURRENT_BINARY_DIR}/libc.so.6
-    COMMAND pad-tls --allow-no-tls ${CMAKE_CURRENT_BINARY_DIR}/libc.so.6
-    DEPENDS pad-tls
+    COMMAND ${CMAKE_BINARY_DIR}/tools/pad-tls/pad-tls --allow-no-tls ${CMAKE_CURRENT_BINARY_DIR}/libc.so.6
+    DEPENDS tools
     COMMENT "Padding TLS segment of libc"
 )
 

From 80b73f8bbabf11ba92ab2efc8398ebb7fb3d5725 Mon Sep 17 00:00:00 2001
From: Ayrton Munoz <adm@immunant.com>
Date: Thu, 14 Mar 2024 17:17:57 -0400
Subject: [PATCH 29/31] CI: Fix target name for WIP aarch64 test

---
 .github/workflows/ci.yml | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index f6d6a16c57..12870e14bc 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -91,8 +91,8 @@ jobs:
             -DCLANG_EXE=`which clang-15`                            \
             -DLLVM_DIR=$LLVM_DIR                                    \
             -G Ninja
-          ninja -v minimal_no_criterion_main_wrapped
+          ninja -v minimal_no_criterion
           qemu-aarch64-static -L /usr/aarch64-linux-gnu/ \
-            tests/minimal_no_criterion/minimal_no_criterion_main_wrapped
+            tests/minimal_no_criterion/minimal_no_criterion
           popd
 

From 7035d9b6034d372378c2071d735a8f93a9cb0487 Mon Sep 17 00:00:00 2001
From: Ayrton Munoz <adm@immunant.com>
Date: Tue, 26 Mar 2024 08:33:20 -0400
Subject: [PATCH 30/31] runtime/libia2: Fix tag used in ia2_mprotect_with_tag

---
 runtime/libia2/ia2.c | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/runtime/libia2/ia2.c b/runtime/libia2/ia2.c
index 2345ef5739..644804a44c 100644
--- a/runtime/libia2/ia2.c
+++ b/runtime/libia2/ia2.c
@@ -133,7 +133,9 @@ int ia2_mprotect_with_tag(void *addr, size_t len, int prot, int tag) {
     /* Assuming we're using st2g. stgm is undefined at EL0 so it's not an option */
     const int granule_sz = 32;
     const int granules_per_page = PAGE_SIZE / 32;
-    size_t tag = ia2_get_tag();
+    /* small sanity check */
+    size_t current_tag = ia2_get_tag();
+    assert(current_tag == tag);
     for (int i = 0; i < granules_per_page; i++) {
         // TODO: It may be possible to simplify this to be more efficient using the addg imm offset
         uint64_t tagged_ptr = insert_tag((uint64_t)addr + (i * granule_sz), tag);

From 203f278dea7542f42a2f2dc225a89eda2b0974ec Mon Sep 17 00:00:00 2001
From: Ayrton Munoz <adm@immunant.com>
Date: Tue, 26 Mar 2024 13:13:55 -0400
Subject: [PATCH 31/31] cmake: Fix dependency between rewrite targets and
 rewriter

CMake has weird rules about what dependencies cause custom commands to be
re-run. See the comment in ia2.cmake for details.
---
 cmake/ia2.cmake | 15 ++++++++++++++-
 1 file changed, 14 insertions(+), 1 deletion(-)

diff --git a/cmake/ia2.cmake b/cmake/ia2.cmake
index eb2f8bca52..d2f64936ac 100644
--- a/cmake/ia2.cmake
+++ b/cmake/ia2.cmake
@@ -172,6 +172,15 @@ execute_process(COMMAND ${CLANG_EXE} -print-file-name=include-fixed
   OUTPUT_STRIP_TRAILING_WHITESPACE)
 message(STATUS "Found Clang fixed headers: ${CLANG_HEADERS_INCLUDE_FIXED}")
 
+file(GLOB REWRITER_SRCS ${CMAKE_SOURCE_DIR}/tools/rewriter/*.cpp)
+# This cannot be in the tools directory CMakeLists.txt because the target is for
+# the top-level CMake project
+add_custom_target(rewriter
+    COMMAND ${CMAKE_COMMAND} --build . -t ia2-rewriter
+    WORKING_DIRECTORY ${CMAKE_BINARY_DIR}/tools
+    # tools dependency is for the CMake config step
+    DEPENDS tools ${REWRITER_SRCS})
+
 # Create call gates for a target
 #
 # Creates call gates for executable target <name> and all of its
@@ -274,7 +283,11 @@ function(add_ia2_call_gates NAME)
         ${ARG_EXTRA_REWRITER_ARGS}
         ${SOURCES}
     WORKING_DIRECTORY ${CMAKE_BINARY_DIR}
-    DEPENDS tools ${SOURCES}
+    # dependencies on custom targets (i.e. the rewriter) do not re-run this
+    # command so we need to add a dependency on the rewriter's sources. We still
+    # need the dependency on the custom rewriter target to make sure it gets
+    # built the first time.
+    DEPENDS ${SOURCES} ${REWRITER_SRCS} rewriter
     VERBATIM
   )