smalloc.cpp

#include <smalloc-cpp.h>

/* heap address => size */
/* static */ std::map<int64_t, int64_t> *tree = NULL;
/* global var address => address of second half in heap */
/* static */ std::map<int64_t, int64_t> *global_tree = NULL;
/* stack address => address of second half in heap, size of stack allocation */
/* static */ std::vector<std::map<int64_t, std::pair<int64_t, int64_t> > *>
    *stack_tree_list = NULL;

struct MemoryBounds {
    std::pair<u_int64_t, u_int64_t> global;
    std::pair<u_int64_t, u_int64_t> stack;
    std::pair<u_int64_t, u_int64_t> heap;
} memory_bounds = {std::make_pair(18446744073709551615, 0), std::make_pair(18446744073709551615, 0), std::make_pair(18446744073709551615, 0)};

/* static */ void /* __attribute__((optnone))  */
updateAddr(std::map<int64_t, int64_t> *tree, int64_t start_addr, int64_t size) {
    assert(tree != NULL && "");
    tree->insert(std::make_pair(start_addr, size));
    // tree->_root = smalloc__insert_node(tree->_root, start_addr, size);
}

/* static */ void /* __attribute__((optnone))  */
updateAddrWithSize(std::map<int64_t, std::pair<int64_t, int64_t> > *tree,
                   int64_t start_addr, int64_t heap_addr, int64_t size) {
    assert(tree != NULL && "");
    tree->insert(std::make_pair(start_addr, std::make_pair(heap_addr, size)));
    // tree->_root = smalloc__insert_node(tree->_root, start_addr, size);
}

/* static */ void * /* __attribute__((optnone))  */ setBit(void *addr) {
    /* assert(checkBit(addr) == 0 && "setBit expect a non secret address"); */
    int64_t int_addr = (int64_t)addr;
    int_addr |= 1UL << smalloc__position;
    return (void *)int_addr;
}

/* static */ void * /* __attribute__((optnone))  */ unsetBit(void *addr) {
    // printf("running unsetBit for %ld ", addr);
    /* assert(checkBit(addr) == 1 && "unsetBit expect a secret address"); */
    /* if (checkBit(addr) != 1) printf("WARN: unsetBit expect a secret address\n"); */

    if (checkBit(addr) == 0) {
        return addr;
    }

    int64_t int_addr = (int64_t)addr;
    int_addr &= ~(1UL << smalloc__position);
    // printf("and returning %ld \n", int_addr);
    return (void *)int_addr;
}

/* static */ int /* __attribute__((optnone))  */ checkBit(void *addr) {
    int64_t int_addr = (int64_t)addr;
    if (memory_bounds.global.first <= int_addr && memory_bounds.global.second >= int_addr) 
        return 1U;
    // secret return 1
    // not secret return 0
    return ((int_addr >> smalloc__position) & 1U);
}

/* static */ void * /* __attribute__((optnone))  */ identity(void *addr) {
    printf("running identity for %ld ", addr);
    int64_t int_addr = (int64_t)addr;
    int_addr |= 1UL << smalloc__position;
    int_addr &= ~(1UL << smalloc__position);
    return (void *)int_addr;
}

/* static */ void * /* __attribute__((optnone))  */ getAddress(void *addr) {
    assert(addr != NULL && "getAddress received null");
    int64_t int_addr = (int64_t)addr;
     // Finally check if it's a global address (lowest addresses)
    if (memory_bounds.global.first <= int_addr && memory_bounds.global.second >= int_addr) 
    if (global_tree) {
        auto res = global_tree->lower_bound(int_addr);
        if (res != global_tree->end() && res->first == int_addr) {
            return (void *)res->second;
        } else if (res != global_tree->begin()) {
            res = std::prev(res, 1);
            int64_t offset = (int64_t)int_addr - (int64_t)res->first;
            int64_t result = res->second + offset;
            return (void *)result;
        }
    }   
    // First check if it's a stack address (highest addresses)
    if (stack_tree_list) {
        for (auto I = stack_tree_list->rbegin(); I != stack_tree_list->rend(); I++) {
            auto res = (*I)->lower_bound(int_addr);
            if (res != stack_tree_list->back()->end() && res->first == int_addr) {
                return (void *)res->second.first;
            } else if (res != (*I)->begin()) {
                res = std::prev(res, 1);
                int64_t start_addr = res->first;
                int64_t size = res->second.second;
                if (int_addr > start_addr + size) {
                    continue;
                } else {
                    void *heap_addr = (void *)(res->second.first + int_addr - start_addr);
                    return heap_addr;
                }
            }
        }
    }
    
    // Then check if it's a heap address (middle addresses)
    if (tree) {
        auto res = tree->lower_bound(int_addr);
        if (res != tree->end() && res->first == int_addr)
            return (void *)(int_addr + res->second / 2);
        else if (res != tree->begin()) {
            res = std::prev(res, 1);
            assert(int_addr <= res->first + res->second / 2 && "out of bound address");
            void *result = (void *)(int_addr + res->second / 2);
            return result;
        }
    }
    
    return NULL;
}

/* static */ void* /* __attribute__((optnone))  */ getBaseAddr(void *addr) {
    int64_t int_addr = (int64_t)addr;
    int64_t base_addr = -1;

    if (memory_bounds.global.first <= int_addr && memory_bounds.global.second >= int_addr) 
    /* printf("getBaseAddr: 0x%lx \n", int_addr); */
    if (global_tree) {
        auto res = global_tree->lower_bound(int_addr);
        if (res != global_tree->end() && res->first == int_addr) {
            base_addr = res->first;
            return (void *)base_addr;
        } else if (res != global_tree->begin()) {
            auto res1 = std::prev(res, 1);
            base_addr = res1->first;
            return (void *)base_addr;
        }
    }

    if (stack_tree_list) {
        for (auto I = stack_tree_list->rbegin(); I != stack_tree_list->rend();
             I++) {
            auto res = (*I)->lower_bound(int_addr);
            if (res != stack_tree_list->back()->end() &&
                res->first == int_addr) {
                base_addr = res->first;
                return (void *)base_addr;
            } else if (res != (*I)->begin()) {
                res = std::prev(res, 1);
                base_addr = res->first;
                return (void *)base_addr;
            }
        }
    }

    if (tree) {
        auto res = tree->lower_bound(int_addr);
        if (res != tree->end() && res->first == int_addr){
            base_addr = res->first;
            return (void *)base_addr;
        }   
        else if (res != tree->begin()) {
            res = std::prev(res, 1);
            base_addr = res->first;
            return (void *)base_addr;
        }
    }

    /* assert(base_addr != -1 && "getBaseAddr returned null"); */
    return nullptr;
    /* return (void *)base_addr; */
}

/* static */ int64_t* /* __attribute__((optnone))  */ getAlignedBaseAddr(void *addr) {
    // cast addr into integer
    int64_t int_addr = (int64_t)addr;
    // get the base address
    int64_t base_addr = (int64_t)getBaseAddr(addr);
    // get the offset
    int64_t offset = int_addr - base_addr;
    // get 8 byte aligned base
    int64_t base8 = base_addr + (offset / 8) * 8;
    return (int64_t*)base8;
}

/* static */ int /* __attribute__((optnone))  */ isShadowAddr(void *addr, int size) {
    void *base_addr = getBaseAddr(addr);
    int64_t int_addr = (int64_t)addr;
    int64_t base_addr_int = (int64_t)base_addr;
    assert(base_addr_int != -1 && "base addr not found");
    int64_t offset = int_addr - base_addr_int;
    int64_t idx = offset / 4;
    return idx % 2;
}

/* static */ int /* __attribute__((optnone))  */ getFirstChunkSize(void *addr, int size) {
    // cast addr into integer
    int64_t int_addr = (int64_t)addr;
    // get 8 byte aligned base
    int64_t base8 = (int64_t)getAlignedBaseAddr(addr);
    // get offset from 8 byte aligned base
    int64_t offset8 = int_addr - base8;
    // get the size of the first chunk
    int chunk_size = 4 - (offset8 % 4);
    int result = chunk_size;
    if (chunk_size > size) 
        result = size;
    return result;
}

/* static */ int /* __attribute__((optnone))  */ getLastChunkSize(void *addr, int size) {
    // get the first chunk size
    int first_chunk_size = getFirstChunkSize(addr, size);
    return (size - first_chunk_size) % 4;
}

/* static */ int /* __attribute__((optnone))  */ getNumFullChunk(void *addr, int size) {
    // get the first chunk size
    int first_chunk_size = getFirstChunkSize(addr, size);
    // get the last chunk size
    int last_chunk_size = getLastChunkSize(addr, size);
    // get the number of full chunks
    return (size - first_chunk_size - last_chunk_size) / 4;
}

/* static */ int64_t /* __attribute__((optnone))  */ getFirstValue(unsigned int Data, int size, int offset) {
    // get lowest size bits of Data
    int64_t temp = (1LL << size * 8) - 1;
    Data &= (int)temp;
    Data = Data << offset * 8;
    int64_t mask = ~0 ^ temp;
    int64_t junk = 0xdeadceef00000000;
    junk = junk & mask;
    return junk | Data;
}

/* static */ int64_t /* __attribute__((optnone))  */ getFirstValueBit(unsigned int Data, int size, int offset) {
    // get lowest size bits of Data
    int64_t temp = (1LL << size) - 1;
    Data &= (int)temp;
    Data = Data << offset * 8;
    int64_t mask = ~0 ^ temp;
    int64_t junk = 0xdeadceef00000000;
    junk = junk & mask;
    return junk | Data;
}

/* static */ int64_t /* __attribute__((optnone))  */ maskLoadedData(int Data, int64_t loadedData, int size, void* addr) {

    // cast addr into integer
    int64_t int_addr = (int64_t)addr;
    // get 8 byte aligned base
    int64_t base8 = (int64_t)getAlignedBaseAddr(addr);
    // get offset from 8 byte aligned base
    int64_t offset8 = (int_addr - base8) % 4;

    // get highest size bits of Data
    // print all the arguments in hex
    /* printf("Data: 0x%x, loadedData: 0x%lx, size: 0x%x\n", Data, loadedData, size); */
    int64_t mask = ~0LL ^ ((1LL << size * 8) - 1);
    mask = (mask << offset8 * 8) | ((1LL << offset8 * 8) - 1);
    int64_t maskedLoadedData = loadedData & (mask & 0xffffffff);
    int64_t dataWithJunk = getFirstValue(Data, size, offset8 % 4);
    return maskedLoadedData | dataWithJunk;
}

/* static */ int64_t /* __attribute__((optnone))  */ maskLoadedDataBit(int Data, int64_t loadedData, int size, void* addr) {

    // cast addr into integer
    int64_t int_addr = (int64_t)addr;
    // get 8 byte aligned base
    int64_t base8 = (int64_t)getAlignedBaseAddr(addr);
    // get offset from 8 byte aligned base
    int64_t offset8 = (int_addr - base8) % 4;

    // get highest size bits of Data
    // print all the arguments in hex
    /* printf("Data: 0x%x, loadedData: 0x%lx, size: 0x%x\n", Data, loadedData, size); */
    int64_t mask = ~0LL ^ ((1LL << size) - 1);
    mask = (mask << offset8 * 8) | ((1LL << offset8 * 8) - 1);
    int64_t maskedLoadedData = loadedData & (mask & 0xffffffff);
    int64_t dataWithJunk = getFirstValueBit(Data, size, offset8 % 4);
    return maskedLoadedData | dataWithJunk;
}

/* static */ int64_t /* __attribute__((optnone))  */ unmaskLoadedData(int64_t loadedData, int size, void* addr, int64_t secret, int64_t oldOffset) {

    int Data = 0;
    // cast addr into integer
    int64_t int_addr = (int64_t)addr;
    // get 8 byte aligned base
    int64_t base8 = (int64_t)getAlignedBaseAddr(addr);
    // get offset from 8 byte aligned base
    int64_t offset8 = (int_addr - base8) % 4;

    // get highest size bits of Data
    // print all the arguments in hex
    /* printf("Data: 0x%x, loadedData: 0x%lx, size: 0x%x\n", Data, loadedData, size); */
    int64_t mask = ~0LL ^ ((1LL << size * 8) - 1);
    mask = (mask << offset8 * 8) | ((1LL << offset8 * 8) - 1);
    mask = ~mask;
    int64_t maskedLoadedData = loadedData & mask;
    maskedLoadedData = maskedLoadedData >> offset8 * 8;
    /* maskedLoadedData = maskedLoadedData << oldOffset; */
    return maskedLoadedData;
}

/* static */ void * /* __attribute__((optnone))  */ incrementAddress(void *addr, int __size) {
    return (void *)((int64_t)addr + __size);
}

/* static */ void * /* __attribute__((optnone))  */ getStoreAddress(void *addr, int __size) {
    // update address map with size
    return getAddress(addr);
}

/* static */ int /* __attribute__((optnone))  */ getLoadSplitSize(void *addr) {
    // return the size fron the address map
    return 0;
}

/* static */ void * /* __attribute__((optnone))  */ smalloc(size_t __size) {
    /* printf("running smalloc "); */
    // align size to 8
    if (__size % 8 != 0) {
        __size += 8 - (__size % 8);
    }
    void *start_addr = malloc(2 * __size);
    int64_t int_addr = (int64_t)start_addr;

    if (memory_bounds.heap.first > int_addr)
        memory_bounds.heap.first = int_addr;
    if (memory_bounds.heap.second < int_addr + 2 * __size)
        memory_bounds.heap.second = int_addr + 2 * __size;

    if (!tree)
        tree = new std::map<int64_t, int64_t>;
    updateAddr(tree, int_addr, 2 * __size);
    void *addr = (void *)int_addr;
    /* printf("returned 0x%lx \n", addr); */
    addr = setBit(addr);
    return addr;
}

/* static */ void * /* __attribute__((optnone))  */
smmap(void *addr, size_t __size, int prot, int flags, int fd, size_t offset) {
    /* printf("running smalloc\n"); */
    // align size to 8
    if (__size % 8 != 0) {
        __size += 8 - (__size % 8);
    }
    void *start_addr = mmap(addr, 2 * __size, prot, flags, fd, offset);
    int64_t int_addr = (int64_t)start_addr;

    if (memory_bounds.heap.first > int_addr)
        memory_bounds.heap.first = int_addr;
    if (memory_bounds.heap.second < int_addr + 2 * __size)
        memory_bounds.heap.second = int_addr + 2 * __size;

    /* printf("Debug: 0x%lx and 0x%lx \n", int_addr, int_addr + 2 * __size); */
    if (!tree)
        tree = new std::map<int64_t, int64_t>;
    updateAddr(tree, int_addr, 2 * __size);
    addr = (void *)int_addr;
    addr = setBit(addr);
    return addr;
}

/* static */ void /* __attribute__((optnone))  */ sfree(void *addr) {
    /* printf("running sfree\n"); */
    assert(addr != NULL && "do not free a null ptr");
    int64_t int_addr = (int64_t)addr;
    // assert(int_addr & (1UL << smalloc__position) &&
    //        "non secret address used with secret free");
    int_addr = (int64_t)unsetBit((void *)int_addr);
    tree->erase(int_addr);
    free((void *)int_addr);
}

/* static */ int /* __attribute__((optnone))  */ smunmap(void *addr,
                                                         size_t __size) {
    /* printf("running sfree\n"); */
    assert(addr != NULL && "do not free a null ptr");
    int64_t int_addr = (int64_t)addr;
    // assert(int_addr & (1UL << smalloc__position) &&
    //        "non secret address used with secret free");
    int_addr = (int64_t)unsetBit((void *)int_addr);
    if(tree) tree->erase(int_addr);
    // align size to 8
    if (__size % 8 != 0) {
        __size += 8 - (__size % 8);
    }
    return munmap((void *)int_addr, 2 * __size);
}

void initStack() {
    /* printf("init stack\n"); */
    std::map<int64_t, std::pair<int64_t, int64_t> > *tree =
        new std::map<int64_t, std::pair<int64_t, int64_t> >;
    if (!stack_tree_list) {
        stack_tree_list =
            new std::vector<std::map<int64_t, std::pair<int64_t, int64_t> > *>;
    }
    stack_tree_list->push_back(tree);
}

void finiStack() {
    /* printf("fini stack\n"); */
    assert(stack_tree_list != NULL && "list must not be null");
    std::map<int64_t, std::pair<int64_t, int64_t> > *heap =
        stack_tree_list->back();
    auto idx = heap->begin();
    while (idx != heap->end()) {
        void *addr = (void *)idx->second.first;
        // free(addr);
        idx = std::next(idx, 1);
    }
    // for(std::pair<int64_t, int64_t> element : *heap) {
    //     free((void*)element.second);
    // }
    stack_tree_list->pop_back();
}

bool isAddrInStackMap(void *addr) {
    assert(stack_tree_list != NULL && "list must not be null");
    std::map<int64_t, std::pair<int64_t, int64_t> > *stackMap =
        stack_tree_list->back();
    if (stackMap->find((int64_t)addr) != stackMap->end())
        return true;
    return false;
}

void *updateStackAddr(void *addr, int64_t size) {
    // expects size in byte
    /* printf("update stack size: %ld \n", size); */
    assert(addr != NULL && "addr must not be null");
    if (size % 8 != 0) {
        size += 8 - (size % 8);
    }
    void *heap_address = malloc(size);
    /* printf("updateStackAddr heap : 0x%lx \n", addr); */
    /* printf("Debug: 0x%lx and ... \n", heap_address); */
    updateAddrWithSize(stack_tree_list->back(), (int64_t)addr,
                       (int64_t)heap_address, size);

    if (memory_bounds.stack.first > (int64_t)addr)
        memory_bounds.stack.first = (int64_t)addr;
    if (memory_bounds.stack.second < (int64_t)addr + size)
        memory_bounds.stack.second = (int64_t)addr + size;

    return setBit(addr);
}

void *updateGlobalAddr(void *addr, int64_t size, int64_t heap_value) {
    addr = unsetBit(addr);
    if (!global_tree)
        global_tree = new std::map<int64_t, int64_t>;
    void *heap_address = malloc(size);
    /* *(int64_t *)heap_address = heap_value; */
    updateAddr(global_tree, (int64_t)addr, (int64_t)heap_address);
    /* printf("updateGlobalAddr Query : 0x%lx and 0x%lx \n", addr, size); */
    /* printf("Debug: 0x%lx and 0x%lx \n", heap_address, heap_value); */

    if (memory_bounds.global.first > (int64_t)addr)
        memory_bounds.global.first = (int64_t)addr;
    if (memory_bounds.global.second < (int64_t)addr + size)
        memory_bounds.global.second = (int64_t)addr + size;

    return setBit(addr);
}

void destroyGlobalAddr() {
    if (!global_tree)
        return;
    for (std::pair<int64_t, int64_t> element : *global_tree) {
        free((void *)element.second);
    }
}

int32_t handle32BitLoad(void* addr) {
    int size = 32;
    int firstSize = 0;
    int oldFirstSize = 0;
    int32_t secret = 0;
    int requiredSize = size / 8;
    addr = unsetBit(addr);
    for (int i = 0; i < 2; i++) {
        addr = (void *)((int64_t)addr + firstSize);
        oldFirstSize += firstSize;
        requiredSize -= firstSize;
        firstSize = getFirstChunkSize(addr, requiredSize);
        if (firstSize <= 0) {
            break;
        }
        void *alignedAddr = getAlignedBaseAddr(addr);
        if (isShadowAddr(addr, size)) {
            alignedAddr = getAddress(alignedAddr);
        }
        int64_t value __attribute__((aligned(8))) =
            *(int64_t __attribute__((aligned(8))) *)alignedAddr;
        int offset = oldFirstSize * 8;
        int32_t secret0 = (int32_t)unmaskLoadedData(value, firstSize, addr, offset, offset);
        secret0 = secret0 << offset;
        secret = secret | secret0;
    }
    return secret;
}

int16_t handle16BitLoad(void* addr) {
    int size = 16;
    int firstSize = 0;
    int oldFirstSize = 0;
    int16_t secret = 0;
    int requiredSize = size / 8;
    addr = unsetBit(addr);
    for (int i = 0; i < 1; i++) {
        addr = (void *)((int64_t)addr + firstSize);
        oldFirstSize += firstSize;
        requiredSize -= firstSize;
        firstSize = getFirstChunkSize(addr, requiredSize);
        if (firstSize <= 0) {
            break;
        }
        void *alignedAddr = getAlignedBaseAddr(addr);
        if (isShadowAddr(addr, size)) {
            alignedAddr = getAddress(alignedAddr);
        }
        int64_t value __attribute__((aligned(8))) =
            *(int64_t __attribute__((aligned(8))) *)alignedAddr;
        int offset = oldFirstSize * 8;
        int16_t secret0 = (int16_t)unmaskLoadedData(value, firstSize, addr, offset, offset);
        secret0 = secret0 << offset;
        secret = secret | secret0;
    }
    return secret;
}

int8_t handle8BitLoad(void* addr) {
    int size = 8;
    int firstSize = 0;
    int oldFirstSize = 0;
    int8_t secret = 0;
    int requiredSize = size / 8;
    addr = unsetBit(addr);
    for (int i = 0; i < 1; i++) {
        addr = (void *)((int64_t)addr + firstSize);
        oldFirstSize += firstSize;
        requiredSize -= firstSize;
        firstSize = getFirstChunkSize(addr, requiredSize);
        if (firstSize <= 0) {
            break;
        }
        void *alignedAddr = getAlignedBaseAddr(addr);
        if (isShadowAddr(addr, size)) {
            alignedAddr = getAddress(alignedAddr);
        }
        int64_t value __attribute__((aligned(8))) =
            *(int64_t __attribute__((aligned(8))) *)alignedAddr;
        int offset = oldFirstSize * 8;
        int8_t secret0 = (int8_t)unmaskLoadedData(value, firstSize, addr, offset, offset);
        secret0 = secret0 << offset;
        secret = secret | secret0;
    }
    return secret;
}

int64_t handle64BitSpillLoad(void* addr) {
    printf("Spill load %p\n", addr);
    int size = 64;
    int firstSize = 0;
    int oldFirstSize = 0;
    int64_t secret = 0;
    int requiredSize = size / 8;
    addr = unsetBit(addr);
    for (int i = 0; i < 3; i++) {
        addr = (void *)((int64_t)addr + firstSize);
        oldFirstSize += firstSize;
        requiredSize -= firstSize;
        firstSize = getFirstChunkSize(addr, requiredSize);
        if (firstSize <= 0) {
            break;
        }
        void *alignedAddr = getAlignedBaseAddr(addr);
        if (isShadowAddr(addr, size)) {
            alignedAddr = getAddress(alignedAddr);
        }
        int64_t value =
            *(int64_t __attribute__((aligned(8))) *)alignedAddr;
        int offset = oldFirstSize * 8;
        int64_t secret0 =
            unmaskLoadedData(value, firstSize, addr, offset, offset);
        secret0 = secret0 << offset;
        secret = secret | secret0;
    }
    printf("load\n");
    return secret;
}

int64_t handle64BitLoad(void* addr) {
    int size = 64;
    int firstSize = 0;
    int oldFirstSize = 0;
    int64_t secret = 0;
    int requiredSize = size / 8;
    addr = unsetBit(addr);
    for (int i = 0; i < 3; i++) {
        addr = (void *)((int64_t)addr + firstSize);
        oldFirstSize += firstSize;
        requiredSize -= firstSize;
        firstSize = getFirstChunkSize(addr, requiredSize);
        if (firstSize <= 0) {
            break;
        }
        void *alignedAddr = getAlignedBaseAddr(addr);
        if (isShadowAddr(addr, size)) {
            alignedAddr = getAddress(alignedAddr);
        }
        int64_t value =
            *(int64_t __attribute__((aligned(8))) *)alignedAddr;
        int offset = oldFirstSize * 8;
        int64_t secret0 =
            unmaskLoadedData(value, firstSize, addr, offset, offset);
        secret0 = secret0 << offset;
        secret = secret | secret0;
    }
    return secret;
}

void handle64BitStore(void *addr, int64_t secret) {
    int size = 64;
    int firstSize = 0;
    int oldFirstSize = 0;
    int requiredSize = size / 8;
    addr = unsetBit(addr);
    for (int i = 0;i < 3;i++){
        addr = (void *)((int64_t)addr + firstSize);
        requiredSize -= firstSize;
        secret = secret >> firstSize * 8;
        firstSize = getFirstChunkSize(addr, requiredSize);
        if (firstSize <= 0) {
            break;
        }
        int32_t secret0 = (int32_t)secret;
        void *alignedAddr = getAlignedBaseAddr(addr);
        if (isShadowAddr(addr, size)) {
            alignedAddr = getAddress(alignedAddr);
        }
        int64_t alignedLoad __attribute__((aligned(8))) =
            *(int64_t __attribute__((aligned(8))) *)alignedAddr;
        int64_t firstValue = maskLoadedData(secret0, alignedLoad, firstSize, addr);
        *(int64_t __attribute__((aligned(8))) *)alignedAddr = firstValue;
    }
}

void handle64BitSpillStore(void *addr, int64_t secret) {
    printf("Spill store %p\n", addr);
    int size = 64;
    int firstSize = 0;
    int oldFirstSize = 0;
    int requiredSize = size / 8;
    addr = unsetBit(addr);

    if(!isAddrInStackMap(addr)){
        updateStackAddr(addr, size);
    }

    for (int i = 0;i < 3;i++){
        addr = (void *)((int64_t)addr + firstSize);
        requiredSize -= firstSize;
        secret = secret >> firstSize * 8;
        firstSize = getFirstChunkSize(addr, requiredSize);
        if (firstSize <= 0) {
            break;
        }
        int32_t secret0 = (int32_t)secret;
        void *alignedAddr = getAlignedBaseAddr(addr);
        if (isShadowAddr(addr, size)) {
            alignedAddr = getAddress(alignedAddr);
        }
        int64_t alignedLoad __attribute__((aligned(8))) =
            *(int64_t __attribute__((aligned(8))) *)alignedAddr;
        int64_t firstValue = maskLoadedData(secret0, alignedLoad, firstSize, addr);
        *(int64_t __attribute__((aligned(8))) *)alignedAddr = firstValue;
    }
    printf("store\n");
}

void handle32BitStore(void *addr, int32_t secret) {
    int size = 32;
    int firstSize = 0;
    int oldFirstSize = 0;
    int requiredSize = size / 8;
    addr = unsetBit(addr);
    for (int i = 0;i < 2;i++){
        addr = (void *)((int64_t)addr + firstSize);
        requiredSize -= firstSize;
        secret = secret >> firstSize * 8;
        firstSize = getFirstChunkSize(addr, requiredSize);
        if (firstSize <= 0) {
            break;
        }
        int32_t secret0 = (int32_t)secret;
        void *alignedAddr = getAlignedBaseAddr(addr);
        if (isShadowAddr(addr, size)) {
            alignedAddr = getAddress(alignedAddr);
        }
        int64_t alignedLoad __attribute__((aligned(8))) =
            *(int64_t __attribute__((aligned(8))) *)alignedAddr;
        int64_t firstValue = maskLoadedData(secret0, alignedLoad, firstSize, addr);
        *(int64_t __attribute__((aligned(8))) *)alignedAddr = firstValue;
    }
}

void handle16BitStore(void *addr, int16_t secret) {
    int size = 16;
    int firstSize = 0;
    int oldFirstSize = 0;
    int requiredSize = size / 8;
    addr = unsetBit(addr);
    for (int i = 0;i < 1;i++){
        addr = (void *)((int64_t)addr + firstSize);
        requiredSize -= firstSize;
        secret = secret >> firstSize * 8;
        firstSize = getFirstChunkSize(addr, requiredSize);
        if (firstSize <= 0) {
            break;
        }
        int16_t secret0 = (int16_t)secret;
        void *alignedAddr = getAlignedBaseAddr(addr);
        if (isShadowAddr(addr, size)) {
            alignedAddr = getAddress(alignedAddr);
        }
        int64_t alignedLoad __attribute__((aligned(8))) =
            *(int64_t __attribute__((aligned(8))) *)alignedAddr;
        int64_t firstValue = maskLoadedData(secret0, alignedLoad, firstSize, addr);
        *(int64_t __attribute__((aligned(8))) *)alignedAddr = firstValue;
    }
}

void handle8BitStore(void *addr, int8_t secret) {
    int size = 8;
    int firstSize = 0;
    int oldFirstSize = 0;
    int requiredSize = size / 8;
    addr = unsetBit(addr);
    for (int i = 0;i < 1;i++){
        addr = (void *)((int64_t)addr + firstSize);
        requiredSize -= firstSize;
        secret = secret >> firstSize * 8;
        firstSize = getFirstChunkSize(addr, requiredSize);
        if (firstSize <= 0) {
            break;
        }
        int8_t secret0 = (int8_t)secret;
        void *alignedAddr = getAlignedBaseAddr(addr);
        if (isShadowAddr(addr, size)) {
            alignedAddr = getAddress(alignedAddr);
        }
        int64_t alignedLoad __attribute__((aligned(8))) =
            *(int64_t __attribute__((aligned(8))) *)alignedAddr;
        int64_t firstValue = maskLoadedData(secret0, alignedLoad, firstSize, addr);
        *(int64_t __attribute__((aligned(8))) *)alignedAddr = firstValue;
    }
}

void handle1BitStore(void *addr, int8_t secret) {
    int size = 1;
    int firstSize = 0;
    int oldFirstSize = 0;
    int requiredSize = 1; // in bits
    addr = unsetBit(addr);
    for (int i = 0;i < 1;i++){
        addr = (void *)((int64_t)addr + firstSize);
        requiredSize -= firstSize;
        secret = secret >> firstSize * 8;
        firstSize = getFirstChunkSize(addr, requiredSize);
        if (firstSize <= 0) {
            break;
        }
        int8_t secret0 = (int8_t)secret;
        void *alignedAddr = getAlignedBaseAddr(addr);
        if (isShadowAddr(addr, size)) {
            alignedAddr = getAddress(alignedAddr);
        }
        int64_t alignedLoad __attribute__((aligned(8))) =
            *(int64_t __attribute__((aligned(8))) *)alignedAddr;
        int64_t firstValue = maskLoadedDataBit(secret0, alignedLoad, firstSize, addr);
        *(int64_t __attribute__((aligned(8))) *)alignedAddr = firstValue;
    }
}