Playing with the stack

#include <windows.h>
#include <string>

#if defined(_MSC_VER)
	#include "..\include\dbghelp.h"
#endif

#define STACKDUMPER_EXPORTS
#define _USE_DISSASEMBLER 

#include "..\include\stackdumper.h"

using namespace std;

void DumpLastError();
#if defined(_MSC_VER)
void getFuncInfo(long addr,string &module,string &funcName);
#endif

///////////////////////////////////////////////////////////////////////////////

BOOL APIENTRY DllMain( HANDLE hModule, DWORD  ul_reason_for_call, LPVOID lpReserved)
{
  switch (ul_reason_for_call)
	{
		case DLL_PROCESS_ATTACH:
		case DLL_THREAD_ATTACH:  
		case DLL_THREAD_DETACH:
		case DLL_PROCESS_DETACH:  
			break;
    }
    return TRUE;
}

ADDR StackDumper::origRetAddr = 0;

StackDumper::StackDumper(bool bUseExitThunk)
{
  mbUseExitThunk = bUseExitThunk;
  file = fopen("stack.log", "wt");
}

void DoTheWork()
{
}

/* static  */
void __declspec(naked) StackDumper::ExitThunk()
{
  __asm
  {
    push ebp
    mov ebp, esp
    sub esp, 4 //make room for one local variable
    pushad 
  }

  DoTheWork();

  long temp;
  temp = origRetAddr;
  __asm{
    popad
    mov esi, temp
    mov esp,ebp
    pop ebp
    jmp esi //I don't have return value on the stack -> can not use ret here.
  }
}

StackDumper::~StackDumper()
{
  fclose(file);

  if (mbUseExitThunk)
  {
    ADDR temp;
    __asm
    {
      mov eax, ebp//ebp = my stack frame
      mov eax, [eax]//get caller stack frame
      mov eax, [eax+4]//get caller ret addr
      mov temp, eax
    }

    if (temp != (ADDR)ExitThunk)
    {
      origRetAddr = temp;

      temp = (ADDR)ExitThunk;

      __asm
      {
        mov eax, ebp
        mov eax, [eax]//get caller stack frame
        mov ecx, temp
        mov [eax+4], ecx//change the ret addr
      }
    }
  }
}

ADDR StackDumper::GetCallerAddr(long _ebp)
{
  long midPtPtr;
  long res = 0;

  __asm{
    mov eax, _ebp
    mov ecx, [eax]//check for end of stack frames list
    test ecx, ecx
    jz label
    mov eax, [eax+4]//get the ret addr
    test eax,eax//there are some cases(Borland C++ RTL) when the function has valid
    //stack frame but the ret addr is 0   
    jz label
    mov midPtPtr, eax
  }

  res = GetFuncAddr(midPtPtr);
label:
  return res;
}


#define PROLOGUE_SIGNATURE 0x00EC8B55
extern "C" DWORD FindAddressOfPrevInstruction(DWORD EIP, PUCHAR instr);

ADDR StackDumper::GetFuncAddr(ADDR midPtPtr)
{
#if defined(_USE_DISSASEMBLER)
  ADDR pKnown = midPtPtr;
  char buff[255];

  long temp;
  do {
    pKnown = FindAddressOfPrevInstruction(pKnown, (PUCHAR)buff);
    temp = (long)(*(long*)pKnown);
  } while((temp&0x00FFFFFF) != PROLOGUE_SIGNATURE);

  return pKnown;

#else
  long temp;
  do {
    temp = (long)(*(long*)midPtPtr);
    if ((temp&0x00FFFFFF) == PROLOGUE_SIGNATURE)
      break;
    midPtPtr--;
  } while(true);

  return midPtPtr;
#endif
}

// !!!!See /Oy option !!!!
void StackDumper::DumpStack(unsigned nDepth)
{
  long m_ebp;
  __asm{
    mov eax, ebp
    mov m_ebp, eax
  }
  
  ADDR callerAddr = GetCallerAddr(m_ebp);

  char buff[1024];
#if defined(_MSC_VER)
  string module,name;
  getFuncInfo(callerAddr,module,name);
  sprintf(buff, "%s!%s (%08X)\n",module.c_str(),name.c_str(),callerAddr);
#else
  sprintf(buff, "(%08X)\n",callerAddr);
#endif
  fwrite(buff, sizeof(char), strlen(buff), file);

  int i = 0;
  while (i < nDepth)
  {
    callerAddr = GetCallerAddr(m_ebp);
    if (callerAddr == 0)
      break;
  #if defined(_MSC_VER)
    string module,name;
    getFuncInfo(callerAddr,module,name);
    sprintf(buff, "%s!%s (%08X)\n",module.c_str(),name.c_str(),callerAddr);
  #else
    sprintf(buff, "(%08X)\n",callerAddr);
  #endif
    fwrite(buff, sizeof(char), strlen(buff), file);
    m_ebp = *((long*)m_ebp);
    i++;
  }
}

bool StackDumper::IsCalledFrom(char* pszFuncName)
{
#if defined(_MSC_VER)
  long m_ebp;
  __asm{
    mov eax, ebp
    mov m_ebp, eax
  }

  ADDR callerAddr;

  HANDLE h = GetCurrentProcess();
	char buf[_MAX_PATH];

  ::GetModuleFileName(GetModuleHandle("stackdumper.dll"),buf, _MAX_PATH);

  char *p = strrchr(buf, '\\');
  if (p)
  	*p = '\0';
  
  if (!::SymInitialize(h,buf,TRUE))
    return false;

  while (true)
  {
    callerAddr = GetCallerAddr(m_ebp);
    if (callerAddr == 0)
      break;

    BYTE szSymbolName[256];
    BYTE buffer[256];
    PSYMBOL_INFO pSymbol = (PSYMBOL_INFO)buffer;

    lstrcpy((char*)szSymbolName, pszFuncName);
    pSymbol->SizeOfStruct = sizeof(SYMBOL_INFO);
    pSymbol->MaxNameLen = sizeof(buffer) - sizeof(SYMBOL_INFO) + 1;

    if (SymFromName(GetCurrentProcess(), (char*)szSymbolName, pSymbol))
    {
      if (pSymbol->Address == callerAddr)
        return true;
    }

    m_ebp = *((long*)m_ebp);
  }

  SymCleanup(GetCurrentProcess());
  return false;
#else
  //map file parsing should be placed here
  return false;
#endif
}

//==============================================================================
// Matt Pietrek, Under the Hood, MSJ April 1997
// Given a linear address, locates the module, section, and offset containing
// that address.
//
// Note: the szModule paramater buffer is an output buffer of length specified
// by the len parameter (in characters!)
//==============================================================================
BOOL GetLogicalAddress(PVOID addr, PTSTR szModule, DWORD len, DWORD& section, DWORD& offset )
{
    MEMORY_BASIC_INFORMATION mbi;

    if ( !VirtualQuery( addr, &mbi, sizeof(mbi) ) )
        return FALSE;

    DWORD hMod = (DWORD)mbi.AllocationBase;

    if ( !GetModuleFileName( (HMODULE)hMod, szModule, len ) )
        return FALSE;

    // Point to the DOS header in memory
    PIMAGE_DOS_HEADER pDosHdr = (PIMAGE_DOS_HEADER)hMod;

    // From the DOS header, find the NT (PE) header
    PIMAGE_NT_HEADERS pNtHdr = (PIMAGE_NT_HEADERS)(hMod + pDosHdr->e_lfanew);

    PIMAGE_SECTION_HEADER pSection = IMAGE_FIRST_SECTION( pNtHdr );

    DWORD rva = (DWORD)addr - hMod; // RVA is offset from module load address

    // Iterate through the section table, looking for the one that encompasses
    // the linear address.
    for (   unsigned i = 0;
            i < pNtHdr->FileHeader.NumberOfSections;
            i++, pSection++ )
    {
        DWORD sectionStart = pSection->VirtualAddress;
        DWORD sectionEnd = sectionStart
                    + max(pSection->SizeOfRawData, pSection->Misc.VirtualSize);

        // Is the address in this section???
        if ( (rva >= sectionStart) && (rva <= sectionEnd) )
        {
            // Yes, address is in the section.  Calculate section and offset,
            // and store in the "section" & "offset" params, which were
            // passed by reference.
            section = i+1;
            offset = rva - sectionStart;
            return TRUE;
        }
    }

    return FALSE;   // Should never get here!
}