E:/PROJECTS/cvsed/mixed/VIRTUA~1/kdclient/tablemgr.h

00001 
00007 #pragma once
00008 
00009 #include <bzscmn/bzscmn.h>
00010 #include <bzscmn/smvec.h>
00011 #include <list>
00012 #include <bzscmn/serializer.h>
00013 
00014 static const int s_EnableLogging = true;
00015 
00016 DECLARE_SERIALIZEABLE_STRUC5(PatchInfoRecord, 
00017                                                          unsigned, TimeStamp,
00018                                                          unsigned, FileSize,
00019                                                          unsigned, FirstSectionSize,
00020 
00021                                                          unsigned, DispatcherTableOffset,
00022                                                          unsigned, EntryCount);
00023 
00024 using namespace BazisLib;
00025 
00026 DECLARE_SERIALIZEABLE_STRUC5_I(PatchingParams,
00027                                                            bool, AllowPatchingAtTableStart, true,
00028                                                            bool, DefaultPatchingAtTableStart, false,
00029 //                                                         String, VMWareCommandToReplace, false,
00030                                                            bool, AllowReplacingFirstCommand, false,
00031                                                            bool, WaitForNonZeroFirstCommand, true,
00032                                                            unsigned, MinRunTimeBeforePatch, 10000);
00033                                                                   
00034         
00036 class PatchInfoDatabase
00037 {
00038 private:
00039         BazisLib::FilePath m_Directory;
00040         BazisLib::SingleMallocVector<PatchInfoRecord> m_Records;
00041 
00042 private:
00043         static bool QueryLoadedImage(void *pPEHeader, PatchInfoRecord &rec);
00044 
00045 public:
00046         PatchInfoDatabase(const BazisLib::FilePath &Directory);
00047         ~PatchInfoDatabase();
00048 
00050 
00055         bool FindRecord(void *pPEHeader, PatchInfoRecord &record);
00056         bool AddRecord(void *pPEHeader, LPVOID lpDispatcherTable, unsigned EntryCount);
00057 };
00058 
00060 
00065 class RPCTableManager
00066 {
00067 private:
00068 #pragma region Auxillary structures
00070         struct AddressRange
00071         {
00072                 char *lpStart;
00073                 char *lpEnd;
00074 
00075                 AddressRange(LPVOID start, LPVOID end)
00076                         : lpStart((char *)start)
00077                         , lpEnd((char *)end)
00078                 {
00079                         ASSERT(lpEnd > lpStart);
00080                 }
00081 
00082                 size_t GetSize()
00083                 {
00084                         return lpEnd - lpStart;
00085                 }
00086         };
00087 
00089         struct StringPointer
00090         {
00091                 char *pszData;
00092                 size_t length;
00093 
00094                 StringPointer(char *psz, size_t len)
00095                         : pszData(psz)
00096                         , length(len)
00097                 {
00098                 }
00099         };
00100 
00102         struct StringReferenceDescriptor
00103         {
00104                 void *pAddress;
00105                 unsigned GroupID;
00106 
00107                 StringReferenceDescriptor(void *addr)
00108                         : pAddress(addr)
00109                         , GroupID(0)
00110                 {
00111 
00112                 }
00113         };
00114 
00116         struct RPCTableInfo
00117         {
00118                 void *pAddr;
00119                 unsigned EntryCount;
00120                 unsigned UsedEntries;
00121                 unsigned FreeEntries;
00122         };
00123 
00125         struct RefGroupDescriptor
00126         {
00127                 void *pAddr;
00128                 unsigned RefCount;
00129                 RPCTableInfo info;
00130 
00131                 RefGroupDescriptor(void *addr)
00132                         : pAddr(addr)
00133                         , RefCount(0)
00134                 {
00135                         memset(&info, 0, sizeof(info));
00136                 }
00137         };
00138 
00139 #pragma endregion
00140 #pragma region RPC handler definition
00141         typedef bool (*GRPCHANDLER)(void *pContext, int ChannelNumber, char *pCommandBody, unsigned CommandBodyLength, char **ppReply, unsigned *pReplyLen);
00143         struct RPCHandlerRecord
00144         {
00145                 char *pCommandPrefix;
00146                 unsigned CommandPrefixLength;
00147                 GRPCHANDLER pHandler;
00148                 void *pHandlerContext;
00149 
00150                 bool Enabled;
00151         };
00152 #pragma endregion
00153 
00154 private:
00155         PatchInfoDatabase m_Database;
00156         PatchingParams m_Params;
00157 
00158         RPCHandlerRecord m_OriginalHandler, m_PatchedEntryCopy;
00159         RPCHandlerRecord *m_pPatchedEntry;
00160 
00161 private:
00162 
00163         enum 
00164         {
00165                 VerifMaxStringLength = 100,
00166         };
00167 
00168         static inline bool PtrInRange(const std::list<AddressRange> &list, void *ptr)
00169         {
00170                 for (std::list<AddressRange>::const_iterator it = list.begin(); it != list.end(); it++)
00171                         if ((ptr >= it->lpStart) && (ptr <= it->lpEnd))
00172                                 return true;
00173                 return false;
00174         }
00175 
00177         static inline bool islen(const char *pStr, size_t expected)
00178         {
00179                 for (size_t i = 0; i < expected; i++)
00180                         if (!pStr[i])
00181                                 return false;
00182                 return !pStr[expected];
00183         }
00184 
00186         static inline bool VerifyEntry(RPCHandlerRecord *pRec, std::list<AddressRange> &dataRanges, std::list<AddressRange> &codeRanges)
00187         {
00188                 switch (*((char *)&pRec->Enabled))
00189                 {
00190                 case 0:
00191 //                      if (pRec->CommandPrefixLength || pRec->pCommandPrefix || pRec->pHandler)
00192 //                              return false;
00193                         break;
00194                 case 1:
00195                         if (pRec->CommandPrefixLength > VerifMaxStringLength)
00196                                 return false;
00197                         if (!PtrInRange(dataRanges, pRec->pCommandPrefix))
00198                                 return false;
00199                         if (!islen(pRec->pCommandPrefix, pRec->CommandPrefixLength))
00200                                 return false;
00201                         if (!PtrInRange(codeRanges, pRec->pHandler))
00202                                 return false;
00203                         return true;
00204                 default:
00205                         return false;
00206                 }
00207                 return true;
00208         }
00209 
00210         static void FindSections(char *lpMainExe, std::list<AddressRange> &dataRanges, std::list<AddressRange> &codeRanges);
00211         static void MakeListOfStrings(const std::list<AddressRange> &dataRanges, BazisLib::SingleMallocVector<StringPointer> &strings, bool FullMode);
00212         static void FindStringRefs(const std::list<AddressRange> &dataRanges,
00213                                                            const BazisLib::SingleMallocVector<StringPointer> &strings,
00214                                                            BazisLib::SingleMallocVector<StringReferenceDescriptor> &stringRefs);
00215         static void GroupStringRefs(BazisLib::SingleMallocVector<StringReferenceDescriptor> &stringRefs,
00216                                                                 BazisLib::SingleMallocVector<RefGroupDescriptor> &groups);
00217         
00218         static bool ScanPotentialRPCTable(void *pAddr,
00219                                                                           std::list<AddressRange> &dataRanges,
00220                                                                           std::list<AddressRange> &codeRanges,
00221                                                                           RPCTableInfo &info);
00222 
00223 private:
00225         bool FindHandlerTable(bool FullMode = false);
00226         void DoPatch(RPCHandlerRecord *pEntry, const char *pszPrefix, size_t prefixLen, GRPCHANDLER pHandler, void *pContext)
00227         {
00228                 m_OriginalHandler = *pEntry;
00229                 pEntry->pCommandPrefix = (char *)pszPrefix;
00230                 pEntry->CommandPrefixLength = (unsigned)prefixLen;
00231                 pEntry->pHandler = pHandler;
00232                 pEntry->pHandlerContext = pContext;
00233                 pEntry->Enabled = true;
00234                 m_pPatchedEntry = pEntry;
00235                 m_PatchedEntryCopy = *pEntry;
00236         }
00237 
00238 public:
00239         RPCTableManager(HINSTANCE hThisDLL);
00240         bool InstallHandler(const char *pszPrefix, size_t prefixLen, GRPCHANDLER pHandler, void *pContext, bool ForceReinstall);
00241         void RestoreOriginalHandler();
00242 
00243         bool IsEntryModified()
00244         {
00245                 if (!m_pPatchedEntry)
00246                         return false;
00247                 //TODO: insert invertion here and test repeated patching bug
00248                 return (memcmp(m_pPatchedEntry, &m_PatchedEntryCopy, sizeof(RPCHandlerRecord)) != 0);
00249         }
00250 
00251 };