18c1441f8SAlexey Samsonov //===-- enable_execute_stack_test.c - Test __enable_execute_stack ----------===//
28c1441f8SAlexey Samsonov //
38c1441f8SAlexey Samsonov //                     The LLVM Compiler Infrastructure
48c1441f8SAlexey Samsonov //
58c1441f8SAlexey Samsonov // This file is dual licensed under the MIT and the University of Illinois Open
68c1441f8SAlexey Samsonov // Source Licenses. See LICENSE.TXT for details.
78c1441f8SAlexey Samsonov //
88c1441f8SAlexey Samsonov //===----------------------------------------------------------------------===//
98c1441f8SAlexey Samsonov 
108c1441f8SAlexey Samsonov 
118c1441f8SAlexey Samsonov #include <stdio.h>
128c1441f8SAlexey Samsonov #include <string.h>
138c1441f8SAlexey Samsonov #include <stdint.h>
148c1441f8SAlexey Samsonov #if defined(_WIN32)
158c1441f8SAlexey Samsonov #include <windows.h>
168c1441f8SAlexey Samsonov void __clear_cache(void* start, void* end)
178c1441f8SAlexey Samsonov {
188c1441f8SAlexey Samsonov     if (!FlushInstructionCache(GetCurrentProcess(), start, end-start))
198c1441f8SAlexey Samsonov         exit(1);
208c1441f8SAlexey Samsonov }
218c1441f8SAlexey Samsonov void __enable_execute_stack(void *addr)
228c1441f8SAlexey Samsonov {
238c1441f8SAlexey Samsonov     MEMORY_BASIC_INFORMATION b;
248c1441f8SAlexey Samsonov 
258c1441f8SAlexey Samsonov     if (!VirtualQuery(addr, &b, sizeof(b)))
268c1441f8SAlexey Samsonov         exit(1);
278c1441f8SAlexey Samsonov     if (!VirtualProtect(b.BaseAddress, b.RegionSize, PAGE_EXECUTE_READWRITE, &b.Protect))
288c1441f8SAlexey Samsonov         exit(1);
298c1441f8SAlexey Samsonov }
308c1441f8SAlexey Samsonov #else
318c1441f8SAlexey Samsonov #include <sys/mman.h>
328c1441f8SAlexey Samsonov extern void __clear_cache(void* start, void* end);
338c1441f8SAlexey Samsonov extern void __enable_execute_stack(void* addr);
348c1441f8SAlexey Samsonov #endif
358c1441f8SAlexey Samsonov 
368c1441f8SAlexey Samsonov typedef int (*pfunc)(void);
378c1441f8SAlexey Samsonov 
388c1441f8SAlexey Samsonov int func1()
398c1441f8SAlexey Samsonov {
408c1441f8SAlexey Samsonov     return 1;
418c1441f8SAlexey Samsonov }
428c1441f8SAlexey Samsonov 
438c1441f8SAlexey Samsonov int func2()
448c1441f8SAlexey Samsonov {
458c1441f8SAlexey Samsonov     return 2;
468c1441f8SAlexey Samsonov }
478c1441f8SAlexey Samsonov 
48*6e6c1c3eSSaleem Abdulrasool void *__attribute__((noinline))
49*6e6c1c3eSSaleem Abdulrasool memcpy_f(void *dst, const void *src, size_t n) {
50*6e6c1c3eSSaleem Abdulrasool // ARM and MIPS nartually align functions, but use the LSB for ISA selection
51*6e6c1c3eSSaleem Abdulrasool // (THUMB, MIPS16/uMIPS respectively).  Ensure that the ISA bit is ignored in
52*6e6c1c3eSSaleem Abdulrasool // the memcpy
53*6e6c1c3eSSaleem Abdulrasool #if defined(__arm__) || defined(__mips__)
54*6e6c1c3eSSaleem Abdulrasool   return (void *)((uintptr_t)memcpy(dst, (void *)((uintptr_t)src & ~1), n) |
55*6e6c1c3eSSaleem Abdulrasool                   ((uintptr_t)src & 1));
56*6e6c1c3eSSaleem Abdulrasool #else
57*6e6c1c3eSSaleem Abdulrasool   return memcpy(dst, (void *)((uintptr_t)src), n);
58*6e6c1c3eSSaleem Abdulrasool #endif
59*6e6c1c3eSSaleem Abdulrasool }
608c1441f8SAlexey Samsonov 
618c1441f8SAlexey Samsonov int main()
628c1441f8SAlexey Samsonov {
638c1441f8SAlexey Samsonov     unsigned char execution_buffer[128];
648c1441f8SAlexey Samsonov     // mark stack page containing execution_buffer to be executable
658c1441f8SAlexey Samsonov     __enable_execute_stack(execution_buffer);
668c1441f8SAlexey Samsonov 
678c1441f8SAlexey Samsonov     // verify you can copy and execute a function
68*6e6c1c3eSSaleem Abdulrasool     pfunc f1 = (pfunc)memcpy_f(execution_buffer, func1, 128);
698c1441f8SAlexey Samsonov     __clear_cache(execution_buffer, &execution_buffer[128]);
708c1441f8SAlexey Samsonov     if ((*f1)() != 1)
718c1441f8SAlexey Samsonov         return 1;
728c1441f8SAlexey Samsonov 
738c1441f8SAlexey Samsonov     // verify you can overwrite a function with another
74*6e6c1c3eSSaleem Abdulrasool     pfunc f2 = (pfunc)memcpy_f(execution_buffer, func2, 128);
758c1441f8SAlexey Samsonov     __clear_cache(execution_buffer, &execution_buffer[128]);
768c1441f8SAlexey Samsonov     if ((*f2)() != 2)
778c1441f8SAlexey Samsonov         return 1;
788c1441f8SAlexey Samsonov 
798c1441f8SAlexey Samsonov     return 0;
808c1441f8SAlexey Samsonov }
81