Reliable clean-up of SafeHandles in Deflater/Inflater

src/libraries/Common/src/System/IO/Compression/ZLibNative.ZStream.cs

@@ -10,7 +10,7 @@ internal static partial class ZLibNative
         /// <summary>
         /// ZLib stream descriptor data structure
         /// Do not construct instances of <see cref="ZStream" /> explicitly.
-        /// Always use <see cref="ZLibNative.CreateZLibStreamForDeflate" /> or <see cref="ZLibNative.CreateZLibStreamForInflate" /> instead.
+        /// Always use <see cref="ZLibNative.ZLibStreamHandle.CreateForDeflate" /> or <see cref="ZLibNative.ZLibStreamHandle.CreateForInflate" /> instead.
         /// Those methods will wrap this structure into a <see cref="ZLibStreamHandle" /> and thus make sure that it is always disposed correctly.
         /// </summary>
         [StructLayout(LayoutKind.Sequential, CharSet = CharSet.Ansi)]

src/libraries/Common/src/System/IO/Compression/ZLibNative.cs

@@ -1,6 +1,7 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
+using System.Diagnostics;
 using System.Runtime.InteropServices;
 using System.Runtime.InteropServices.Marshalling;
 using System.Security;
@@ -173,6 +174,41 @@ public ZLibStreamHandle()
                 SetHandle(IntPtr.Zero);
             }
 
+            public static ZLibStreamHandle CreateForDeflate(CompressionLevel level, int windowBits,
+                int memLevel, CompressionStrategy strategy)
+            {
+                ZLibStreamHandle zLibStreamHandle = new ZLibStreamHandle();
+
+                try
+                {
+                    zLibStreamHandle.DeflateInit2_(level, windowBits, memLevel, strategy);
+                }
+                catch (Exception)
+                {
+                    zLibStreamHandle.Dispose();
+                    throw;
+                }
+
+                return zLibStreamHandle;
+            }
+
+            public static ZLibStreamHandle CreateForInflate(int windowBits)
+            {
+                ZLibStreamHandle zLibStreamHandle = new ZLibStreamHandle();
+
+                try
+                {
+                    zLibStreamHandle.InflateInit2_(windowBits);
+                }
+                catch (Exception)
+                {
+                    zLibStreamHandle.Dispose();
+                    throw;
+                }
+
+                return zLibStreamHandle;
+            }
+
             public override bool IsInvalid
             {
                 get { return handle == new IntPtr(-1); }
@@ -228,20 +264,53 @@ private void EnsureState(State requiredState)
                     throw new InvalidOperationException("InitializationState != " + requiredState.ToString());
             }
 
-            public unsafe ErrorCode DeflateInit2_(CompressionLevel level, int windowBits, int memLevel, CompressionStrategy strategy)
+            private void EnsureNativeHandleInitialized(ErrorCode zlibErrorCode, string zlibErrorContext)
             {
-                EnsureNotDisposed();
-                EnsureState(State.NotInitialized);
+                Debug.Assert(zlibErrorContext is "deflateInit2_" or "inflateInit2_");
 
-                fixed (ZStream* stream = &_zStream)
+                if (zlibErrorCode is not ErrorCode.Ok)
                 {
-                    ErrorCode errC = Interop.ZLib.DeflateInit2_(stream, level, CompressionMethod.Deflated, windowBits, memLevel, strategy);
-                    _initializationState = State.InitializedForDeflate;
+                    string zlibErrorMessage = GetErrorMessage();
+                    string exceptionMessage = zlibErrorCode switch
+                    {
+                        // Not enough memory
+                        ErrorCode.MemError => SR.ZLibErrorNotEnoughMemory,
 
-                    return errC;
+                        // zlib library is incompatible with the version assumed
+                        ErrorCode.VersionError => SR.ZLibErrorVersionMismatch,
+
+                        // Parameters are invalid
+                        ErrorCode.StreamError => SR.ZLibErrorIncorrectInitParameters,
+
+                        _ => SR.Format(SR.ZLibErrorUnexpected, (int)zlibErrorCode)
+                    };
+
+                    throw new ZLibException(exceptionMessage, zlibErrorContext, (int)zlibErrorCode, zlibErrorMessage);
                 }
             }
 
+            private unsafe void DeflateInit2_(CompressionLevel level, int windowBits, int memLevel, CompressionStrategy strategy)
+            {
+                Debug.Assert(InitializationState == State.NotInitialized);
+
+                ErrorCode errC;
+
+                try
+                {
+                    fixed (ZStream* stream = &_zStream)
+                    {
+                        errC = Interop.ZLib.DeflateInit2_(stream, level, CompressionMethod.Deflated, windowBits, memLevel, strategy);
+                    }
+                }
+                catch (Exception e) // Could not load the ZLib dll
+                {
+                    throw new ZLibException(SR.ZLibErrorDLLLoadError, e);
+                }
+
+                EnsureNativeHandleInitialized(errC, "deflateInit2_");
+                _initializationState = State.InitializedForDeflate;
+            }
+
             public unsafe ErrorCode Deflate(FlushCode flush)
             {
                 EnsureNotDisposed();
@@ -267,18 +336,26 @@ public unsafe ErrorCode DeflateEnd()
                 }
             }
 
-            public unsafe ErrorCode InflateInit2_(int windowBits)
+            private unsafe void InflateInit2_(int windowBits)
             {
-                EnsureNotDisposed();
-                EnsureState(State.NotInitialized);
+                Debug.Assert(InitializationState == State.NotInitialized);
 
-                fixed (ZStream* stream = &_zStream)
-                {
-                    ErrorCode errC = Interop.ZLib.InflateInit2_(stream, windowBits);
-                    _initializationState = State.InitializedForInflate;
+                ErrorCode errC;
 
-                    return errC;
+                try
+                {
+                    fixed (ZStream* stream = &_zStream)
+                    {
+                        errC = Interop.ZLib.InflateInit2_(stream, windowBits);
+                    }
                 }
+                catch (Exception e) // Could not load the ZLib dll
+                {
+                    throw new ZLibException(SR.ZLibErrorDLLLoadError, e);
+                }
+
+                EnsureNativeHandleInitialized(errC, "inflateInit2_");
+                _initializationState = State.InitializedForInflate;
             }
 
             public unsafe ErrorCode InflateReset2_(int windowBits)
@@ -317,21 +394,8 @@ public unsafe ErrorCode InflateEnd()
                 }
             }
 
-            // This can work even after XxflateEnd().
+            // This can work even after XxflateEnd(). Gets the error message from the native library.
             public unsafe string GetErrorMessage() => Utf8StringMarshaller.ConvertToManaged(_zStream.msg) ?? string.Empty;
         }
-
-        public static ErrorCode CreateZLibStreamForDeflate(out ZLibStreamHandle zLibStreamHandle, CompressionLevel level,
-            int windowBits, int memLevel, CompressionStrategy strategy)
-        {
-            zLibStreamHandle = new ZLibStreamHandle();
-            return zLibStreamHandle.DeflateInit2_(level, windowBits, memLevel, strategy);
-        }
-
-        public static ErrorCode CreateZLibStreamForInflate(out ZLibStreamHandle zLibStreamHandle, int windowBits)
-        {
-            zLibStreamHandle = new ZLibStreamHandle();
-            return zLibStreamHandle.InflateInit2_(windowBits);
-        }
     }
 }

src/libraries/System.IO.Compression/src/Resources/Strings.resx

@@ -168,7 +168,7 @@
     <value>The version of the underlying compression routine does not match expected version.</value>
   </data>
   <data name="ZLibErrorUnexpected" xml:space="preserve">
-    <value>The underlying compression routine returned an unexpected error code.</value>
+    <value>The underlying compression routine returned an unexpected error code: '{0}'.</value>
   </data>
   <data name="CDCorrupt" xml:space="preserve">
     <value>Central Directory corrupt.</value>

src/libraries/System.IO.Compression/src/System/IO/Compression/DeflateZLib/DeflateStream.cs

@@ -79,7 +79,7 @@ internal DeflateStream(Stream stream, CompressionMode mode, bool leaveOpen, int
                     if (!stream.CanRead)
                         throw new ArgumentException(SR.NotSupported_UnreadableStream, nameof(stream));
 
-                    _inflater = new Inflater(windowBits, uncompressedSize);
+                    _inflater = Inflater.CreateInflater(windowBits, uncompressedSize);
                     _stream = stream;
                     _mode = CompressionMode.Decompress;
                     _leaveOpen = leaveOpen;
@@ -114,7 +114,7 @@ internal void InitializeDeflater(Stream stream, ZLibNative.CompressionLevel comp
             if (!stream.CanWrite)
                 throw new ArgumentException(SR.NotSupported_UnwritableStream, nameof(stream));
 
-            _deflater = new Deflater(compressionLevel, strategy, windowBits, GetMemLevel(compressionLevel));
+            _deflater = Deflater.CreateDeflater(compressionLevel, strategy, windowBits, GetMemLevel(compressionLevel));
 
             _stream = stream;
             _mode = CompressionMode.Compress;

src/libraries/System.IO.Compression/src/System/IO/Compression/DeflateZLib/Deflater.cs

@@ -3,7 +3,6 @@
 
 using System.Buffers;
 using System.Diagnostics;
-using System.Security;
 
 using ZErrorCode = System.IO.Compression.ZLibNative.ErrorCode;
 using ZFlushCode = System.IO.Compression.ZLibNative.FlushCode;
@@ -28,37 +27,9 @@ internal sealed class Deflater : IDisposable
         // on the stream explicitly.
         private object SyncLock => this;
 
-        internal Deflater(ZLibNative.CompressionLevel compressionLevel, ZLibNative.CompressionStrategy strategy, int windowBits, int memLevel)
+        private Deflater(ZLibNative.ZLibStreamHandle zlibStream)
         {
-            Debug.Assert(windowBits >= minWindowBits && windowBits <= maxWindowBits);
-
-            ZErrorCode errC;
-            try
-            {
-                errC = ZLibNative.CreateZLibStreamForDeflate(out _zlibStream, compressionLevel, windowBits, memLevel, strategy);
-            }
-            catch (Exception cause)
-            {
-                throw new ZLibException(SR.ZLibErrorDLLLoadError, cause);
-            }
-
-            switch (errC)
-            {
-                case ZErrorCode.Ok:
-                    return;
-
-                case ZErrorCode.MemError:
-                    throw new ZLibException(SR.ZLibErrorNotEnoughMemory, "deflateInit2_", (int)errC, _zlibStream.GetErrorMessage());
-
-                case ZErrorCode.VersionError:
-                    throw new ZLibException(SR.ZLibErrorVersionMismatch, "deflateInit2_", (int)errC, _zlibStream.GetErrorMessage());
-
-                case ZErrorCode.StreamError:
-                    throw new ZLibException(SR.ZLibErrorIncorrectInitParameters, "deflateInit2_", (int)errC, _zlibStream.GetErrorMessage());
-
-                default:
-                    throw new ZLibException(SR.ZLibErrorUnexpected, "deflateInit2_", (int)errC, _zlibStream.GetErrorMessage());
-            }
+            _zlibStream = zlibStream;
         }
 
         ~Deflater()
@@ -77,9 +48,12 @@ private void Dispose(bool disposing)
             if (!_isDisposed)
             {
                 if (disposing)
+                {
                     _zlibStream.Dispose();
+                }
 
-                DeallocateInputBufferHandle();
+                // Unpin the input buffer, but avoid modifying the ZLibStreamHandle (which may have been disposed of).
+                DeallocateInputBufferHandle(resetStreamHandle: false);
                 _isDisposed = true;
             }
         }
@@ -129,7 +103,7 @@ internal int GetDeflateOutput(byte[] outputBuffer)
                 // Before returning, make sure to release input buffer if necessary:
                 if (0 == _zlibStream.AvailIn)
                 {
-                    DeallocateInputBufferHandle();
+                    DeallocateInputBufferHandle(resetStreamHandle: true);
                 }
             }
         }
@@ -179,12 +153,16 @@ internal bool Flush(byte[] outputBuffer, out int bytesRead)
             return ReadDeflateOutput(outputBuffer, ZFlushCode.SyncFlush, out bytesRead) == ZErrorCode.Ok;
         }
 
-        private void DeallocateInputBufferHandle()
+        private void DeallocateInputBufferHandle(bool resetStreamHandle)
         {
             lock (SyncLock)
             {
-                _zlibStream.AvailIn = 0;
-                _zlibStream.NextIn = ZLibNative.ZNullPtr;
+                if (resetStreamHandle)
+                {
+                    _zlibStream.AvailIn = 0;
+                    _zlibStream.NextIn = ZLibNative.ZNullPtr;
+                }
+
                 _inputBufferHandle.Dispose();
             }
         }
@@ -217,5 +195,14 @@ private ZErrorCode Deflate(ZFlushCode flushCode)
                     throw new ZLibException(SR.ZLibErrorUnexpected, "deflate", (int)errC, _zlibStream.GetErrorMessage());
             }
         }
+
+        public static Deflater CreateDeflater(ZLibNative.CompressionLevel compressionLevel, ZLibNative.CompressionStrategy strategy, int windowBits, int memLevel)
+        {
+            Debug.Assert(windowBits >= minWindowBits && windowBits <= maxWindowBits);
+
+            ZLibNative.ZLibStreamHandle zlibStream = ZLibNative.ZLibStreamHandle.CreateForDeflate(compressionLevel, windowBits, memLevel, strategy);
+
+            return new Deflater(zlibStream);
+        }
     }
 }