Support newlines in key names

stormshield · Feb 11, 2020 · 038c610 · 038c610
1 parent fb165dc
commit 038c610
Show file tree

Hide file tree

Showing 6 changed files with 32 additions and 21 deletions.
diff --git a/Binx64/HiveSwarming.exe b/Binx64/HiveSwarming.exe
diff --git a/Binx64/HiveSwarming.pdb b/Binx64/HiveSwarming.pdb
diff --git a/CommonFunctions.cpp b/CommonFunctions.cpp
@@ -57,4 +57,18 @@ void ReportError
     }
 
     std::wcerr << std::endl << std::endl;
-}
+}
+
+VOID GlobalStringSubstitute(
+    _Inout_ std::wstring& String,
+    _In_ const std::wstring& Pattern,
+    _In_ const std::wstring& Replacement
+)
+{
+    size_t Position = String.find(Pattern);
+    while (Position != String.npos)
+    {
+        String.replace(Position, Pattern.length(), Replacement);
+        Position = String.find(Pattern, Position + Replacement.length());
+    }
+}
diff --git a/CommonFunctions.h b/CommonFunctions.h
@@ -19,4 +19,14 @@ void ReportError
 (
     _In_ const HRESULT ErrorCode,
     _In_ const std::wstring& Context = std::wstring{}
-);
+);
+
+/// @brief Perform global replacement of substring in a std::wstring
+/// @param[in,out] String String on which to perform substitutions
+/// @param[in] Pattern Substring that will be replaced by #Replacement in #String
+/// @param[in] Replacement What is inserted in place of #Substring in #String
+VOID GlobalStringSubstitute(
+    _Inout_ std::wstring& String,
+    _In_ const std::wstring& Pattern,
+    _In_ const std::wstring& Replacement
+);
diff --git a/InternalToRegfile.cpp b/InternalToRegfile.cpp
@@ -8,24 +8,6 @@
 #include <sstream>
 #include <iomanip>
 
-/// @brief Perform global replacement of substring in a std::wstring
-/// @param[in,out] String String on which to perform substitutions
-/// @param[in] Pattern Substring that will be replaced by #Replacement in #String
-/// @param[in] Replacement What is inserted in place of #Substring in #String
-static VOID GlobalStringSubstitute(
-    _Inout_ std::wstring& String,
-    _In_ const std::wstring& Pattern,
-    _In_ const std::wstring& Replacement
-)
-{
-    size_t Position = String.find(Pattern);
-    while (Position != String.npos)
-    {
-        String.replace(Position, Pattern.length(), Replacement);
-        Position = String.find(Pattern, Position + Replacement.length());
-    }
-}
-
 /// @brief Append binary contents of a std::wstring to an opened file
 /// @param[in] OutFileHandle Handle to the file
 /// @param[in] String Input string
@@ -316,9 +298,12 @@ HRESULT RenderRegistryKey
         NewPath = PathSoFar + Constants::RegFiles::PathSeparator + RegKey.Name;
     }
 
+    std::wstring EscapedPath { NewPath };
+    GlobalStringSubstitute(EscapedPath, L"\n", L"\r\n");
+
     {
         std::wostringstream KeySpecStream;
-        KeySpecStream << Constants::RegFiles::KeyOpening << NewPath << Constants::RegFiles::KeyClosing << Constants::RegFiles::NewLines;
+        KeySpecStream << Constants::RegFiles::KeyOpening << EscapedPath << Constants::RegFiles::KeyClosing << Constants::RegFiles::NewLines;
 
         Result = WriteStringBufferToFile(OutFileHandle, KeySpecStream.str());
         if (FAILED(Result))

diff --git a/RegfileToInternal.cpp b/RegfileToInternal.cpp
@@ -386,8 +386,10 @@ static HRESULT RegListToInternal
             break;
         }
         const std::wstring_view KeyName{ &KeyPath[PathPrefix.length()], KeyPath.length() - PathPrefix.length() };
+
         RegistryKey NewKey;
         NewKey.Name = KeyName;
+        GlobalStringSubstitute(NewKey.Name, L"\r\n", L"\n");
 
         RegList.remove_prefix(EndKeyPos + 1);
         RegList.remove_prefix(Constants::RegFiles::NewLines.length());