VOID HookTCP( const IMG img, const AFUNPTR replacement, const char *functionName )
{
fprintf(LogFile, "Found %s...\r\n", functionName );
std::cerr << "Found " << functionName << "..." << endl;
// hook
RTN rtn = RTN_FindByName( img, functionName );
PROTO proto =
PROTO_Allocate( PIN_PARG(WINDOWS::INT),
CALLINGSTD_STDCALL,
functionName,
PIN_PARG(WINDOWS::SOCKET), // s,
PIN_PARG(WINDOWS::CHAR *), // buf,
PIN_PARG(WINDOWS::INT), // len
PIN_PARG(WINDOWS::INT), // flags
PIN_PARG_END()
);
RTN_ReplaceSignature( rtn, replacement,
IARG_PROTOTYPE, proto,
IARG_ORIG_FUNCPTR,
IARG_FUNCARG_ENTRYPOINT_VALUE, 0,
IARG_FUNCARG_ENTRYPOINT_VALUE, 1,
IARG_FUNCARG_ENTRYPOINT_VALUE, 2,
IARG_FUNCARG_ENTRYPOINT_VALUE, 3,
IARG_CONTEXT,
IARG_PTR, functionName,
IARG_END
);
PROTO_Free( proto );
}
VOID Image(IMG img, void *v)
{
RTN rtn = RTN_FindByName(img, "ReplacedX87Regs");
if (RTN_Valid(rtn))
{
PROTO proto = PROTO_Allocate(PIN_PARG(int), CALLINGSTD_DEFAULT, "ReplacedX87Regs", PIN_PARG_END());
RTN_ReplaceSignature(rtn, AFUNPTR(REPLACE_ReplacedX87Regs),
IARG_PROTOTYPE, proto,
(KnobUseIargConstContext)?IARG_CONST_CONTEXT:IARG_CONTEXT,
IARG_THREAD_ID,
IARG_ORIG_FUNCPTR,
IARG_END);
PROTO_Free(proto);
printf ("TOOL found and replaced ReplacedX87Regs\n");
fflush (stdout);
RTN rtn = RTN_FindByName(img, "ExecutedAtFunc");
if (RTN_Valid(rtn))
{
executeAtAddr = RTN_Address(rtn);
printf ("TOOL found ExecutedAtFunc for later PIN_ExecuteAt\n");
fflush (stdout);
}
rtn = RTN_FindByName(img, "DumpX87RegsAtException");
if (RTN_Valid(rtn))
{
dumpX87RegsAtExceptionAddr = RTN_Address(rtn);
printf ("TOOL found DumpX87RegsAtException for later Exception\n");
fflush (stdout);
}
}
}
// Hooks
VOID HookUDP( const IMG img, const AFUNPTR replacement, const char *functionName )
{
std::cerr << "Found " << functionName << "..." << endl;
RTN rtn = RTN_FindByName(img, functionName );
PROTO proto =
PROTO_Allocate( PIN_PARG(WINDOWS::INT), // retval
CALLINGSTD_STDCALL,
functionName,
PIN_PARG(WINDOWS::SOCKET), // s,
PIN_PARG(WINDOWS::CHAR *), // buf,
PIN_PARG(WINDOWS::INT), // len
PIN_PARG(WINDOWS::INT), // flags
PIN_PARG(WINDOWS::SOCKADDR *), // from|to,
PIN_PARG(WINDOWS::INT *), // fromlen|tolen,
PIN_PARG_END()
);
RTN_ReplaceSignature(rtn, replacement,
IARG_PROTOTYPE, proto,
IARG_ORIG_FUNCPTR,
IARG_FUNCARG_ENTRYPOINT_VALUE, 0,
IARG_FUNCARG_ENTRYPOINT_VALUE, 1,
IARG_FUNCARG_ENTRYPOINT_VALUE, 2,
IARG_FUNCARG_ENTRYPOINT_VALUE, 3,
IARG_FUNCARG_ENTRYPOINT_VALUE, 4,
IARG_FUNCARG_ENTRYPOINT_VALUE, 5,
IARG_CONTEXT,
IARG_PTR, functionName,
IARG_END
);
std::cerr << "Replaced " << RTN_Name( rtn ) << " and stuff. " << endl;
PROTO_Free(proto);
}
void DoReplacementFunc( IMG img, char * funcName)
{
RTN rtnToReplace;
printf ("Image %s\n", IMG_Name(img).c_str());
for (SEC sec = IMG_SecHead(img); SEC_Valid(sec); sec = SEC_Next(sec))
{
for (RTN rtn = SEC_RtnHead(sec); RTN_Valid(rtn); rtn = RTN_Next(rtn))
{
printf (" Rtn: %s %s\n", RTN_Name(rtn).c_str(), funcName);
if (strstr( RTN_Name(rtn).c_str(), funcName))
{
//printf (" found\n");
foundFunc = true;
rtnToReplace = rtn;
break;
}
}
if (foundFunc)
{
break;
}
}
if (!foundFunc)
{
return;
}
printf ( "Found %s %x\n", funcName, RTN_Address(rtnToReplace));
// commented out so that absence of pdb file will not cause failure
if (RTN_IsSafeForProbedReplacement(rtnToReplace))
{
printf ( "RTN_ReplaceSignatureProbed on %s\n", funcName);
foundFunc = true;
PROTO protoOfStdCallFunction1ToBeReplacedByPin
= PROTO_Allocate( PIN_PARG(void *),
CALLINGSTD_STDCALL,
"protoOfStdCallFunction1ToBeReplacedByPin",
PIN_PARG(char),
PIN_PARG(int),
PIN_PARG(char),
PIN_PARG(int),
PIN_PARG_END() );
RTN_ReplaceSignatureProbed(rtnToReplace, AFUNPTR(ReplacementFunc),
IARG_PROTOTYPE, protoOfStdCallFunction1ToBeReplacedByPin,
IARG_ORIG_FUNCPTR,
IARG_FUNCARG_ENTRYPOINT_VALUE, 0,
IARG_FUNCARG_ENTRYPOINT_VALUE, 1,
IARG_FUNCARG_ENTRYPOINT_VALUE, 2,
IARG_FUNCARG_ENTRYPOINT_VALUE, 3,
IARG_END);
PROTO_Free( protoOfStdCallFunction1ToBeReplacedByPin );
}
}
VOID ImageLoad(IMG img, VOID *v)
{
RTN rtn = RTN_FindByName(img, "TestIargPreserveInReplacement");
if (RTN_Valid(rtn))
{
printf ("found TestIargPreserveInReplacement\n");
fflush (stdout);
PROTO proto = PROTO_Allocate( PIN_PARG(void), CALLINGSTD_DEFAULT,
"TestIargPreserveInReplacementProto", PIN_PARG_END() );
REGSET regsFP;
REGSET_Clear(regsFP);
REGSET_Insert (regsFP, REG_X87);
RTN_ReplaceSignature(
rtn, AFUNPTR(TestIargPreserveReplacement),
IARG_PROTOTYPE, proto,
IARG_CONTEXT,
IARG_ORIG_FUNCPTR,
IARG_PRESERVE, ®sFP,
IARG_END);
PROTO_Free( proto );
}
rtn = RTN_FindByName(img, "TestIargPreserveInReplacement1");
if (RTN_Valid(rtn))
{
printf ("found TestIargPreserveInReplacement1\n");
fflush (stdout);
PROTO proto = PROTO_Allocate( PIN_PARG(void), CALLINGSTD_DEFAULT,
"TestIargPreserveInReplacementProto", PIN_PARG_END() );
RTN_ReplaceSignature(
rtn, AFUNPTR(TestIargPreserveReplacement1),
IARG_PROTOTYPE, proto,
IARG_CONTEXT,
IARG_ORIG_FUNCPTR,
IARG_END);
PROTO_Free( proto );
}
}
void ImageLoad (IMG img, void *context)
{
fprintf (stderr, "Notified of load of %s at [%p,%p]\n",
IMG_Name(img).c_str(),
(char *)IMG_LowAddress(img), (char *)IMG_HighAddress(img));
// See if this is ntdll.dll
char szName[_MAX_FNAME];
char szExt[_MAX_EXT];
_splitpath_s (IMG_Name(img).c_str(),
NULL, 0,
NULL, 0,
szName, _MAX_FNAME,
szExt, _MAX_EXT);
strcat_s (szName, _MAX_FNAME, szExt);
if (0 != _stricmp ("ntdll.dll", szName))
return;
RTN rtn = RTN_FindByName (img, "RtlAllocateHeap");
if (RTN_Invalid() == rtn)
{
fprintf (stderr, "Failed to find RtlAllocateHeap in %s\n",
IMG_Name(img).c_str());
return;
}
fprintf(stderr,"Replacing\n");
PROTO protoRtlAllocateHeap =
PROTO_Allocate (PIN_PARG(void *),
CALLINGSTD_STDCALL,
"RtlAllocateHeap",
PIN_PARG(WINDOWS::PVOID), // HeapHandle
PIN_PARG(WINDOWS::ULONG), // Flags
PIN_PARG(WINDOWS::SIZE_T), // Size
PIN_PARG_END());
RTN_ReplaceSignature (rtn, (AFUNPTR)replacement_RtlAllocateHeap,
IARG_PROTOTYPE, protoRtlAllocateHeap,
IARG_ORIG_FUNCPTR,
IARG_FUNCARG_ENTRYPOINT_VALUE, 0,
IARG_FUNCARG_ENTRYPOINT_VALUE, 1,
IARG_FUNCARG_ENTRYPOINT_VALUE, 2,
IARG_CONTEXT,
IARG_END);
PROTO_Free (protoRtlAllocateHeap);
}
/* ===================================================================== */
VOID ImageLoad(IMG img, VOID *v)
{
if (!IMG_IsMainExecutable(img))
{
return;
}
const char * name = "check_sp_value";
RTN rtn = RTN_FindByName(img, name);
if (RTN_Valid(rtn))
{
PROTO proto = PROTO_Allocate(PIN_PARG(int),
CALLINGSTD_DEFAULT, name,
PIN_PARG(void*),
PIN_PARG_END() );
if ( ! PIN_IsProbeMode() )
{
RTN_ReplaceSignature(
rtn, AFUNPTR(check_sp_value),
IARG_PROTOTYPE, proto,
IARG_FUNCARG_ENTRYPOINT_VALUE, 0,
IARG_REG_VALUE, REG_STACK_PTR,
IARG_REG_VALUE, REG_ESP,
IARG_REG_VALUE, REG_SP,
IARG_END);
}
else if (RTN_IsSafeForProbedReplacement(rtn))
{
RTN_ReplaceSignatureProbed(
rtn, AFUNPTR(check_sp_value),
IARG_PROTOTYPE, proto,
IARG_FUNCARG_ENTRYPOINT_VALUE, 0,
IARG_REG_VALUE, REG_STACK_PTR,
IARG_REG_VALUE, REG_ESP,
IARG_REG_VALUE, REG_SP,
IARG_END);
}
PROTO_Free( proto);
}
}
// WSASendTo HOOK Untested, I think. I forget :/.
VOID HookWSASendTo( const IMG img, AFUNPTR replacement, const char *functionName )
{
std::cerr << "Found " << functionName << "..." << endl;
RTN rtn = RTN_FindByName(img, functionName );
PROTO proto =
PROTO_Allocate( PIN_PARG(WINDOWS::INT), // retval
CALLINGSTD_STDCALL,
functionName,
PIN_PARG(WINDOWS::SOCKET), // s,
PIN_PARG(WINDOWS::LPWSABUF), // lpBuffers,
PIN_PARG(WINDOWS::DWORD), // dwBufferCount
PIN_PARG(WINDOWS::LPDWORD), // lpNumberOfBytesSend
PIN_PARG(WINDOWS::DWORD), // dwFlags
PIN_PARG(WINDOWS::SOCKADDR *), // lpTo,
PIN_PARG(WINDOWS::INT), // iToLen,
PIN_PARG(WINDOWS::LPWSAOVERLAPPED), // lpOverlapped
PIN_PARG(WINDOWS::LPWSAOVERLAPPED_COMPLETION_ROUTINE), // lpCompletionRoutine
PIN_PARG_END()
);
RTN_ReplaceSignature(rtn, replacement,
IARG_PROTOTYPE, proto,
IARG_ORIG_FUNCPTR,
IARG_FUNCARG_ENTRYPOINT_VALUE, 0,
IARG_FUNCARG_ENTRYPOINT_VALUE, 1,
IARG_FUNCARG_ENTRYPOINT_VALUE, 2,
IARG_FUNCARG_ENTRYPOINT_VALUE, 3,
IARG_FUNCARG_ENTRYPOINT_VALUE, 4,
IARG_FUNCARG_ENTRYPOINT_VALUE, 5,
IARG_FUNCARG_ENTRYPOINT_VALUE, 6,
IARG_FUNCARG_ENTRYPOINT_VALUE, 7,
IARG_FUNCARG_ENTRYPOINT_VALUE, 8,
IARG_CONTEXT,
IARG_PTR, functionName,
IARG_END
);
std::cerr << "Replaced " << RTN_Name( rtn ) << endl;
PROTO_Free(proto);
}
void InsertProbe( IMG img, char * funcName)
{
/*
printf ("Image %s\n", IMG_Name(img).c_str());
for (SEC sec = IMG_SecHead(img); SEC_Valid(sec); sec = SEC_Next(sec))
{
for (RTN rtn = SEC_RtnHead(sec); RTN_Valid(rtn); rtn = RTN_Next(rtn))
{
printf (" Rtn: %s %s\n", RTN_Name(rtn).c_str(), funcName);
if (strstr( RTN_Name(rtn).c_str(), funcName))
{
printf (" found\n");
}
}
}
*/
RTN allocRtn = RTN_FindByName(img, funcName);
if (RTN_Valid(allocRtn) && RTN_IsSafeForProbedReplacement(allocRtn))
{
fprintf (fp, "RTN_ReplaceSignatureProbed on %s\n", funcName);
PROTO protoHeapAlloc = PROTO_Allocate( PIN_PARG(void *), CALLINGSTD_STDCALL,
"protoHeapAlloc", PIN_PARG(WINDOWS::HANDLE),
PIN_PARG(WINDOWS::DWORD),PIN_PARG(WINDOWS::DWORD), PIN_PARG_END() );
RTN_ReplaceSignatureProbed(allocRtn, AFUNPTR(ReplacementFunc),
IARG_PROTOTYPE, protoHeapAlloc,
IARG_ORIG_FUNCPTR,
IARG_FUNCARG_ENTRYPOINT_VALUE, 0,
IARG_FUNCARG_ENTRYPOINT_VALUE, 1,
IARG_FUNCARG_ENTRYPOINT_VALUE, 2,
IARG_CONTEXT,
IARG_RETURN_IP,
IARG_END);
PROTO_Free( protoHeapAlloc );
}
VOID Image(IMG img, void *v)
{
RTN rtn = RTN_FindByName(img, "Replaced");
if (RTN_Valid(rtn))
{
PROTO proto = PROTO_Allocate(PIN_PARG(int), CALLINGSTD_DEFAULT, "Replaced", PIN_PARG_END());
RTN_ReplaceSignature(rtn, AFUNPTR(REPLACE_Replaced),
IARG_PROTOTYPE, proto,
(KnobUseIargConstContext)?IARG_CONST_CONTEXT:IARG_CONTEXT,
IARG_THREAD_ID,
IARG_ORIG_FUNCPTR,
IARG_END);
PROTO_Free(proto);
}
rtn = RTN_FindByName(img, "Inner");
if (RTN_Valid(rtn))
{
RTN_Open(rtn);
RTN_InsertCall(rtn, IPOINT_BEFORE, AFUNPTR(AtInner), IARG_REG_VALUE, scratchReg, IARG_END);
RTN_Close(rtn);
}
}
static VOID
HookHeapFunctions(IMG img)
{
RTN rtn;
// check this image actually has the heap functions.
if ((rtn = RTN_FindByName(img, "RtlAllocateHeap")) == RTN_Invalid())
return;
// hook RtlAllocateHeap
RTN rtlAllocate = RTN_FindByName(img, "RtlAllocateHeap");
PROTO protoRtlAllocateHeap = \
PROTO_Allocate( PIN_PARG(void *),
CALLINGSTD_STDCALL,
"RtlAllocateHeap",
PIN_PARG(WINDOWS::PVOID), // HeapHandle
PIN_PARG(WINDOWS::ULONG), // Flags
PIN_PARG(WINDOWS::SIZE_T), // Size
PIN_PARG_END()
);
RTN_ReplaceSignature(rtlAllocate,(AFUNPTR)replacementRtlAllocateHeap,
IARG_PROTOTYPE, protoRtlAllocateHeap,
IARG_ORIG_FUNCPTR,
IARG_FUNCARG_ENTRYPOINT_VALUE, 0,
IARG_FUNCARG_ENTRYPOINT_VALUE, 1,
IARG_FUNCARG_ENTRYPOINT_VALUE, 2,
IARG_CONTEXT,
IARG_END
);
PROTO_Free(protoRtlAllocateHeap);
// replace RtlReAllocateHeap()
RTN rtlReallocate = RTN_FindByName(img, "RtlReAllocateHeap");
PROTO protoRtlReAllocateHeap = \
PROTO_Allocate( PIN_PARG(void *), CALLINGSTD_STDCALL,
"RtlReAllocateHeap",
PIN_PARG(WINDOWS::PVOID), // HeapHandle
PIN_PARG(WINDOWS::ULONG), // Flags
PIN_PARG(WINDOWS::PVOID), // MemoryPtr
PIN_PARG(WINDOWS::SIZE_T),// Size
PIN_PARG_END()
);
RTN_ReplaceSignature(rtlReallocate,(AFUNPTR)replacementRtlReAllocateHeap,
IARG_PROTOTYPE, protoRtlReAllocateHeap,
IARG_ORIG_FUNCPTR,
IARG_FUNCARG_ENTRYPOINT_VALUE, 0,
IARG_FUNCARG_ENTRYPOINT_VALUE, 1,
IARG_FUNCARG_ENTRYPOINT_VALUE, 2,
IARG_FUNCARG_ENTRYPOINT_VALUE, 3,
IARG_CONTEXT,
IARG_END
);
PROTO_Free(protoRtlReAllocateHeap);
// replace RtlFreeHeap
RTN rtlFree = RTN_FindByName(img, "RtlFreeHeap");
PROTO protoRtlFreeHeap = \
PROTO_Allocate( PIN_PARG(void *), CALLINGSTD_STDCALL,
"RtlFreeHeap",
PIN_PARG(WINDOWS::PVOID), // HeapHandle
PIN_PARG(WINDOWS::ULONG), // Flags
PIN_PARG(WINDOWS::PVOID), // MemoryPtr
PIN_PARG_END()
);
RTN_ReplaceSignature(rtlFree,(AFUNPTR)replacementRtlFreeHeap,
IARG_PROTOTYPE, protoRtlFreeHeap,
IARG_ORIG_FUNCPTR,
IARG_FUNCARG_ENTRYPOINT_VALUE, 0,
IARG_FUNCARG_ENTRYPOINT_VALUE, 1,
IARG_FUNCARG_ENTRYPOINT_VALUE, 2,
IARG_CONTEXT,
IARG_END
);
PROTO_Free(protoRtlAllocateHeap);
}
