1+ //
2+ // Created by xia__mc on 2024/12/10.
3+ //
4+
5+ #include " ASM.h"
6+ #include " utils/memory/AlignedAllocator.h"
7+ #include " utils/memory/FastMemcpy.h"
8+ #include " Compat.h"
9+
10+ #ifdef WINDOWS
11+
12+ #include " windows.h"
13+
14+ #endif
15+
16+ extern " C"
17+
18+ static PyTypeObject ASMType = {
19+ PyVarObject_HEAD_INIT (&PyType_Type, 0 )
20+ };
21+
22+ static int ASM_init ([[maybe_unused]] ASM *self,
23+ [[maybe_unused]] PyObject *args, [[maybe_unused]] PyObject *kwargs) {
24+ return 0 ;
25+ }
26+
27+ static void ASM_dealloc (ASM *self) {
28+ Py_TYPE (self)->tp_free ((PyObject *) self);
29+ }
30+
31+ static PyObject *ASM_enter (PyObject *self, [[maybe_unused]] PyObject *args) {
32+ Py_INCREF (self);
33+ return self;
34+ }
35+
36+ static PyObject *ASM_exit ([[maybe_unused]] PyObject *self,
37+ [[maybe_unused]] PyObject *const *args, [[maybe_unused]] Py_ssize_t nargs) {
38+ Py_RETURN_NONE;
39+ }
40+
41+ static PyObject *ASM_run ([[maybe_unused]] PyObject *__restrict self,
42+ PyObject *const *__restrict args, Py_ssize_t nargs) noexcept {
43+ #ifdef WINDOWS
44+ if (nargs != 1 ) {
45+ PyErr_SetString (PyExc_TypeError, " Function takes exactly 1 arguments (__code)"
46+ return nullptr ;
47+ }
48+
49+ if (!PyBytes_Check (args[0 ])) {
50+ PyErr_SetString (PyExc_TypeError, " Expected a bytes object."
51+ return nullptr ;
52+ }
53+
54+ const auto *__restrict code = (const unsigned char *) PyBytes_AsString (args[0 ]);
55+ if (code == nullptr ) {
56+ PyErr_SetString (PyExc_ValueError, " Failed to convert argument to c str."
57+ return nullptr ;
58+ }
59+
60+ const auto size = (size_t ) PyBytes_Size (args[0 ]) * sizeof (unsigned char );
61+
62+ unsigned char *__restrict target;
63+ const bool aligned = (uintptr_t ) code % 16 == 0 ;
64+ if (!aligned) {
65+ try {
66+ target = (unsigned char *) alignedAlloc (size, 16 );
67+ } catch (const std::exception &e) {
68+ PyErr_SetString (PyExc_MemoryError, e.what ());
69+ return nullptr ;
70+ }
71+ fast_memcpy (target, code, size);
72+ } else {
73+ target = (unsigned char *) code;
74+ }
75+
76+ DWORD oldProtect;
77+ if (!VirtualProtect ((LPVOID) target, size,
78+ PAGE_EXECUTE_READWRITE, &oldProtect)) {
79+ PyErr_SetString (PyExc_OSError, " Failed to make ASM executable."
80+
81+ if (!aligned) {
82+ alignedFree (target);
83+ }
84+ return nullptr ;
85+ }
86+
87+ ((void (*)()) target)();
88+
89+ if (!VirtualProtect ((LPVOID) target, size, oldProtect, &oldProtect)) {
90+ PyErr_SetString (PyExc_OSError, " Failed to restore memory protection."
91+
92+ if (!aligned) {
93+ alignedFree (target);
94+ }
95+ return nullptr ;
96+ }
97+
98+ if (!aligned) {
99+ alignedFree (target);
100+ }
101+
102+ Py_RETURN_NONE;
103+ #else
104+ PyErr_SetString (PyExc_NotImplementedError, " ASM is not supported on this architecture."
105+ return nullptr ;
106+ #endif
107+ }
108+
109+ static PyObject *ASM_runFast ([[maybe_unused]] PyObject *__restrict self,
110+ PyObject *const *__restrict args, [[maybe_unused]] Py_ssize_t nargs) noexcept {
111+ #ifdef WINDOWS
112+ const auto *__restrict code = (const unsigned char *) PyBytes_AS_STRING (args[0 ]);
113+ const auto size = (size_t ) PyBytes_GET_SIZE (args[0 ]) * sizeof (unsigned char );
114+
115+ if ((uintptr_t ) code % 16 != 0 ) {
116+ unsigned char *__restrict target;
117+ try {
118+ target = (unsigned char *) alignedAlloc (size, 16 );
119+ } catch (const std::exception &e) {
120+ PyErr_SetString (PyExc_MemoryError, e.what ());
121+ return nullptr ;
122+ }
123+ fast_memcpy (target, code, size);
124+
125+ DWORD oldProtect;
126+ if (!VirtualProtect ((LPVOID) target, size,
127+ PAGE_EXECUTE_READWRITE, &oldProtect)) {
128+ PyErr_SetString (PyExc_OSError, " Failed to make ASM executable."
129+ alignedFree (target);
130+ return nullptr ;
131+ }
132+
133+ ((void (*)()) target)();
134+
135+ if (!VirtualProtect ((LPVOID) target, size, oldProtect, &oldProtect)) {
136+ PyErr_SetString (PyExc_OSError, " Failed to restore memory protection."
137+ alignedFree (target);
138+ return nullptr ;
139+ }
140+
141+ alignedFree (target);
142+ } else {
143+ DWORD oldProtect;
144+ if (!VirtualProtect ((LPVOID) code, size,
145+ PAGE_EXECUTE_READWRITE, &oldProtect)) {
146+ PyErr_SetString (PyExc_OSError, " Failed to make ASM executable."
147+ return nullptr ;
148+ }
149+
150+ ((void (*)()) code)();
151+
152+ if (!VirtualProtect ((LPVOID) code, size, oldProtect, &oldProtect)) {
153+ PyErr_SetString (PyExc_OSError, " Failed to restore memory protection."
154+ return nullptr ;
155+ }
156+ }
157+
158+ Py_RETURN_NONE;
159+ #else
160+ PyErr_SetString (PyExc_NotImplementedError, " ASM is not supported on this architecture."
161+ return nullptr ;
162+ #endif
163+ }
164+
165+ static PyMethodDef ASM_methods[] = {
166+ {" __enter__" nullptr },
167+ {" __exit__" nullptr },
168+ {" run" nullptr },
169+ {" runFast" nullptr },
170+ {nullptr , nullptr , 0 , nullptr }
171+ };
172+
173+ void initializeASMType (PyTypeObject &type) {
174+ type.tp_name = " ASM"
175+ type.tp_basicsize = sizeof (ASM);
176+ type.tp_itemsize = 0 ;
177+ type.tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE;
178+ type.tp_methods = ASM_methods;
179+ type.tp_init = (initproc) ASM_init;
180+ type.tp_new = PyType_GenericNew;
181+ type.tp_dealloc = (destructor) ASM_dealloc;
182+ type.tp_alloc = PyType_GenericAlloc;
183+ type.tp_free = PyObject_Del;
184+ }
185+
186+ static struct PyModuleDef ASM_module = {
187+ PyModuleDef_HEAD_INIT,
188+ " __pyfastutil.ASM"
189+ " Allow access to native asm."
190+ -1 ,
191+ nullptr , nullptr , nullptr , nullptr , nullptr
192+ };
193+
194+ #pragma clang diagnostic push
195+ #pragma ide diagnostic ignored "OCUnusedGlobalDeclarationInspection"
196+ PyMODINIT_FUNC PyInit_ASM () {
197+ initializeASMType (ASMType);
198+
199+ PyObject *object = PyModule_Create (&ASM_module);
200+ if (object == nullptr )
201+ return nullptr ;
202+
203+ Py_INCREF (&ASMType);
204+ if (PyModule_AddObject (object, " ASM" 0 ) {
205+ Py_DECREF (&ASMType);
206+ Py_DECREF (object);
207+ return nullptr ;
208+ }
209+
210+ return object;
211+ }
212+ #pragma clang diagnostic pop
213+
214+ }
0 commit comments