Clear buffers before returning them to the ArrayPool (#124)

It is possible that sensitive data such as PII is being compressed or
decompressed using Snappier. We don't want to return such data to the
ArrayPool without zeroing it first, as it could create a security
vulnerability if the buffer is reused by some other portion of an
application that isn't properly handling the buffer.

In some cases, we request the clear as part of the return. However, in
others we know we've only used a subset of the buffer so we can optimize
by only clearing the portion we've used.

This change also removes some unnecessary try..finally blocks to return
arrays to the pool during compression. Compression doesn't typically
throw exceptions, and in any extreme corner cases we'll simply not
return the array to the pool. This simplifies the code and provides a
minor performance improvement.

Author: brantburnett

Snappier/Internal/ByteArrayPoolMemoryOwner.cs

@@ -45,6 +45,9 @@ public void Dispose()
             byte[]? innerArray = _innerArray;
             if (innerArray is not null)
             {
+                // Clear the used portion of the array before returning it to the pool
+                Memory.Span.Clear();
+
                 _innerArray = null;
                 Memory = default;
                 ArrayPool<byte>.Shared.Return(innerArray);

Snappier/Internal/Helpers.cs

@@ -1,8 +1,10 @@
 using System;
+using System.Buffers;
 using System.Buffers.Binary;
 using System.Diagnostics;
 using System.Runtime.CompilerServices;
 using System.Runtime.InteropServices;
+
 #if NET6_0_OR_GREATER
 using System.Numerics;
 using System.Runtime.Intrinsics;
@@ -45,6 +47,23 @@ public static int MaxCompressedLength(int sourceBytes)
             return 32 + sourceBytes + sourceBytes / 6 + 1;
         }
 
+        // Constant for MaxCompressedLength when passed Constants.BlockSize, keep this in sync with the above method
+        public const int MaxBlockCompressedLength = (int)(32 + Constants.BlockSize + Constants.BlockSize / 6 + 1);
+
+        /// <summary>
+        /// Clears the array and returns it to the pool, clearing only the used portion of the array.
+        /// This is a minor performance optimization to avoid clearing the entire array.
+        /// </summary>
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        public static void ClearAndReturn(byte[] array, int usedLength)
+        {
+            Debug.Assert(array is not null);
+            Debug.Assert(usedLength >= 0);
+
+            array.AsSpan(0, usedLength).Clear();
+            ArrayPool<byte>.Shared.Return(array);
+        }
+
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         public static bool LeftShiftOverflows(byte value, int shift)
         {

Snappier/Internal/SnappyCompressor.cs

@@ -60,21 +60,18 @@ public bool TryCompress(ReadOnlySpan<byte> input, Span<byte> output, out int byt
                     // compress to a temporary buffer and copy the compressed data to the output span.
 
                     byte[] scratch = ArrayPool<byte>.Shared.Rent(maxOutput);
-                    try
+                    written = CompressFragment(fragment, scratch.AsSpan(), hashTable);
+                    if (output.Length < written)
                     {
-                        written = CompressFragment(fragment, scratch.AsSpan(), hashTable);
-                        if (output.Length < written)
-                        {
-                            bytesWritten = 0;
-                            return false;
-                        }
-
-                        scratch.AsSpan(0, written).CopyTo(output);
-                    }
-                    finally
-                    {
-                        ArrayPool<byte>.Shared.Return(scratch);
+                        Helpers.ClearAndReturn(scratch, written);
+                        bytesWritten = 0;
+                        return false;
                     }
+
+                    Span<byte> writtenScratch = scratch.AsSpan(0, written);
+                    writtenScratch.CopyTo(output);
+                    writtenScratch.Clear();
+                    ArrayPool<byte>.Shared.Return(scratch);
                 }
 
                 output = output.Slice(written);
@@ -132,19 +129,17 @@ public void Compress(ReadOnlySequence<byte> input, IBufferWriter<byte> bufferWri
 
                     int fragmentLength = (int)fragment.Length;
                     byte[] scratch = ArrayPool<byte>.Shared.Rent(fragmentLength);
-                    try
-                    {
-                        fragment.CopyTo(scratch);
 
-                        CompressFragment(scratch.AsSpan(0, fragmentLength), bufferWriter);
+                    Span<byte> usedScratch = scratch.AsSpan(0, fragmentLength);
+                    fragment.CopyTo(usedScratch);
 
-                        // Advance the length of the entire fragment
-                        input = input.Slice(position);
-                    }
-                    finally
-                    {
-                        ArrayPool<byte>.Shared.Return(scratch);
-                    }
+                    CompressFragment(usedScratch, bufferWriter);
+
+                    usedScratch.Clear();
+                    ArrayPool<byte>.Shared.Return(scratch);
+
+                    // Advance the length of the entire fragment
+                    input = input.Slice(position);
                 }
             }
         }

Snappier/Internal/SnappyDecompressor.cs

@@ -526,7 +526,8 @@ private int? ExpectedLength
                 {
                     if (_lookbackBufferArray is not null)
                     {
-                        ArrayPool<byte>.Shared.Return(_lookbackBufferArray);
+                        // Clear the used portion of the lookback buffer before returning
+                        Helpers.ClearAndReturn(_lookbackBufferArray, _lookbackPosition);
                     }
 
                     if (BufferWriter is not null)
@@ -731,7 +732,9 @@ public void Dispose()
         {
             if (_lookbackBufferArray is not null)
             {
-                ArrayPool<byte>.Shared.Return(_lookbackBufferArray);
+                // Clear the used portion of the lookback buffer before returning
+                Helpers.ClearAndReturn(_lookbackBufferArray, _lookbackPosition);
+
                 _lookbackBufferArray = null;
                 _lookbackBuffer = default;
             }

Snappier/Internal/SnappyStreamCompressor.cs

@@ -276,7 +276,7 @@ private void EnsureBuffer()
         {
             // Allocate enough room for the stream header and block headers
             _outputBuffer ??=
-                ArrayPool<byte>.Shared.Rent(Helpers.MaxCompressedLength((int) Constants.BlockSize) + 8 + SnappyHeader.Length);
+                ArrayPool<byte>.Shared.Rent(Helpers.MaxBlockCompressedLength + 8 + SnappyHeader.Length);
 
             // Allocate enough room for the stream header and block headers
             _inputBuffer ??= ArrayPool<byte>.Shared.Rent((int) Constants.BlockSize);
@@ -289,12 +289,12 @@ public void Dispose()
 
             if (_outputBuffer is not null)
             {
-                ArrayPool<byte>.Shared.Return(_outputBuffer);
+                ArrayPool<byte>.Shared.Return(_outputBuffer, clearArray: true);
                 _outputBuffer = null;
             }
             if (_inputBuffer is not null)
             {
-                ArrayPool<byte>.Shared.Return(_inputBuffer);
+                ArrayPool<byte>.Shared.Return(_inputBuffer, clearArray: true);
                 _inputBuffer = null;
             }
         }

Snappier/Snappy.cs

@@ -102,21 +102,16 @@ public static IMemoryOwner<byte> CompressToMemory(ReadOnlySpan<byte> input)
         {
             byte[] buffer = ArrayPool<byte>.Shared.Rent(GetMaxCompressedLength(input.Length));
 
-            try
+            if (!TryCompress(input, buffer, out int length))
             {
-                if (!TryCompress(input, buffer, out int length))
-                {
-                    // Should be unreachable since we're allocating a buffer of the correct size.
-                    ThrowHelper.ThrowInvalidOperationException();
-                }
+                // The amount of data written is unknown, so clear the entire buffer when returning
+                ArrayPool<byte>.Shared.Return(buffer, clearArray: true);
 
-                return new ByteArrayPoolMemoryOwner(buffer, length);
-            }
-            catch
-            {
-                ArrayPool<byte>.Shared.Return(buffer);
-                throw;
+                // Should be unreachable since we're allocating a buffer of the correct size.
+                ThrowHelper.ThrowInvalidOperationException();
             }
+
+            return new ByteArrayPoolMemoryOwner(buffer, length);
         }
 
         /// <summary>

Snappier/SnappyStream.cs

@@ -494,7 +494,7 @@ protected override void Dispose(bool disposing)
                             _buffer = null;
                             if (!AsyncOperationIsActive)
                             {
-                                ArrayPool<byte>.Shared.Return(buffer);
+                                ArrayPool<byte>.Shared.Return(buffer, clearArray: true);
                             }
                         }
 
@@ -545,7 +545,7 @@ public override async ValueTask DisposeAsync()
                             _buffer = null;
                             if (!AsyncOperationIsActive)
                             {
-                                ArrayPool<byte>.Shared.Return(buffer);
+                                ArrayPool<byte>.Shared.Return(buffer, clearArray: true);
                             }
                         }
                     }