x64下动态代码构建RUNTIME_FUNCTION

  1 #include "pch.h"
  2 #include <iostream>
  3 #include <vector>
  4 
  5 #include <stdio.h>
  6 #include <stdlib.h>
  7 #include <stdint.h>
  8 
  9 #include <windows.h>
 10 
 11 typedef uint8_t UBYTE;
 12 typedef uint16_t USHORT;
 13 
 14 typedef void(__cdecl* CallTestFunc)(void);
 15 
 16 typedef enum _UNWIND_OP_CODES {
 17     UWOP_PUSH_NONVOL = 0,
 18     UWOP_ALLOC_LARGE,       // 1
 19     UWOP_ALLOC_SMALL,       // 2
 20     UWOP_SET_FPREG,         // 3
 21     UWOP_SAVE_NONVOL,       // 4
 22     UWOP_SAVE_NONVOL_FAR,   // 5
 23     UWOP_SPARE_CODE1,       // 6
 24     UWOP_SPARE_CODE2,       // 7
 25     UWOP_SAVE_XMM128,       // 8
 26     UWOP_SAVE_XMM128_FAR,   // 9
 27     UWOP_PUSH_MACHFRAME     // 10
 28 } UNWIND_OP_CODES, *PUNWIND_OP_CODES;
 29 
 30 typedef union _UNWIND_CODE {
 31     struct {
 32         UBYTE CodeOffset;
 33         UBYTE UnwindOp : 4;
 34         UBYTE OpInfo : 4;
 35     };
 36     USHORT FrameOffset;
 37 } UNWIND_CODE, *PUNWIND_CODE;
 38 
 39 typedef struct _UNWIND_INFO {
 40     UBYTE Version : 3;
 41     UBYTE Flags : 5;
 42     UBYTE SizeOfProlog;
 43     UBYTE CountOfCodes;
 44     UBYTE FrameRegister : 4;
 45     UBYTE FrameOffset : 4;
 46     UNWIND_CODE UnwindCode[1];
 47     /*    UNWIND_CODE MoreUnwindCode[((CountOfCodes + 1) & ~1) - 1];
 48      *    OPTIONAL ULONG ExceptionHandler;
 49      *    OPTIONAL ULONG ExceptionData[]; */
 50 } UNWIND_INFO, *PUNWIND_INFO;
 51 
 52 typedef struct {
 53     uint8_t code[0x1000];
 54     RUNTIME_FUNCTION function_table[1];
 55     UNWIND_INFO unwind_info[1];
 56 } DYNSECTION;
 57 
 58 
 59 //    FILTER    HANDLER        FINALLY_HANDLER
 60 #define UNW_FLAG_NHANDLER 0x0        //    NO      NO           
 61 #define UNW_FLAG_EHANDLER 0x1        //    YES      YES
 62 #define UNW_FLAG_UHANDLER 0x2        //                            YES
 63 #define UNW_FLAG_CHAININFO 0x4        //    multi UNWIND_INFO
 64 
 65 struct STACK_FRAME_INFO {
 66     PVOID returnAddress;
 67     PVOID functionAddress;
 68 };
 69 
 70 struct CALL_STACK_INFO {
 71     PVOID stackBottom;
 72     PVOID stackTop;
 73     DWORD dwFrameCount;
 74     STACK_FRAME_INFO* pFrameList;
 75 };
 76 
 77 typedef ULONG(WINAPI *RTLWALKFRAMECHAIN)(OUT PVOID *Callers, IN ULONG Count, IN ULONG Flags);
 78 
 79 void GetCallStackInfo(CALL_STACK_INFO& callStack, DWORD dwMaxFrame)
 80 {
 81     std::vector<STACK_FRAME_INFO> vecFrames;
 82 #ifdef _M_IX86
 83     STACK_FRAME_INFO StackFrame;
 84     std::vector<PVOID> vecRetAddr;
 85     ULONG StackCount = 0;
 86     RTLWALKFRAMECHAIN pRtlWalkFrameChain =
 87         (RTLWALKFRAMECHAIN)GetProcAddress(GetModuleHandleW(L"ntdll.dll"), "RtlWalkFrameChain");
 88 
 89     vecRetAddr.resize(50);
 90     StackCount = pRtlWalkFrameChain(&vecRetAddr[0], vecRetAddr.size(), 0);
 91 
 92     for (ULONG i = 0;i < StackCount;i++)
 93     {
 94         ZeroMemory(&StackFrame, sizeof(StackFrame));
 95 
 96         if (vecRetAddr[i] == nullptr)
 97         {
 98             break;
 99         }
100 
101         StackFrame.returnAddress = vecRetAddr[i];
102 
103         vecFrames.push_back(StackFrame);
104     }
105 
106 #elif (_M_X64)
107     CONTEXT                       Context;
108     KNONVOLATILE_CONTEXT_POINTERS NvContext;
109     UNWIND_HISTORY_TABLE          UnwindHistoryTable;
110     PRUNTIME_FUNCTION             RuntimeFunction;
111     PVOID                         HandlerData;
112     ULONG64                       EstablisherFrame;
113     ULONG64                       ImageBase;
114     STACK_FRAME_INFO              StackFrame;
115 
116     RtlCaptureContext(&Context);
117 
118     ZeroMemory(&UnwindHistoryTable, sizeof(UNWIND_HISTORY_TABLE));
119 
120     for (ULONG Frame = 0; Frame <= dwMaxFrame; Frame++)
121     {
122         RuntimeFunction = RtlLookupFunctionEntry(
123             Context.Rip,
124             &ImageBase,
125             &UnwindHistoryTable
126         );
127 
128         ZeroMemory(&NvContext, sizeof(KNONVOLATILE_CONTEXT_POINTERS));
129 
130         if (!RuntimeFunction)
131         {
132             Context.Rip = (ULONG64)(*(PULONG64)Context.Rsp);
133             Context.Rsp += 8;
134         }
135         else
136         {
137             RtlVirtualUnwind(
138                 UNW_FLAG_NHANDLER,
139                 ImageBase,
140                 Context.Rip,
141                 RuntimeFunction,
142                 &Context,
143                 &HandlerData,
144                 &EstablisherFrame,
145                 &NvContext);
146         }
147 
148         if (!Context.Rip)
149         {
150             break;
151         }
152         
153         ZeroMemory(&StackFrame, sizeof(StackFrame));
154 
155         StackFrame.returnAddress = (PVOID)Context.Rip;
156         StackFrame.functionAddress = (PVOID)(ImageBase + RuntimeFunction->BeginAddress);
157 
158         vecFrames.push_back(StackFrame);
159     }
160 #endif
161     
162     callStack.pFrameList = (STACK_FRAME_INFO*)malloc(vecFrames.size() * sizeof(STACK_FRAME_INFO));
163     if (!callStack.pFrameList)
164     {
165         callStack.dwFrameCount = 0;
166         return;
167     }
168 
169     callStack.dwFrameCount = vecFrames.size();
170     for (DWORD dwFrame = 0; dwFrame < vecFrames.size(); dwFrame++)
171     {
172         CopyMemory(&callStack.pFrameList[dwFrame], &vecFrames[dwFrame], sizeof(STACK_FRAME_INFO));
173     }
174 }
175 
176 void Func1()
177 {
178     CALL_STACK_INFO callStack;
179     ZeroMemory(&callStack, sizeof(callStack));
180     GetCallStackInfo(callStack, 100);
181 
182     for (DWORD dwFrame = 0; dwFrame < callStack.dwFrameCount; dwFrame++)
183     {
184         printf(
185             "FRAME %02x: FuncAddrss=%p ReturnAddress=%p \n",
186             dwFrame,
187             callStack.pFrameList[dwFrame].functionAddress,
188             callStack.pFrameList[dwFrame].returnAddress);
189     }
190 
191 }
192 
193 void DumpCallStack()
194 {
195     Func1();
196 }
197 
198 DWORD TestForSehAddFunctionTable();
199 DWORD TestForUnwindFunctionTable();
200 
201 int main()
202 {
203     TestForUnwindFunctionTable();
204 
205     ::system("pause");
206 }
207 
208 #ifdef _M_X64
209 
210 DWORD TestForUnwindFunctionTable()
211 {
212     DYNSECTION *pDynData = (DYNSECTION*)VirtualAlloc(NULL, 0x2000, MEM_COMMIT, PAGE_EXECUTE_READWRITE);
213     uint8_t *code = pDynData->code;
214     ULONG_PTR p = 0;
215     code[p++] = 0x56; //push rsi
216     code[p++] = 0x57; //push rdi
217     code[p++] = 0x48;
218     code[p++] = 0x81;
219     code[p++] = 0xec;
220     code[p++] = 0x88;
221     code[p++] = 0x07;
222     code[p++] = 0x00;
223     code[p++] = 0x00; //sub rsp,788h
224     ULONG_PTR prologSize = p;
225     code[p++] = 0x48;
226     code[p++] = 0xb8; //mov rax, xxxxxxxh
227     ULONG_PTR returnAddress = p;
228     for (int i = 0; i < sizeof(ULONG_PTR); i++)
229     {
230         code[p++] = 0x00;
231     }
232     code[p++] = 0x50; //push rax
233     code[p++] = 0x48;
234     code[p++] = 0xb8; //mov rax, xxxxxxxh
235     ULONG_PTR targetAddress = p;
236     for (int i = 0; i < sizeof(ULONG_PTR); i++)
237     {
238         code[p++] = 0x00;
239     }
240     code[p++] = 0x50; //push rax
241     code[p++] = 0xc3; //ret
242     ULONG_PTR offsetRip = p;
243     code[p++] = 0x48;
244     code[p++] = 0x81;
245     code[p++] = 0xc4;
246     code[p++] = 0x88;
247     code[p++] = 0x07;
248     code[p++] = 0x00;
249     code[p++] = 0x00; //add rsp,788h
250     code[p++] = 0x5f; //pop rdi
251     code[p++] = 0x5e; //pop rsi
252     code[p++] = 0xc3; //ret
253     ULONG_PTR trampoline = p;
254 
255     UNWIND_INFO *pUnwindInfo = pDynData->unwind_info;
256     pUnwindInfo[0].Version = 1;
257     pUnwindInfo[0].Flags = UNW_FLAG_NHANDLER;
258     pUnwindInfo[0].SizeOfProlog = (UBYTE)prologSize;
259     pUnwindInfo[0].CountOfCodes = 4;
260     pUnwindInfo[0].FrameRegister = 0;
261     pUnwindInfo[0].FrameOffset = 0;
262 
263     DWORD64 dyn_base = (DWORD64)pDynData;    
264 
265     UNWIND_CODE *pUnwindCode = pUnwindInfo->UnwindCode;
266 
267     //pUnwindCode[0].CodeOffset = 9;
268     //pUnwindCode[0].UnwindOp = UWOP_ALLOC_LARGE;
269     //pUnwindCode[0].OpInfo = 0;
270     //
271     //pUnwindCode[1].CodeOffset = 2;
272     //pUnwindCode[1].UnwindOp = UWOP_PUSH_NONVOL;
273     //pUnwindCode[1].OpInfo = 7;
274     //
275     //pUnwindCode[2].CodeOffset = 1;
276     //pUnwindCode[2].UnwindOp = UWOP_PUSH_NONVOL;
277     //pUnwindCode[2].OpInfo = 6;
278     pUnwindCode[0].FrameOffset = 0x0109;
279     pUnwindCode[1].FrameOffset = 0x7002;
280     pUnwindCode[2].FrameOffset = 0x6001;
281 
282     RUNTIME_FUNCTION *function_table = pDynData->function_table;
283     function_table[0].BeginAddress = (DWORD64)&code[0] - dyn_base; // set RVA of dynamic code start
284     function_table[0].EndAddress = (DWORD64)&code[trampoline] - dyn_base; // RVA of dynamic code end
285     function_table[0].UnwindInfoAddress = (DWORD64)pUnwindCode - dyn_base; // RVA of unwind info
286 
287     *(DWORD64*)&code[returnAddress] = (DWORD64)code + offsetRip;
288     *(DWORD64*)&code[targetAddress] = (DWORD64)DumpCallStack; // VA of target
289 
290     if (!RtlAddFunctionTable(function_table, 1, dyn_base)) {
291         printf("RtlAddFunctionTable() failed, exit.\n");
292         exit(EXIT_FAILURE);
293     }
294 
295     void(*call)() = (void(*)())code;
296     (*call)();
297 
298     return 0;
299 }
300 
301 static EXCEPTION_DISPOSITION handler(
302     PEXCEPTION_RECORD ExceptionRecord, 
303     ULONG64 EstablisherFrame, 
304     PCONTEXT ContextRecord, 
305     PDISPATCHER_CONTEXT DispatcherContext
306 ) 
307 {
308     printf("handler!\n");
309     ContextRecord->Rip += 3;
310     return ExceptionContinueExecution;
311 }
312 
313 DWORD TestForSehAddFunctionTable()
314 {
315     int ret;
316     RUNTIME_FUNCTION *q;
317     DYNSECTION *dynsection = (DYNSECTION*)VirtualAlloc(NULL, 0x2000, MEM_COMMIT, PAGE_EXECUTE_READWRITE);
318 
319     uint8_t *code = dynsection->code;
320     size_t p = 0;
321     code[p++] = 0xb8; // mov rax, 42
322     code[p++] = 0x2a;
323     code[p++] = 0x00;
324     code[p++] = 0x00;
325     code[p++] = 0x00;
326     code[p++] = 0xc6; // mov byte [rax], 0  -- raises exception!
327     code[p++] = 0x00;
328     code[p++] = 0x00;
329     code[p++] = 0xc3; // ret
330 
331     size_t trampoline = p;
332     code[p++] = 0x48; // mov rax, 
333     code[p++] = 0xb8;
334     size_t patch_handler_address = p;
335     code[p++] = 0x00; // address to handler patched here
336     code[p++] = 0x00;
337     code[p++] = 0x00;
338     code[p++] = 0x00;
339     code[p++] = 0x00;
340     code[p++] = 0x00;
341     code[p++] = 0x00;
342     code[p++] = 0x00;
343     code[p++] = 0xff; // jmp rax
344     code[p++] = 0xe0;
345 
346     DWORD64 dyn_base = 0;
347     q = RtlLookupFunctionEntry((DWORD64)code, &dyn_base, NULL);
348     printf("lookup ‘code‘ %p %llx\n", q, dyn_base); // no function table entry
349 
350     DWORD64 image_base = 0;
351     q = RtlLookupFunctionEntry((DWORD64)main, &image_base, NULL);
352     printf("lookup ‘main‘ %p %llx\n", q, image_base); // there is a function table entry
353 
354     dyn_base = (DWORD64)dynsection;
355     UNWIND_INFO *unwind_info = dynsection->unwind_info;
356     unwind_info[0].Version = 1;
357     unwind_info[0].Flags = UNW_FLAG_EHANDLER;
358     unwind_info[0].SizeOfProlog = 0;
359     unwind_info[0].CountOfCodes = 0;
360     unwind_info[0].FrameRegister = 0;
361     unwind_info[0].FrameOffset = 0;
362     *(DWORD *)&unwind_info[0].UnwindCode = (DWORD64)&code[trampoline] - dyn_base;
363 
364     RUNTIME_FUNCTION *function_table = dynsection->function_table;
365     function_table[0].BeginAddress = (DWORD64)&code[0] - dyn_base; // set RVA of dynamic code start
366     function_table[0].EndAddress = (DWORD64)&code[trampoline] - dyn_base; // RVA of dynamic code end
367     function_table[0].UnwindInfoAddress = (DWORD64)unwind_info - dyn_base; // RVA of unwind info
368 
369     *(DWORD64 *)&code[patch_handler_address] = (DWORD64)handler; // VA of handler
370 
371     printf("main VA %016llx\n", (DWORD64)main);
372     printf("code VA %016llx\n", (DWORD64)code);
373     printf("function table VA %016llx\n", (DWORD64)function_table);
374     printf("unwind info VA %016llx\n", (DWORD64)unwind_info);
375     printf("handler VA %016llx\n", (DWORD64)handler);
376     printf("RUNTIME_FUNCTION begin RVA %08x, end RVA %08x, unwind RVA %08x\n",
377         function_table[0].BeginAddress, function_table[0].EndAddress,
378         function_table[0].UnwindInfoAddress);
379     printf("UNWIND_INFO handler RVA %08x\n", *(DWORD *)&unwind_info[0].UnwindCode);
380 
381     if (!RtlAddFunctionTable(function_table, 1, dyn_base)) {
382         printf("RtlAddFunctionTable() failed, exit.\n");
383         exit(EXIT_FAILURE);
384     }
385 
386     q = RtlLookupFunctionEntry((DWORD64)code, &dyn_base, NULL);
387     printf("lookup ‘code‘ %p %llx\n", q, dyn_base); // should return address of function table entry
388 
389     uint64_t(*call)() = (uint64_t(*)()) code;
390     uint64_t result = (*call)();
391     printf("result = %llx\n", result);
392 
393     if (!RtlDeleteFunctionTable(function_table)) {
394         printf("RtlDeleteFunctionTable() failed, exit.\n");
395         exit(EXIT_FAILURE);
396     }
397 
398     return EXIT_SUCCESS;
399 }
400 #endif