#include <arm_acle.h>
#include <asm/hwcap.h>
#include <linux/mman.h>
#include <stdio.h>
#include <sys/auxv.h>
#include <sys/mman.h>
#include <sys/prctl.h>
#include <unistd.h>

int print_result(char *ptr) {
  // Page size allows the test to try reading off of the end of the page
  printf("buffer: %p page_size: 0x%x\n", ptr, sysconf(_SC_PAGESIZE));

  // Exit after some time, so we don't leave a zombie process
  // if the test framework lost track of us.
  sleep(60);
  return 0;
}

int main(int argc, char const *argv[]) {
  if (prctl(PR_SET_TAGGED_ADDR_CTRL,
            PR_TAGGED_ADDR_ENABLE | PR_MTE_TCF_SYNC |
                // Allow all tags to be generated by the addg
                // instruction __arm_mte_increment_tag produces.
                (0xffff << PR_MTE_TAG_SHIFT),
            0, 0, 0)) {
    return print_result(NULL);
  }

  size_t page_size = sysconf(_SC_PAGESIZE);
  char *buf = mmap(0, page_size, PROT_READ | PROT_WRITE | PROT_MTE,
                   MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
  if (buf == MAP_FAILED)
    return print_result(NULL);

  // Set incrementing tags until end of the page
  char *tagged_ptr = buf;
  // This intrinsic treats the addresses as if they were untagged
  while (__arm_mte_ptrdiff(tagged_ptr, buf) < page_size) {
    // This sets the allocation tag
    __arm_mte_set_tag(tagged_ptr);
    // Set the tag of the next granule (hence +16) to the next
    // tag value. Returns a new pointer with the new logical tag.
    // Tag values wrap at 0xF so it'll cycle.
    tagged_ptr = __arm_mte_increment_tag(tagged_ptr + 16, 1);
  }

  // lldb-server should be removing the top byte from addresses passed
  // to ptrace. So put some random bits in there.
  // ptrace expects you to remove them but it can still succeed if you
  // don't. So this isn't proof that we're removing them, it's just a
  // smoke test in case something didn't account for them.
  buf = (char *)((size_t)buf | ((size_t)0xAA << 56));
  return print_result(buf);
}
