Ensure that MimeMessage.WriteTo() always ends with a new-line

Fixes jstedfast/MailKit#251
jstedfast · Oct 16, 2015 · 9b01f59 · 9b01f59
1 parent 937a54d
commit 9b01f59
Show file tree

Hide file tree

Showing 4 changed files with 30 additions and 14 deletions.
diff --git a/MimeKit/FormatOptions.cs b/MimeKit/FormatOptions.cs
@@ -117,13 +117,13 @@ public NewLineFormat NewLineFormat {
 			}
 		}
 
-		internal IMimeFilter CreateNewLineFilter ()
+		internal IMimeFilter CreateNewLineFilter (bool ensureNewLine = false)
 		{
 			switch (NewLineFormat) {
 			case NewLineFormat.Unix:
-				return new Dos2UnixFilter ();
+				return new Dos2UnixFilter (ensureNewLine);
 			default:
-				return new Unix2DosFilter ();
+				return new Unix2DosFilter (ensureNewLine);
 			}
 		}
 

diff --git a/MimeKit/IO/Filters/Dos2UnixFilter.cs b/MimeKit/IO/Filters/Dos2UnixFilter.cs
@@ -33,6 +33,7 @@ namespace MimeKit.IO.Filters {
 	/// </remarks>
 	public class Dos2UnixFilter : MimeFilterBase
 	{
+		bool ensureNewLine;
 		byte pc;
 
 		/// <summary>
@@ -41,11 +42,13 @@ public class Dos2UnixFilter : MimeFilterBase
 		/// <remarks>
 		/// Creates a new <see cref="Dos2UnixFilter"/>.
 		/// </remarks>
-		public Dos2UnixFilter ()
+		/// <param name="ensureNewLine">Ensure that the stream ends with a new line.</param>
+		public Dos2UnixFilter (bool ensureNewLine = false)
 		{
+			this.ensureNewLine = ensureNewLine;
 		}
 
-		unsafe int Filter (byte* inbuf, int length, byte* outbuf)
+		unsafe int Filter (byte* inbuf, int length, byte* outbuf, bool flush)
 		{
 			byte* inend = inbuf + length;
 			byte* outptr = outbuf;
@@ -65,6 +68,9 @@ unsafe int Filter (byte* inbuf, int length, byte* outbuf)
 				pc = *inptr++;
 			}
 
+			if (flush && ensureNewLine && pc != (byte) '\n')
+				*outptr++ = (byte) '\n';
+
 			return (int) (outptr - outbuf);
 		}
 
@@ -85,15 +91,15 @@ unsafe int Filter (byte* inbuf, int length, byte* outbuf)
 		protected override byte[] Filter (byte[] input, int startIndex, int length, out int outputIndex, out int outputLength, bool flush)
 		{
 			if (pc == (byte) '\r')
-				EnsureOutputSize (length + 1, false);
+				EnsureOutputSize (length + (flush && ensureNewLine ? 2 : 1), false);
 			else
-				EnsureOutputSize (length, false);
+				EnsureOutputSize (length + (flush && ensureNewLine ? 1 : 0), false);
 
 			outputIndex = 0;
 
 			unsafe {
 				fixed (byte* inptr = input, outptr = OutputBuffer) {
-					outputLength = Filter (inptr + startIndex, length, outptr);
+					outputLength = Filter (inptr + startIndex, length, outptr, flush);
 				}
 			}
 

diff --git a/MimeKit/IO/Filters/Unix2DosFilter.cs b/MimeKit/IO/Filters/Unix2DosFilter.cs
@@ -33,6 +33,7 @@ namespace MimeKit.IO.Filters {
 	/// </remarks>
 	public class Unix2DosFilter : MimeFilterBase
 	{
+		bool ensureNewLine;
 		byte pc;
 
 		/// <summary>
@@ -41,11 +42,13 @@ public class Unix2DosFilter : MimeFilterBase
 		/// <remarks>
 		/// Creates a new <see cref="Unix2DosFilter"/>.
 		/// </remarks>
-		public Unix2DosFilter ()
+		/// <param name="ensureNewLine">Ensure that the stream ends with a new line.</param>
+		public Unix2DosFilter (bool ensureNewLine = false)
 		{
+			this.ensureNewLine = ensureNewLine;
 		}
 
-		unsafe int Filter (byte* inbuf, int length, byte* outbuf)
+		unsafe int Filter (byte* inbuf, int length, byte* outbuf, bool flush)
 		{
 			byte* inend = inbuf + length;
 			byte* outptr = outbuf;
@@ -65,6 +68,11 @@ unsafe int Filter (byte* inbuf, int length, byte* outbuf)
 				pc = *inptr++;
 			}
 
+			if (flush && ensureNewLine && pc != (byte) '\n') {
+				*outptr++ = (byte) '\r';
+				*outptr++ = (byte) '\n';
+			}
+
 			return (int) (outptr - outbuf);
 		}
 
@@ -84,13 +92,13 @@ unsafe int Filter (byte* inbuf, int length, byte* outbuf)
 		/// <param name="flush">If set to <c>true</c>, all internally buffered data should be flushed to the output buffer.</param>
 		protected override byte[] Filter (byte[] input, int startIndex, int length, out int outputIndex, out int outputLength, bool flush)
 		{
-			EnsureOutputSize (length * 2, false);
+			EnsureOutputSize (length * 2 + (flush && ensureNewLine ? 2 : 0), false);
 
 			outputIndex = 0;
 
 			unsafe {
 				fixed (byte* inptr = input, outptr = OutputBuffer) {
-					outputLength = Filter (inptr + startIndex, length, outptr);
+					outputLength = Filter (inptr + startIndex, length, outptr, flush);
 				}
 			}
 

diff --git a/MimeKit/MimePart.cs b/MimeKit/MimePart.cs
@@ -572,7 +572,7 @@ public override void Prepare (EncodingConstraint constraint, int maxLineLength =
 					filtered.Add (EncoderFilter.Create (encoding));
 
 					if (encoding != ContentEncoding.Binary)
-						filtered.Add (options.CreateNewLineFilter ());
+						filtered.Add (options.CreateNewLineFilter (true));
 
 					ContentObject.DecodeTo (filtered, cancellationToken);
 					filtered.Flush (cancellationToken);
@@ -592,7 +592,9 @@ public override void Prepare (EncodingConstraint constraint, int maxLineLength =
 				}
 			} else if (encoding != ContentEncoding.Binary) {
 				using (var filtered = new FilteredStream (stream)) {
-					filtered.Add (options.CreateNewLineFilter ());
+					// Note: if we are writing the top-level MimePart, make sure it ends with a new-line so that
+					// MimeMessage.WriteTo() *always* ends with a new-line.
+					filtered.Add (options.CreateNewLineFilter (Headers.Suppress));
 					ContentObject.WriteTo (filtered, cancellationToken);
 					filtered.Flush (cancellationToken);
 				}