[置顶] ATL窗口thunk机制的剖析与实现

ATL的窗口过程采用了一种称为thunk的调用机制，这是为了解决如下问题:
UI线程创建窗口的必要步骤是，注册一个窗口类->创建窗口->显示窗口。对应的API函数为RegisterClass ，CreateWindowEx，ShowWindow。注册窗口的阶段，需要传送一个回调函数的地址，消息循环中的DispatchMessage负责将消息分发给这个函数。每个窗口都有一个消息处理函数，这个函数就是注册窗口阶段传入的回调函数地址。通过函数GetWindowLong可以获取这个指针。ATL的窗口类中，这个函数是一个静态函数。不同的窗口应该有不同的窗口过程，因此建立窗口类消息循环的关键在于：这个公共的回调函数必须能够取得窗口对应的窗口对象，窗口对象包含了真正的消息处理代码。
如何从一个公共的类静态函数取得对象地址呢？
一个简单的方案是处理WM_NCCREATE消息，从lParam参数获取创建窗口时置入的对象地址，调用SetWindowLong将其放入内核里面的窗口对象中，以后在公共的回调函数中，调用GetWindowLong即可取得对象地址，之后调用对象的函数。
这种方法比较简单但是比较平庸，效率上，每次处理一个消息，都要一次系统调用，一次系统调用花费几百个CUP周期，效率不够好。
ATL的实现：
ATL采用一种称为thunk的机制，简单来说，公共的窗口过程依然是类的一个静态函数，但只负责窗口的第一个消息WM_NCCREATE，其目的是在堆上开辟一小块数据区，称为thunk，这小块数据区其实是一段机器码，把这个地址换成窗口过程函数地址，通过调用SetWindowLongPtr实现。因此以后真正的窗口过程是thunk，thunk负责把hwnd替换为对象的地址，然后jump到另一个公共的类静态函数。
jump到的函数虽然是一个静态函数，不含有this指针，但是堆栈上的第一个参数已经被替换为对象的地址，用得到的对象地址就可以调用真正的消息处理函数，根据消息映射宏调用对应消息的处理过程。窗口句柄和指针都是32位，或者64位（64为平台）。
实际的实现需要考虑
1 数据执行保护 (DEP)
2 WM_NCCREATE处理中，如何得到对象地址。
ATL采用了独特的方法解决第二个问题，其实比较简单的方法仍然是在创建窗口的时候传递一个对象地址的参数。此外，一个更好的选择是借助TLS来存储需要的参数。但ATL的方案从效率上没有问题，因为在整个窗口的存在期间，这部分的代码总共执行一次。
thunk定义为：

#pragma pack(push,1)
struct _stdcallthunk
{
DWORD   m_mov;          // mov dword ptr [esp+0x4], pThis (esp+0x4 is hWnd)
DWORD   m_this;         //
BYTE    m_jmp;          // jmp WndProc
DWORD   m_relproc;      // relative jmp
BOOL Init(DWORD_PTR proc, void* pThis)
{
m_mov = 0x042444C7;  //C7 44 24 0C
m_this = PtrToUlong(pThis);
m_jmp = 0xe9;
m_relproc = DWORD((INT_PTR)proc - ((INT_PTR)this+sizeof(_stdcallthunk)));
// write block from data cache and
//  flush from instruction cache
FlushInstructionCache(GetCurrentProcess(), this, sizeof(_stdcallthunk));
return TRUE;
}
//some thunks will dynamically allocate the memory for the code
void* GetCodeAddress()
{
return this;
}
void* operator new(size_t)
{
return __AllocStdCallThunk();
}
void operator delete(void* pThunk)
{
__FreeStdCallThunk(pThunk);
}
};
#pragma pack(pop)

以上代码是ATL的，下面我重新定义的一个beta_thunk，就是从这个结构来的。因为是机器码，所以必须字节对齐。

ATL重载了new操作符，实际上是调用了__AllocStdCallThunk函数分配，很可惜这个函数只有声明。但我们可以猜测它的实现。
thunk的大小只需要13个字节，就算是16个字节吧。因为要执行堆上的代码，因此必须可读可写可执行，防止数据执行保护。关键点在于，这部分数据必须与通常的内存分配空间分开。举个例子，假如我不重载new,而是直接new出一块空间作为thunk并初始化，那么就会出现问题。首先，DEP这关过不了。其次，假如代码的其他地方有一个数据比方说数组溢出错误，不小心覆盖了thunk，这必然引起崩溃，并且还难以定位。
一个简单的解决之道是直接调用低层内存管理函数VirtualAlloc，分配一块可读可写可执行的区域，作为一个thunk区域，并初始化。但这样做有一个问题，在r3下，这个函数的分配粒度是64kb，效率太差。
我认为最好的方案是，在程序初始化时就分配64kb,利用56kb作为thunk堆。这样，56kb/16byte,理论上可以支持同时存在3584个窗口。64kb的最开始页面和最后一个页面设置为不可访问，这样预防溢出错误。

这个方案其实就是自己实现一个thunk堆管理器。以下是代码部分：

class CYYCriticalSection
{
public:
CYYCriticalSection()throw()
{
memset(&m_sec,0,sizeof(CRITICAL_SECTION));
}
~CYYCriticalSection()throw()
{
}
HRESULT Init()throw()
{
HRESULT hRes = E_FAIL;
__try
{
InitializeCriticalSection(&m_sec);
hRes = S_OK;
}
// structured exception may be raised in low memory situations
__except(STATUS_NO_MEMORY == GetExceptionCode())
{
hRes = E_OUTOFMEMORY;
}
return hRes;
}
HRESULT Term() throw()
{
DeleteCriticalSection(&m_sec);
return S_OK;
}
HRESULT Lock()throw()
{
EnterCriticalSection(&m_sec);
return S_OK;
}
HRESULT Unlock()throw()
{
LeaveCriticalSection(&m_sec);
return S_OK;
}
private:
CRITICAL_SECTION m_sec;
};

template<typename TLock>
class beta_CComCritSecLock
{
public:
beta_CComCritSecLock( TLock& cs);
~beta_CComCritSecLock() throw();
HRESULT Lock() throw();
void Unlock() throw();
// Implementation
private:
TLock& m_cs;
bool m_bLocked;
};
template<typename TLock >
inline beta_CComCritSecLock< TLock >::beta_CComCritSecLock( TLock& cs) :
m_cs( cs ),	m_bLocked( false )
{
}
template< class TLock >
inline beta_CComCritSecLock< TLock >::~beta_CComCritSecLock() throw()
{
if( m_bLocked )
{
Unlock();
m_bLocked = false;
}
}
template< class TLock >
inline HRESULT beta_CComCritSecLock< TLock >::Lock() throw()
{
HRESULT hr;
ATLASSERT( !m_bLocked );
hr = m_cs.Lock();
if( FAILED( hr ) )
{
return( hr );
}
m_bLocked = true;
return( S_OK );
}
template< class TLock >
inline void beta_CComCritSecLock< TLock >::Unlock() throw()
{
ATLASSUME( m_bLocked );
m_cs.Unlock();
m_bLocked = false;
}

#define PAGE_SIZE	0x1000
#define BETA_64KB       0x10000
#define PAGE_ROUND_UP(x)                                    \
((((ULONG_PTR)(x)) + PAGE_SIZE-1)  & (~(PAGE_SIZE-1)))

#define ROUND_DOWN(n, align) \
(((ULONG_PTR)(n)) & ~((align) - 1))

#define ROUND_UP(n, align) \
ROUND_DOWN(((ULONG_PTR)(n)) + (align) - 1, (align))

static unsigned char bitmapMask[8] =
{
0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80
};

#pragma pack(push,1)
struct beta_stdcallthunk
{
DWORD   m_mov;          // mov dword ptr [esp+0x4], pThis (esp+0x4 is hWnd)//
DWORD   m_this;         //
BYTE    m_jmp;          // jmp WndProc
DWORD   m_relproc;      // relative jmp
void* GetCodeAddress()
{
return this;
}
BOOL Init(DWORD_PTR proc, void* pThis)
{
m_mov = 0x042444C7;  //C7 44 24 0C
m_this = PtrToUlong(pThis);
m_jmp = 0xe9;
m_relproc = DWORD((INT_PTR)proc - ((INT_PTR)this+sizeof(beta_stdcallthunk)));
// write block from data cache and
//  flush from instruction cache
FlushInstructionCache(GetCurrentProcess(), this, sizeof(beta_stdcallthunk));
return TRUE;
}
};
#pragma pack(pop)
//对大页面特殊处理//
class MyHeap
{
public:
MyHeap():m_heap(NULL),m_heapbk(NULL),_handle(NULL),m_bf(true),_bLPage(false),
_bitcount(448),_pagesize(0x400000)
{
memset(_bitmap,0,448);
SYSTEM_INFO sysinfo;
::GetSystemInfo(&sysinfo);
_pagesize = sysinfo.dwPageSize;
if(sysinfo.dwPageSize > PAGE_SIZE)
{
_bLPage = true;
}
}
char *getHeap()
{
if(m_bf  && NULL == m_heapbk)
{
m_heap =(char *) ::VirtualAlloc(NULL,BETA_64KB,MEM_COMMIT|MEM_RESERVE,PAGE_EXECUTE_READWRITE);
if(m_heap)
{
#ifdef _DEBUG
//protection code 防御性代码，保证任何情况下都正确分配到，但可以去掉//
assert(PAGE_ROUND_UP(m_heap) == (ULONG_PTR)m_heap);
#endif
m_heapbk = m_heap;
//fangyuxing daima
m_bf = false;
char *tem = m_heap;
char *tpend = (char *)(ROUND_UP((m_heap + PAGE_SIZE),BETA_64KB));
tpend -= PAGE_SIZE;
m_heap = (char *)(PAGE_ROUND_UP(m_heap));
if(m_heap > tem )
{// never coding//防御性代码，不应执行到这里
tpend = tem + BETA_64KB;
tpend = (char *)(ROUND_DOWN(tpend,PAGE_SIZE));
tpend -= PAGE_SIZE;
}
//protectMemory
DWORD lpflOldProtect;
::VirtualProtect(m_heap,PAGE_SIZE,PAGE_NOACCESS,&lpflOldProtect);
::VirtualProtect(tpend,PAGE_SIZE,PAGE_NOACCESS,&lpflOldProtect);
m_heap += PAGE_SIZE;
_bitcount = ((long)(tpend - m_heap))/(_Granularity*8);
#ifdef _DEBUG
//protection code
assert(_bitcount == 448);
#endif
}
}
return m_heap;
}
char *MyHeapAlloc();
void MyHeapFree(char *pv);
~MyHeap()
{
if(m_heapbk)
{
::VirtualFree(m_heapbk,0,MEM_RELEASE|MEM_DECOMMIT);
m_heapbk = NULL;
m_heap = NULL;
}
if(_handle)
{
::HeapDestroy(_handle);
_handle = NULL;
}
m_bf = true;
}
//如果是大页面，改用这个函数
HANDLE LgetHeap()
{
if(m_bf  && NULL == _handle)
{
_handle = ::HeapCreate(0,_pagesize,_pagesize);
if(_handle)
m_bf = false;
}
return _handle;
}
bool IsLPage()
{
return _bLPage;
}
private:
bool m_bf;
bool _bLPage;
char *m_heap;
char *m_heapbk;
unsigned char _bitmap[448];
static int _Granularity;
int _bitcount;// bytecount;
HANDLE _handle;
int _pagesize;
};
//<=448
char *MyHeap::MyHeapAlloc()
{
beta_CComCritSecLock<CYYCriticalSection> lock(beta_WinModule.m_csWindowCreate);
int i;
int j = 0;
for(i = 0;i<448;++i)
{
if(_bitmap[i] == 0xff)
{}
else
break;
}
if(i >= _bitcount)
return NULL;
for(j=0;j<8;++j)
{
if(0 == (_bitmap[i] & bitmapMask[j]))
{
break;
}
}
assert(j<8);
_bitmap[i] = _bitmap[i] | bitmapMask[j];
return (i*8 + j)*_Granularity + m_heap;
}

void MyHeap::MyHeapFree(char *pv)
{
ATLASSERT(unsigned long(pv - m_heap) >= 0);
unsigned long  region = static_cast<unsigned long>(pv - m_heap);
unsigned long t = region/_Granularity;
assert((region%_Granularity) == 0);
int bytepos = t/8;
int bitpos = t%8;
//同步
beta_CComCritSecLock<CYYCriticalSection> lock(beta_WinModule.m_csWindowCreate);
if(_bitmap[bytepos]&bitmapMask[bitpos])
{
_bitmap[bytepos] = _bitmap[bytepos]&(~bitmapMask[bitpos]);
return;
}
assert(0);
}

分析：以上代码自建了一个thunk堆管理器。必须考虑线程同步，多个线程同时访问的时候才能正确。同时这个堆管理器考虑到大页面。什么是大页面？极个别的windows系统才会将页面粒度设为1Mb或者4Mb，针对奇葩windows，做了特殊处理。

在程序中可以设一个全局变量，如果是大页面，那么调用LgetHeap，得到句柄后调用HeapAlloc分配thunk，再设置访问方式，可读可写可执行。