1 files changed, 469 insertions, 0 deletions
diff --git a/src/ToolBox/superpmi/superpmi-shared/methodcontextreader.cpp b/src/ToolBox/superpmi/superpmi-shared/methodcontextreader.cpp
new file mode 100644
index 0000000000..1e375e0d6a
--- /dev/null
+++ b/src/ToolBox/superpmi/superpmi-shared/methodcontextreader.cpp
@@ -0,0 +1,469 @@
+//
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+//
+
+//----------------------------------------------------------
+// MethodContextReader.cpp - Abstraction for reading MethodContexts
+//                         Should eventually support multithreading
+//----------------------------------------------------------
+
+#include "standardpch.h"
+#include "tocfile.h"
+#include "methodcontextreader.h"
+#include "methodcontext.h"
+#include "logging.h"
+#include "runtimedetails.h"
+
+// Just a helper...
+HANDLE MethodContextReader::OpenFile(const char *inputFile, DWORD flags)
+{
+    HANDLE fileHandle = CreateFileA(inputFile, GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL | flags, NULL);
+    if (fileHandle == INVALID_HANDLE_VALUE)
+    {
+        LogError("Failed to open file '%s'. GetLastError()=%u", inputFile, GetLastError());
+    }
+    return fileHandle;
+}
+
+static std::string to_lower(const std::string &input)
+{
+    std::string res = input;
+    std::transform(input.cbegin(), input.cend(), res.begin(), tolower);
+    return res;
+}
+
+// Looks for a file named foo.origSuffix.newSuffix or foo.newSuffix
+// but only if foo.origSuffix exists.
+//
+// Note: filename extensions must be lower-case, even on case-sensitive file systems!
+std::string MethodContextReader::CheckForPairedFile(const std::string &fileName, const std::string &origSuffix, const std::string &newSuffix)
+{
+    std::string tmp = to_lower(origSuffix);
+
+    // First, check to see if foo.origSuffix exists
+    size_t suffix_offset = fileName.find_last_of('.');
+    if ((SSIZE_T)suffix_offset <= 0
+        || (tmp != to_lower(fileName.substr(suffix_offset)))
+        || (GetFileAttributesA(fileName.c_str()) == INVALID_FILE_ATTRIBUTES))
+        return std::string();
+
+    // next, check foo.orig.new from foo.orig
+    tmp = fileName + newSuffix;
+    if (GetFileAttributesA(tmp.c_str()) != INVALID_FILE_ATTRIBUTES)
+        return tmp;
+
+    // Now let's check for foo.new from foo.new.orig
+    tmp = fileName.substr(0, suffix_offset);
+    if (GetFileAttributesA(tmp.c_str()) != INVALID_FILE_ATTRIBUTES)
+        return tmp;
+
+    // Finally, lets try foo.orig from foo.new
+    tmp += newSuffix;
+    if (GetFileAttributesA(tmp.c_str()) != INVALID_FILE_ATTRIBUTES)
+        return tmp;
+
+    return std::string();
+}
+
+
+MethodContextReader::MethodContextReader(const char *inputFileName, const int *indexes, int indexCount, char *hash, int offset, int increment)
+    : fileHandle(INVALID_HANDLE_VALUE)
+    , fileSize(0)
+    , curMCIndex(0)
+    , Indexes(indexes)
+    , IndexCount(indexCount)
+    , curIndexPos(0)
+    , Hash(hash)
+    , curTOCIndex(0)
+    , Offset(offset)
+    , Increment(increment)
+{
+    this->mutex = CreateMutexA(NULL, FALSE, nullptr);
+
+    std::string tocFileName, mchFile;
+
+    // First, check to see if they passed an MCH file (look for a paired MCT file)
+    tocFileName = MethodContextReader::CheckForPairedFile(inputFileName, ".mch", ".mct");
+    if (!tocFileName.empty())
+    {
+        mchFile = inputFileName;
+    }
+    else
+    {
+        // Okay, it wasn't an MCH file, let's check to see if it was an MCT file
+        // so check for a paired MCH file instead
+        mchFile = MethodContextReader::CheckForPairedFile(inputFileName, ".mct", ".mch");
+        if (!mchFile.empty())
+        {
+            tocFileName = inputFileName;
+        }
+        else
+        {
+            mchFile = inputFileName;
+        }
+    }
+
+    if (!tocFileName.empty())
+        this->tocFile.LoadToc(tocFileName.c_str());
+
+    // we'll get here even if we don't have a valid index file
+    this->fileHandle = OpenFile(mchFile.c_str(), (this->hasTOC() && this->hasIndex()) ? FILE_ATTRIBUTE_NORMAL : FILE_FLAG_SEQUENTIAL_SCAN);
+    if (this->fileHandle != INVALID_HANDLE_VALUE)
+    {
+        GetFileSizeEx(this->fileHandle, (PLARGE_INTEGER)&this->fileSize);
+    }
+}
+
+MethodContextReader::~MethodContextReader()
+{
+    if (fileHandle != INVALID_HANDLE_VALUE)
+    {
+        CloseHandle(this->fileHandle);
+    }
+
+    CloseHandle(this->mutex);
+}
+
+bool MethodContextReader::AcquireLock()
+{
+    DWORD res = WaitForSingleObject(this->mutex, INFINITE);
+    return (res == WAIT_OBJECT_0);
+}
+
+void MethodContextReader::ReleaseLock()
+{
+    ReleaseMutex(this->mutex);
+}
+
+bool MethodContextReader::atEof()
+{
+    __int64 pos = 0;
+    SetFilePointerEx(this->fileHandle, *(PLARGE_INTEGER)&pos, (PLARGE_INTEGER)&pos, FILE_CURRENT); // LARGE_INTEGER is a crime against humanity
+    return pos == this->fileSize;
+}
+
+MethodContextBuffer MethodContextReader::ReadMethodContextNoLock(bool justSkip)
+{
+    DWORD bytesRead;
+    char buff[512];
+    unsigned int totalLen = 0;
+    if (atEof())
+    {
+        return MethodContextBuffer();
+    }
+    Assert(ReadFile(this->fileHandle, buff, 2 + sizeof(unsigned int), &bytesRead, NULL) == TRUE);
+    AssertMsg((buff[0] == 'm') && (buff[1] == 'c'), "Didn't find magic number");
+    memcpy(&totalLen, &buff[2], sizeof(unsigned int));
+    if (justSkip)
+    {
+        __int64 pos = totalLen + 2;
+        // Just move the file pointer ahead the correct number of bytes
+        AssertMsg(SetFilePointerEx(this->fileHandle, *(PLARGE_INTEGER)&pos, (PLARGE_INTEGER)&pos, FILE_CURRENT) == TRUE, "SetFilePointerEx failed (Error %X)", GetLastError());
+
+        //Increment curMCIndex as we advanced the file pointer by another MC
+        ++curMCIndex;
+
+        return MethodContextBuffer(0);
+    }
+    else
+    {
+        unsigned char *buff2 = new unsigned char[totalLen + 2]; //total + End Canary
+        Assert(ReadFile(this->fileHandle, buff2, totalLen + 2, &bytesRead, NULL) == TRUE);
+
+        //Increment curMCIndex as we read another MC
+        ++curMCIndex;
+
+        return MethodContextBuffer(buff2, totalLen);
+    }
+}
+
+MethodContextBuffer MethodContextReader::ReadMethodContext(bool acquireLock, bool justSkip)
+{
+    if (acquireLock && !this->AcquireLock())
+    {
+        LogError("Can't acquire the reader lock!");
+        return MethodContextBuffer(-1);
+    }
+
+    struct Param {
+        MethodContextReader* pThis;
+        MethodContextBuffer ret;
+        bool justSkip;
+    } param;
+    param.pThis = this;
+    param.ret = MethodContextBuffer(-2);
+    param.justSkip = justSkip;
+
+    PAL_TRY(Param*, pParam, &param)
+    {
+        pParam->ret = pParam->pThis->ReadMethodContextNoLock(pParam->justSkip);
+    }
+    PAL_FINALLY
+    {
+        this->ReleaseLock();
+    }
+    PAL_ENDTRY
+
+    return param.ret;
+}
+
+// Read a method context buffer from the ContextCollection
+// (either a hive [single] or an index)
+MethodContextBuffer MethodContextReader::GetNextMethodContext()
+{
+    struct Param : FilterSuperPMIExceptionsParam_CaptureException
+    {
+        MethodContextReader*    pThis;
+        MethodContextBuffer     mcb;
+    } param;
+    param.pThis = this;
+
+    PAL_TRY(Param*, pParam, &param)
+    {
+        pParam->mcb = pParam->pThis->GetNextMethodContextHelper();
+    }
+    PAL_EXCEPT_FILTER(FilterSuperPMIExceptions_CaptureExceptionAndStop)
+    {
+        LogError("Method %d is of low integrity.", GetMethodContextIndex());
+        param.mcb = MethodContextBuffer(-1);
+    }
+    PAL_ENDTRY
+
+    return param.mcb;
+}
+
+MethodContextBuffer MethodContextReader::GetNextMethodContextHelper()
+{
+    //If we have an offset/increment combo
+    if (this->Offset > 0 && this->Increment > 0)
+        return GetNextMethodContextFromOffsetIncrement();
+
+    //If we have an index
+    if (this->hasIndex())
+    {
+        if (this->curIndexPos < this->IndexCount)
+        {
+            //If we are not done with all of them
+            return GetNextMethodContextFromIndexes();
+        }
+        else //We are done with all of them, return
+            return MethodContextBuffer();
+    }
+
+    //If we have a hash
+    if (this->Hash != nullptr)
+        return GetNextMethodContextFromHash();
+
+    //If we don't have any of these options return all MCs one by one
+    return this->ReadMethodContext(true);
+}
+
+// Read a method context buffer from the ContextCollection using Indexes
+MethodContextBuffer MethodContextReader::GetNextMethodContextFromIndexes()
+{
+    //Assert if we don't have an Index or we are done with all the indexes
+    Assert(this->hasIndex() && this->curIndexPos < this->IndexCount);
+
+    if (this->hasTOC())
+    {
+        // If we have an index & we have a TOC, we can just jump to that method!
+        return this->GetSpecificMethodContext(this->Indexes[this->curIndexPos++]);
+    }
+    else
+    {
+        // Find the current method (either #0, or the previous index)
+        int curMethod = this->curIndexPos ? this->Indexes[this->curIndexPos - 1] : 0;
+        // Get the next method
+        int nextMethod = this->Indexes[this->curIndexPos++];
+        // Skip over methods until we get to the right now
+        while (++curMethod < nextMethod)
+        {
+            // Skip a method context
+            MethodContextBuffer mcb = this->ReadMethodContext(true, true);
+            if (mcb.allDone() || mcb.Error())
+                return mcb;
+        }
+    }
+    return this->ReadMethodContext(true);
+}
+
+// Read a method context buffer from the ContextCollection using Hash
+MethodContextBuffer MethodContextReader::GetNextMethodContextFromHash()
+{
+    //Assert if we don't have a valid hash
+    Assert(this->Hash != nullptr);
+
+    if (this->hasTOC())
+    {
+        //We have a TOC so lets go through the TOCElements
+        //one-by-one till we find a matching hash
+        for (; curTOCIndex < (int)this->tocFile.GetTocCount(); curTOCIndex++)
+        {
+            if (_strnicmp(this->Hash, this->tocFile.GetElementPtr(curTOCIndex)->Hash, MD5_HASH_BUFFER_SIZE) == 0)
+            {
+                //We found a match, return this specific method
+                return this->GetSpecificMethodContext(this->tocFile.GetElementPtr(curTOCIndex++)->Number);
+            }
+        }
+
+        //No more matches in the TOC for our hash value
+        return MethodContextBuffer();
+    }
+    else
+    {
+        // Keep reading all MCs until we hit a match
+        //or we reach the end or hit an error
+        while (true)
+        {
+            // Read a method context
+            // we can't skip because we need to calculate hashes
+            MethodContextBuffer mcb = this->ReadMethodContext(true, false);
+            if (mcb.allDone() || mcb.Error())
+                return mcb;
+
+            char mcHash[MD5_HASH_BUFFER_SIZE];
+
+            //Create a temporary copy of mcb.buff plus ending 2-byte canary
+            //this will get freed up by MethodContext constructor
+            unsigned char *buff = new unsigned char[mcb.size + 2];
+            memcpy(buff, mcb.buff, mcb.size + 2);
+
+            MethodContext *mc;
+
+            if (!MethodContext::Initialize(-1, buff, mcb.size, &mc))
+                return MethodContextBuffer(-1);
+
+            mc->dumpMethodMD5HashToBuffer(mcHash, MD5_HASH_BUFFER_SIZE);
+            delete mc;
+
+            if (_strnicmp(this->Hash, mcHash, MD5_HASH_BUFFER_SIZE) == 0)
+            {
+                //We found a match, return this specific method
+                return mcb;
+            }
+        }
+    }
+
+    //We should never get here under normal conditions
+    AssertMsg(true, "Unexpected condition hit while reading input file.");
+    return MethodContextBuffer(-1);
+}
+
+// Read a method context buffer from the ContextCollection using offset/increment
+MethodContextBuffer MethodContextReader::GetNextMethodContextFromOffsetIncrement()
+{
+    //Assert if we don't have a valid increment/offset combo
+    Assert(this->Offset > 0 && this->Increment > 0);
+
+    int methodNumber = this->curMCIndex > 0 ? this->curMCIndex + this->Increment : this->Offset;
+
+    if (this->hasTOC())
+    {
+        //Check if we are within the TOC
+        if ((int)this->tocFile.GetTocCount() >= methodNumber)
+        {
+            //We have a TOC so we can request a specific method context
+            return this->GetSpecificMethodContext(methodNumber);
+        }
+        else
+            return MethodContextBuffer();
+    }
+    else
+    {
+        // Keep skipping MCs until we get to the one we need to return
+        while (this->curMCIndex + 1 < methodNumber)
+        {
+            //skip over a method
+            MethodContextBuffer mcb = this->ReadMethodContext(true, true);
+            if (mcb.allDone() || mcb.Error())
+                return mcb;
+        }
+    }
+    return this->ReadMethodContext(true);
+}
+
+bool MethodContextReader::hasIndex()
+{
+    return this->IndexCount > 0;
+}
+
+bool MethodContextReader::hasTOC()
+{
+    return this->tocFile.GetTocCount() > 0;
+}
+
+bool MethodContextReader::isValid()
+{
+    return this->fileHandle != INVALID_HANDLE_VALUE && this->mutex != INVALID_HANDLE_VALUE;
+}
+
+double MethodContextReader::PercentComplete()
+{
+    if (this->hasIndex() && this->hasTOC())
+    {
+        // Best estimate I can come up with...
+        return 100.0 * (double)this->curIndexPos / (double)this->IndexCount;
+    }
+    this->AcquireLock();
+    __int64 pos = 0;
+    SetFilePointerEx(this->fileHandle, *(PLARGE_INTEGER)&pos, (PLARGE_INTEGER)&pos, FILE_CURRENT);
+    this->ReleaseLock();
+    return 100.0 * (double)pos / (double)this->fileSize;
+}
+
+// Binary search to get this method number from the index
+// Returns -1 for not found, or -2 for not indexed
+// Interview question alert: hurray for CLR headers incompatibility with STL :-(
+// Note that TOC is 0 based and MC# are 1 based!
+__int64 MethodContextReader::GetOffset(unsigned int methodNumber)
+{
+    if (!this->hasTOC())
+        return -2;
+    size_t high = this->tocFile.GetTocCount() - 1;
+    size_t low = 0;
+    while (low <= high)
+    {
+        size_t pos = (high + low) / 2;
+        unsigned int num = this->tocFile.GetElementPtr(pos)->Number;
+        if (num == methodNumber)
+            return this->tocFile.GetElementPtr(pos)->Offset;
+        if (num > methodNumber)
+            high = pos - 1;
+        else
+            low = pos + 1;
+    }
+    return -1;
+}
+
+MethodContextBuffer MethodContextReader::GetSpecificMethodContext(unsigned int methodNumber)
+{
+    __int64 pos = this->GetOffset(methodNumber);
+    if (pos < 0)
+    {
+        return MethodContextBuffer(-3);
+    }
+
+    // Take the IO lock before we set the file pointer, so we can do this on multiple threads
+    if (!this->AcquireLock())
+    {
+        return MethodContextBuffer(-2);
+    }
+    if (SetFilePointerEx(this->fileHandle, *(PLARGE_INTEGER)&pos, (PLARGE_INTEGER)&pos, FILE_BEGIN))
+    {
+        // ReadMethodContext will release the lock, but we already acquired it
+        MethodContextBuffer mcb = this->ReadMethodContext(false);
+
+        //The curMCIndex value updated by ReadMethodContext() is incorrect
+        //since we are repositioning the file pointer we need to update it
+        curMCIndex = methodNumber;
+
+        return mcb;
+    }
+    else
+    {
+        // Don't forget to release the lock!
+        this->ReleaseLock();
+        return MethodContextBuffer(-4);
+    }
+}