More Span-based Stream methods

3 files changed, 142 insertions, 32 deletions

diff --git a/crypto/src/bcpg/ArmoredInputStream.cs b/crypto/src/bcpg/ArmoredInputStream.cs
index 1c5ebd7c5..4fbb8baae 100644
--- a/crypto/src/bcpg/ArmoredInputStream.cs
+++ b/crypto/src/bcpg/ArmoredInputStream.cs
@@ -1,3 +1,4 @@
+using System;
 using System.Collections.Generic;
 using System.IO;
 using System.Text;
@@ -330,6 +331,26 @@ namespace Org.BouncyCastle.Bcpg
             return pos;
         }
 
+#if NETCOREAPP2_1_OR_GREATER || NETSTANDARD2_1_OR_GREATER
+        public override int Read(Span<byte> buffer)
+        {
+            /*
+             * TODO Currently can't return partial data when exception thrown (breaking test case), so we don't inherit
+             * the base class implementation. Probably the reason is that throws don't mark this instance as 'failed'.
+             */
+            int pos = 0;
+            while (pos < buffer.Length)
+            {
+                int b = ReadByte();
+                if (b < 0)
+                    break;
+
+                buffer[pos++] = (byte)b;
+            }
+            return pos;
+        }
+#endif
+
         public override int ReadByte()
         {
             if (start)
diff --git a/crypto/src/bcpg/BcpgInputStream.cs b/crypto/src/bcpg/BcpgInputStream.cs
index 7a19a90dd..3b6f61bbc 100644
--- a/crypto/src/bcpg/BcpgInputStream.cs
+++ b/crypto/src/bcpg/BcpgInputStream.cs
@@ -57,6 +57,21 @@ namespace Org.BouncyCastle.Bcpg
             return 1;
         }
 
+#if NETCOREAPP2_1_OR_GREATER || NETSTANDARD2_1_OR_GREATER
+        public override int Read(Span<byte> buffer)
+        {
+			if (!next)
+				return m_in.Read(buffer);
+
+			if (nextB < 0)
+				return 0;
+
+            buffer[0] = (byte)nextB;
+            next = false;
+            return 1;
+        }
+#endif
+
         public byte[] ReadAll()
         {
 			return Streams.ReadAll(this);
@@ -312,9 +327,8 @@ namespace Org.BouncyCastle.Bcpg
 						int readLen = (dataLength > count || dataLength < 0) ? count : dataLength;
 						int len = m_in.Read(buffer, offset, readLen);
 						if (len < 1)
-						{
 							throw new EndOfStreamException("Premature end of stream in PartialInputStream");
-						}
+
 						dataLength -= len;
 						return len;
 					}
@@ -324,6 +338,29 @@ namespace Org.BouncyCastle.Bcpg
 				return 0;
 			}
 
+#if NETCOREAPP2_1_OR_GREATER || NETSTANDARD2_1_OR_GREATER
+            public override int Read(Span<byte> buffer)
+            {
+				do
+				{
+					if (dataLength != 0)
+					{
+                        int count = buffer.Length;
+						int readLen = (dataLength > count || dataLength < 0) ? count : dataLength;
+						int len = m_in.Read(buffer[..readLen]);
+						if (len < 1)
+							throw new EndOfStreamException("Premature end of stream in PartialInputStream");
+
+						dataLength -= len;
+						return len;
+					}
+				}
+				while (partial && ReadPartialDataLength() >= 0);
+
+				return 0;
+            }
+#endif
+
             private int ReadPartialDataLength()
             {
                 int l = m_in.ReadByte();
diff --git a/crypto/src/bcpg/BcpgOutputStream.cs b/crypto/src/bcpg/BcpgOutputStream.cs
index fbdd75bff..690686d88 100644
--- a/crypto/src/bcpg/BcpgOutputStream.cs
+++ b/crypto/src/bcpg/BcpgOutputStream.cs
@@ -164,7 +164,7 @@ namespace Org.BouncyCastle.Bcpg
 
             if (partialBuffer != null)
             {
-                PartialFlush(true);
+                PartialFlushLast();
                 partialBuffer = null;
             }
 
@@ -215,19 +215,26 @@ namespace Org.BouncyCastle.Bcpg
             }
         }
 
-        private void PartialFlush(bool isLast)
+        private void PartialFlush()
         {
-            if (isLast)
-            {
-                WriteNewPacketLength(partialOffset);
-                outStr.Write(partialBuffer, 0, partialOffset);
-            }
-            else
-            {
-                outStr.WriteByte((byte)(0xE0 | partialPower));
-                outStr.Write(partialBuffer, 0, partialBufferLength);
-            }
+            outStr.WriteByte((byte)(0xE0 | partialPower));
+            outStr.Write(partialBuffer, 0, partialBufferLength);
+            partialOffset = 0;
+        }
+
+#if NETCOREAPP2_1_OR_GREATER || NETSTANDARD2_1_OR_GREATER
+        private void PartialFlush(ref ReadOnlySpan<byte> buffer)
+        {
+            outStr.WriteByte((byte)(0xE0 | partialPower));
+            outStr.Write(buffer[..partialBufferLength]);
+            buffer = buffer[partialBufferLength..];
+        }
+#endif
 
+        private void PartialFlushLast()
+        {
+            WriteNewPacketLength(partialOffset);
+            outStr.Write(partialBuffer, 0, partialOffset);
             partialOffset = 0;
         }
 
@@ -235,40 +242,71 @@ namespace Org.BouncyCastle.Bcpg
         {
             Streams.ValidateBufferArguments(buffer, offset, count);
 
+#if NETCOREAPP2_1_OR_GREATER || NETSTANDARD2_1_OR_GREATER
+            PartialWrite(buffer.AsSpan(offset, count));
+#else
             if (partialOffset == partialBufferLength)
             {
-                PartialFlush(false);
+                PartialFlush();
             }
 
             if (count <= (partialBufferLength - partialOffset))
             {
                 Array.Copy(buffer, offset, partialBuffer, partialOffset, count);
                 partialOffset += count;
+                return;
             }
-            else
+
+            int diff = partialBufferLength - partialOffset;
+            Array.Copy(buffer, offset, partialBuffer, partialOffset, diff);
+            offset += diff;
+            count -= diff;
+            PartialFlush();
+            while (count > partialBufferLength)
             {
-                int diff = partialBufferLength - partialOffset;
-                Array.Copy(buffer, offset, partialBuffer, partialOffset, diff);
-                offset += diff;
-                count -= diff;
-                PartialFlush(false);
-                while (count > partialBufferLength)
-                {
-                    Array.Copy(buffer, offset, partialBuffer, 0, partialBufferLength);
-                    offset += partialBufferLength;
-                    count -= partialBufferLength;
-                    PartialFlush(false);
-                }
-                Array.Copy(buffer, offset, partialBuffer, 0, count);
-                partialOffset += count;
+                Array.Copy(buffer, offset, partialBuffer, 0, partialBufferLength);
+                offset += partialBufferLength;
+                count -= partialBufferLength;
+                PartialFlush();
             }
+            Array.Copy(buffer, offset, partialBuffer, 0, count);
+            partialOffset = count;
+#endif
         }
 
+#if NETCOREAPP2_1_OR_GREATER || NETSTANDARD2_1_OR_GREATER
+        private void PartialWrite(ReadOnlySpan<byte> buffer)
+        {
+            if (partialOffset == partialBufferLength)
+            {
+                PartialFlush();
+            }
+
+            if (buffer.Length <= (partialBufferLength - partialOffset))
+            {
+                buffer.CopyTo(partialBuffer.AsSpan(partialOffset));
+                partialOffset += buffer.Length;
+                return;
+            }
+
+            int diff = partialBufferLength - partialOffset;
+            buffer[..diff].CopyTo(partialBuffer.AsSpan(partialOffset));
+            buffer = buffer[diff..];
+            PartialFlush();
+            while (buffer.Length > partialBufferLength)
+            {
+                PartialFlush(ref buffer);
+            }
+            buffer.CopyTo(partialBuffer);
+            partialOffset = buffer.Length;
+        }
+#endif
+
         private void PartialWriteByte(byte value)
         {
             if (partialOffset == partialBufferLength)
             {
-                PartialFlush(false);
+                PartialFlush();
             }
 
             partialBuffer[partialOffset++] = value;
@@ -286,6 +324,20 @@ namespace Org.BouncyCastle.Bcpg
             }
         }
 
+#if NETCOREAPP2_1_OR_GREATER || NETSTANDARD2_1_OR_GREATER
+        public override void Write(ReadOnlySpan<byte> buffer)
+        {
+            if (partialBuffer != null)
+            {
+                PartialWrite(buffer);
+            }
+            else
+            {
+                outStr.Write(buffer);
+            }
+        }
+#endif
+
         public override void WriteByte(byte value)
         {
             if (partialBuffer != null)
@@ -370,7 +422,7 @@ namespace Org.BouncyCastle.Bcpg
         {
             if (partialBuffer != null)
             {
-                PartialFlush(true);
+                PartialFlushLast();
                 Array.Clear(partialBuffer, 0, partialBuffer.Length);
                 partialBuffer = null;
             }