Check max length when writing MemoryStream (#52)

Author: patagonaa

System.IO.Streams/MemoryStream.cs

@@ -328,6 +328,9 @@ public override long Seek(
         }
 
         /// <inheritdoc/>
+        /// <exception cref="ObjectDisposedException"></exception>
+        /// <exception cref="NotSupportedException">The stream buffer does not have the capacity to hold the data and is not expandable.</exception>
+        /// <exception cref="ArgumentOutOfRangeException">The MemoryStream max size was exceeded</exception>
         public override void SetLength(long value)
         {
             EnsureOpen();
@@ -364,6 +367,12 @@ public virtual byte[] ToArray()
         }
 
         /// <inheritdoc/>
+        /// <exception cref="ObjectDisposedException"></exception>
+        /// <exception cref="NotSupportedException">The stream buffer does not have the capacity to hold the data and is not expandable.</exception>
+        /// <exception cref="ArgumentOutOfRangeException">
+        /// The MemoryStream max size was exceeded or
+        /// <paramref name="offset"/> or <paramref name="count"/> are less than 0
+        /// </exception>
         public override void Write(byte[] buffer, int offset, int count)
         {
             EnsureOpen();
@@ -401,6 +410,9 @@ public override void Write(byte[] buffer, int offset, int count)
         }
 
         /// <inheritdoc/>
+        /// <exception cref="ObjectDisposedException"></exception>
+        /// <exception cref="NotSupportedException">The stream buffer does not have the capacity to hold the data and is not expandable.</exception>
+        /// <exception cref="ArgumentOutOfRangeException">The MemoryStream max size was exceeded</exception>
         public override void WriteByte(byte value)
         {
             EnsureOpen();
@@ -452,7 +464,8 @@ private void EnsureOpen()
         /// </summary>
         /// <param name="value">Value for the new capacity.</param>
         /// <returns><see langword="true"/> if allocation for a new array was successful. <see langword="false"/> otherwise.</returns>
-        /// <exception cref="NotSupportedException"></exception>
+        /// <exception cref="NotSupportedException">The stream buffer does not have the capacity to hold the data and is not expandable.</exception>
+        /// <exception cref="ArgumentOutOfRangeException">The MemoryStream max size would be exceeded by resizing</exception>
         private bool EnsureCapacity(int value)
         {
             if (value > _capacity)
@@ -464,11 +477,21 @@ private bool EnsureCapacity(int value)
                     newCapacity = CapacityDefaultSize;
                 }
 
+                if (value > MemStreamMaxLength)
+                {
+                    throw new ArgumentOutOfRangeException();
+                }
+
                 if (newCapacity < _capacity * 2)
                 {
                     newCapacity = _capacity * 2;
                 }
 
+                if (newCapacity > MemStreamMaxLength)
+                {
+                    newCapacity = MemStreamMaxLength;
+                }
+
                 if (!_expandable && newCapacity > _capacity)
                 {
                     throw new NotSupportedException();

UnitTests/MemoryStreamUnitTests/Write.cs

@@ -385,6 +385,30 @@ public void BoundaryCheck()
             }
         }
 
+        [TestMethod]
+        public void Write_AtMaxLength_Works()
+        {
+            var maxLength = 0xFFFF;
+            using (var ms = new MemoryStream())
+            {
+                ms.Position = maxLength - 1;
+                // should resize the buffer to the max size
+                ms.WriteByte(0);
+            }
+        }
+
+        [TestMethod]
+        public void Write_AboveMaxLength_Throws()
+        {
+            var maxLength = 0xFFFF;
+            using (var ms = new MemoryStream())
+            {
+                ms.Position = maxLength;
+                // should throw as the max size is exceeded
+                Assert.ThrowsException(typeof(ArgumentOutOfRangeException), () => ms.WriteByte(0));
+            }
+        }
+
         #endregion Test Cases
     }
 }