Add diagnostics IPC set environment variable command (#52567)

src/coreclr/vm/eventing/eventpipe/ds-rt-coreclr.h

@@ -6,6 +6,7 @@
 
 #ifdef ENABLE_PERFTRACING
 #include "ep-rt-coreclr.h"
+#include "ds-process-protocol.h"
 #include "ds-profiler-protocol.h"
 #include "ds-dump-protocol.h"
 
@@ -320,6 +321,13 @@ ds_rt_profiler_startup (DiagnosticsStartupProfilerCommandPayload *payload)
 }
 #endif // PROFILING_SUPPORTED
 
+static
+uint32_t
+ds_rt_set_environment_variable (const ep_char16_t *name, const ep_char16_t *value)
+{
+	return SetEnvironmentVariableW(reinterpret_cast<LPCWSTR>(name), reinterpret_cast<LPCWSTR>(value)) ? S_OK : HRESULT_FROM_WIN32(GetLastError());
+}
+
 /*
 * DiagnosticServer.
 */

src/mono/mono/eventpipe/ds-rt-mono.h

@@ -244,6 +244,25 @@ ds_rt_profiler_startup (DiagnosticsStartupProfilerCommandPayload *payload)
 	return DS_IPC_E_NOTSUPPORTED;
 }
 
+/*
+* Environment variables
+*/
+
+static
+uint32_t
+ds_rt_set_environment_variable (const ep_char16_t *name, const ep_char16_t *value)
+{
+	gchar *nameNarrow = ep_rt_utf16_to_utf8_string (name, ep_rt_utf16_string_len (name));
+	gchar *valueNarrow = ep_rt_utf16_to_utf8_string (value, ep_rt_utf16_string_len (value));
+
+	gboolean success = g_setenv(nameNarrow, valueNarrow, true);
+
+	g_free (nameNarrow);
+	g_free (valueNarrow);
+
+	return success ? DS_IPC_S_OK : DS_IPC_E_FAIL;
+}
+
 /*
 * DiagnosticServer.
 */

src/native/eventpipe/ds-process-protocol.c

@@ -65,6 +65,12 @@ process_protocol_helper_resume_runtime_startup (
 	DiagnosticsIpcMessage *message,
 	DiagnosticsIpcStream *stream);
 
+static
+bool
+process_protocol_helper_set_environment_variable (
+	DiagnosticsIpcMessage *message,
+	DiagnosticsIpcStream *stream);
+
 static
 bool
 process_protocol_helper_unknown_command (
@@ -443,6 +449,89 @@ process_protocol_helper_resume_runtime_startup (
 	return result;
 }
 
+DiagnosticsSetEnvironmentVariablePayload *
+ds_set_environment_variable_payload_alloc (void)
+{
+	return ep_rt_object_alloc (DiagnosticsSetEnvironmentVariablePayload);
+}
+
+void
+ds_set_environment_variable_payload_free (DiagnosticsSetEnvironmentVariablePayload *payload)
+{
+	ep_return_void_if_nok (payload != NULL);
+	ep_rt_byte_array_free (payload->incoming_buffer);
+	ep_rt_object_free (payload);
+}
+
+static
+uint8_t *
+set_environment_variable_command_try_parse_payload (
+	uint8_t *buffer,
+	uint16_t buffer_len)
+{
+	EP_ASSERT (buffer != NULL);
+
+	uint8_t * buffer_cursor = buffer;
+	uint32_t buffer_cursor_len = buffer_len;
+
+	DiagnosticsSetEnvironmentVariablePayload *instance = ds_set_environment_variable_payload_alloc ();
+	ep_raise_error_if_nok (instance != NULL);
+
+	instance->incoming_buffer = buffer;
+
+	if (!ds_ipc_message_try_parse_string_utf16_t (&buffer_cursor, &buffer_cursor_len, &instance->name) ||
+		!ds_ipc_message_try_parse_string_utf16_t (&buffer_cursor, &buffer_cursor_len, &instance->value))
+		ep_raise_error ();
+
+ep_on_exit:
+	return (uint8_t *)instance;
+
+ep_on_error:
+	ds_set_environment_variable_payload_free (instance);
+	instance = NULL;
+	ep_exit_error_handler ();
+}
+
+static
+bool
+process_protocol_helper_set_environment_variable (
+	DiagnosticsIpcMessage *message,
+	DiagnosticsIpcStream *stream)
+{
+	EP_ASSERT (message != NULL);
+	EP_ASSERT (stream != NULL);
+
+	if (!stream)
+		return false;
+
+    bool result = false;
+    DiagnosticsSetEnvironmentVariablePayload *payload = (DiagnosticsSetEnvironmentVariablePayload *)ds_ipc_message_try_parse_payload (message, set_environment_variable_command_try_parse_payload);
+	if (!payload) {
+		ds_ipc_message_send_error (stream, DS_IPC_E_BAD_ENCODING);
+		ep_raise_error ();
+	}
+
+	ds_ipc_result_t ipc_result;
+	ipc_result = ds_rt_set_environment_variable (payload->name, payload->value);
+	if (ipc_result != DS_IPC_S_OK) {
+		ds_ipc_message_send_error (stream, ipc_result);
+		ep_raise_error ();
+	} else {
+		ds_ipc_message_send_success (stream, ipc_result);
+	}
+
+	result = true;
+
+ep_on_exit:
+	ds_set_environment_variable_payload_free (payload);
+    ds_ipc_stream_free (stream);
+	return result;
+
+ep_on_error:
+	EP_ASSERT (!result);
+	ep_exit_error_handler ();
+}
+
 static
 bool
 process_protocol_helper_unknown_command (
@@ -475,6 +564,9 @@ ds_process_protocol_helper_handle_ipc_message (
 	case DS_PROCESS_COMMANDID_GET_PROCESS_ENV:
 		result = process_protocol_helper_get_process_env (message, stream);
 		break;
+	case DS_PROCESS_COMMANDID_SET_ENV_VAR:
+		result = process_protocol_helper_set_environment_variable (message, stream);
+		break;
 	default:
 		result = process_protocol_helper_unknown_command (message, stream);
 		break;

src/native/eventpipe/ds-process-protocol.h

@@ -87,6 +87,33 @@ ds_env_info_payload_init (DiagnosticsEnvironmentInfoPayload *payload);
 void
 ds_env_info_payload_fini (DiagnosticsEnvironmentInfoPayload *payload);
 
+/*
+* DiagnosticsSetEnvironmentVariablePayload
+*/
+
+#if defined(DS_INLINE_GETTER_SETTER) || defined(DS_IMPL_PROCESS_PROTOCOL_GETTER_SETTER)
+struct _DiagnosticsSetEnvironmentVariablePayload {
+#else
+struct _DiagnosticsSetEnvironmentVariablePayload_Internal {
+#endif
+	uint8_t * incoming_buffer;
+
+	const ep_char16_t *name;
+	const ep_char16_t *value;
+};
+
+#if !defined(DS_INLINE_GETTER_SETTER) && !defined(DS_IMPL_PROCESS_PROTOCOL_GETTER_SETTER)
+struct _DiagnosticsSetEnvironmentVariablePayload {
+	uint8_t _internal [sizeof (struct _DiagnosticsSetEnvironmentVariablePayload_Internal)];
+};
+#endif
+
+DiagnosticsSetEnvironmentVariablePayload *
+ds_set_environment_variable_payload_alloc (void);
+
+void
+ds_set_environment_variable_payload_free (DiagnosticsSetEnvironmentVariablePayload *payload);
+
 /*
  * DiagnosticsProcessProtocolHelper.
  */
@@ -98,3 +125,4 @@ ds_process_protocol_helper_handle_ipc_message (
 
 #endif /* ENABLE_PERFTRACING */
 #endif /* __DIAGNOSTICS_PROCESS_PROTOCOL_H__ */
+

src/native/eventpipe/ds-rt.h

@@ -139,6 +139,21 @@ static
 uint32_t
 ds_rt_profiler_startup (DiagnosticsStartupProfilerCommandPayload *payload);
 
+/*
+* Environment variables
+*/
+
+static
+uint32_t
+ds_rt_set_environment_variable (const ep_char16_t *name, const ep_char16_t *value);
+
+static
+uint32_t
+ds_rt_get_environment_variable (const ep_char16_t *name,
+								uint32_t valueBufferLength,
+								uint32_t *valueLengthOut,
+								ep_char16_t *valueBuffer);
+
 /*
 * DiagnosticServer.
 */

src/native/eventpipe/ds-types.h

@@ -20,6 +20,8 @@ typedef struct _DiagnosticsStartupProfilerCommandPayload DiagnosticsStartupProfi
 typedef struct _DiagnosticsConnectPort DiagnosticsConnectPort;
 typedef struct _DiagnosticsEnvironmentInfoPayload DiagnosticsEnvironmentInfoPayload;
 typedef struct _DiagnosticsGenerateCoreDumpCommandPayload DiagnosticsGenerateCoreDumpCommandPayload;
+typedef struct _DiagnosticsSetEnvironmentVariablePayload DiagnosticsSetEnvironmentVariablePayload;
+typedef struct _DiagnosticsGetEnvironmentVariablePayload DiagnosticsGetEnvironmentVariablePayload;
 typedef struct _DiagnosticsIpcHeader DiagnosticsIpcHeader;
 typedef struct _DiagnosticsIpcMessage DiagnosticsIpcMessage;
 typedef struct _DiagnosticsListenPort DiagnosticsListenPort;
@@ -64,6 +66,7 @@ typedef enum {
 	DS_PROCESS_COMMANDID_GET_PROCESS_INFO = 0x00,
 	DS_PROCESS_COMMANDID_RESUME_RUNTIME = 0x01,
 	DS_PROCESS_COMMANDID_GET_PROCESS_ENV = 0x02,
+	DS_PROCESS_COMMANDID_SET_ENV_VAR = 0x03,
 	// future
 } DiagnosticsProcessCommandId;
 
@@ -127,6 +130,8 @@ typedef int32_t ds_ipc_result_t;
 #define DS_IPC_E_NOT_YET_AVAILABLE ((ds_ipc_result_t)(0x8013135bL))
 #define DS_IPC_E_RUNTIME_UNINITIALIZED ((ds_ipc_result_t)(0x80131371L))
 #define DS_IPC_E_INVALIDARG ((ds_ipc_result_t)(0x80070057L))
+#define DS_IPC_E_INSUFFICIENT_BUFFER ((ds_ipc_result_t)(0x8007007A))
+#define DS_IPC_E_ENVVAR_NOT_FOUND ((ds_ipc_result_t)(0x800000CB))
 
 #endif /* ENABLE_PERFTRACING */
 #endif /* __DIAGNOSTICS_TYPES_H__ */

src/tests/profiler/common/DiagnosticsIPCWorkaround.cs

@@ -6,6 +6,7 @@
 using System.Text;
 using System.Reflection;
 using Microsoft.Diagnostics.NETCore.Client;
+using Tracing.Tests.Common;
 
 // This is to work around having to wait for an update to the DiagnosticsClient nuget before adding
 // a test. I really hope this isn't permanent
@@ -41,36 +42,11 @@ public void SetStartupProfiler(Guid profilerGuid, string profilerPath)
             MethodInfo startupProfiler = typeof(DiagnosticsClient).GetMethod("SetStartupProfiler", BindingFlags.Public);
             if (startupProfiler != null)
             {
-                throw new Exception("You updated DiagnosticsClient to a version that supports SetStartupProfiler, please remove this nonsense and use the real code.");
+                throw new Exception("You updated DiagnosticsClient to a version that supports SetStartupProfiler, please remove this and use the real code.");
             }
-
-            DiagnosticsClient client = new DiagnosticsClient(_processId);
-
-            Console.WriteLine("Sending startup profiler message.");            
-            // Send StartupProfiler command
-            object ipcMessage = MakeStartupProfilerMessage(profilerGuid, profilerPath);
-            SendMessage(_processId, ipcMessage);
-        }
-
-        private static object MakeHeader(byte commandSet, byte commandId)
-        {
-            Type commandEnumType = GetPrivateType("Microsoft.Diagnostics.NETCore.Client.DiagnosticsServerCommandSet");
-            object enumCommandSet = Enum.ToObject(commandEnumType, commandSet);
-            Type ipcHeaderType = GetPrivateType("Microsoft.Diagnostics.NETCore.Client.IpcHeader");
-            object ipcHeader = Activator.CreateInstance(ipcHeaderType, new object[] { enumCommandSet, commandId });
-            return ipcHeader;
-        }
-
-        private static object MakeMessage(object ipcHeader, byte[] payload)
-        {
-                Type ipcMessageType = GetPrivateType("Microsoft.Diagnostics.NETCore.Client.IpcMessage");
-                object ipcMessage = Activator.CreateInstance(ipcMessageType, new object[] { ipcHeader, payload });
-                return ipcMessage;
-        }
-
-        private static object MakeStartupProfilerMessage(Guid profilerGuid, string profilerPath)
-        {
-            // public IpcMessage(IpcHeader header, byte[] payload = null)
+
+            Console.WriteLine("Sending startup profiler message.");
+
             using (var stream = new MemoryStream())
             using (var writer = new BinaryWriter(stream))
             {
@@ -80,51 +56,63 @@ private static object MakeStartupProfilerMessage(Guid profilerGuid, string profi
                 writer.Flush();
                 byte[] payload = stream.ToArray();
 
-                object ipcHeader = MakeHeader(3, 2);
-                return MakeMessage(ipcHeader, payload);
+                var message = new IpcMessage(0x03, 0x02, payload);
+                Console.WriteLine($"Sent: {message.ToString()}");
+                IpcMessage response = IpcClient.SendMessage(ConnectionHelper.GetStandardTransport(_processId), message);
+                Console.WriteLine($"Received: {response.ToString()}");
             }
-        }
 
-        private static object MakeResumeRuntimeMessage()
-        {
-            object ipcHeader = MakeHeader(4, 1);
-            return MakeMessage(ipcHeader, new byte[0]);
+            Console.WriteLine("Finished sending startup profiler message.");
         }
 
-        private static Type GetPrivateType(string typeName)
+        public bool SetEnvironmentVariable(string name, string val)
         {
-            return typeof(DiagnosticsClient).Assembly.GetType(typeName);
-        }
+            MethodInfo setEnvironmentVariable = typeof(DiagnosticsClient).GetMethod("SetEnvironmentVariable", BindingFlags.Public);
+            if (setEnvironmentVariable != null)
+            {
+                throw new Exception("You updated DiagnosticsClient to a version that supports SetEnvironmentVariable, please remove this and use the real code.");
+            }
 
-        private static int GetResponseCommandSet(object response)
-        {
-            PropertyInfo header = response.GetType().GetProperty("Header");
-            object ipcHeader = header.GetValue(response);
+            Console.WriteLine($"Sending SetEnvironmentVariable message name={name} value={val ?? "NULL"}.");
+
+            using (var stream = new MemoryStream())
+            using (var writer = new BinaryWriter(stream))
+            {
+                writer.WriteString(name);
+                writer.WriteString(val);
 
-            FieldInfo commandId = ipcHeader.GetType().GetField("CommandSet");
-            byte id = (byte)commandId.GetValue(ipcHeader);
-            return id;
-        }
+                writer.Flush();
+                byte[] payload = stream.ToArray();
+
+                var message = new IpcMessage(0x04, 0x03, payload);
+                Console.WriteLine($"Sent: {message.ToString()}");
+                IpcMessage response = IpcClient.SendMessage(ConnectionHelper.GetStandardTransport(_processId), message);
+                Console.WriteLine($"Received: {response.ToString()}");
+
+                if (response.Header.CommandSet != 255 || response.Header.CommandId != 0)
+                {
+                    Console.WriteLine($"SetEnvironmentVariable failed.");
+                    return false;
+                }
+            }
 
-        private static int GetResponseCommandId(object response)
-        {
-            PropertyInfo header = response.GetType().GetProperty("Header");
-            object ipcHeader = header.GetValue(response);
+            Console.WriteLine($"Finished sending SetEnvironmentVariable message.");
 
-            FieldInfo commandId = ipcHeader.GetType().GetField("CommandId");
-            byte id = (byte)commandId.GetValue(ipcHeader);
-            return id;
+            return true;
         }
 
-        private static void SendMessage(object processId, object message)
+        private static string ReadString(byte[] buffer)
         {
-            Type ipcClientType = GetPrivateType("Microsoft.Diagnostics.NETCore.Client.IpcClient");
-            Type ipcMessageType = GetPrivateType("Microsoft.Diagnostics.NETCore.Client.IpcMessage");
-            MethodInfo clientSendMessage = ipcClientType.GetMethod("SendMessage", new Type[] { typeof(int), ipcMessageType } );
-            object response = clientSendMessage.Invoke(null, new object[] { processId, message });
-            int responseCommandSet = GetResponseCommandSet(response);
-            int responseCommandId = GetResponseCommandId(response);
-            Console.WriteLine($"SendMessage response CommandSet={responseCommandSet} CommandId={responseCommandId}");
+            int index = 0;
+            // Length of the string of UTF-16 characters
+            int length = (int)BitConverter.ToUInt32(buffer, index);
+            index += sizeof(UInt32);
+
+            int size = (int)length * sizeof(char);
+            // The string contains an ending null character; remove it before returning the value
+            string value = Encoding.Unicode.GetString(buffer, index, size).Substring(0, length - 1);
+            index += size;
+            return value;
         }
     }
 }