Fix Base64.DecodeFromUtf8 consuming whitespace in partial final quantum when isFinalBlock=false

src/libraries/System.Memory/tests/Base64/Base64DecoderUnitTests.cs

@@ -1004,6 +1004,89 @@ public void DecodingWithEmbeddedWhiteSpaceIntoSmallDestination_ActualDestination
             Assert.Equal(new byte[] { 1, 2, 3, 4 }, destination4);
         }
 
+        [Theory]
+        [InlineData("AQ\r\nQ=")]
+        [InlineData("AQ\r\nQ=\r\n")]
+        [InlineData("AQ Q=")]
+        [InlineData("AQ\tQ=")]
+        public void DecodingWithWhiteSpaceSplitFinalQuantumAndIsFinalBlockFalse(string base64String)
+        {
+            // When a final quantum (containing padding) is split by whitespace and isFinalBlock=false,
+            // the decoder should not consume any bytes, allowing the caller to retry with isFinalBlock=true
+            ReadOnlySpan<byte> base64Data = Encoding.ASCII.GetBytes(base64String);
+            var output = new byte[10];
+
+            // First call with isFinalBlock=false should consume 0 bytes
+            OperationStatus status = Base64.DecodeFromUtf8(base64Data, output, out int bytesConsumed, out int bytesWritten, isFinalBlock: false);
+            Assert.Equal(0, bytesConsumed);
+            Assert.Equal(0, bytesWritten);
+            Assert.Equal(OperationStatus.InvalidData, status);
+
+            // Second call with isFinalBlock=true should succeed
+            status = Base64.DecodeFromUtf8(base64Data, output, out bytesConsumed, out bytesWritten, isFinalBlock: true);
+            Assert.Equal(OperationStatus.Done, status);
+            Assert.Equal(base64Data.Length, bytesConsumed);
+            Assert.Equal(2, bytesWritten); // "AQQ=" decodes to 2 bytes: {1, 4}
+            Assert.Equal(new byte[] { 1, 4 }, output[..2]);
+        }
+
+        [Fact]
+        public void DecodingCompleteQuantumWithIsFinalBlockFalse()
+        {
+            // Complete quantum without padding should be decoded even when isFinalBlock=false
+            ReadOnlySpan<byte> base64Data = "AAAA"u8;
+            var output = new byte[10];
+
+            OperationStatus status = Base64.DecodeFromUtf8(base64Data, output, out int bytesConsumed, out int bytesWritten, isFinalBlock: false);
+            Assert.Equal(OperationStatus.Done, status);
+            Assert.Equal(4, bytesConsumed);
+            Assert.Equal(3, bytesWritten);
+        }
+
+        [Fact]
+        public void DecodingPaddedQuantumWithIsFinalBlockFalse()
+        {
+            // Quantum with padding should not be decoded when isFinalBlock=false
+            ReadOnlySpan<byte> base64Data = "AAA="u8;
+            var output = new byte[10];
+
+            OperationStatus status = Base64.DecodeFromUtf8(base64Data, output, out int bytesConsumed, out int bytesWritten, isFinalBlock: false);
+            Assert.Equal(OperationStatus.InvalidData, status);
+            Assert.Equal(0, bytesConsumed);
+            Assert.Equal(0, bytesWritten);
+        }
+
+        [Theory]
+        [InlineData("AQIDBAUG AQ\r\nQ=", 9, 6, "AQ\r\nQ=")]          // Two complete blocks, then whitespace-split final quantum
+        [InlineData("AQID BAUG AQ\r\nQ=", 10, 6, "AQ\r\nQ=")]        // Two blocks with space, then whitespace-split final quantum
+        [InlineData("AQIDBAUG\r\nAQID AQ\r\nQ=", 15, 9, "AQ\r\nQ=")] // Multiple blocks with various whitespace patterns
+        public void DecodingWithValidDataBeforeWhiteSpaceSplitFinalQuantum(string base64String, int expectedBytesConsumedFirstCall, int expectedBytesWrittenFirstCall, string expectedRemainingAfterFirstCall)
+        {
+            // When there's valid data before a whitespace-split final quantum and isFinalBlock=false,
+            // verify the streaming scenario works correctly
+            ReadOnlySpan<byte> base64Data = Encoding.ASCII.GetBytes(base64String);
+            var output = new byte[100];
+
+            // First call with isFinalBlock=false should decode the valid complete blocks and stop before the incomplete final quantum
+            OperationStatus status = Base64.DecodeFromUtf8(base64Data, output, out int bytesConsumed, out int bytesWritten, isFinalBlock: false);
+
+            Assert.Equal(OperationStatus.InvalidData, status);
+            Assert.Equal(expectedBytesConsumedFirstCall, bytesConsumed);
+            Assert.Equal(expectedBytesWrittenFirstCall, bytesWritten);
+
+            // Verify that only the final block remains
+            ReadOnlySpan<byte> remaining = base64Data.Slice(bytesConsumed);
+            string remainingString = Encoding.ASCII.GetString(remaining);
+            Assert.Equal(expectedRemainingAfterFirstCall, remainingString);
+
+            // Verify we can complete decoding by retrying with the FULL input and isFinalBlock=true
+            Array.Clear(output, 0, output.Length);
+            status = Base64.DecodeFromUtf8(base64Data, output, out bytesConsumed, out bytesWritten, isFinalBlock: true);
+            Assert.Equal(OperationStatus.Done, status);
+            Assert.Equal(base64Data.Length, bytesConsumed);
+            Assert.True(bytesWritten > 0, "Should have decoded data");
+        }
+
         [Fact]
         public void DecodingWithEmbeddedWhiteSpaceIntoSmallDestination_TrailingWhiteSpacesAreConsumed()
         {
@@ -1020,5 +1103,88 @@ public void DecodingWithEmbeddedWhiteSpaceIntoSmallDestination_TrailingWhiteSpac
             Assert.Equal(destination.Length, written);
             Assert.Equal(new byte[] { 240, 159, 141, 137, 240, 159 }, destination);
         }
+
+        [Theory]
+        [InlineData("AQ\r\nQ=")]
+        [InlineData("AQ\r\nQ=\r\n")]
+        [InlineData("AQ Q=")]
+        [InlineData("AQ\tQ=")]
+        public void DecodingFromCharsWithWhiteSpaceSplitFinalQuantumAndIsFinalBlockFalse(string base64String)
+        {
+            // When a final quantum (containing padding) is split by whitespace and isFinalBlock=false,
+            // the decoder should not consume any bytes, allowing the caller to retry with isFinalBlock=true
+            ReadOnlySpan<char> base64Data = base64String.AsSpan();
+            var output = new byte[10];
+
+            // First call with isFinalBlock=false should consume 0 bytes
+            OperationStatus status = Base64.DecodeFromChars(base64Data, output, out int bytesConsumed, out int bytesWritten, isFinalBlock: false);
+            Assert.Equal(0, bytesConsumed);
+            Assert.Equal(0, bytesWritten);
+            Assert.Equal(OperationStatus.InvalidData, status);
+
+            // Second call with isFinalBlock=true should succeed
+            status = Base64.DecodeFromChars(base64Data, output, out bytesConsumed, out bytesWritten, isFinalBlock: true);
+            Assert.Equal(OperationStatus.Done, status);
+            Assert.Equal(base64Data.Length, bytesConsumed);
+            Assert.Equal(2, bytesWritten); // "AQQ=" decodes to 2 bytes: {1, 4}
+            Assert.Equal(new byte[] { 1, 4 }, output[..2]);
+        }
+
+        [Fact]
+        public void DecodingFromCharsCompleteQuantumWithIsFinalBlockFalse()
+        {
+            // Complete quantum without padding should be decoded even when isFinalBlock=false
+            ReadOnlySpan<char> base64Data = "AAAA".AsSpan();
+            var output = new byte[10];
+
+            OperationStatus status = Base64.DecodeFromChars(base64Data, output, out int bytesConsumed, out int bytesWritten, isFinalBlock: false);
+            Assert.Equal(OperationStatus.Done, status);
+            Assert.Equal(4, bytesConsumed);
+            Assert.Equal(3, bytesWritten);
+        }
+
+        [Fact]
+        public void DecodingFromCharsPaddedQuantumWithIsFinalBlockFalse()
+        {
+            // Quantum with padding should not be decoded when isFinalBlock=false
+            ReadOnlySpan<char> base64Data = "AAA=".AsSpan();
+            var output = new byte[10];
+
+            OperationStatus status = Base64.DecodeFromChars(base64Data, output, out int bytesConsumed, out int bytesWritten, isFinalBlock: false);
+            Assert.Equal(OperationStatus.InvalidData, status);
+            Assert.Equal(0, bytesConsumed);
+            Assert.Equal(0, bytesWritten);
+        }
+
+        [Theory]
+        [InlineData("AQIDBAUG AQ\r\nQ=", 9, 6, "AQ\r\nQ=")]          // Two complete blocks, then whitespace-split final quantum
+        [InlineData("AQID BAUG AQ\r\nQ=", 10, 6, "AQ\r\nQ=")]        // Two blocks with space, then whitespace-split final quantum
+        [InlineData("AQIDBAUG\r\nAQID AQ\r\nQ=", 15, 9, "AQ\r\nQ=")] // Multiple blocks with various whitespace patterns
+        public void DecodingFromCharsWithValidDataBeforeWhiteSpaceSplitFinalQuantum(string base64String, int expectedBytesConsumedFirstCall, int expectedBytesWrittenFirstCall, string expectedRemainingAfterFirstCall)
+        {
+            // When there's valid data before a whitespace-split final quantum and isFinalBlock=false,
+            // verify the streaming scenario works correctly
+            ReadOnlySpan<char> base64Data = base64String.AsSpan();
+            var output = new byte[100];
+
+            // First call with isFinalBlock=false should decode the valid complete blocks and stop before the incomplete final quantum
+            OperationStatus status = Base64.DecodeFromChars(base64Data, output, out int bytesConsumed, out int bytesWritten, isFinalBlock: false);
+
+            Assert.Equal(OperationStatus.InvalidData, status);
+            Assert.Equal(expectedBytesConsumedFirstCall, bytesConsumed);
+            Assert.Equal(expectedBytesWrittenFirstCall, bytesWritten);
+
+            // Verify that only the final block remains
+            ReadOnlySpan<char> remaining = base64Data.Slice(bytesConsumed);
+            string remainingString = new string(remaining);
+            Assert.Equal(expectedRemainingAfterFirstCall, remainingString);
+
+            // Verify we can complete decoding by retrying with the FULL input and isFinalBlock=true
+            Array.Clear(output, 0, output.Length);
+            status = Base64.DecodeFromChars(base64Data, output, out bytesConsumed, out bytesWritten, isFinalBlock: true);
+            Assert.Equal(OperationStatus.Done, status);
+            Assert.Equal(base64Data.Length, bytesConsumed);
+            Assert.True(bytesWritten > 0, "Should have decoded data");
+        }
     }
 }

src/libraries/System.Private.CoreLib/src/System/Buffers/Text/Base64Helper/Base64DecoderHelper.cs

@@ -467,6 +467,14 @@ internal static OperationStatus DecodeWithWhiteSpaceBlockwise<TBase64Decoder>(TB
 
             while (!source.IsEmpty)
             {
+                // Skip over any leading whitespace
+                if (IsWhiteSpace(source[0]))
+                {
+                    source = source.Slice(1);
+                    bytesConsumed++;
+                    continue;
+                }
+
                 int encodedIdx = 0;
                 int bufferIdx = 0;
                 int skipped = 0;
@@ -485,12 +493,7 @@ internal static OperationStatus DecodeWithWhiteSpaceBlockwise<TBase64Decoder>(TB
                 }
 
                 source = source.Slice(encodedIdx);
-                bytesConsumed += skipped;
-
-                if (bufferIdx == 0)
-                {
-                    continue;
-                }
+                Debug.Assert(bufferIdx > 0);
 
                 bool hasAnotherBlock;
 
@@ -522,14 +525,17 @@ internal static OperationStatus DecodeWithWhiteSpaceBlockwise<TBase64Decoder>(TB
                 }
 
                 status = DecodeFrom<TBase64Decoder, byte>(decoder, buffer.Slice(0, bufferIdx), bytes, out int localConsumed, out int localWritten, localIsFinalBlock, ignoreWhiteSpace: false);
-                bytesConsumed += localConsumed;
-                bytesWritten += localWritten;
 
                 if (status != OperationStatus.Done)
                 {
+                    Debug.Assert(localConsumed == 0 && localWritten == 0, "On failure, should not have consumed or written any bytes");
                     return status;
                 }
 
+                bytesConsumed += skipped;
+                bytesConsumed += localConsumed;
+                bytesWritten += localWritten;
+
                 // The remaining data must all be whitespace in order to be valid.
                 if (!hasAnotherBlock)
                 {
@@ -551,6 +557,7 @@ internal static OperationStatus DecodeWithWhiteSpaceBlockwise<TBase64Decoder>(TB
                 }
 
                 bytes = bytes.Slice(localWritten);
+                Debug.Assert(!source.IsEmpty);
             }
 
             return status;
@@ -565,6 +572,14 @@ internal static OperationStatus DecodeWithWhiteSpaceBlockwise<TBase64Decoder>(TB
 
             while (!source.IsEmpty)
             {
+                // Skip over any leading whitespace
+                if (IsWhiteSpace(source[0]))
+                {
+                    source = source.Slice(1);
+                    bytesConsumed++;
+                    continue;
+                }
+
                 int encodedIdx = 0;
                 int bufferIdx = 0;
                 int skipped = 0;
@@ -583,12 +598,7 @@ internal static OperationStatus DecodeWithWhiteSpaceBlockwise<TBase64Decoder>(TB
                 }
 
                 source = source.Slice(encodedIdx);
-                bytesConsumed += skipped;
-
-                if (bufferIdx == 0)
-                {
-                    continue;
-                }
+                Debug.Assert(bufferIdx > 0);
 
                 bool hasAnotherBlock;
 
@@ -620,14 +630,17 @@ internal static OperationStatus DecodeWithWhiteSpaceBlockwise<TBase64Decoder>(TB
                 }
 
                 status = DecodeFrom(decoder, buffer.Slice(0, bufferIdx), bytes, out int localConsumed, out int localWritten, localIsFinalBlock, ignoreWhiteSpace: false);
-                bytesConsumed += localConsumed;
-                bytesWritten += localWritten;
 
                 if (status != OperationStatus.Done)
                 {
+                    Debug.Assert(localConsumed == 0 && localWritten == 0, "On failure, should not have consumed or written any bytes");
                     return status;
                 }
 
+                bytesConsumed += skipped;
+                bytesConsumed += localConsumed;
+                bytesWritten += localWritten;
+
                 // The remaining data must all be whitespace in order to be valid.
                 if (!hasAnotherBlock)
                 {
@@ -648,6 +661,7 @@ internal static OperationStatus DecodeWithWhiteSpaceBlockwise<TBase64Decoder>(TB
                 }
 
                 bytes = bytes.Slice(localWritten);
+                Debug.Assert(!source.IsEmpty);
             }
 
             return status;