From 28fc03a7b97538845e38042a83d38585e5a757d0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=C3=9Eorvaldur=20Hafdal?= Date: Wed, 15 Jul 2026 11:51:11 +0000 Subject: [PATCH 1/2] SS-4 Add streaming DATA message store path --- src/SmtpServer.Tests/PipeReaderTests.cs | 62 ++++++++ src/SmtpServer.Tests/SmtpServerTests.cs | 62 ++++++++ src/SmtpServer/IO/PipeReaderExtensions.cs | 135 ++++++++++++++++++ src/SmtpServer/Protocol/DataCommand.cs | 48 +++++-- .../Storage/IStreamingMessageStore.cs | 23 +++ 5 files changed, 320 insertions(+), 10 deletions(-) create mode 100644 src/SmtpServer/Storage/IStreamingMessageStore.cs diff --git a/src/SmtpServer.Tests/PipeReaderTests.cs b/src/SmtpServer.Tests/PipeReaderTests.cs index b9a085a7..79756f6e 100644 --- a/src/SmtpServer.Tests/PipeReaderTests.cs +++ b/src/SmtpServer.Tests/PipeReaderTests.cs @@ -3,6 +3,7 @@ using System.Text; using System.Threading.Tasks; using SmtpServer.IO; +using SmtpServer.Protocol; using SmtpServer.Text; using Xunit; @@ -93,5 +94,66 @@ await reader.ReadDotBlockAsync( // assert Assert.Equal("abcd\r\n.1234", text); } + + [Fact] + public async Task CanStreamBlockWithDotStuffingRemoved() + { + // arrange + var reader = CreatePipeReader("abcd\r\n..1234\r\n.\r\n"); + var writer = new Pipe(); + + var maxMessageSizeOptions = new MaxMessageSizeOptions(); + + // act + await reader.ReadDotBlockAsync(writer.Writer, maxMessageSizeOptions); + var text = await ReadAllAsync(writer.Reader); + + // assert + Assert.Equal("abcd\r\n.1234", text); + } + + [Fact] + public async Task CanEnforceMaxMessageSizeWhenStreamingBlock() + { + // arrange + var reader = CreatePipeReader("abcd\r\n1234\r\n.\r\n"); + var writer = new Pipe(); + + var maxMessageSizeOptions = new MaxMessageSizeOptions(MaxMessageSizeHandling.Strict, 5); + + // act + var exception = await Assert.ThrowsAsync( + async () => await reader.ReadDotBlockAsync(writer.Writer, maxMessageSizeOptions)); + + // assert + Assert.True(exception.IsQuitRequested); + } + + static async Task ReadAllAsync(PipeReader reader) + { + using var stream = new MemoryStream(); + + while (true) + { + var result = await reader.ReadAsync(); + var buffer = result.Buffer; + + foreach (var segment in buffer) + { + stream.Write(segment.Span); + } + + reader.AdvanceTo(buffer.End); + + if (result.IsCompleted) + { + break; + } + } + + reader.Complete(); + + return Encoding.ASCII.GetString(stream.ToArray()); + } } } diff --git a/src/SmtpServer.Tests/SmtpServerTests.cs b/src/SmtpServer.Tests/SmtpServerTests.cs index 292cdaa1..ab35b8e9 100644 --- a/src/SmtpServer.Tests/SmtpServerTests.cs +++ b/src/SmtpServer.Tests/SmtpServerTests.cs @@ -8,8 +8,10 @@ using SmtpServer.Storage; using SmtpServer.Tests.Mocks; using System; +using System.Buffers; using System.Diagnostics; using System.IO; +using System.IO.Pipelines; using System.Linq; using System.Net; using System.Net.Security; @@ -52,6 +54,21 @@ public void CanReceiveMessage() } } + [Fact] + public void CanReceiveMessageUsingStreamingMessageStore() + { + var streamingMessageStore = new StreamingMockMessageStore(); + + using (CreateServer(services => services.Add(streamingMessageStore))) + { + MailClient.Send(MailClient.Message(from: "test1@test.com", to: "test2@test.com", text: "streamed body")); + } + + Assert.True(streamingMessageStore.StreamingSaveCalled); + Assert.False(streamingMessageStore.BufferedSaveCalled); + Assert.Contains("streamed body", streamingMessageStore.Message); + } + [Theory] [InlineData("Assunto teste acento çãõáéíóú", "utf-8")] [InlineData("שלום שלום שלום", "windows-1255")] @@ -655,5 +672,50 @@ SmtpServerDisposable CreateServer( /// The cancellation token source for the test. /// public CancellationTokenSource CancellationTokenSource { get; } + + sealed class StreamingMockMessageStore : MessageStore, IStreamingMessageStore + { + public override Task SaveAsync(ISessionContext context, IMessageTransaction transaction, ReadOnlySequence buffer, CancellationToken cancellationToken) + { + BufferedSaveCalled = true; + + return Task.FromResult(SmtpResponse.Ok); + } + + public async Task SaveAsync(ISessionContext context, IMessageTransaction transaction, PipeReader reader, CancellationToken cancellationToken) + { + StreamingSaveCalled = true; + + using var stream = new MemoryStream(); + + while (true) + { + var result = await reader.ReadAsync(cancellationToken); + var buffer = result.Buffer; + + foreach (var segment in buffer) + { + stream.Write(segment.Span); + } + + reader.AdvanceTo(buffer.End); + + if (result.IsCompleted) + { + break; + } + } + + Message = Encoding.UTF8.GetString(stream.ToArray()); + + return SmtpResponse.Ok; + } + + public bool BufferedSaveCalled { get; private set; } + + public bool StreamingSaveCalled { get; private set; } + + public string Message { get; private set; } + } } } diff --git a/src/SmtpServer/IO/PipeReaderExtensions.cs b/src/SmtpServer/IO/PipeReaderExtensions.cs index c0c4f470..6c56b712 100644 --- a/src/SmtpServer/IO/PipeReaderExtensions.cs +++ b/src/SmtpServer/IO/PipeReaderExtensions.cs @@ -15,6 +15,7 @@ internal static class PipeReaderExtensions static readonly byte[] CRLF = { 13, 10 }; static readonly byte[] DotBlock = { 13, 10, 46, 13, 10 }; static readonly byte[] DotBlockStuffing = { 13, 10, 46, 46 }; + const int DotBlockTailLength = 4; /// /// Read from the reader until the sequence is found. @@ -177,5 +178,139 @@ static ReadOnlySequence Unstuff(ReadOnlySequence buffer) return segments.Build(); } } + + /// + /// Reads a dot block from the reader and writes the unstuffed content to the writer. + /// + /// The reader to read from. + /// The writer to stream the message content to. + /// Handling of MaxMessageSize. + /// The cancellation token. + /// The value that was read from the buffer. + internal static async ValueTask ReadDotBlockAsync(this PipeReader reader, PipeWriter writer, IMaxMessageSizeOptions maxMessageSizeOptions, CancellationToken cancellationToken = default) + { + if (reader == null) + { + throw new ArgumentNullException(nameof(reader)); + } + + if (writer == null) + { + throw new ArgumentNullException(nameof(writer)); + } + + Exception error = null; + + try + { + await ReadDotBlockAsync(reader, writer, maxMessageSizeOptions, cancellationToken, DotBlockTailLength).ConfigureAwait(false); + } + catch (Exception exception) + { + error = exception; + throw; + } + finally + { + writer.Complete(error); + } + } + + static async ValueTask ReadDotBlockAsync(PipeReader reader, PipeWriter writer, IMaxMessageSizeOptions maxMessageSizeOptions, CancellationToken cancellationToken, int tailLength) + { + long consumedLength = 0; + + while (true) + { + var read = await reader.ReadAsync(cancellationToken).ConfigureAwait(false); + var buffer = read.Buffer; + + if (buffer.IsEmpty && read.IsCompleted) + { + return; + } + + var head = buffer.GetPosition(0); + if (buffer.TryFind(DotBlock, ref head, out var tail)) + { + var body = buffer.Slice(buffer.Start, head); + + EnsureMessageSize(maxMessageSizeOptions, consumedLength + body.Length); + + await WriteUnstuffedAsync(writer, body, cancellationToken).ConfigureAwait(false); + + reader.AdvanceTo(tail); + return; + } + + if (read.IsCompleted) + { + reader.AdvanceTo(buffer.End); + return; + } + + if (buffer.Length > tailLength) + { + var safeLength = buffer.Length - tailLength; + var safeBuffer = buffer.Slice(0, safeLength); + + EnsureMessageSize(maxMessageSizeOptions, consumedLength + safeBuffer.Length); + + await WriteUnstuffedAsync(writer, safeBuffer, cancellationToken).ConfigureAwait(false); + + consumedLength += safeBuffer.Length; + + reader.AdvanceTo(buffer.GetPosition(safeLength), buffer.End); + continue; + } + + reader.AdvanceTo(buffer.Start, buffer.End); + } + } + + static void EnsureMessageSize(IMaxMessageSizeOptions maxMessageSizeOptions, long length) + { + if (maxMessageSizeOptions.Handling == MaxMessageSizeHandling.Strict && length > maxMessageSizeOptions.Length) + { + throw new SmtpResponseException(SmtpResponse.MaxMessageSizeExceeded, true); + } + } + + static async ValueTask WriteUnstuffedAsync(PipeWriter writer, ReadOnlySequence buffer, CancellationToken cancellationToken) + { + var head = buffer.GetPosition(0); + var start = head; + + while (buffer.TryFind(DotBlockStuffing, ref head, out var tail)) + { + var slice = buffer.Slice(start, buffer.GetPosition(3, head)); + + Write(writer, slice); + + start = tail; + head = tail; + } + + Write(writer, buffer.Slice(start)); + + await writer.FlushAsync(cancellationToken).ConfigureAwait(false); + } + + static void Write(PipeWriter writer, ReadOnlySequence buffer) + { + var position = buffer.GetPosition(0); + + while (buffer.TryGet(ref position, out var memory)) + { + if (memory.Length == 0) + { + continue; + } + + var span = writer.GetSpan(memory.Length); + memory.Span.CopyTo(span); + writer.Advance(memory.Length); + } + } } } diff --git a/src/SmtpServer/Protocol/DataCommand.cs b/src/SmtpServer/Protocol/DataCommand.cs index b30ce23c..7df894a8 100644 --- a/src/SmtpServer/Protocol/DataCommand.cs +++ b/src/SmtpServer/Protocol/DataCommand.cs @@ -1,4 +1,5 @@ using System; +using System.IO.Pipelines; using System.Threading; using System.Threading.Tasks; using SmtpServer.ComponentModel; @@ -45,16 +46,9 @@ internal override async Task ExecuteAsync(SmtpSessionContext context, Canc { using var container = new DisposableContainer(messageStore); - SmtpResponse response = null; - - await context.Pipe.Input.ReadDotBlockAsync( - async buffer => - { - // ReSharper disable once AccessToDisposedClosure - response = await container.Instance.SaveAsync(context, context.Transaction, buffer, cancellationToken).ConfigureAwait(false); - }, - context.ServerOptions.MaxMessageSizeOptions, - cancellationToken).ConfigureAwait(false); + var response = container.Instance is IStreamingMessageStore streamingMessageStore + ? await SaveAsync(streamingMessageStore, context, cancellationToken).ConfigureAwait(false) + : await SaveAsync(container.Instance, context, cancellationToken).ConfigureAwait(false); await context.Pipe.Output.WriteReplyAsync(response, cancellationToken).ConfigureAwait(false); } @@ -70,5 +64,39 @@ await context.Pipe.Input.ReadDotBlockAsync( return true; } + + static async Task SaveAsync(IMessageStore messageStore, SmtpSessionContext context, CancellationToken cancellationToken) + { + SmtpResponse response = null; + + await context.Pipe.Input.ReadDotBlockAsync( + async buffer => + { + response = await messageStore.SaveAsync(context, context.Transaction, buffer, cancellationToken).ConfigureAwait(false); + }, + context.ServerOptions.MaxMessageSizeOptions, + cancellationToken).ConfigureAwait(false); + + return response; + } + + static async Task SaveAsync(IStreamingMessageStore messageStore, SmtpSessionContext context, CancellationToken cancellationToken) + { + var pipe = new Pipe(); + + try + { + var readTask = context.Pipe.Input.ReadDotBlockAsync(pipe.Writer, context.ServerOptions.MaxMessageSizeOptions, cancellationToken).AsTask(); + var saveTask = messageStore.SaveAsync(context, context.Transaction, pipe.Reader, cancellationToken); + + await readTask.ConfigureAwait(false); + + return await saveTask.ConfigureAwait(false); + } + finally + { + pipe.Reader.Complete(); + } + } } } diff --git a/src/SmtpServer/Storage/IStreamingMessageStore.cs b/src/SmtpServer/Storage/IStreamingMessageStore.cs new file mode 100644 index 00000000..996949ee --- /dev/null +++ b/src/SmtpServer/Storage/IStreamingMessageStore.cs @@ -0,0 +1,23 @@ +using System.IO.Pipelines; +using System.Threading; +using System.Threading.Tasks; +using SmtpServer.Protocol; + +namespace SmtpServer.Storage +{ + /// + /// Streaming Message Store Interface + /// + public interface IStreamingMessageStore : IMessageStore + { + /// + /// Save the given message to the underlying storage system. + /// + /// The session level context. + /// The SMTP message transaction to store. + /// The reader that streams the message content. + /// The cancellation token. + /// The response code to return that indicates the result of the message being saved. + Task SaveAsync(ISessionContext context, IMessageTransaction transaction, PipeReader reader, CancellationToken cancellationToken); + } +} From 5df5db2b1f04db8007acf6c7e068f0cef26937be Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=C3=9Eorvaldur=20Hafdal?= Date: Wed, 15 Jul 2026 11:57:45 +0000 Subject: [PATCH 2/2] SS-4 Add DATA streaming benchmark --- .../DataStoreBenchmarks.cs | 113 ++++++++++++++++++ src/SmtpServer.Benchmarks/Program.cs | 18 ++- 2 files changed, 120 insertions(+), 11 deletions(-) create mode 100644 src/SmtpServer.Benchmarks/DataStoreBenchmarks.cs diff --git a/src/SmtpServer.Benchmarks/DataStoreBenchmarks.cs b/src/SmtpServer.Benchmarks/DataStoreBenchmarks.cs new file mode 100644 index 00000000..1756c3ea --- /dev/null +++ b/src/SmtpServer.Benchmarks/DataStoreBenchmarks.cs @@ -0,0 +1,113 @@ +using System; +using System.Buffers; +using System.IO; +using System.IO.Pipelines; +using System.Threading; +using System.Threading.Tasks; +using BenchmarkDotNet.Attributes; +using MimeKit; +using SmtpServer.ComponentModel; +using SmtpServer.Protocol; +using SmtpServer.Storage; +using SmtpClient = MailKit.Net.Smtp.SmtpClient; + +namespace SmtpServer.Benchmarks +{ + [MemoryDiagnoser] + [ShortRunJob] + public class DataStoreBenchmarks + { + const int Port = 9026; + + SmtpServer _smtpServer; + CancellationTokenSource _smtpServerCancellationTokenSource; + SmtpClient _smtpClient; + MimeMessage _message; + + public enum StoreMode + { + BufferedMaterializing, + StreamingDrain + } + + [Params(StoreMode.BufferedMaterializing, StoreMode.StreamingDrain)] + public StoreMode Mode { get; set; } + + [GlobalSetup] + public void SmtpServerSetup() + { + _message = MimeMessage.Load(typeof(DataStoreBenchmarks).Assembly.GetManifestResourceStream("SmtpServer.Benchmarks.Test3.eml")); + _smtpServerCancellationTokenSource = new CancellationTokenSource(); + + var serviceProvider = new ServiceProvider(); + serviceProvider.Add(Mode == StoreMode.StreamingDrain + ? (IMessageStore)new StreamingDrainMessageStore() + : new BufferedMaterializingMessageStore()); + + _smtpServer = new SmtpServer( + new SmtpServerOptionsBuilder() + .Port(Port, false) + .Build(), + serviceProvider); + + _ = _smtpServer.StartAsync(_smtpServerCancellationTokenSource.Token); + + _smtpClient = new SmtpClient(); + _smtpClient.Connect("localhost", Port); + } + + [GlobalCleanup] + public Task SmtpServerCleanupAsync() + { + _smtpClient.Disconnect(true); + _smtpClient.Dispose(); + + _smtpServerCancellationTokenSource.Cancel(); + _smtpServerCancellationTokenSource.Dispose(); + + return _smtpServer.ShutdownTask; + } + + [Benchmark] + public void SendMessage() + { + _smtpClient.Send(_message); + } + + sealed class BufferedMaterializingMessageStore : MessageStore + { + public override Task SaveAsync(ISessionContext context, IMessageTransaction transaction, ReadOnlySequence buffer, CancellationToken cancellationToken) + { + _ = buffer.ToArray(); + + return Task.FromResult(SmtpResponse.Ok); + } + } + + sealed class StreamingDrainMessageStore : MessageStore, IStreamingMessageStore + { + public override Task SaveAsync(ISessionContext context, IMessageTransaction transaction, ReadOnlySequence buffer, CancellationToken cancellationToken) + { + throw new NotSupportedException(); + } + + public async Task SaveAsync(ISessionContext context, IMessageTransaction transaction, PipeReader reader, CancellationToken cancellationToken) + { + while (true) + { + var result = await reader.ReadAsync(cancellationToken).ConfigureAwait(false); + var buffer = result.Buffer; + + reader.AdvanceTo(buffer.End); + + if (result.IsCompleted) + { + break; + } + } + + return SmtpResponse.Ok; + } + } + } +} diff --git a/src/SmtpServer.Benchmarks/Program.cs b/src/SmtpServer.Benchmarks/Program.cs index e8c71cb9..10523351 100644 --- a/src/SmtpServer.Benchmarks/Program.cs +++ b/src/SmtpServer.Benchmarks/Program.cs @@ -8,17 +8,13 @@ public class Program { public static void Main(string[] args) { - //var summary = BenchmarkRunner.Run( - // ManualConfig - // .Create(DefaultConfig.Instance) - // .With(ConfigOptions.DisableOptimizationsValidator)); - - //var summary = BenchmarkRunner.Run(); - - var summary = BenchmarkRunner.Run( - ManualConfig - .Create(DefaultConfig.Instance) - .With(ConfigOptions.DisableOptimizationsValidator)); + BenchmarkSwitcher + .FromAssembly(typeof(Program).Assembly) + .Run( + args, + ManualConfig + .Create(DefaultConfig.Instance) + .WithOptions(ConfigOptions.DisableOptimizationsValidator)); } } }