test_mte_customWorkload.cpp

/*
 * To be compiled with -march=armv8.5-a+memtag
 */
#include <errno.h>
#include <stdint.h>
#include <stdio.h>
#include <cstdlib>
#include <unistd.h>
#include <sys/auxv.h>
#include <sys/mman.h>
#include <sys/prctl.h>
// #include <fstream>
// #include <vector>
// #include <string>

/*
 * From arch/arm64/include/uapi/asm/hwcap.h
 */
#define HWCAP2_MTE              (1 << 18)

/*
 * From arch/arm64/include/uapi/asm/mman.h
 */
#define PROT_MTE                 0x20

/*
 * From include/uapi/linux/prctl.h
 */
#define PR_SET_TAGGED_ADDR_CTRL 55
#define PR_GET_TAGGED_ADDR_CTRL 56
# define PR_TAGGED_ADDR_ENABLE  (1UL << 0)
# define PR_MTE_TCF_SHIFT       1
# define PR_MTE_TCF_NONE        (0UL << PR_MTE_TCF_SHIFT)
# define PR_MTE_TCF_SYNC        (1UL << PR_MTE_TCF_SHIFT)
# define PR_MTE_TCF_ASYNC       (2UL << PR_MTE_TCF_SHIFT)
# define PR_MTE_TCF_MASK        (3UL << PR_MTE_TCF_SHIFT)
# define PR_MTE_TAG_SHIFT       3
# define PR_MTE_TAG_MASK        (0xffffUL << PR_MTE_TAG_SHIFT)

#define GRANULE_SZ 16
#define MALLOC_HDR_SZ 64

/*
 * Insert a random logical tag into the given pointer.
 */
#define insert_random_tag(ptr) ({                       \
        uint64_t __val;                                 \
        asm("irg %0, %1" : "=r" (__val) : "r" (ptr));   \
        printf("IN insert_random_tag: tagged ptr: %llu \n", __val);               \
        __val;                                          \
})

/*
 * Set the allocation tag of 1granule on the destination address.
 */
        // printf("IN set_tag: for addr %p \n", tagged_addr);             
        // printf("Done with set_tag!\n");                                
#define set_tag(tagged_addr) do {                                      \
        asm volatile("stg %0, [%0]" : : "r" (tagged_addr) : "memory"); \
} while (0)

void printIntArr(int* arr, int len)
{
        for (int x = 0; x < len; x++)
                printf("arr[%i] : %i, ", x, arr[x]);
        printf("\n");
}

int main(int argc, char const *argv[])
{
        // Make malloc_sizes a LARGE array!
        int total_mallocs = 0;
        long total_mallocs_sz = 0;
        int malloc_sizes[5000]; // +ve elems: new malloc, -ve elems: del (-1*elem)th malloc, i.e. tag=0
        int malloc_ids[5000];
        int malloc_orig_offsets[5000];

        int useless_arr[1000];

        // long long int first_malloc_addr = 0;
        // // open file:
        FILE *fp;
        fp = fopen("apache_workload_Logv4_readValg_19351.txt", "r");
        // // read file: put malloc sizes & freed mallocs in malloc_sizes arr.
        if (fp)
        {
                char buf[256];
                while (true)
                {
                        if (fscanf(fp, "%s", buf) == EOF)
                                break;
                        if (buf[0] == 'm')
                        {
                                int msz; // char msz[256];
                                long moffset; // char moffset[256];
                                int mid; // char mid[256];
                                
                                if ((fscanf(fp, "%i", &msz) == 0) || (fscanf(fp, "%li", &moffset) == 0) || (fscanf(fp, "%i", &mid) == 0))
                                        printf("Parse: ERRORRRRRR!!!! malloc: sz / offset / id unavailable!!! \n");
                                malloc_sizes[total_mallocs] = msz;
                                malloc_orig_offsets[total_mallocs] = moffset;
                                malloc_ids[total_mallocs] = mid;

                                total_mallocs += 1;
                                total_mallocs_sz = (total_mallocs_sz > (moffset+msz)) ? (total_mallocs_sz) : (moffset+msz);

                                printf("Parse: GOT malloc: sz %i, offset: %li, id: %i, total_mallocs: %i, total_mallocs_sz: %li \n", (msz), (moffset), (mid), total_mallocs, total_mallocs_sz);
                        }
                        else if (buf[0] == 'f')
                        {
                                long foffset; // [256];
                                int fid; // [256];

                                if ( (fscanf(fp, "%li", &foffset) == 0) || (fscanf(fp, "%i", &fid) == 0) )
                                        printf("Parse: ERRORRRRRR!!!! free: offset / id unavailable!!!! \n");
                                
                                malloc_sizes[total_mallocs] = -1 * fid;
                                total_mallocs += 1;

                                if (malloc_orig_offsets[fid] != foffset)
                                        printf("Parse: ERRORRRR!!!! free command offset not matching!!! \n");
                                printf("Parse: GOT free: id: %i, offset: %li, malloc_orig_offsets: %i \n", (fid), (foffset), malloc_orig_offsets[fid]);
                        }
                }
        }
        // {
        //         char buf [1024];
        //         while (true)
        //         {       
        //                 // fscanf: reads until first space, fgets: reads until ? chars or until first \n | EOF
        //                 if (fscanf(buf, 1024, fp) == NULL)
        // Line0: has global base addr.
        // if "malloc" : get id, sz, ptr. [add to malloc_sizes, malloc_ids, malloc_orig_offsets]
        // if "malloc": get id, offset, sz.
        // if "free" : get id
        //         }
        // }
        // TODO: Use malloc last addr - first_malloc_addr as total_mallocs_sz.

        // For testing: Test1: first 4 allocs of apache:
        // total_mallocs = 4;
        // malloc_sizes[0] = 208;
        // malloc_orig_offsets[0] = 64;
        // malloc_sizes[1] = 4096;
        // malloc_orig_offsets[1] = 336;
        // malloc_sizes[2] = 8192;
        // malloc_orig_offsets[2] = 4496;
        // malloc_sizes[3] = 8192;
        // malloc_orig_offsets[3] = 12752;
        // total_mallocs_sz = 20944;

        // Test2:
        // total_mallocs = 10;
        // for (int i = 0; i < total_mallocs; i++)
        // {
        //         if ((i < 3) || (i == 6) || (i == 8))
        //         {
        //                 malloc_sizes[i] = 16 * (i+1) * (i/2 + 1);
        //                 malloc_orig_offsets[i] = total_mallocs_sz + MALLOC_HDR_SZ;
        //                 total_mallocs_sz += malloc_sizes[i] + MALLOC_HDR_SZ;
        //         }
        //         else if ((i >= 3) && (i < 6))
        //                 malloc_sizes[i] = -1*(i-3); // free-ing prev mallocs
        // }
        // malloc_sizes[7] = -6;
        // malloc_sizes[9] = -8;

        char *page_ptr;
        unsigned long page_sz = sysconf(_SC_PAGESIZE);
        unsigned long hwcap2 = getauxval(AT_HWCAP2);

        /* check if MTE is present */
        if (!(hwcap2 & HWCAP2_MTE))
                return EXIT_FAILURE;

        /*
         * Enable the tagged address ABI, synchronous or asynchronous MTE
         * tag check faults (based on per-CPU preference) and allow all
         * non-zero tags in the randomly generated set.
         */
        if (prctl(PR_SET_TAGGED_ADDR_CTRL,
                  PR_TAGGED_ADDR_ENABLE | PR_MTE_TCF_SYNC | PR_MTE_TCF_ASYNC |
                  (0xfffe << PR_MTE_TAG_SHIFT),
                  0, 0, 0)) {
                perror("prctl() failed");
                return EXIT_FAILURE;
        }

        // need to explicitly cast since c++ doesn't allow casting void* to unsigned char* [that'd be okay in C]
        // addr (arg0), offset (arg4) need to be a multiple of page sz, but not length.
        long total_mmap_sz = page_sz * (total_mallocs_sz/page_sz + 1);
        
        
        page_ptr = static_cast<char*>(mmap(0, total_mmap_sz, PROT_READ | PROT_WRITE,
                 MAP_PRIVATE | MAP_ANONYMOUS, -1, 0));
        if (page_ptr == MAP_FAILED) {
                perror("mmap() failed");
                return EXIT_FAILURE;
        }

        // need to enable MTE protection for the mem starting from page_ptr, of page_sz bytes.
        if (mprotect(page_ptr, total_mmap_sz, PROT_READ | PROT_WRITE | PROT_MTE)) {
                perror("mprotect() failed");
                return EXIT_FAILURE;
        }   

        printf("TOTAL sz of mallocs: %li, did mmap sz %li, page_ptr: %p, Starting to allocate now! \n", total_mallocs_sz, total_mmap_sz, page_ptr);

        // START allocating mem for malloc_sizes arr:
        unsigned int malloc_start_inds[total_mallocs]; // stores start_ind offset OF actual_malloc ptr (i.e. end of Metadata) from page_ptr.
        // char *malloc_start_ptrs[total_mallocs]; NOT USING THIS cuz storing the actual start_ptr in any local variable is causing segFaults.
        unsigned long malloc_tags[total_mallocs];
        int first_avail = 0;

        for (int i = 0; i < total_mallocs; i++)
        {
                if (malloc_sizes[i] > 0)
                {
                        int ith_offset = (first_avail + MALLOC_HDR_SZ); // counting metadata space
                        if (malloc_orig_offsets[i] >= (ith_offset) )
                        {
                                // malloc - using new space
                                malloc_start_inds[i] = ith_offset;
                                first_avail = (ith_offset + malloc_sizes[i]);
                                printf("MALLOC: i=%i Allcoated fresh! at offset: %i, orig_offset: %i, sz: %i \n", i, ith_offset, malloc_orig_offsets[i], malloc_sizes[i]);
                        }
                        else
                        {
                                malloc_start_inds[i] = malloc_orig_offsets[i];
                                printf("MALLOC: i=%i RE-using FREED memory! first_avail: %i, org_offset: %i, sz: %i \n", i, first_avail, malloc_orig_offsets[i], malloc_sizes[i]);
                        }
                        printf("page + offset: %p, i: %i, total_mallocs: %i, ", (page_ptr + malloc_start_inds[i]), i, total_mallocs );
                        // char* start_ptr = (char*)(page_ptr + malloc_start_inds[i]); // seg-fault!!!
                        // printf("Untagged start_ptr: %p, ", start_ptr); // malloc_start_ptrs[i]

                        // tag first 4 granules with tag=0 (NOTE, we have metadata for every alloc)
                        // if (i > 0)
                        for (int gmi = 0; gmi < MALLOC_HDR_SZ/GRANULE_SZ; gmi++)
                                set_tag(page_ptr + malloc_start_inds[i] - MALLOC_HDR_SZ + gmi*16);

                        printf("Finished set_tag for %i bytes, tag=0, ", MALLOC_HDR_SZ);

                        // generate new random tag
                        malloc_tags[i] = (((char*)insert_random_tag(page_ptr + malloc_start_inds[i])) - page_ptr) >> 56;
                        printf("ith Tag: %lu , Tagged start_ptr: %p \n", malloc_tags[i], (page_ptr + malloc_start_inds[i] + (malloc_tags[i] << 56)) );

                        // tag all granules with new tag
                        for (int gi = 0; gi < (malloc_sizes[i]/GRANULE_SZ); gi++)
                        {
                                if (gi % 50 == 0)
                                        printf("MALLOC=%i ABOUT to set tag for gi = %i ", i, gi);
                                set_tag( ((page_ptr + malloc_start_inds[i] + (malloc_tags[i] << 56))) + gi*GRANULE_SZ);
                        }

                        // check access working: for all bytes
                        // printf("Checking access now: \n");
                        for (int mi = 0; mi < MALLOC_HDR_SZ; mi++)
                                useless_arr[mi%1000] = *((page_ptr + malloc_start_inds[i] + 0 - MALLOC_HDR_SZ) + mi); // tag=0

                        for (int bi = 0; bi < malloc_sizes[i]; bi++)
                        {
                                // char* i_bi_byte = ((page_ptr + malloc_start_inds[i] + (malloc_tags[i] << 56)))+bi+MALLOC_HDR_SZ;
                                // *i_bi_byte = bi + 7;
                                // printf(" for bi = %i \n", bi);
                                useless_arr[bi%1000] = *((page_ptr + malloc_start_inds[i] + (malloc_tags[i] << 56)) + bi) ;
                                // printf("mi[%i] = %i, ", bi, *((page_ptr + malloc_start_inds[i] + (malloc_tags[i] << 56)) + MALLOC_HDR_SZ + bi) );
                        }
                        // printf("Trying to access MALLOC METADATA: (should segfault) \n");
                        // printf("malloc_start_ptrs[%i][0]: %i \n", i, malloc_start_ptrs[i][0]); // THIS DID SEGFAULT!!! YAYY!!!
                }
                else
                {
                        int free_index = -1*malloc_sizes[i];
                        printf("FREE-ing index %i at offset: %i (tag: %lu) \n !!!", free_index, malloc_start_inds[free_index], malloc_tags[free_index]);

                        // tag metadata & mem with 0
                        malloc_tags[free_index] = 0;
                        // malloc_start_ptrs[free_index] = page_ptr + malloc_start_inds[free_index]; NOT USING THIS ARRAY [segFault]

                        for (int gi = 0; gi < (malloc_sizes[free_index] + MALLOC_HDR_SZ)/GRANULE_SZ; gi++)
                                set_tag( (page_ptr + malloc_start_inds[free_index]) + gi*GRANULE_SZ - MALLOC_HDR_SZ);
                }
        }        

        return 0;
}

// stg count:
// 46*4 md-t0
//