luaur_code_gen/methods/
code_allocator_allocate.rs

1use crate::functions::flush_instruction_cache_code_allocator::flush_instruction_cache;
2use crate::functions::make_pages_executable_code_allocator::make_pages_executable;
3use crate::functions::make_pages_read_only_code_allocator::make_pages_read_only;
4use crate::macros::codegen_assert::CODEGEN_ASSERT;
5use crate::records::code_allocation_data::CodeAllocationData;
6use crate::records::code_allocator::CodeAllocator;
7use core::ffi::c_void;
8use core::ptr;
9
10impl CodeAllocator {
11    pub fn allocate(
12        &mut self,
13        data: *const u8,
14        data_size: usize,
15        code: *const u8,
16        code_size: usize,
17    ) -> CodeAllocationData {
18        use luaur_common::FFlag;
19        CODEGEN_ASSERT!(FFlag::LuauCodegenFreeBlocks.get());
20
21        let mut start_offset = 0;
22        let code_offset: usize;
23        let data_offset: usize;
24        let page_aligned_size: usize;
25        let total_size: usize;
26
27        if FFlag::LuauCodegenProtectData.get() {
28            if data_size != 0 {
29                if CodeAllocator::align_to_page_size(Self::kMaxReservedDataSize + data_size)
30                    + code_size
31                    > self.block_size
32                {
33                    return CodeAllocationData::default();
34                }
35
36                if CodeAllocator::align_to_page_size(data_size) + code_size
37                    > (self.block_end as usize - self.block_pos as usize)
38                {
39                    if !self.allocate_new_block(&mut start_offset) {
40                        return CodeAllocationData::default();
41                    }
42                    CODEGEN_ASSERT!(
43                        CodeAllocator::align_to_page_size(start_offset + data_size) + code_size
44                            <= (self.block_end as usize - self.block_pos as usize)
45                    );
46                }
47
48                code_offset = CodeAllocator::align_to_page_size(start_offset + data_size);
49                data_offset = code_offset - data_size;
50                total_size = CodeAllocator::align_to_page_size(data_size) + code_size;
51                page_aligned_size = CodeAllocator::align_to_page_size(code_offset + code_size);
52            } else {
53                let ts = code_size;
54                if ts > self.block_size - Self::kMaxReservedDataSize {
55                    return CodeAllocationData::default();
56                }
57
58                if ts > (self.block_end as usize - self.block_pos as usize) {
59                    if !self.allocate_new_block(&mut start_offset) {
60                        return CodeAllocationData::default();
61                    }
62                    CODEGEN_ASSERT!(ts <= (self.block_end as usize - self.block_pos as usize));
63                }
64
65                data_offset = start_offset;
66                code_offset = start_offset;
67                page_aligned_size = CodeAllocator::align_to_page_size(start_offset + ts);
68                total_size = ts;
69            }
70        } else {
71            let k_code_alignment = 32;
72            let aligned_data_size = (data_size + (k_code_alignment - 1)) & !(k_code_alignment - 1);
73            let ts = aligned_data_size + code_size;
74
75            if ts > self.block_size - Self::kMaxReservedDataSize {
76                return CodeAllocationData::default();
77            }
78
79            if ts > (self.block_end as usize - self.block_pos as usize) {
80                if !self.allocate_new_block(&mut start_offset) {
81                    return CodeAllocationData::default();
82                }
83                CODEGEN_ASSERT!(ts <= (self.block_end as usize - self.block_pos as usize));
84            }
85
86            data_offset = start_offset + aligned_data_size - data_size;
87            code_offset = start_offset + aligned_data_size;
88            page_aligned_size = CodeAllocator::align_to_page_size(start_offset + ts);
89            total_size = ts;
90        }
91
92        CODEGEN_ASSERT!(
93            CodeAllocator::align_to_page_size(self.block_pos as usize) == self.block_pos as usize
94        );
95
96        if data_size != 0 {
97            unsafe {
98                ptr::copy_nonoverlapping(data, self.block_pos.add(data_offset), data_size);
99            }
100        }
101        if code_size != 0 {
102            unsafe {
103                ptr::copy_nonoverlapping(code, self.block_pos.add(code_offset), code_size);
104            }
105        }
106
107        if FFlag::LuauCodegenProtectData.get() {
108            if data_size != 0 {
109                if !make_pages_read_only(self.block_pos, code_offset) {
110                    return CodeAllocationData::default();
111                }
112                if !make_pages_executable(
113                    unsafe { self.block_pos.add(code_offset) },
114                    page_aligned_size - code_offset,
115                ) {
116                    return CodeAllocationData::default();
117                }
118            } else if !make_pages_executable(self.block_pos, page_aligned_size) {
119                return CodeAllocationData::default();
120            }
121        } else if !make_pages_executable(self.block_pos, page_aligned_size) {
122            return CodeAllocationData::default();
123        }
124
125        self.live_allocations += 1;
126        flush_instruction_cache(unsafe { self.block_pos.add(code_offset) }, code_size);
127
128        let result = CodeAllocationData {
129            start: unsafe { self.block_pos.add(start_offset) },
130            size: total_size,
131            code_start: unsafe { self.block_pos.add(code_offset) },
132            allocation_start: self.block_pos,
133            allocation_size: page_aligned_size,
134        };
135
136        if page_aligned_size <= (self.block_end as usize - self.block_pos as usize) {
137            self.block_pos = unsafe { self.block_pos.add(page_aligned_size) };
138            CODEGEN_ASSERT!(
139                CodeAllocator::align_to_page_size(self.block_pos as usize)
140                    == self.block_pos as usize
141            );
142            CODEGEN_ASSERT!(self.block_pos <= self.block_end);
143        } else {
144            self.block_pos = self.block_end;
145        }
146
147        result
148    }
149}
luaur_code_gen/methods/code_allocator_allocate.rs

luaur_code_gen/methods/
code_allocator_allocate.rs
