dmp-rt/smalloc-vector.cpp at 64bit-split · dmp-mitigation/dmp-rt · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
#include <smalloc-cpp-vec.h>

/* static */ std::vector<std::pair<int64_t, int64_t>> *tree = NULL;

bool sortPair(const std::pair<int, int> &A, const std::pair<int, int> &B) {
    return (A.first < B.first);
}

/* static */ void /* __attribute__((optnone))  */
updateAddr(int64_t start_addr, int64_t size) {
    if (!tree)
         tree = new std::vector<std::pair<int64_t, int64_t>>;
    tree->push_back(std::make_pair(start_addr, size));
    std::sort(tree->begin(), tree->end(), sortPair);
    // 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");
    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;
    // secret return 1
    // not secret return 0
    return ((int_addr >> smalloc__position) & 1U);
}

/* static */ void * /* __attribute__((optnone))  */ identity(void *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) {
    // printf("running getAddress for 0x%lx \n", addr);
    assert(addr != NULL && "getAddress recieved null");
    assert(checkBit(addr) == 0 && "getAddress recieved a secret address");
    int64_t int_addr = (int64_t)addr;
    auto res = std::lower_bound(tree->begin(), tree->end(), std::make_pair(int_addr, 0), sortPair);
    if(res == tree->end() || (res != tree->end() && res->first > int_addr)){
        if(res == tree->begin())
            assert(false && "no address found");
        res = std::prev(res, 1);
    }
    assert(res->first <= int_addr && "wrong result");
    assert(int_addr <= res->first + res->second && "out of bound addr");
    assert(int_addr <= res->first + res->second / 2 &&
           "addr in shadow mem");
    void* result = (void *)((int64_t)addr + (int64_t)addr - res->first +
                    res->second / 2);
    assert(result != NULL && "getAddress returned null");
    return result;
}

/* static */ void * /* __attribute__((optnone))  */ smalloc(size_t __size) {
    printf("running smalloc ");
    void *start_addr = malloc(2 * __size);
    int64_t int_addr = (int64_t)start_addr;
    updateAddr(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, off_t offset) {
    printf("running smalloc\n");
    void *start_addr = mmap(addr, 2 * __size, prot, flags, fd, offset);
    int64_t int_addr = (int64_t)start_addr;
    printf("Debug: 0x%lx and 0x%lx \n", int_addr, int_addr + 2 * __size);
    updateAddr(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);
    // tree->erase(int_addr);
    return munmap((void *)int_addr, 2 * __size);
}