diff --git a/src/NServiceBus.AcceptanceTests/ApprovalFiles/AcceptanceTestsShippedSourceFilesApproval.ApproveShippedSourceFiles.approved.txt b/src/NServiceBus.AcceptanceTests/ApprovalFiles/AcceptanceTestsShippedSourceFilesApproval.ApproveShippedSourceFiles.approved.txt index aa3251ab29b..a489f49089a 100644 --- a/src/NServiceBus.AcceptanceTests/ApprovalFiles/AcceptanceTestsShippedSourceFilesApproval.ApproveShippedSourceFiles.approved.txt +++ b/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 diff --git a/src/NServiceBus.AcceptanceTests/Audit/When_headers_are_mutated.cs b/src/NServiceBus.AcceptanceTests/Audit/When_headers_are_mutated.cs new file mode 100644 index 00000000000..94e1b6200aa --- /dev/null +++ b/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() + .WithEndpoint(b => b.When(session => + { + var options = new SendOptions(); + options.RouteToThisEndpoint(); + options.SetHeader("AuditTestHeader", "original"); + return session.Send(new MessageToBeAudited(), options); + })) + .WithEndpoint() + .Run(); + + Assert.That(context.AuditedHeader, Is.EqualTo("original")); + } + + public class Context : ScenarioContext + { + public string AuditedHeader { get; set; } + } + + public class MutatingEndpoint : EndpointConfigurationBuilder + { + public MutatingEndpoint() => EndpointSetup(c => + { + c.AuditProcessedMessagesTo(); + c.RegisterMessageMutator(new HeaderMutator()); + }); + + class HeaderMutator : IMutateIncomingTransportMessages + { + public Task MutateIncoming(MutateIncomingTransportMessageContext context) + { + context.Headers["AuditTestHeader"] = "mutated"; + return Task.CompletedTask; + } + } + + [Handler] + public class MessageHandler : IHandleMessages + { + public Task Handle(MessageToBeAudited message, IMessageHandlerContext context) => Task.CompletedTask; + } + } + + public class AuditSpyEndpoint : EndpointConfigurationBuilder + { + public AuditSpyEndpoint() => EndpointSetup(); + + [Handler] + public class MessageHandler(Context testContext) : IHandleMessages + { + public Task Handle(MessageToBeAudited message, IMessageHandlerContext context) + { + testContext.AuditedHeader = context.MessageHeaders["AuditTestHeader"]; + testContext.MarkAsCompleted(); + return Task.CompletedTask; + } + } + } + + public class MessageToBeAudited : IMessage; +} diff --git a/src/NServiceBus.Core.Tests/Transports/IncomingMessageTests.cs b/src/NServiceBus.Core.Tests/Transports/IncomingMessageTests.cs index ba070119845..d98a1794dac 100644 --- a/src/NServiceBus.Core.Tests/Transports/IncomingMessageTests.cs +++ b/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 + { + { Headers.MessageId, "id" }, + { Headers.ContentType, "text/plain" }, + { "RemovedHeader", "original" } + }; + var message = new IncomingMessage("id", headers, System.Array.Empty()); + + 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")); + } + } } \ No newline at end of file diff --git a/src/NServiceBus.Core/Audit/InvokeAuditPipelineBehavior.cs b/src/NServiceBus.Core/Audit/InvokeAuditPipelineBehavior.cs index a7ec2fc198b..5aaaf85505f 100644 --- a/src/NServiceBus.Core/Audit/InvokeAuditPipelineBehavior.cs +++ b/src/NServiceBus.Core/Audit/InvokeAuditPipelineBehavior.cs @@ -21,6 +21,7 @@ public async Task Invoke(IIncomingPhysicalMessageContext context, Func(context.Message.Headers), context.Message.Body); diff --git a/src/NServiceBus.Core/Transports/IncomingMessage.cs b/src/NServiceBus.Core/Transports/IncomingMessage.cs index fd8ebcdbeb1..65d9e17c314 100644 --- a/src/NServiceBus.Core/Transports/IncomingMessage.cs +++ b/src/NServiceBus.Core/Transports/IncomingMessage.cs @@ -35,6 +35,7 @@ public IncomingMessage(string nativeMessageId, Dictionary header NativeMessageId = nativeMessageId; MessageId = GetOrSetMessageIdFromHeaders(headers, nativeMessageId); Headers = headers; + originalHeaders = new Dictionary(headers); Body = body; ReceiveProperties = receiveProperties; } @@ -108,5 +109,18 @@ internal void RevertToOriginalBodyIfNeeded() } } + /// + /// Makes sure that the headers are reset to the exact state as they were when the message was created. + /// + internal void RevertToOriginalHeadersIfNeeded() + { + Headers.Clear(); + foreach (var header in originalHeaders) + { + Headers.Add(header.Key, header.Value); + } + } + ReadOnlyMemory? originalBody; + readonly Dictionary originalHeaders; } \ No newline at end of file