home *** CD-ROM | disk | FTP | other *** search
/ Computer Shopper 275 / DPCS0111DVD.ISO / Toolkit / Audio-Visual / VirtualDub / Source / VirtualDub-1.9.10-src.7z / src / system / source / thunk.cpp < prev    next >
Encoding:
C/C++ Source or Header  |  2009-09-14  |  7.8 KB  |  307 lines
  1. #include "stdafx.h"
  2. #include <windows.h>
  3. #include <map>
  4. #include <vd2/system/atomic.h>
  5. #include <vd2/system/refcount.h>
  6. #include <vd2/system/thunk.h>
  7. #include <vd2/system/binary.h>
  8.  
  9. class IVDJITAllocator {};
  10.  
  11. class VDJITAllocator : public vdrefcounted<IVDJITAllocator> {
  12. public:
  13.     VDJITAllocator();
  14.     ~VDJITAllocator();
  15.  
  16.     void *Allocate(size_t len);
  17.     void Free(void *p, size_t len);
  18.  
  19.     void EndUpdate(void *p, size_t len);
  20.  
  21. protected:
  22.     typedef std::map<void *, size_t> FreeChunks;
  23.     FreeChunks mFreeChunks;
  24.     FreeChunks::iterator mNextChunk;
  25.  
  26.     typedef std::map<void *, size_t> Allocations;
  27.     Allocations mAllocations;
  28.  
  29.     uintptr        mAllocationGranularity;
  30. };
  31.  
  32. VDJITAllocator::VDJITAllocator()
  33.     : mNextChunk(mFreeChunks.end())
  34. {
  35.     SYSTEM_INFO si;
  36.     GetSystemInfo(&si);
  37.  
  38.     mAllocationGranularity = si.dwAllocationGranularity;
  39. }
  40.  
  41. VDJITAllocator::~VDJITAllocator() {
  42.     for(Allocations::iterator it(mAllocations.begin()), itEnd(mAllocations.end()); it!=itEnd; ++it) {
  43.         VirtualFree(it->first, 0, MEM_RELEASE);
  44.     }
  45. }
  46.  
  47. void *VDJITAllocator::Allocate(size_t len) {
  48.     len = (len + 15) & ~(size_t)15;
  49.  
  50.     FreeChunks::iterator itMark(mNextChunk), itEnd(mFreeChunks.end()), it(itMark);
  51.  
  52.     if (it == itEnd)
  53.         it = mFreeChunks.begin();
  54.  
  55.     for(;;) {
  56.         for(; it!=itEnd; ++it) {
  57.             if (it->second >= len) {
  58.                 it->second -= len;
  59.  
  60.                 void *p = (char *)it->first + it->second;
  61.  
  62.                 if (!it->second) {
  63.                     if (mNextChunk == it)
  64.                         ++mNextChunk;
  65.  
  66.                     mFreeChunks.erase(it);
  67.                 }
  68.  
  69.                 return p;
  70.             }
  71.         }
  72.  
  73.         if (itEnd == itMark)
  74.             break;
  75.  
  76.         it = mFreeChunks.begin();
  77.         itEnd = itMark;
  78.     }
  79.  
  80.     size_t alloclen = (len + mAllocationGranularity - 1) & ~(mAllocationGranularity - 1);
  81.  
  82.     void *p = VirtualAlloc(NULL, alloclen, MEM_COMMIT, PAGE_EXECUTE_READWRITE);
  83.     if (p) {
  84.         try {
  85.             Allocations::iterator itA(mAllocations.insert(Allocations::value_type(p, alloclen)).first);
  86.  
  87.             try {
  88.                 if (len < alloclen)
  89.                     mFreeChunks.insert(FreeChunks::value_type((char *)p + len, alloclen - len));
  90.  
  91.             } catch(...) {
  92.                 mAllocations.erase(itA);
  93.                 throw;
  94.             }
  95.         } catch(...) {
  96.             VirtualFree(p, 0, MEM_RELEASE);
  97.             p = NULL;
  98.         }
  99.     }
  100.  
  101.     return p;
  102. }
  103.  
  104. void VDJITAllocator::Free(void *p, size_t len) {
  105.     VDASSERT(p);
  106.     VDASSERT(len < 0x10000);
  107.  
  108.     FreeChunks::iterator cur(mFreeChunks.lower_bound(p));
  109.     if (cur != mFreeChunks.end() && (char *)p + len == cur->first) {
  110.         len += cur->second;
  111.         if (mNextChunk == cur)
  112.             ++mNextChunk;
  113.         cur = mFreeChunks.erase(cur);
  114.     }
  115.  
  116.     if (cur != mFreeChunks.begin()) {
  117.         FreeChunks::iterator prev(cur);
  118.  
  119.         --prev;
  120.         if ((char *)prev->first + prev->second == p) {
  121.             p = prev->first;
  122.             len += prev->second;
  123.             if (mNextChunk == prev)
  124.                 ++mNextChunk;
  125.             mFreeChunks.erase(prev);
  126.         }
  127.     }
  128.  
  129.     uintptr start = (size_t)p;
  130.     uintptr end = start + len;
  131.  
  132.     if (!((start | end) & (mAllocationGranularity - 1))) {
  133.         Allocations::iterator it(mAllocations.find(p));
  134.  
  135.         if (it != mAllocations.end()) {
  136.             VirtualFree((void *)start, 0, MEM_RELEASE);
  137.             mAllocations.erase(it);
  138.             return;
  139.         }
  140.     }
  141.  
  142.     mFreeChunks.insert(FreeChunks::value_type((void *)start, end-start));
  143. }
  144.  
  145. void VDJITAllocator::EndUpdate(void *p, size_t len) {
  146.     FlushInstructionCache(GetCurrentProcess(), p, len);
  147. }
  148.  
  149. ///////////////////////////////////////////////////////////////////////////
  150.  
  151. VDJITAllocator *g_pVDJITAllocator;
  152. VDAtomicInt g_VDJITAllocatorLock;
  153.  
  154. bool VDInitThunkAllocator() {
  155.     bool success = true;
  156.  
  157.     while(g_VDJITAllocatorLock.xchg(1))
  158.         ::Sleep(1);
  159.  
  160.     if (!g_pVDJITAllocator) {
  161.         g_pVDJITAllocator = new_nothrow VDJITAllocator;
  162.         if (!g_pVDJITAllocator)
  163.             success = false;
  164.     }
  165.  
  166.     if (success)
  167.         g_pVDJITAllocator->AddRef();
  168.  
  169.     VDVERIFY(1 == g_VDJITAllocatorLock.xchg(0));
  170.  
  171.     return success;
  172. }
  173.  
  174. void VDShutdownThunkAllocator() {
  175.     while(g_VDJITAllocatorLock.xchg(1))
  176.         ::Sleep(1);
  177.  
  178.     VDASSERT(g_pVDJITAllocator);
  179.  
  180.     if (!g_pVDJITAllocator->Release())
  181.         g_pVDJITAllocator = NULL;
  182.  
  183.     VDVERIFY(1 == g_VDJITAllocatorLock.xchg(0));
  184. }
  185.  
  186. void *VDAllocateThunkMemory(size_t len) {
  187.     return g_pVDJITAllocator->Allocate(len);
  188. }
  189.  
  190. void VDFreeThunkMemory(void *p, size_t len) {
  191.     g_pVDJITAllocator->Free(p, len);
  192. }
  193.  
  194. void VDSetThunkMemory(void *p, const void *src, size_t len) {
  195.     memcpy(p, src, len);
  196.     g_pVDJITAllocator->EndUpdate(p, len);
  197. }
  198.  
  199. void VDFlushThunkMemory(void *p, size_t len) {
  200.     g_pVDJITAllocator->EndUpdate(p, len);
  201. }
  202.  
  203. ///////////////////////////////////////////////////////////////////////////
  204.  
  205. #ifdef _M_AMD64
  206.     extern "C" void VDMethodToFunctionThunk64();
  207. #else
  208.     extern "C" void VDMethodToFunctionThunk32();
  209.     extern "C" void VDMethodToFunctionThunk32_4();
  210.     extern "C" void VDMethodToFunctionThunk32_8();
  211.     extern "C" void VDMethodToFunctionThunk32_12();
  212.     extern "C" void VDMethodToFunctionThunk32_16();
  213. #endif
  214.  
  215. VDFunctionThunk *VDCreateFunctionThunkFromMethod(void *method, void *pThis, size_t argbytes, bool stdcall_thunk) {
  216. #if defined(_M_IX86)
  217.     void *pThunk = VDAllocateThunkMemory(16);
  218.  
  219.     if (!pThunk)
  220.         return NULL;
  221.  
  222.     if (stdcall_thunk || !argbytes) {    // thiscall -> stdcall (easy case)
  223.         uint8 thunkbytes[16]={
  224.             0xB9, 0x00, 0x00, 0x00, 0x00,                // mov ecx, this
  225.             0xE9, 0x00, 0x00, 0x00, 0x00                // jmp fn
  226.         };
  227.  
  228.  
  229.         VDWriteUnalignedLEU32(thunkbytes+1, (uint32)(uintptr)pThis);
  230.         VDWriteUnalignedLEU32(thunkbytes+6, (uint32)method - ((uint32)pThunk + 10));
  231.  
  232.         VDSetThunkMemory(pThunk, thunkbytes, 15);
  233.     } else {                // thiscall -> cdecl (hard case)
  234.         uint8 thunkbytes[16]={
  235.             0xE8, 0x00, 0x00, 0x00, 0x00,                // call VDFunctionThunk32
  236.             0xC3,                                        // ret
  237.             argbytes,                                    // db argbytes
  238.             0,                                            // db 0
  239.             0x00, 0x00, 0x00, 0x00,                        // dd method
  240.             0x00, 0x00, 0x00, 0x00,                        // dd this
  241.         };
  242.  
  243.         void *adapter;
  244.  
  245.         switch(argbytes) {
  246.         case 4:        adapter = VDMethodToFunctionThunk32_4;    break;
  247.         case 8:        adapter = VDMethodToFunctionThunk32_8;    break;
  248.         case 12:    adapter = VDMethodToFunctionThunk32_12;    break;
  249.         case 16:    adapter = VDMethodToFunctionThunk32_16;    break;
  250.         default:    adapter = VDMethodToFunctionThunk32;    break;
  251.         }
  252.  
  253.         VDWriteUnalignedLEU32(thunkbytes+1, (uint32)(uintptr)adapter - ((uint32)pThunk + 5));
  254.         VDWriteUnalignedLEU32(thunkbytes+8, (uint32)(uintptr)method);
  255.         VDWriteUnalignedLEU32(thunkbytes+12, (uint32)(uintptr)pThis);
  256.  
  257.         VDSetThunkMemory(pThunk, thunkbytes, 16);
  258.     }
  259.  
  260.     return (VDFunctionThunk *)pThunk;
  261. #elif defined(_M_AMD64)
  262.     void *pThunk = VDAllocateThunkMemory(44);
  263.     if (!pThunk)
  264.         return NULL;
  265.  
  266.     uint8 thunkbytes[44]={
  267.         0x48, 0x8D, 0x04, 0x25, 0x10, 0x00, 0x00,    // lea rax, [eip+16]
  268.         0x00,
  269.         0xFF, 0x24, 0x25, 0x08, 0x00, 0x00, 0x00,    // jmp qword ptr [rip+8]
  270.         0x90,                                        // nop
  271.         0, 0, 0, 0, 0, 0, 0, 0,                        // dq VDFunctionThunk64
  272.         0, 0, 0, 0, 0, 0, 0, 0,                        // dq method
  273.         0, 0, 0, 0, 0, 0, 0, 0,                        // dq this
  274.         0, 0, 0, 0                                    // dd argspillbytes
  275.     };
  276.  
  277.     VDWriteUnalignedLEU64(thunkbytes+16, (uint64)(uintptr)VDMethodToFunctionThunk64);
  278.     VDWriteUnalignedLEU64(thunkbytes+24, (uint64)(uintptr)method);
  279.     VDWriteUnalignedLEU64(thunkbytes+32, (uint64)(uintptr)pThis);
  280.  
  281.     // The stack must be aligned to a 16 byte boundary when the CALL
  282.     // instruction occurs. On entry to VDFunctionThunk64(), the stack is misaligned
  283.     // to 16n+8. Therefore, the number of argbytes must be 16m+8 and the number of
  284.     // argspillbytes must be 16m+8-24.
  285.     VDWriteUnalignedLEU32(thunkbytes+40, argbytes < 32 ? 0 : ((argbytes - 16 + 15) & ~15));
  286.  
  287.     VDSetThunkMemory(pThunk, thunkbytes, 44);
  288.  
  289.     return (VDFunctionThunk *)pThunk;
  290. #else
  291.     return NULL;
  292. #endif
  293. }
  294.  
  295. void VDDestroyFunctionThunk(VDFunctionThunk *pFnThunk) {
  296.     // validate thunk
  297. #if defined(_M_IX86)
  298.     VDASSERT(((const uint8 *)pFnThunk)[0] == 0xB9 || ((const uint8 *)pFnThunk)[0] == 0xE8);
  299.     VDFreeThunkMemory(pFnThunk, 16);
  300. #elif defined(_M_AMD64)
  301.     VDFreeThunkMemory(pFnThunk, 44);
  302. #else
  303.     VDASSERT(false);
  304. #endif
  305.  
  306. }
  307.