Fix mixed elevation integration issue by using MemoryStream (#4413) (#…

…4414)
microsoft · Apr 26, 2024 · f18924d · f18924d
1 parent 83ec285
commit f18924d
Show file tree

Hide file tree

Showing 4 changed files with 23 additions and 24 deletions.
diff --git a/src/AppInstallerCLICore/ConfigurationDynamicRuntimeFactory.cpp b/src/AppInstallerCLICore/ConfigurationDynamicRuntimeFactory.cpp
@@ -153,18 +153,16 @@ namespace AppInstaller::CLI::ConfigurationRemoting
                 highIntegritySet.Serialize(memoryStream);
 
                 Streams::DataReader reader(memoryStream.GetInputStreamAt(0));
-                reader.UnicodeEncoding(Streams::UnicodeEncoding::Utf8);
-                reader.LoadAsync((uint32_t)memoryStream.Size());
-
-                winrt::hstring result;
-                uint32_t bytesToRead = reader.UnconsumedBufferLength();
-
-                if (bytesToRead > 0)
-                {
-                    result = reader.ReadString(bytesToRead);
-                }
-
-                return winrt::to_string(result);
+                THROW_HR_IF(E_UNEXPECTED, memoryStream.Size() > std::numeric_limits<uint32_t>::max());
+                uint32_t streamSize = (uint32_t)memoryStream.Size();
+                std::vector<uint8_t> bytes;
+                bytes.resize(streamSize);
+                reader.LoadAsync(streamSize);
+                reader.ReadBytes(bytes);
+                reader.DetachStream();
+                memoryStream.Close();
+
+                return { bytes.begin(), bytes.end() };
             }
 
             ProcessorMap::iterator CreateSetProcessorForIntegrityLevel(Security::IntegrityLevel integrityLevel)

diff --git a/src/ConfigurationRemotingServer/Program.cs b/src/ConfigurationRemotingServer/Program.cs
@@ -8,7 +8,6 @@
 using System.Text.Json.Serialization;
 using Microsoft.Management.Configuration;
 using Microsoft.Management.Configuration.Processor;
-using Windows.Storage.Streams;
 using WinRT;
 
 namespace ConfigurationRemotingServer
@@ -118,14 +117,14 @@ static int Main(string[] args)
 
                     // Parse limitation set.
                     byte[] limitationSetBytes = Encoding.UTF8.GetBytes(commandStr.Substring(secondSeparatorIndex + CommandLineSectionSeparator.Length));
-                    InMemoryRandomAccessStream limitationSetStream = new InMemoryRandomAccessStream();
-                    DataWriter streamWriter = new DataWriter(limitationSetStream);
-                    streamWriter.WriteBytes(limitationSetBytes);
-                    streamWriter.StoreAsync().GetAwaiter().GetResult();
-                    streamWriter.DetachStream();
-                    limitationSetStream.Seek(0);
+                    MemoryStream memoryStream = new MemoryStream();
+                    memoryStream.Write(limitationSetBytes);
+                    memoryStream.Flush();
+                    memoryStream.Seek(0, SeekOrigin.Begin);
                     ConfigurationProcessor processor = new ConfigurationProcessor(factory);
-                    var limitationSetResult = processor.OpenConfigurationSet(limitationSetStream);
+                    var limitationSetResult = processor.OpenConfigurationSet(memoryStream.AsInputStream());
+                    memoryStream.Close();
+
                     if (limitationSetResult.ResultCode != null)
                     {
                         throw limitationSetResult.ResultCode;

diff --git a/src/Microsoft.Management.Configuration/ConfigurationProcessor.cpp b/src/Microsoft.Management.Configuration/ConfigurationProcessor.cpp
@@ -232,8 +232,9 @@ namespace winrt::Microsoft::Management::Configuration::implementation
             // This is done here to enable easy cancellation propagation to the stream reads.
             uint32_t bufferSize = 1 << 20;
             Windows::Storage::Streams::Buffer buffer(bufferSize);
-            Windows::Storage::Streams::InputStreamOptions readOptions = 
-                Windows::Storage::Streams::InputStreamOptions::Partial | Windows::Storage::Streams::InputStreamOptions::ReadAhead;
+
+            // Memory stream in mixed elevation does not support InputStreamOptions as flags.
+            Windows::Storage::Streams::InputStreamOptions readOptions = Windows::Storage::Streams::InputStreamOptions::Partial;
             std::string inputString;
 
             for (;;)

diff --git a/src/Microsoft.Management.Configuration/ConfigurationSet.cpp b/src/Microsoft.Management.Configuration/ConfigurationSet.cpp
@@ -128,10 +128,11 @@ namespace winrt::Microsoft::Management::Configuration::implementation
     {
         std::unique_ptr<ConfigurationSetSerializer> serializer = ConfigurationSetSerializer::CreateSerializer(m_schemaVersion);
         hstring result = serializer->Serialize(this);
+        auto resultUtf8 = winrt::to_string(result);
+        std::vector<uint8_t> bytes(resultUtf8.begin(), resultUtf8.end());
 
         Windows::Storage::Streams::DataWriter dataWriter{ stream };
-        dataWriter.UnicodeEncoding(Windows::Storage::Streams::UnicodeEncoding::Utf8);
-        dataWriter.WriteString(result);
+        dataWriter.WriteBytes(bytes);
         dataWriter.StoreAsync().get();
         dataWriter.DetachStream();
     }