[Tracer] Fix Rabbitmq7 header injection overwrites user supplied basic properties [from external PR #6730] (#6753)

Based on #6730 by @johang88.

## Summary of changes
Fixes the header injection so that basic properties are preserved from
the supplied argument list in case `PopulateBasicPropertiesHeaders`
returns null which happens if the basic properties argument is writable
or if no changes had to be made to it.

## Reason for change
Fixes a bug that would remove all user supplied basic properties.

## Implementation details
Stores the basicproperties argument in the active scope and retrieves it
a null return value is encountered. Changed
`CachedBasicPropertiesHelper` to call the copy constructor that takes a
`IReadOnlyBasicProperties` as argument.

## Test coverage
Added basic properties to one of the rabbitmq7 test methods as well as
assertions.

## Other details
Fixes #6723

---------

Co-authored-by: Andrew Lock <[email protected]>
Co-authored-by: Johan Gustafsson <[email protected]>

Author: 3 people

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>();
@@ -61,6 +97,6 @@ public class PopulateBasicPropertiesHeadersIntegration
             duckType.Headers,
             int.TryParse(tags.MessageSize, out var bodyLength) ? bodyLength : 0);
 
-        return new CallTargetReturn<TReturn?>(returnValue);
+        return new CallTargetReturn<TReturn?>(basicProperties);
     }
 }

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>

tracer/test/test-applications/integrations/Samples.RabbitMQ/Program.cs

@@ -11,9 +11,11 @@
 #if RABBITMQ_7_0
 using IRabbitChannel = RabbitMQ.Client.IChannel;
 using IRabbitConsumer = RabbitMQ.Client.IAsyncBasicConsumer;
+using RabbitProperties = RabbitMQ.Client.BasicProperties;
 #else
 using IRabbitChannel = RabbitMQ.Client.IModel;
 using IRabbitConsumer = RabbitMQ.Client.IBasicConsumer;
+using RabbitProperties = RabbitMQ.Client.IBasicProperties;
 #endif
 
 namespace Samples.RabbitMQ
@@ -26,6 +28,7 @@ public static class Program
         private static readonly string exchangeName = "test-exchange-name";
         private static readonly string routingKey = "test-routing-key";
         private static readonly string queueName = "test-queue-name";
+        private static readonly string customHeaderName = "x-custom-header";
 
         private static string Host()
         {
@@ -77,6 +80,8 @@ private static async Task PublishAndGet(string consumerType, bool useDefaultQueu
                 string publishExchangeName;
                 string publishQueueName;
                 string publishRoutingKey;
+                string messageId = Guid.NewGuid().ToString();
+                string headerValue = Guid.NewGuid().ToString();
 
                 using (SampleHelpers.CreateScope(messagePrefix))
                 {
@@ -103,13 +108,20 @@ private static async Task PublishAndGet(string consumerType, bool useDefaultQueu
                     // Test an empty BasicGetResult
                     await Helper.BasicGetAsync(channel, publishQueueName);
 
+                    // Setup basic properties to verify instrumentation preserves properties and headers.
+                    var properties = Helper.CreateBasicProperties(channel);
+                    properties.MessageId = messageId;
+                    properties.Headers = new Dictionary<string, object>
+                    {
+                        { customHeaderName, headerValue }
+                    };
+
                     // Send message to the default exchange and use new queue as the routingKey
                     string message = $"{messagePrefix} - Message";
                     var body = Encoding.UTF8.GetBytes(message);
 
-                    await Helper.BasicPublishAsync(channel, exchange: publishExchangeName,
-                                                   routingKey: publishRoutingKey,
-                                                   body: body);
+                    await Helper.BasicPublishAsync(channel, publishExchangeName, publishRoutingKey, body, properties);
+
                     Console.WriteLine($"BasicPublish - Sent message: {message}");
                 }
 
@@ -123,6 +135,19 @@ await Helper.BasicPublishAsync(channel, exchange: publishExchangeName,
                 var resultMessage = Encoding.UTF8.GetString(result.Body);
 #endif
                 Console.WriteLine($"[Program.PublishAndGetDefault] BasicGet - Received message: {resultMessage}");
+
+                if (result.BasicProperties.MessageId != messageId)
+                {
+                    throw new Exception("MessageId was not preserved in BasicProperties");
+                }
+
+                if (result.BasicProperties.Headers is null ||
+                    !result.BasicProperties.Headers.TryGetValue(customHeaderName, out var receivedHeaderValue) ||
+                    receivedHeaderValue is not byte[] receivedHeaderValueString ||
+                    Encoding.UTF8.GetString(receivedHeaderValueString) != headerValue)
+                {
+                    throw new Exception("Custom header was not preserved in BasicProperties");
+                }
             }
         }
 
@@ -466,13 +491,45 @@ public static Task BasicPublishAsync(IRabbitChannel channel, string exchange, st
 #endif
         }
 
+#if RABBITMQ_6_0 || RABBITMQ_7_0
+        public static Task BasicPublishAsync(IRabbitChannel channel, string exchange, string routingKey, ReadOnlyMemory<byte> body, RabbitProperties basicProperties = null)
+#else
+        public static Task BasicPublishAsync(IRabbitChannel channel, string exchange, string routingKey, byte[] body, RabbitProperties basicProperties = null)
+#endif
+        {
+#if RABBITMQ_7_0
+            return channel.BasicPublishAsync(
+                               exchange: exchange,
+                               routingKey: routingKey,
+                               mandatory: false,
+                               basicProperties: basicProperties,
+                               body: body)
+                          .AsTask();
+#else
+            channel.BasicPublish(exchange: exchange,
+                                 routingKey: routingKey,
+                                 basicProperties: basicProperties,
+                                 body: body);
+            return Task.CompletedTask;
+#endif
+        }
+
         public static Task BasicConsumeAsync(IRabbitChannel channel, string queue, IRabbitConsumer consumer)
         {
 #if RABBITMQ_7_0
             return channel.BasicConsumeAsync(queue, autoAck: true, consumer);
 #else
             channel.BasicConsume(queue, true, consumer);
             return Task.CompletedTask;
+#endif
+        }
+
+        public static RabbitProperties CreateBasicProperties(IRabbitChannel channel)
+        {
+#if RABBITMQ_7_0
+            return new BasicProperties();
+#else
+            return channel.CreateBasicProperties();
 #endif
         }
     }