Add GetProcessEnvironment command to diagnostics server (#40556)

src/coreclr/src/vm/processdiagnosticsprotocolhelper.cpp

@@ -11,7 +11,7 @@
 
 #ifdef FEATURE_PERFTRACING
 
-static inline uint32_t GetStringLength(const WCHAR *&value)
+static inline uint32_t GetStringLength(const WCHAR *value)
 {
     return static_cast<uint32_t>(wcslen(value) + 1);
 }
@@ -35,6 +35,28 @@ static bool TryWriteString(uint8_t * &bufferCursor, uint16_t &bufferLen, const W
     return true;
 }
 
+static bool TryWriteString(IpcStream *pStream, const WCHAR *value)
+{
+    uint32_t stringLen = GetStringLength(value);
+    uint32_t stringLenInBytes = stringLen * sizeof(WCHAR);
+    uint32_t totalStringSizeInBytes = stringLenInBytes + sizeof(uint32_t);
+
+    uint32_t totalBytesWritten = 0;
+    uint32_t nBytesWritten = 0;
+
+    bool fSuccess = pStream->Write(&stringLen, sizeof(stringLen), nBytesWritten);
+    totalBytesWritten += nBytesWritten;
+
+    if (fSuccess)
+    {
+        fSuccess &= pStream->Write(value, stringLenInBytes, nBytesWritten);
+        totalBytesWritten += nBytesWritten;
+    }
+
+    ASSERT(totalBytesWritten == totalStringSizeInBytes);
+    return fSuccess && totalBytesWritten == totalStringSizeInBytes;
+}
+
 uint16_t ProcessInfoPayload::GetSize()
 {
     LIMITED_METHOD_CONTRACT;
@@ -114,6 +136,134 @@ bool ProcessInfoPayload::Flatten(BYTE * &lpBuffer, uint16_t &cbSize)
     return fSuccess;
 }
 
+void EnvironmentHelper::PopulateEnvironment()
+{
+    CONTRACTL
+    {
+        NOTHROW;
+        GC_TRIGGERS;
+        MODE_PREEMPTIVE;
+    }
+    CONTRACTL_END;
+
+    if (Environment != nullptr)
+        return;
+
+    // env block is an array of strings of the form "key=value" delimited by null and terminated by null
+    // e.g., "key=value\0key=value\0\0";
+    LPWSTR envBlock = GetEnvironmentStringsW();
+    LPWSTR envCursor = envBlock;
+    // first calculate the buffer size
+    uint32_t nWchars = 0;
+    uint32_t nEntries = 0;
+    while(*envCursor != 0) 
+    {
+        uint32_t len = static_cast<uint32_t>(wcslen(envCursor) + 1);
+        nEntries++;
+        nWchars += len;
+        envCursor += len;
+    }
+
+    // serialized size value can't be larger than uint32_t
+    S_UINT32 size(sizeof(uint32_t) + (nEntries * sizeof(uint32_t)) + (nWchars * sizeof(WCHAR)));
+    // if the envblock size value is larger than a uint32_t then we should bail
+    if (!size.IsOverflow())
+    {
+        // copy out the Environment block
+        LPWSTR tmpEnv = new (nothrow) WCHAR[nWchars + 1];
+        if (tmpEnv != nullptr)
+        {
+            memcpy(tmpEnv, envBlock, (nWchars + 1) * sizeof(WCHAR));
+            Environment = tmpEnv;
+        }
+        _nEnvEntries = nEntries;
+        _nWchars = nWchars;
+    }
+
+    FreeEnvironmentStringsW(envBlock);
+
+    return;
+}
+
+uint32_t EnvironmentHelper::GetEnvironmentBlockSize()
+{
+    LIMITED_METHOD_CONTRACT;
+
+    S_UINT32 size(sizeof(uint32_t) + (_nEnvEntries * sizeof(uint32_t)) + (_nWchars * sizeof(WCHAR)));
+    return size.IsOverflow() ? sizeof(uint32_t) : size.Value();
+}
+
+bool EnvironmentHelper::WriteToStream(IpcStream *pStream)
+{
+    CONTRACTL
+    {
+        NOTHROW;
+        GC_TRIGGERS;
+        MODE_PREEMPTIVE;
+        PRECONDITION(pStream != nullptr);
+        PRECONDITION(Environment != nullptr);
+    }
+    CONTRACTL_END;
+
+    // Array<Array<WCHAR>>
+    uint32_t nBytesWritten = 0;
+    bool fSuccess = pStream->Write(&_nEnvEntries, sizeof(_nEnvEntries), nBytesWritten);
+
+    LPCWSTR cursor = Environment;
+    for (uint32_t i = 0; i < _nEnvEntries; i++)
+    {
+        if (!fSuccess || cursor == nullptr)
+            break;
+
+        uint32_t len = static_cast<uint32_t>(wcslen(cursor) + 1);
+        fSuccess &= TryWriteString(pStream, cursor);
+        cursor += len;
+    }
+
+    return fSuccess;
+}
+
+void ProcessDiagnosticsProtocolHelper::GetProcessEnvironment(DiagnosticsIpc::IpcMessage& message, IpcStream *pStream)
+{
+    CONTRACTL
+    {
+        THROWS;
+        GC_TRIGGERS;
+        MODE_PREEMPTIVE;
+        PRECONDITION(pStream != nullptr);
+    }
+    CONTRACTL_END;
+
+    // GetProcessEnvironment responds with a Diagnostics IPC message
+    // who's payload is a uint32_t that specifies the number of bytes
+    // that will follow.  The optional continuation will contain
+    // the Environemnt streamed in the standard Diagnostics IPC format (length-prefixed arrays).
+
+    struct EnvironmentHelper helper = {};
+    helper.PopulateEnvironment();
+
+    struct EnvironmentHelper::InitialPayload payload = { helper.GetEnvironmentBlockSize(), 0 };
+
+    DiagnosticsIpc::IpcMessage ProcessEnvironmentResponse;
+    const bool fSuccess = ProcessEnvironmentResponse.Initialize(DiagnosticsIpc::GenericSuccessHeader, payload) ?
+        ProcessEnvironmentResponse.Send(pStream) :
+        DiagnosticsIpc::IpcMessage::SendErrorMessage(pStream, E_FAIL);
+
+    if (!fSuccess)
+    {
+        STRESS_LOG0(LF_DIAGNOSTICS_PORT, LL_ERROR, "Failed to send DiagnosticsIPC response");
+    }
+    else
+    {
+        // send the environment
+        const bool fSuccessWriteEnv = helper.WriteToStream(pStream);
+        if (!fSuccessWriteEnv)
+            STRESS_LOG0(LF_DIAGNOSTICS_PORT, LL_ERROR, "Failed to stream environment");
+    }
+
+    delete pStream;
+}
+
 void ProcessDiagnosticsProtocolHelper::GetProcessInfo(DiagnosticsIpc::IpcMessage& message, IpcStream *pStream)
 {
     CONTRACTL
@@ -202,6 +352,9 @@ void ProcessDiagnosticsProtocolHelper::HandleIpcMessage(DiagnosticsIpc::IpcMessa
     case ProcessCommandId::ResumeRuntime:
         ProcessDiagnosticsProtocolHelper::ResumeRuntimeStartup(message, pStream);
         break;
+    case ProcessCommandId::GetProcessEnvironment:
+        ProcessDiagnosticsProtocolHelper::GetProcessEnvironment(message, pStream);
+        break;
 
     default:
         STRESS_LOG1(LF_DIAGNOSTICS_PORT, LL_WARNING, "Received unknown request type (%d)\n", message.GetHeader().CommandSet);

src/coreclr/src/vm/processdiagnosticsprotocolhelper.h

@@ -18,8 +18,9 @@ class IpcStream;
 // see diagnosticsipc.h and diagnosticserver.h for more details
 enum class ProcessCommandId : uint8_t
 {
-    GetProcessInfo = 0x00,
-    ResumeRuntime  = 0x01,
+    GetProcessInfo        = 0x00,
+    ResumeRuntime         = 0x01,
+    GetProcessEnvironment = 0x02
     // future
 };
 
@@ -33,7 +34,6 @@ struct ProcessInfoPayload
     // GUID = 16 little endian bytes
     // wchar = 2 little endian bytes, UTF16 encoding
     // array<T> = uint length, length # of Ts
-    // string = (array<char> where the last char must = 0) or (length = 0)
 
     // ProcessInfo = long pid, string cmdline, string OS, string arch, GUID runtimeCookie
     uint64_t ProcessId;
@@ -45,12 +45,44 @@ struct ProcessInfoPayload
     bool Flatten(BYTE * &lpBuffer, uint16_t& cbSize);
 };
 
+struct EnvironmentHelper
+{
+    // The environemnt is sent back as an optional continuation stream of data.
+    // It is encoded in the typical length-prefixed array format as defined in
+    // the Diagnostics IPC Spec: https://github.com/dotnet/diagnostics/blob/master/documentation/design-docs/ipc-protocol.md
+
+    struct InitialPayload
+    {
+        uint32_t continuationSizeInBytes;
+        uint16_t future;
+    };
+
+    // sent as: Array<Array<WCHAR>>
+    NewArrayHolder<const WCHAR> Environment = nullptr;
+
+    void PopulateEnvironment();
+    uint32_t GetNumberOfElements() { PopulateEnvironment(); return _nEnvEntries; }
+
+    // Write the environment block to the stream
+    bool WriteToStream(IpcStream *pStream);
+
+    // The size in bytes of the Diagnostic IPC Protocol encoded Environment Block
+    // It is encoded as Array<Array<WCHAR>> so this will return at least sizeof(uint32_t)
+    // if the env block is empty or failed to be snapshotted since the stream will
+    // just contain 0 for the array length.
+    uint32_t GetEnvironmentBlockSize();
+private:
+    uint32_t _nEnvEntries = 0;
+    uint32_t _nWchars = 0;
+};
+
 class ProcessDiagnosticsProtocolHelper
 {
 public:
     // IPC event handlers.
     static void HandleIpcMessage(DiagnosticsIpc::IpcMessage& message, IpcStream *pStream);
     static void GetProcessInfo(DiagnosticsIpc::IpcMessage& message, IpcStream *pStream);
+    static void GetProcessEnvironment(DiagnosticsIpc::IpcMessage& message, IpcStream *pStream);
     static void ResumeRuntimeStartup(DiagnosticsIpc::IpcMessage& message, IpcStream *pStream);
 };
 

src/tests/tracing/eventpipe/processenvironment/processenvironment.cs

@@ -0,0 +1,96 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+using System;
+using System.Collections.Generic;
+using System.Diagnostics.Tracing;
+using System.Diagnostics;
+using System.IO;
+using System.Linq;
+using System.Runtime.InteropServices;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.Diagnostics.Tools.RuntimeClient;
+using Microsoft.Diagnostics.Tracing;
+using Tracing.Tests.Common;
+
+namespace Tracing.Tests.ProcessEnvironmentValidation
+{
+    public class ProcessEnvironmentValidation
+    {
+        public static int Main(string[] args)
+        {
+
+            Process currentProcess = Process.GetCurrentProcess();
+            int pid = currentProcess.Id;
+            Logger.logger.Log($"Test PID: {pid}");
+            var testEnvPairs = new Dictionary<string, string>
+            {
+                { "TESTKEY1", "TESTVAL1" },
+                { "TESTKEY2", "TESTVAL2" },
+                { "TESTKEY3", "__TEST__VAL=;--3" }
+            };
+
+            foreach (var (key, val) in testEnvPairs)
+                System.Environment.SetEnvironmentVariable(key, val);
+
+            Stream stream = ConnectionHelper.GetStandardTransport(pid);
+
+            // 0x04 = ProcessCommandSet, 0x02 = ProcessInfo
+            var processInfoMessage = new IpcMessage(0x04, 0x02);
+            Logger.logger.Log($"Wrote: {processInfoMessage}");
+            Stream continuationStream = IpcClient.SendMessage(stream, processInfoMessage, out IpcMessage response);
+            Logger.logger.Log($"Received: {response}");
+
+            Utils.Assert(response.Header.CommandSet == 0xFF, $"Response must have Server command set. Expected: 0xFF, Received: 0x{response.Header.CommandSet:X2}"); // server
+            Utils.Assert(response.Header.CommandId == 0x00, $"Response must have OK command id. Expected: 0x00, Received: 0x{response.Header.CommandId:X2}"); // OK
+
+            UInt32 continuationSizeInBytes = BitConverter.ToUInt32(response.Payload[0..4]);
+            Logger.logger.Log($"continuation size: {continuationSizeInBytes} bytes");
+            UInt16 future = BitConverter.ToUInt16(response.Payload[4..]);
+            Logger.logger.Log($"future value: {future}");
+
+            using var memoryStream = new MemoryStream();
+            Logger.logger.Log($"Starting to copy continuation");
+            continuationStream.CopyTo(memoryStream);
+            Logger.logger.Log($"Finished copying continuation");
+            byte[] envBlock = memoryStream.ToArray();
+            Logger.logger.Log($"Total bytes in continuation: {envBlock.Length}");
+
+            Utils.Assert(envBlock.Length == continuationSizeInBytes, $"Continuation size must equal the reported size in the payload response.  Expected: {continuationSizeInBytes} bytes, Received: {envBlock.Length} bytes");
+
+            // VALIDATE ENV
+            // env block is sent as Array<LPCWSTR> (length-prefixed array of length-prefixed wchar strings)
+            int start = 0;
+            int end = start + 4 /* sizeof(uint32_t) */;
+            UInt32 envCount = BitConverter.ToUInt32(envBlock[start..end]);
+            Logger.logger.Log($"envCount: {envCount}");
+
+            var env = new Dictionary<string,string>();
+            for (int i = 0; i < envCount; i++)
+            {
+                start = end;
+                end = start + 4 /* sizeof(uint32_t) */;
+                UInt32 pairLength = BitConverter.ToUInt32(envBlock[start..end]);
+
+                start = end;
+                end = start + ((int)pairLength * sizeof(char));
+                Utils.Assert(end <= envBlock.Length, $"String end can't exceed payload size. Expected: <{envBlock.Length}, Received: {end} (decoded length: {pairLength})");
+                string envPair = System.Text.Encoding.Unicode.GetString(envBlock[start..end]).TrimEnd('\0');
+                int equalsIndex = envPair.IndexOf('=');
+                env[envPair[0..equalsIndex]] = envPair[(equalsIndex+1)..];
+            }
+            Logger.logger.Log($"finished parsing env");
+
+
+            foreach (var (key, val) in testEnvPairs)
+                Utils.Assert(env.ContainsKey(key) && env[key].Equals(val), $"Did not find test environment pair in the environment block: '{key}' = '{val}'");
+            Logger.logger.Log($"Saw test values in env");
+
+            Utils.Assert(end == envBlock.Length, $"Full payload should have been read. Expected: {envBlock.Length}, Received: {end}");
+
+            return 100;
+        }
+    }
+}

src/tests/tracing/eventpipe/processenvironment/processenvironment.csproj

@@ -0,0 +1,16 @@
+<Project Sdk="Microsoft.NET.Sdk">
+  <PropertyGroup>
+    <TargetFrameworkIdentifier>.NETCoreApp</TargetFrameworkIdentifier>
+    <OutputType>exe</OutputType>
+    <CLRTestKind>BuildAndRun</CLRTestKind>
+    <AllowUnsafeBlocks>true</AllowUnsafeBlocks>
+    <CLRTestPriority>0</CLRTestPriority>
+    <UnloadabilityIncompatible>true</UnloadabilityIncompatible>
+    <JitOptimizationSensitive>true</JitOptimizationSensitive>
+    <IlasmRoundTripIncompatible>true</IlasmRoundTripIncompatible>
+  </PropertyGroup>
+  <ItemGroup>
+    <Compile Include="$(MSBuildProjectName).cs" />
+    <ProjectReference Include="../common/common.csproj" />
+  </ItemGroup>
+</Project>

src/tests/tracing/eventpipe/processinfo/processinfo.cs

@@ -82,7 +82,7 @@ public static int Main(string[] args)
             var processInfoMessage = new IpcMessage(0x04, 0x00);
             Logger.logger.Log($"Wrote: {processInfoMessage}");
             IpcMessage response = IpcClient.SendMessage(stream, processInfoMessage);
-            Logger.logger.Log($"Received: {response}");
+            Logger.logger.Log($"Received: <omitted>");
 
             Utils.Assert(response.Header.CommandSet == 0xFF, $"Response must have Server command set. Expected: 0xFF, Received: 0x{response.Header.CommandSet:X2}"); // server
             Utils.Assert(response.Header.CommandId == 0x00, $"Response must have OK command id. Expected: 0x00, Received: 0x{response.Header.CommandId:X2}"); // OK
@@ -95,7 +95,7 @@ public static int Main(string[] args)
             // LPCWSTR Arch;
 
             int totalSize = response.Payload.Length;
-            Logger.logger.Log($"Total size of Payload == {totalSize} b");
+            Logger.logger.Log($"Total size of Payload = {totalSize} bytes");
 
             // VALIDATE PID
             int start = 0;