[go: up one dir, main page]

gg-alloc 1.0.0

A custom allocator that only returns pointers above 2G
Documentation 
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
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146

use std::alloc::{GlobalAlloc, Layout};

pub struct GgAlloc<A> {
    alloc: A,
}

impl<A> GgAlloc<A> {
    pub const fn new(alloc: A) -> Self {
        Self { alloc }
    }
}

fn pointer_above_2g(ptr: *mut u8) -> bool {
    //(ptr as usize) > (isize::MAX as usize)
    (ptr as u32) > (i32::MAX as u32)
}

fn alloc_fully_below_2g(ptr: *mut u8, layout: Layout) -> bool {
    !pointer_above_2g(ptr) && !pointer_above_2g(unsafe { ptr.add(layout.size() - 1) })
}

unsafe impl<A> GlobalAlloc for GgAlloc<A>
where
    A: GlobalAlloc,
{
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        let mut ret = self.alloc.alloc(layout);

        loop {
            if ret.is_null() {
                // Failed to allocate, return null pointer
                break;
            }

            if pointer_above_2g(ret) {
                // Allocated above 2G, return this pointer
                break;
            }

            // Here the allocation was successful but below 2G, so fill memory a bit
            // in chunks of 128 MB
            let mut size = 1 << 27;
            loop {
                let test_layout = Layout::from_size_align(size, 1).unwrap();
                let fill_ptr = self.alloc.alloc(test_layout);
                if !fill_ptr.is_null() && alloc_fully_below_2g(fill_ptr, test_layout) {
                    // leak allocation and continue
                } else {
                    // Free and try again with half size
                    if !fill_ptr.is_null() {
                        self.alloc.dealloc(fill_ptr, test_layout);
                    }
                    size /= 2;

                    if size < 1 {
                        // Exit fill loop if size is small
                        break;
                    }
                }
            }

            if !alloc_fully_below_2g(ret, layout) {
                // Free and allocate something with half size
                self.alloc.dealloc(ret, layout);
            }

            // fill memory a bit
            // in chunks of same size as the requested allocation
            loop {
                let test_layout = layout;
                let fill_ptr = self.alloc.alloc(test_layout);
                if !fill_ptr.is_null() && alloc_fully_below_2g(fill_ptr, test_layout) {
                    // leak allocation and continue
                } else {
                    // this allocation is either above 2GB or right at the 2GB edge, we are done
                    if pointer_above_2g(fill_ptr) {
                        // above 2GB, we can use it as the return pointer
                        return fill_ptr;
                    } else {
                        // right at the edge, leak this allocation and try again in next iteration
                        // of the outer loop
                        break;
                    }
                }
            }

            ret = self.alloc.alloc(layout);
        }

        ret
    }
    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        self.alloc.dealloc(ptr, layout);
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::alloc::System;

    #[global_allocator]
    static A: GgAlloc<System> = GgAlloc::new(System);

    fn assert_above_2g(ptr: *mut u8) {
        assert!(pointer_above_2g(ptr), "{:p}", ptr,);
    }

    #[test]
    fn alloc_string() {
        let s = format!("allocating a string!");
        assert_above_2g(s.as_str() as *const _ as *mut _);
    }

    #[test]
    fn alloc_one_byte() {
        let s = format!("1");
        assert!(
            pointer_above_2g(s.as_str() as *const _ as *mut _),
            "{:p}",
            s.as_str()
        );
    }

    #[test]
    fn alloc_many_bytes() {
        let s = format!("1");
        assert_above_2g(s.as_str() as *const _ as *mut _);
        let s = format!("12");
        assert_above_2g(s.as_str() as *const _ as *mut _);
        let s = format!("123");
        assert_above_2g(s.as_str() as *const _ as *mut _);
        let s = format!("1234");
        assert_above_2g(s.as_str() as *const _ as *mut _);
        let s = format!("12345");
        assert_above_2g(s.as_str() as *const _ as *mut _);
        let s = format!("123456");
        assert_above_2g(s.as_str() as *const _ as *mut _);
        let s = format!("1234567");
        assert_above_2g(s.as_str() as *const _ as *mut _);
        let s = format!("12345678");
        assert_above_2g(s.as_str() as *const _ as *mut _);
        let s = format!("123456789");
        assert_above_2g(s.as_str() as *const _ as *mut _);
    }
}