关于HOOK系统API时的递归调用（或叫重入）问题的解决办法

原贴：http://bbs.driverdevelop.com/htm_data/39/0303/36274.html

关于重入的问题很容易解决的，用一个TLS变量就可以了。由于只在调用线程的范围内改变TLS变量的值，所以改变当前线程的TLS变量不会改变其它线程中该TLS变量的值，ZwQueryDirectoryFile就可以根据该TLS变量的值来判断是否之前已经进入了一个ZwQueryDirectroyFile调用而且未返回。
而且TlsGetValue、TlsSetValue我跟踪过，其只是简单地利用段寄存器来区分对不同线程中该TLS变量的存取，并没有调用任何API，所以可以完全不用担心调用TlsGetValue、TlsSetValue会产生其它重入问题。

驱动中可以使用TLS吗？
DWORD
APIENTRY
TlsAlloc(
    VOID
    )

/*++

Routine Description:

    A TLS index may be allocated using TlsAlloc.  Win32 garuntees a
    minimum number of TLS indexes are available in each process.  The
    constant TLS_MINIMUM_AVAILABLE defines the minimum number of
    available indexes.  This minimum is at least 64 for all Win32
    systems.

Arguments:

    None.

Return Value:

    Not-0xffffffff - Returns a TLS index that may be used in a
        subsequent call to TlsFree, TlsSetValue, or TlsGetValue.  The
        storage associated with the index is initialized to NULL.

    0xffffffff - The operation failed. Extended error status is available
        using GetLastError.


--*/

{
    PPEB Peb;
    PTEB Teb;
    DWORD Index;

    Peb = NtCurrentPeb();
    Teb = NtCurrentTeb();

    RtlAcquirePebLock();
    try {

        Index = RtlFindClearBitsAndSet((PRTL_BITMAP)Peb->TlsBitmap,1,0);
        if ( Index == 0xffffffff ) {
            Index = RtlFindClearBitsAndSet((PRTL_BITMAP)Peb->TlsExpansionBitmap,1,0);
            if ( Index == 0xffffffff ) {
                BaseSetLastNTError(STATUS_NO_MEMORY);
                }
            else {
                if ( !Teb->TlsExpansionSlots ) {
                    Teb->TlsExpansionSlots = RtlAllocateHeap(
                                                RtlProcessHeap(),
                                                MAKE_TAG( TMP_TAG ) | HEAP_ZERO_MEMORY,
                                                TLS_EXPANSION_SLOTS * sizeof(PVOID)
                                                );
                    if ( !Teb->TlsExpansionSlots ) {
                        RtlClearBits((PRTL_BITMAP)Peb->TlsExpansionBitmap,Index,1);
                        Index = 0xffffffff;
                        BaseSetLastNTError(STATUS_NO_MEMORY);
                        return Index;
                        }
                    }
                Teb->TlsExpansionSlots[Index] = NULL;
                Index += TLS_MINIMUM_AVAILABLE;
                }
            }
        else {
            Teb->TlsSlots[Index] = NULL;
            }
        }
    finally {
        RtlReleasePebLock();
        }
#if DBG
    Index |= TLS_MASK;
#endif
    return Index;
}
这个函数是一个纯AP的函数吧！

我一般采用“阻塞”的概念，一个TLS变量的值为0表示无阻塞状态，大于0表示阻塞状态。下面是示例

// 程序开始
DWORD g_tls_isBlock = ::TlsAlloc();

// 程序结束
::TlsFree(g_tls_isBlock);

// 定义阻塞支持
#define BLOCKCOUNT()    ((LONG)::TlsGetValue(g_tls_isBlock))
#defien ISBLOCK()    (BLOCKCOUNT() > 0)
#define BLOCK()    (::TlsSetValue(g_tls_isBlock, (LPVOID)(BLOCKCOUNT() + 1))
#define UNBLOCK()    (::TlsSetValue(g_tls_isBlock, (LPVOID)(BLOCKCOUNT() - 1))


// 使用示例
void MyAPI()
{
    if (ISBLOCK())
    {
        return;
    }

    ...

    BLOCK();
    MyApi2();
    UNBLOCK();
}


void MyApi2()
{
    ...
    MyAPI();
    ...
}

引用第1楼dreamsity于2008-09-10 17:05发表的  :
驱动中可以使用TLS吗？
我知道IoSetTopOfIrp是用的TLS，但不知道怎么在驱动中分配自己的TLS。

我的是在Rang 3下的，Rang 0下的有待其它高手研究。。。

驱动下应该是可以使用的，最多就自己直接存取线程TED里的TLS数组。。。当然最好是有其它更安全稳定的方法去在驱动里存取TLS变量。。。

不过这个的确是一个思路，IRP_MJ_CREATE是必须同步完成的，
也就是说有可能通过TLS来避免IRP_MJ_CREATE重入。
虽然内核内本身没有TLS的支持，但是否可以通过ROOKIT的技术增加一个TLS接口？
这样做是很有商业价值的，呵呵。值得研究！

引用第1楼dreamsity于2008-09-10 17:05发表的  :
驱动中可以使用TLS吗？
DWORD
APIENTRY
TlsAlloc(
    VOID
.......

TlsAlloc和TlsFree只在程序开始和结束(或其它适合的位置)去调用的;

在HOOK的函数里只用到TlsGetValue和TlsSetValue。。。

http://www.diybl.com/course/4_webprogram/asp.net/netjs/200865/122185.html
好像有说到内核模式下使用tls的...

其实TLS是次要的，关键是TLS的思想。
在驱动中是完全有可能去组织一个数据结构来实践（代替）TLS的。。。

我在某个开发包中看过某牛实现的tls驱动实现方案. 有一定的技巧.

看起来好复杂啊，越简单才越不会出问题