Revert message headers to their original state before auditing

src/NServiceBus.AcceptanceTests/ApprovalFiles/AcceptanceTestsShippedSourceFilesApproval.ApproveShippedSourceFiles.approved.txt

@@ -4,6 +4,7 @@ Audit/When_audit_is_overridden_in_code.cs
 Audit/When_audit_is_overridden_in_environment.cs
 Audit/When_auditing.cs
 Audit/When_auditing_message_with_TimeToBeReceived.cs
+Audit/When_headers_are_mutated.cs
 BehaviorBuilderRegistrationExtensions.cs
 ConfigureEndpointAcceptanceTestingPersistence.cs
 ConfigureEndpointAcceptanceTestingTransport.cs

src/NServiceBus.AcceptanceTests/Audit/When_headers_are_mutated.cs

@@ -0,0 +1,75 @@
+namespace NServiceBus.AcceptanceTests.Audit;
+
+using System.Threading.Tasks;
+using AcceptanceTesting;
+using AcceptanceTesting.Customization;
+using EndpointTemplates;
+using MessageMutator;
+using NUnit.Framework;
+
+public class When_headers_are_mutated : NServiceBusAcceptanceTest
+{
+    [Test]
+    public async Task Should_audit_original_headers()
+    {
+        var context = await Scenario.Define<Context>()
+            .WithEndpoint<MutatingEndpoint>(b => b.When(session =>
+            {
+                var options = new SendOptions();
+                options.RouteToThisEndpoint();
+                options.SetHeader("AuditTestHeader", "original");
+                return session.Send(new MessageToBeAudited(), options);
+            }))
+            .WithEndpoint<AuditSpyEndpoint>()
+            .Run();
+
+        Assert.That(context.AuditedHeader, Is.EqualTo("original"));
+    }
+
+    public class Context : ScenarioContext
+    {
+        public string AuditedHeader { get; set; }
+    }
+
+    public class MutatingEndpoint : EndpointConfigurationBuilder
+    {
+        public MutatingEndpoint() => EndpointSetup<DefaultServer>(c =>
+        {
+            c.AuditProcessedMessagesTo<AuditSpyEndpoint>();
+            c.RegisterMessageMutator(new HeaderMutator());
+        });
+
+        class HeaderMutator : IMutateIncomingTransportMessages
+        {
+            public Task MutateIncoming(MutateIncomingTransportMessageContext context)
+            {
+                context.Headers["AuditTestHeader"] = "mutated";
+                return Task.CompletedTask;
+            }
+        }
+
+        [Handler]
+        public class MessageHandler : IHandleMessages<MessageToBeAudited>
+        {
+            public Task Handle(MessageToBeAudited message, IMessageHandlerContext context) => Task.CompletedTask;
+        }
+    }
+
+    public class AuditSpyEndpoint : EndpointConfigurationBuilder
+    {
+        public AuditSpyEndpoint() => EndpointSetup<DefaultServer>();
+
+        [Handler]
+        public class MessageHandler(Context testContext) : IHandleMessages<MessageToBeAudited>
+        {
+            public Task Handle(MessageToBeAudited message, IMessageHandlerContext context)
+            {
+                testContext.AuditedHeader = context.MessageHeaders["AuditTestHeader"];
+                testContext.MarkAsCompleted();
+                return Task.CompletedTask;
+            }
+        }
+    }
+
+    public class MessageToBeAudited : IMessage;
+}

src/NServiceBus.Core.Tests/Transports/IncomingMessageTests.cs

@@ -52,4 +52,29 @@ public void Should_assign_transport_message_id_when_NServiceBus_message_id_heade
             Assert.That(message.MessageId, Is.EqualTo("nativeId"));
         }
     }
+
+    [Test]
+    public void RevertToOriginalHeadersIfNeeded_should_restore_original_headers()
+    {
+        var headers = new Dictionary<string, string>
+        {
+            { Headers.MessageId, "id" },
+            { Headers.ContentType, "text/plain" },
+            { "RemovedHeader", "original" }
+        };
+        var message = new IncomingMessage("id", headers, System.Array.Empty<byte>());
+
+        message.Headers[Headers.ContentType] = "application/json";
+        message.Headers["AddedHeader"] = "added";
+        message.Headers.Remove("RemovedHeader");
+
+        message.RevertToOriginalHeadersIfNeeded();
+
+        using (Assert.EnterMultipleScope())
+        {
+            Assert.That(message.Headers[Headers.ContentType], Is.EqualTo("text/plain"));
+            Assert.That(message.Headers.ContainsKey("AddedHeader"), Is.False);
+            Assert.That(message.Headers["RemovedHeader"], Is.EqualTo("original"));
+        }
+    }
 }

src/NServiceBus.Core/Audit/InvokeAuditPipelineBehavior.cs

@@ -21,6 +21,7 @@ public async Task Invoke(IIncomingPhysicalMessageContext context, Func<IIncoming
         await next(context).ConfigureAwait(false);
 
         context.Message.RevertToOriginalBodyIfNeeded();
+        context.Message.RevertToOriginalHeadersIfNeeded();
 
         var processedMessage = new OutgoingMessage(context.Message.MessageId, new Dictionary<string, string>(context.Message.Headers), context.Message.Body);
 

src/NServiceBus.Core/Transports/IncomingMessage.cs

@@ -35,6 +35,7 @@ public IncomingMessage(string nativeMessageId, Dictionary<string, string> header
         NativeMessageId = nativeMessageId;
         MessageId = GetOrSetMessageIdFromHeaders(headers, nativeMessageId);
         Headers = headers;
+        originalHeaders = new Dictionary<string, string>(headers);
         Body = body;
         ReceiveProperties = receiveProperties;
     }
@@ -108,5 +109,18 @@ internal void RevertToOriginalBodyIfNeeded()
         }
     }
 
+    /// <summary>
+    /// Makes sure that the headers are reset to the exact state as they were when the message was created.
+    /// </summary>
+    internal void RevertToOriginalHeadersIfNeeded()
+    {
+        Headers.Clear();
+        foreach (var header in originalHeaders)
+        {
+            Headers.Add(header.Key, header.Value);
+        }
+    }
+
     ReadOnlyMemory<byte>? originalBody;
+    readonly Dictionary<string, string> originalHeaders;
 }