taskswch.cpp

ICEExt for Driver Studio3.2的sourcecode

/*++
    This file is mostly derived from mamaich's debugger codename BlindStudio.

    Copyleft  (c) 2002 Sten
    Contact information:
        mail: [e-mail]

    Copyright (c) 2001 mamaich
    Contact information:
    	web: www.reversing.net/mamaich
        mail: [e-mail]
        ICQ# 70241285

    This program is free software; you can redistribute it and/or
    modify it under the terms of the GNU General Public License
    as published by the Free Software Foundation; either version 2
    of the License, or (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program; if not, write to the Free Software
    Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.

 
Module Name:
    taskswch.cpp

Abstract:  Task switch hook. Clears my breakpoint on task switching. 

Revision History:

 Sten        05/06/2002
      The file is fully rewritten for better understanding.
 mamaich     14/06/2001
      Initial release

--*/

extern "C"{
#pragma warning ( push, 3 )
#include <ntddk.h>
#pragma warning ( pop )
}

#pragma warning ( disable: 4514 ) // unreferenced inline function has been removed

#include "defs.h"
#include "search.h"
#include "bpr.h"
#include "ntoskrnl.h"
#include "undoc.h"

void MySwapContextHook_NT();
void MySwapContextHook_W3K();
void MySwapContextEndHook();

static ULONG_PTR SwapContextPtr=0;
static ULONG_PTR SwapContextCnt=0;

static PVOID pSwapContextHook      =0;
static PVOID pSwapContextSavedBytes=0;

// The code "mov byte ptr es:[esi+2Dh], 2" is near the start of 
// SwitchContext function in NTOSKRNL.EXE. It is the only place we see
// these bytes in NT4 & 2K and the first place in XP
UCHAR SwapContextStart_NT[] =
{
    0x26, 0xc6, 0x46, 0x2d, 0x02          // mov byte ptr es:[esi+2Dh], 2
};

// MS changed SwapContext routine implementation in Windows 2003 Server
// So we should use other template to search it
UCHAR SwapContextStart_W3K[] =
{
    0x51,                                 // push    ecx
                               
                                          //loc_407F51:                             
    0x80, 0x7E, 0x5D, 0x00,               // cmp     byte ptr [esi+5Dh], 0
    0x74, 0x04,                           // jz      short loc_407F5B
    0xF3, 0x90,                           // pause
    0xEB                                  // jmp     short loc_407F51
};


////////////////////////////////////////////////////////////////////////////
//
// InitSwapContextHook
//
//   Hooks the SwapContext function
//
//   Note: InitNTosKernel MUST be called before InitSwapContextHook
//
////////////////////////////////////////////////////////////////////////////

NTSTATUS InitSwapContextHook()
{
    SwapContextPtr = (ULONG_PTR)((ULONG_PTR)RabSearch(SwapContextStart_NT, sizeof(SwapContextStart_NT),
                                              (unsigned char *)hNTOSKRNL, NTOS_TEXT_SIZE) + (ULONG)hNTOSKRNL); 

    if ((ULONG_PTR)SwapContextPtr < (ULONG_PTR)hNTOSKRNL)
    {
        // try to locate Windows 2003 Server SwapContext routine
        SwapContextPtr = (ULONG_PTR)((ULONG_PTR)RabSearch(SwapContextStart_W3K, sizeof(SwapContextStart_W3K),
                                               (unsigned char *)hNTOSKRNL, NTOS_TEXT_SIZE) + (ULONG_PTR)hNTOSKRNL); 

        if ((ULONG_PTR)SwapContextPtr < (ULONG_PTR)hNTOSKRNL)
        {
            DbgPrint("InitSwapContext: can't find SwapContext function %08X.\n",SwapContextPtr);
            SwapContextPtr = 0;
            return STATUS_PROCEDURE_NOT_FOUND;
        }

        // install W3K specific SwapContext hook
        pSwapContextHook = (PVOID)&MySwapContextHook_W3K;
        pSwapContextSavedBytes = (PVOID)&SwapContextStart_W3K;
    }
    else
    {
        // install NT+ generic SwapContext hook
        pSwapContextHook = (PVOID)&MySwapContextHook_NT;
        pSwapContextSavedBytes = (PVOID)&SwapContextStart_NT;
    }

    DbgPrint("SwapContextPtr:                               %08X\n", SwapContextPtr);

    SwapContextCnt = SwapContextPtr + 5;

    __asm cli; // Disable task switching
    
    ULONG Addr=(ULONG_PTR)pSwapContextHook - (SwapContextPtr+5);
    *(PUCHAR)SwapContextPtr=(UCHAR)0xe9;
    *(PULONG)(SwapContextPtr+1)=Addr;

    __asm sti; 

    return STATUS_SUCCESS;
}

////////////////////////////////////////////////////////////////////////////
//
// RemoveSwapContextHook
//
//   Uninstall swap context hook
//
////////////////////////////////////////////////////////////////////////////

void RemoveSwapContextHook()
{
    if (SwapContextPtr)
    {
         __asm cli;  // Disable task switching

         memcpy((void*)SwapContextPtr, pSwapContextSavedBytes,5);
       
         __asm sti;
             SwapContextPtr=0;
    }
}

////////////////////////////////////////////////////////////////////////////
//
// HandleContextSwap_NT
//
//   This code is called when windows decides to swap context of 2 threads.
//   I should deactivate breakpoints on range in the old process
//   NT+ version
//  
////////////////////////////////////////////////////////////////////////////

void __stdcall HandleContextSwap_NT(ETHREAD *New, ETHREAD *Old)
{
	UNREFERENCED_PARAMETER(Old);

    bpr_DeactivateAll(); // Deativate all BPRs in the old process

    PULONG Stack=(PULONG)New->Tcb.KernelStack;

    // Stack of SwapContext in NT4-XP looks like:
    // esp+8 -> return address | Stack[2]
    // esp+4 -> EFLAGS         | Stack[1]
    // esp   -> ECX            | Stack[0]

    // I'll insert address of my hook function between EFLAGS and return addreses

    New->Tcb.KernelStack=((char*)New->Tcb.KernelStack)-4;
    memmove(New->Tcb.KernelStack,Stack,8);
    Stack[1]=(ULONG)MySwapContextEndHook;
}

////////////////////////////////////////////////////////////////////////////
//
// HandleContextSwap_W3K
//
//   This code is called when windows decides to swap context of 2 threads.
//   I should deactivate breakpoints on range in the old process
//   Windows 2003 Server version
//
////////////////////////////////////////////////////////////////////////////

void __stdcall HandleContextSwap_W3K(ETHREAD *New, ETHREAD *Old)
{
	UNREFERENCED_PARAMETER(Old);

    bpr_DeactivateAll(); // Deativate all BPRs in the old process

    PULONG Stack=*(PULONG*)((ULONG)New+0x20); //->Tcb.KernelStack;
                                              // ETHREAD structure was changed
                                              // in W3K

    // Stack of SwapContext in W3K looks like:
    // esp+8 -> return address | Stack[2]
    // esp+4 -> ???            | Stack[1]
    // esp   -> ECX            | Stack[0]

    // I'll insert address of my hook function between EFLAGS and return addreses
    *(PULONG*)((ULONG)New+0x20) -= 1;
    memmove(*(PULONG*)((ULONG)New+0x20),Stack,8);
    Stack[1]=(ULONG)MySwapContextEndHook;
}

/////////////////////////////////////////////////////////////////////////////
//
// HandleContextSwapEnd
//
//   This code is called after context is swapped. It reenables all my
//   breakpoints in the new process
//
/////////////////////////////////////////////////////////////////////////////

void __stdcall HandleContextSwapEnd()
{
     bpr_ActivateAll(); // Activate all BPRs in the new process
}

/////////////////////////////////////////////////////////////////////////////
//
// MySwapContextHook_NT
//
//   Hook. Called from NTOSKRNL!SwapContext
//
/////////////////////////////////////////////////////////////////////////////

void __declspec(naked) MySwapContextHook_NT()
{
    __asm
    {
          mov          byte ptr es:[esi+2Dh], 2
          pushfd
          pushad
          push         edi
          push         esi
          call         HandleContextSwap_NT
          popad
          popfd
          jmp          [SwapContextCnt]
    }
}

/////////////////////////////////////////////////////////////////////////////
//
// MySwapContextHook_W3K
//
//   Hook. Called from NTOSKRNL!SwapContext
//
/////////////////////////////////////////////////////////////////////////////

void __declspec(naked) MySwapContextHook_W3K()
{
    __asm
    {
          push         ecx

swap_wait:
          ; this loop stolen directly from W3K ntoskrnl.
          cmp          byte ptr [esi+5Dh], 0 ; _KTHREAD.SwapBusy
          jz           short swap_ok
          pause
          jmp          short swap_wait

swap_ok:
          pushfd
          pushad
          push         edi
          push         esi
          
          call         HandleContextSwap_W3K

          popad
          popfd

          cmp          byte ptr [esi+5Dh], 0 ; _KTHREAD.SwapBusy
          jmp          [SwapContextCnt]
    }
}

/////////////////////////////////////////////////////////////////////////////
//
// MySwapContextEndHook
//
//   Hook. Called on return from NTOSKRNL!SwapContext
//
/////////////////////////////////////////////////////////////////////////////

void __declspec(naked) MySwapContextEndHook()
{
    __asm
    {
          pushfd
          pushad
          call    HandleContextSwapEnd
          popad
          popfd
          retn
    }
}