Fix: Rabbitmq7 header injection overwrites user supplied basic properties [from external PR #6730]

tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/RabbitMQ/CachedBasicPropertiesHelper.cs

@@ -1,23 +1,51 @@
-// <copyright file="CachedBasicPropertiesHelper.cs" company="Datadog">
+// <copyright file="CachedBasicPropertiesHelper.cs" company="Datadog">
 // Unless explicitly stated otherwise all files in this repository are licensed under the Apache 2 License.
 // This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2017 Datadog, Inc.
 // </copyright>
 
 #nullable enable
-using Datadog.Trace.Util;
+using System;
+using System.Reflection.Emit;
+using Datadog.Trace.DuckTyping;
+using Datadog.Trace.Logging;
 
 namespace Datadog.Trace.ClrProfiler.AutoInstrumentation.RabbitMQ;
 
 internal static class CachedBasicPropertiesHelper<TBasicProperties>
 {
-    // ReSharper disable once StaticMemberInGenericType
-    private static readonly ActivatorHelper Activator;
+    private static readonly Func<object, TBasicProperties>? Activator;
 
     static CachedBasicPropertiesHelper()
     {
-        Activator = new ActivatorHelper(typeof(TBasicProperties).Assembly.GetType("RabbitMQ.Client.BasicProperties")!);
+        try
+        {
+            var targetType = typeof(TBasicProperties).Assembly.GetType("RabbitMQ.Client.BasicProperties")!;
+            var parameterType = typeof(TBasicProperties).Assembly.GetType("RabbitMQ.Client.IReadOnlyBasicProperties")!;
+
+            var constructor = targetType.GetConstructor([parameterType])!;
+
+            var createBasicPropertiesMethod = new DynamicMethod(
+                $"TypeActivator_{targetType.Name}_{parameterType.Name}",
+                targetType,
+                [typeof(object)],
+                typeof(DuckType).Module,
+                skipVisibility: true);
+
+            var il = createBasicPropertiesMethod.GetILGenerator();
+            il.Emit(OpCodes.Ldarg_0);                  // Load first argument (as object).
+            il.Emit(OpCodes.Castclass, parameterType); // Cast to IReadOnlyBasicProperties
+            il.Emit(OpCodes.Newobj, constructor);      // Call constructor
+            il.Emit(OpCodes.Ret);                      // Return new instance
+
+            Activator = (Func<object, TBasicProperties>)createBasicPropertiesMethod.CreateDelegate(typeof(Func<object, TBasicProperties>));
+        }
+        catch (Exception ex)
+        {
+            var log = DatadogLogging.GetLoggerFor(typeof(CachedBasicPropertiesHelper<TBasicProperties>));
+            log.Error(ex, "Failed to create activator for {TBasicProperties}", typeof(TBasicProperties).FullName);
+        }
     }
 
-    public static TBasicProperties CreateHeaders()
-        => (TBasicProperties)Activator.CreateInstance();
+    public static TBasicProperties? CreateHeaders(object readonlyBasicProperties)
+        => Activator is null ? default : Activator(readonlyBasicProperties);
 }

tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/RabbitMQ/PopulateBasicPropertiesHeadersIntegration.cs

@@ -1,4 +1,4 @@
-// <copyright file="PopulateBasicPropertiesHeadersIntegration.cs" company="Datadog">
+// <copyright file="PopulateBasicPropertiesHeadersIntegration.cs" company="Datadog">
 // Unless explicitly stated otherwise all files in this repository are licensed under the Apache 2 License.
 // This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2017 Datadog, Inc.
 // </copyright>
@@ -10,6 +10,7 @@
 using System.ComponentModel;
 using Datadog.Trace.ClrProfiler.CallTarget;
 using Datadog.Trace.DuckTyping;
+using Datadog.Trace.Logging;
 using Datadog.Trace.Propagators;
 using Datadog.Trace.Tagging;
 
@@ -31,6 +32,13 @@ namespace Datadog.Trace.ClrProfiler.AutoInstrumentation.RabbitMQ;
 [EditorBrowsable(EditorBrowsableState.Never)]
 public class PopulateBasicPropertiesHeadersIntegration
 {
+    private static readonly IDatadogLogger Log = DatadogLogging.GetLoggerFor<PopulateBasicPropertiesHeadersIntegration>();
+
+    internal static CallTargetState OnMethodBegin<TTarget, TBasicProperties, TActivity>(TTarget instance, TBasicProperties basicProperties, TActivity sendActivity, ulong publishSequenceNumber)
+    {
+        return new CallTargetState(null, basicProperties);
+    }
+
     internal static CallTargetReturn<TReturn?> OnMethodEnd<TTarget, TReturn>(TTarget instance, TReturn returnValue, Exception? exception, in CallTargetState state)
     {
         var tracer = Tracer.Instance;
@@ -45,8 +53,36 @@ public class PopulateBasicPropertiesHeadersIntegration
             return new CallTargetReturn<TReturn?>(returnValue);
         }
 
-        returnValue ??= CachedBasicPropertiesHelper<TReturn>.CreateHeaders();
-        var duckType = returnValue.DuckCast<IBasicProperties>()!;
+        // TReturn is type RabbitMQ.Client.BasicProperties
+        TReturn? basicProperties = default;
+
+        // PopulateBasicPropertiesHeaders returns null if the supplied IReadOnlyBasicProperties
+        // does not have to be modified or if it's a writable instance.
+        // If that is the case then we have to fetch IReadOnlyBasicProperties from the argument
+        // list instead of creating a new instance that overwrites the supplied properties.
+        if (returnValue is null)
+        {
+            // state.State is of type RabbitMQ.Client.IReadOnlyBasicProperties
+            if (state.State is TReturn writable)
+            {
+                // Use the existing BasicProperties if it's already of type RabbitMQ.Client.BasicProperties
+                basicProperties = writable;
+            }
+            else if (state.State is null)
+            {
+                // This case cannot happen as argument is not nullable.
+                Log.Warning("Invalid state: PopulateBasicPropertiesHeaders() is returning null and CallTargetState.State has type {Type}", state.State?.GetType().FullName ?? "null");
+                return new CallTargetReturn<TReturn?>(returnValue);
+            }
+            else
+            {
+                // create new BasicProperties using the BasicProperties(IReadOnlyBasicProperties) copy constructor
+                basicProperties = CachedBasicPropertiesHelper<TReturn>.CreateHeaders(state.State!);
+            }
+        }
+
+        // duck cast so we can access the Headers property
+        var duckType = basicProperties.DuckCast<IBasicProperties>()!;
 
         // add distributed tracing headers to the message
         duckType.Headers ??= new Dictionary<string, object>();

tracer/test/Datadog.Trace.ClrProfiler.IntegrationTests/RabbitMQTests.cs

@@ -106,6 +106,7 @@ private string GetPackageVersionSuffix(string packageVersion)
             => packageVersion switch
             {
                 null or "" => "6_x", // the default version in the csproj
+                _ when new Version(packageVersion) >= new Version("7.0.0") => "7_x",
                 _ when new Version(packageVersion) >= new Version("6.0.0") => "6_x",
                 _ when new Version(packageVersion) >= new Version("5.0.0") => "5_x",
                 _ => "3_x",

tracer/test/Datadog.Trace.ClrProfiler.Managed.Tests/AutoInstrumentation/RabbitMQ/CachedBasicPropertiesHelperTests.cs

@@ -0,0 +1,66 @@
+// <copyright file="CachedBasicPropertiesHelperTests.cs" company="Datadog">
+// Unless explicitly stated otherwise all files in this repository are licensed under the Apache 2 License.
+// This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2017 Datadog, Inc.
+// </copyright>
+
+using System;
+using System.Collections.Generic;
+using Datadog.Trace.ClrProfiler.AutoInstrumentation.RabbitMQ;
+using FluentAssertions;
+using RabbitMQ.Client;
+using Xunit;
+
+namespace Datadog.Trace.ClrProfiler.Managed.Tests.AutoInstrumentation.RabbitMQ;
+
+public class CachedBasicPropertiesHelperTests
+{
+    [Fact]
+    public void CreateHeaders_With_Null_BasicProperties_Throws()
+    {
+        // the BasicProperties(IReadOnlyBasicProperties) constructor throws NullReferenceException if parameter is null
+        var func = () => CachedBasicPropertiesHelper<BasicProperties>.CreateHeaders(null!);
+        func.Should().ThrowExactly<NullReferenceException>();
+    }
+
+    [Fact]
+    public void CreateHeaders_BasicProperties_With_Null_Headers()
+    {
+        var originalProperties = new BasicProperties();
+
+        var newProperties = CachedBasicPropertiesHelper<BasicProperties>.CreateHeaders(originalProperties);
+
+        newProperties.Should().NotBeNull();
+        newProperties.Headers.Should().BeNull();
+    }
+
+    [Fact]
+    public void CreateHeaders_BasicProperties_With_Empty_Headers()
+    {
+        var headers = new Dictionary<string, object>();
+
+        var originalProperties = new BasicProperties
+        {
+            Headers = headers
+        };
+
+        var newProperties = CachedBasicPropertiesHelper<BasicProperties>.CreateHeaders(originalProperties);
+
+        newProperties.Should().NotBeNull();
+        newProperties.Headers.Should().BeSameAs(originalProperties.Headers)
+                     .And.Subject.Should().BeEmpty();
+    }
+
+    [Fact]
+    public void CreateHeaders_BasicProperties_With_Headers()
+    {
+        var originalProperties = new BasicProperties();
+        originalProperties.Headers ??= new Dictionary<string, object>();
+        originalProperties.Headers["key1"] = "value1";
+
+        var newProperties = CachedBasicPropertiesHelper<BasicProperties>.CreateHeaders(originalProperties);
+
+        newProperties.Should().NotBeNull();
+        newProperties.Headers.Should().BeSameAs(originalProperties.Headers)
+                     .And.Subject.Should().Contain("key1", "value1");
+    }
+}

tracer/test/Datadog.Trace.ClrProfiler.Managed.Tests/Datadog.Trace.ClrProfiler.Managed.Tests.csproj

@@ -26,6 +26,7 @@
     <PackageReference Include="log4net" Version="$(Log4NetVersion)" />
     <PackageReference Include="Newtonsoft.Json" Version="13.0.2" />
     <PackageReference Include="NLog" Version="$(NLogVersion)" Condition=" $(TargetFramework.StartsWith('net4'))" />
+    <PackageReference Include="RabbitMQ.Client" Version="7.1.1" />
   </ItemGroup>
 
   <ItemGroup>