Fix memory corruption in GetResultFileName in SuperPMI (#24537)

4 files changed, 64 insertions, 72 deletions

diff --git a/src/ToolBox/superpmi/superpmi-shared/spmiutil.cpp b/src/ToolBox/superpmi/superpmi-shared/spmiutil.cpp
index 986c6d1f5e..79caf6020d 100644
--- a/src/ToolBox/superpmi/superpmi-shared/spmiutil.cpp
+++ b/src/ToolBox/superpmi/superpmi-shared/spmiutil.cpp
@@ -138,98 +138,94 @@ bool LoadRealJitLib(HMODULE& jitLib, WCHAR* jitLibPath)
     return true;
 }
 
-void cleanupExecutableName(WCHAR* executableName)
+void ReplaceIllegalCharacters(WCHAR* fileName)
 {
     WCHAR* quote = nullptr;
 
     // If there are any quotes in the file name convert them to spaces.
-    while ((quote = wcsstr(executableName, W("\""))) != nullptr)
+    while ((quote = wcsstr(fileName, W("\""))) != nullptr)
     {
         *quote = W(' ');
     }
 
     // Remove any illegal or annoying characters from the file name by converting them to underscores.
-    while ((quote = wcspbrk(executableName, W("=<>:\"/\\|?! *.,"))) != nullptr)
+    while ((quote = wcspbrk(fileName, W("=<>:\"/\\|?! *.,"))) != nullptr)
     {
         *quote = W('_');
     }
 }
 
-void generateRandomSuffix(WCHAR* buffer, size_t length)
-{
-    unsigned randNumber = 0;
-#ifdef FEATURE_PAL
-    PAL_Random(&randNumber, sizeof(randNumber));
-#else  // !FEATURE_PAL
-    rand_s(&randNumber);
-#endif // !FEATURE_PAL
-
-    swprintf_s(buffer, length, W("%08X"), randNumber);
-}
-
 // All lengths in this function exclude the terminal NULL.
-WCHAR* getResultFileName(const WCHAR* folderPath, WCHAR* executableName, const WCHAR* extension)
+WCHAR* GetResultFileName(const WCHAR* folderPath, const WCHAR* fileName, const WCHAR* extension)
 {
-    cleanupExecutableName(executableName);
-
-    size_t executableNameLength = wcslen(executableName);
-    size_t extensionLength      = wcslen(extension);
-    size_t folderPathLength     = wcslen(folderPath);
+    const size_t folderPathLength   = wcslen(folderPath);
+    const size_t fileNameLength     = wcslen(fileName);
+    const size_t extensionLength    = wcslen(extension);
+    const size_t maxPathLength      = MAX_PATH - 50; // subtract 50 because excel doesn't like paths longer then 230.
+    const size_t randomStringLength = 8;
 
-    size_t dataFileNameLength = folderPathLength + 1 + executableNameLength + 1 + extensionLength;
+    size_t fullPathLength   = folderPathLength + 1 + extensionLength;
+    bool appendRandomString = false;
 
-    const size_t maxAcceptablePathLength =
-        MAX_PATH - 50; // subtract 50 because excel doesn't like paths longer then 230.
+    if (fileNameLength > 0)
+    {
+        fullPathLength += fileNameLength;
+    }
+    else
+    {
+        fullPathLength += randomStringLength;
+        appendRandomString = true;
+    }
 
-#ifdef FEATURE_PAL
-        assert(executableNameLength == 0);
-#endif // FEATURE_PAL
+    size_t charsToDelete = 0;
 
-    if (dataFileNameLength > maxAcceptablePathLength || executableNameLength == 0)
+    if (fullPathLength > maxPathLength)
     {
         // The path name is too long; creating the file will fail. This can happen because we use the command line,
         // which for ngen includes lots of environment variables, for example.
-        // Shorten the executable file name and add a random string to it to avoid collisions.
+        // Shorten the file name and add a random string to the end to avoid collisions.
 
-        const size_t randStringLength = 8;
+        charsToDelete = fullPathLength - maxPathLength + randomStringLength;
 
-#ifndef FEATURE_CORECLR
-
-        size_t lengthToBeDeleted = (dataFileNameLength - maxAcceptablePathLength) + randStringLength;
-
-        if (executableNameLength <= lengthToBeDeleted)
+        if (fileNameLength >= charsToDelete)
+        {
+            appendRandomString = true;
+            fullPathLength = maxPathLength;
+        }
+        else
         {
-            LogError("getResultFileName - path to the output file is too long '%ws\\%ws.%ws(%d)'", folderPath,
-                     executableName, extension, dataFileNameLength);
+            LogError("GetResultFileName - path to the output file is too long '%ws\\%ws.%ws(%d)'", folderPath, fileName, extension, fullPathLength);
             return nullptr;
         }
+    }
 
-        executableNameLength -= lengthToBeDeleted;
-        executableName[executableNameLength] = 0;
+    WCHAR* fullPath = new WCHAR[fullPathLength + 1];
+    fullPath[0] = W('\0');
+    wcsncat_s(fullPath, fullPathLength + 1, folderPath, folderPathLength);
+    wcsncat_s(fullPath, fullPathLength + 1, DIRECTORY_SEPARATOR_STR_W, 1);
 
-#endif // FEATURE_CORECLR
+    if (fileNameLength > charsToDelete)
+    {
+        wcsncat_s(fullPath, fullPathLength + 1, fileName, fileNameLength - charsToDelete);
+        ReplaceIllegalCharacters(fullPath + folderPathLength + 1);
+    }
 
-        executableNameLength += randStringLength;
-        WCHAR randNumberString[randStringLength + 1];
-        generateRandomSuffix(randNumberString, randStringLength + 1);
-        wcsncat_s(executableName, executableNameLength + 1, randNumberString, randStringLength);
+    if (appendRandomString)
+    {
+       unsigned randomNumber = 0;
 
-        executableNameLength = wcslen(executableName);
+#ifdef FEATURE_PAL
+       PAL_Random(&randomNumber, sizeof(randomNumber));
+#else  // !FEATURE_PAL
+       rand_s(&randomNumber);
+#endif // !FEATURE_PAL
 
-        dataFileNameLength = folderPathLength + 1 + executableNameLength + 1 + extensionLength;
-        if (dataFileNameLength > maxAcceptablePathLength)
-        {
-            LogError("getResultFileName - were not able to short the result file name.", folderPath, executableName,
-                     extension, dataFileNameLength);
-            return nullptr;
-        }
+       WCHAR randomString[randomStringLength + 1];
+       swprintf_s(randomString, randomStringLength + 1, W("%08X"), randomNumber);
+       wcsncat_s(fullPath, fullPathLength + 1, randomString, randomStringLength);
     }
 
-    WCHAR* dataFileName = new WCHAR[dataFileNameLength + 1];
-    dataFileName[0]     = 0;
-    wcsncat_s(dataFileName, dataFileNameLength + 1, folderPath, folderPathLength);
-    wcsncat_s(dataFileName, dataFileNameLength + 1, DIRECTORY_SEPARATOR_STR_W, 1);
-    wcsncat_s(dataFileName, dataFileNameLength + 1, executableName, executableNameLength);
-    wcsncat_s(dataFileName, dataFileNameLength + 1, extension, extensionLength);
-    return dataFileName;
+    wcsncat_s(fullPath, fullPathLength + 1, extension, extensionLength);
+
+    return fullPath;
 }
diff --git a/src/ToolBox/superpmi/superpmi-shared/spmiutil.h b/src/ToolBox/superpmi/superpmi-shared/spmiutil.h
index 55573ae74d..5ae422ffc4 100644
--- a/src/ToolBox/superpmi/superpmi-shared/spmiutil.h
+++ b/src/ToolBox/superpmi/superpmi-shared/spmiutil.h
@@ -25,6 +25,6 @@ LPSTR GetCommandLineA();
 
 bool LoadRealJitLib(HMODULE& realJit, WCHAR* realJitPath);
 
-WCHAR* getResultFileName(const WCHAR* folderPath, WCHAR* executableName, const WCHAR* extension);
+WCHAR* GetResultFileName(const WCHAR* folderPath, const WCHAR* fileName, const WCHAR* extension);
 
 #endif // !_SPMIUtil
diff --git a/src/ToolBox/superpmi/superpmi-shim-collector/superpmi-shim-collector.cpp b/src/ToolBox/superpmi/superpmi-shim-collector/superpmi-shim-collector.cpp
index 4068dac931..f937ce6edf 100644
--- a/src/ToolBox/superpmi/superpmi-shim-collector/superpmi-shim-collector.cpp
+++ b/src/ToolBox/superpmi/superpmi-shim-collector/superpmi-shim-collector.cpp
@@ -65,15 +65,11 @@ void SetLogPath()
 void SetLogPathName()
 {
     // NOTE: under PAL, we don't get the command line, so we depend on the random number generator to give us a unique
-    // filename.
-    const WCHAR* originalExecutableName = GetCommandLineW();
-    size_t       executableNameLength   = wcslen(originalExecutableName);
-    WCHAR*       executableName         = new WCHAR[executableNameLength + 1];
-    wcscpy_s(executableName, executableNameLength + 1, originalExecutableName);
-    executableName[executableNameLength] = W('\0');
-
-    const WCHAR* DataFileExtension = W(".mc");
-    g_dataFileName                 = getResultFileName(g_logPath, executableName, DataFileExtension);
+    // filename
+    const WCHAR* fileName  = GetCommandLineW();
+    const WCHAR* extension = W(".mc");
+
+    g_dataFileName = GetResultFileName(g_logPath, fileName, extension);
 }
 
 // TODO: this only works for ANSI file paths...
diff --git a/src/ToolBox/superpmi/superpmi-shim-counter/methodcallsummarizer.cpp b/src/ToolBox/superpmi/superpmi-shim-counter/methodcallsummarizer.cpp
index ca1e191c08..3f41973d16 100644
--- a/src/ToolBox/superpmi/superpmi-shim-counter/methodcallsummarizer.cpp
+++ b/src/ToolBox/superpmi/superpmi-shim-counter/methodcallsummarizer.cpp
@@ -14,10 +14,10 @@ MethodCallSummarizer::MethodCallSummarizer(WCHAR* logPath)
     names    = nullptr;
     counts   = nullptr;
 
-    WCHAR*       executableName    = GetCommandLineW();
-    const WCHAR* dataFileExtension = W(".csv");
+    const WCHAR* fileName  = GetCommandLineW();
+    const WCHAR* extension = W(".csv");
 
-    dataFileName = getResultFileName(logPath, executableName, dataFileExtension);
+    dataFileName = GetResultFileName(logPath, fileName, extension);
 }
 
 // lots of ways will be faster.. this happens to be decently simple and good enough for the task at hand and nicely