// mtetoyctl.c - userland CLI for the "MTE" mac_policy syscall toybox
//
// Build:
//   clang -O2 -Wall -Wextra -std=c11 -o mtetoyctl mtetoyctl.c
//
// Examples:
//   sudo ./mtetoyctl gettoymem
//   sudo ./mtetoyctl malloc 0x1000
//   sudo ./mtetoyctl ldg 0xfffffe0000000000
//   sudo ./mtetoyctl stg 0xfffffe0000000000 0xabfffffe0000000000
//   sudo ./mtetoyctl irg 0xfffffe0000000000 0x0
//   sudo ./mtetoyctl gmi 0xabfffffe0000000000 0x0
//   sudo ./mtetoyctl read 8  0xabfffffe0000000000
//   sudo ./mtetoyctl read 16 0xabfffffe0000000000
//   sudo ./mtetoyctl write 8  0xabfffffe0000000000 0x1122334455667788
//   sudo ./mtetoyctl write 16 0xabfffffe0000000000 0x1122334455667788 0x99aabbccddeeff00
//
// Notes:
// - This talks to your kernel MAC policy (mpc_name "MTE") via __mac_syscall.
// - Addresses returned by GETTOYMEM/MALLOC are *kernel* addresses in your policy implementation.

#include <errno.h>
#include <inttypes.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#if !defined(__APPLE__)
#error "This tool is intended for macOS."
#endif

// Userland prototype for the MAC policy syscall wrapper.
extern int __mac_syscall(char *policy, int call, uintptr_t arg);

#define POLICY_NAME "MTE"

// Call numbers from your policy.
#define MTE_TOYBOX_GETTOYMEM    0
#define MTE_TOYBOX_MALLOC       1
#define MTE_TOYBOX_LDG          2
#define MTE_TOYBOX_STG          3
#define MTE_TOYBOX_IRG          4
#define MTE_TOYBOX_GMI          5
#define MTE_TOYBOX_READ         6
#define MTE_TOYBOX_WRITE        7

// Structs matching mtetoybox.c exactly.
typedef struct {
    uintptr_t addr;
} mte_toybox_gettoymem_t;

typedef struct {
    uint64_t  size;
    uintptr_t addr;
} mte_toybox_malloc_t;

typedef struct {
    uintptr_t addr;
    uintptr_t tagged_addr;
} mte_toybox_ldg_t;

typedef struct {
    uintptr_t addr;
    uintptr_t tagged_addr;
} mte_toybox_stg_t;

typedef struct {
    uintptr_t addr;
    uintptr_t tagged_addr;
    uint64_t  exclude;
} mte_toybox_irg_t;

typedef struct {
    uintptr_t tagged_addr;
    uint64_t  exclude_in;
    uint64_t  exclude_out;
} mte_toybox_gmi_t;

typedef struct {
    uint64_t  size;
    uintptr_t tagged_addr;
    uint64_t  low;
    uint64_t  high;
} mte_toybox_read_t;

typedef mte_toybox_read_t mte_toybox_write_t;

_Static_assert(sizeof(mte_toybox_gettoymem_t) == 8,  "size mismatch gettoymem");
_Static_assert(sizeof(mte_toybox_malloc_t)    == 16, "size mismatch malloc");
_Static_assert(sizeof(mte_toybox_ldg_t)       == 16, "size mismatch ldg");
_Static_assert(sizeof(mte_toybox_stg_t)       == 16, "size mismatch stg");
_Static_assert(sizeof(mte_toybox_irg_t)       == 24, "size mismatch irg");
_Static_assert(sizeof(mte_toybox_gmi_t)       == 24, "size mismatch gmi");
_Static_assert(sizeof(mte_toybox_read_t)      == 32, "size mismatch read/write");

static uint64_t parse_u64(const char *s) {
    if (!s || !*s) {
        fprintf(stderr, "error: missing numeric argument\n");
        exit(2);
    }
    errno = 0;
    char *end = NULL;
    unsigned long long v = strtoull(s, &end, 0);
    if (errno || !end || *end != '\0') {
        fprintf(stderr, "error: invalid number: %s\n", s);
        exit(2);
    }
    return (uint64_t)v;
}

static uintptr_t parse_uptr(const char *s) {
    return (uintptr_t)parse_u64(s);
}

static void die_errno(const char *what) {
    fprintf(stderr, "error: %s: %s (errno=%d)\n", what, strerror(errno), errno);
    exit(1);
}

static int mte_call(int call, void *buf) {
    int rc = __mac_syscall((char *)POLICY_NAME, call, (uintptr_t)buf);
    if (rc < 0) {
        // __mac_syscall returns -1 and sets errno on failure.
        return -1;
    }
    return rc;
}

static void print_ptr(const char *name, uintptr_t v) {
    printf("%s = 0x%016" PRIxPTR "\n", name, v);
    printf("  tag = 0x%02" PRIxPTR ", untagged = 0x%016" PRIxPTR "\n",
           (v >> 56) & 0xFF, v | 0xFF00000000000000ULL);
}

static void usage(const char *argv0) {
    fprintf(stderr,
        "Usage:\n"
        "  %s gettoymem\n"
        "  %s malloc <size>\n"
        "  %s ldg <addr>\n"
        "  %s stg <addr> <tagged_addr>\n"
        "  %s irg <addr> [exclude]\n"
        "  %s gmi <tagged_addr> [exclude_in]\n"
        "  %s read  <size:{1,2,4,8,16}> <tagged_addr>\n"
        "  %s write <size:{1,2,4,8,16}> <tagged_addr> <low> [high_if_16]\n"
        "\n"
        "Notes:\n"
        "  - Talks to mac_policy syscall policy \"%s\".\n",
        argv0, argv0, argv0, argv0, argv0, argv0, argv0, argv0, POLICY_NAME
    );
    exit(2);
}

int main(int argc, char **argv) {
    if (argc < 2) usage(argv[0]);

    const char *cmd = argv[1];

    if (strcmp(cmd, "gettoymem") == 0) {
        mte_toybox_gettoymem_t d = {0};
        int rc = mte_call(MTE_TOYBOX_GETTOYMEM, &d);
        if (rc < 0) die_errno("__mac_syscall(GETTOYMEM)");
        printf("rc=%d\n", rc);
        print_ptr("toymem", d.addr);
        return 0;
    }

    if (strcmp(cmd, "malloc") == 0) {
        if (argc != 3) usage(argv[0]);
        mte_toybox_malloc_t d = {0};
        d.size = parse_u64(argv[2]);

        int rc = mte_call(MTE_TOYBOX_MALLOC, &d);
        if (rc < 0) die_errno("__mac_syscall(MALLOC)");
        printf("rc=%d\n", rc);
        printf("size = 0x%016" PRIx64 "\n", d.size);
        print_ptr("addr", d.addr);
        return 0;
    }

    if (strcmp(cmd, "ldg") == 0) {
        if (argc != 3) usage(argv[0]);
        mte_toybox_ldg_t d = {0};
        d.addr = parse_uptr(argv[2]);

        int rc = mte_call(MTE_TOYBOX_LDG, &d);
        if (rc < 0) die_errno("__mac_syscall(LDG)");
        printf("rc=%d\n", rc);
        print_ptr("addr", d.addr);
        print_ptr("tagged_addr", d.tagged_addr);
        return 0;
    }

    if (strcmp(cmd, "stg") == 0) {
        if (argc != 4) usage(argv[0]);
        mte_toybox_stg_t d = {0};
        d.addr = parse_uptr(argv[2]);
        d.tagged_addr = parse_uptr(argv[3]);

        int rc = mte_call(MTE_TOYBOX_STG, &d);
        if (rc < 0) die_errno("__mac_syscall(STG)");
        printf("rc=%d\n", rc);
        print_ptr("addr", d.addr);
        print_ptr("tagged_addr", d.tagged_addr);
        return 0;
    }

    if (strcmp(cmd, "irg") == 0) {
        if (argc != 3 && argc != 4) usage(argv[0]);
        mte_toybox_irg_t d = {0};
        d.addr = parse_uptr(argv[2]);
        d.exclude = (argc == 4) ? parse_u64(argv[3]) : 0;

        int rc = mte_call(MTE_TOYBOX_IRG, &d);
        if (rc < 0) die_errno("__mac_syscall(IRG)");
        printf("rc=%d\n", rc);
        print_ptr("addr", d.addr);
        printf("exclude = 0x%016" PRIx64 "\n", d.exclude);
        print_ptr("tagged_addr", d.tagged_addr);
        return 0;
    }

    if (strcmp(cmd, "gmi") == 0) {
        if (argc != 3 && argc != 4) usage(argv[0]);
        mte_toybox_gmi_t d = {0};
        d.tagged_addr = parse_uptr(argv[2]);
        d.exclude_in  = (argc == 4) ? parse_u64(argv[3]) : 0;
        d.exclude_out = 0;

        int rc = mte_call(MTE_TOYBOX_GMI, &d);
        if (rc < 0) die_errno("__mac_syscall(GMI)");
        printf("rc=%d\n", rc);
        print_ptr("tagged_addr", d.tagged_addr);
        printf("exclude_in  = 0x%016" PRIx64 "\n", d.exclude_in);
        printf("exclude_out = 0x%016" PRIx64 "\n", d.exclude_out);
        return 0;
    }

    if (strcmp(cmd, "read") == 0) {
        if (argc != 4) usage(argv[0]);
        mte_toybox_read_t d = {0};
        d.size = parse_u64(argv[2]);
        d.tagged_addr = parse_uptr(argv[3]);

        if (!(d.size == 1 || d.size == 2 || d.size == 4 || d.size == 8 || d.size == 16)) {
            fprintf(stderr, "error: read size must be 1,2,4,8,16\n");
            return 2;
        }

        int rc = mte_call(MTE_TOYBOX_READ, &d);
        if (rc < 0) die_errno("__mac_syscall(READ)");
        printf("rc=%d\n", rc);
        printf("size = %" PRIu64 "\n", d.size);
        print_ptr("tagged_addr", d.tagged_addr);

        if (d.size == 16) {
            printf("value = 0x%016" PRIx64 "%016" PRIx64 "  (high:low)\n", d.high, d.low);
        } else {
            printf("value = 0x%016" PRIx64 "\n", d.low);
        }
        return 0;
    }

    if (strcmp(cmd, "write") == 0) {
        if (argc != 5 && argc != 6) usage(argv[0]);
        mte_toybox_write_t d = {0};
        d.size = parse_u64(argv[2]);
        d.tagged_addr = parse_uptr(argv[3]);
        d.low = parse_u64(argv[4]);
        d.high = (argc == 6) ? parse_u64(argv[5]) : 0;

        if (!(d.size == 1 || d.size == 2 || d.size == 4 || d.size == 8 || d.size == 16)) {
            fprintf(stderr, "error: write size must be 1,2,4,8,16\n");
            return 2;
        }
        if (d.size == 16 && argc != 6) {
            fprintf(stderr, "error: write size 16 requires <high>\n");
            return 2;
        }

        int rc = mte_call(MTE_TOYBOX_WRITE, &d);
        if (rc < 0) die_errno("__mac_syscall(WRITE)");
        printf("rc=%d\n", rc);
        printf("size = %" PRIu64 "\n", d.size);
        print_ptr("tagged_addr", d.tagged_addr);
        if (d.size == 16) {
            printf("wrote = 0x%016" PRIx64 "%016" PRIx64 "  (high:low)\n", d.high, d.low);
        } else {
            printf("wrote = 0x%016" PRIx64 "\n", d.low);
        }
        return 0;
    }

    usage(argv[0]);
}