SafeDeflateStream

Introduction

With .NET 2.0, Microsoft finally decided to integrate the deflate algorithm into the framework. Great! This means we no longer need to buy the Component One library and modify it to true asynchronous usage, or to hack ICSharpCode's for similar functionality. Unfortunately, Microsoft's implementation suffers a severe problem. If you are decompressing a stream using async methods (BeginRead) and the stream is corrupt, the framework throws an exception that crashes your program - you cannot catch and suppress this exception. Yuck! My application, Swapper.NET, can often get such streams (when receiving compressed traffic from other P2P nodes), this was unacceptable.

Using the code

Simply put, any place you would ordinarily construct a System.IO.Compression.DeflateStream, you instead construct a RevolutionaryStuff.JBT.Compression.SafeDeflateStream class. The problem outlined in the introduction only manifests itself when you are using BeginRead on a DeflateStream where mode=Decompress and the input data may be corrupt, so you could use my wrapper only in that case. I use it wherever, just in case I find a new code in Microsoft's implementation.

Given the compactness of this code, I felt it was simpler to post it in its entirety here rather than attach a zip. It is commented, but if you think you need more explanations, please speak up so you can influence my forthcoming articles.

using System;
using System.Diagnostics;
using System.Reflection;
using System.IO;
using System.IO.Compression;

namespace RevolutionaryStuff.JBT.Compression
{
    /// <summary>
    /// This is a small wrapper of DeflateStream.
    /// 
    /// If you use BeginRead on a DeflateStream
    /// with corrupt data, an exception will occur in the inflater.
    /// Microsoft did not catch this exception,
    /// and it is not in a place where you app can handle it.
    /// So when this occurs, a nasty exception
    /// will be thrown from which you cannot gracefully recover.
    /// This class fixes this issue by trapping
    /// the underlying exception when it is thrown, and re-throwing
    /// it in the correct place, EndRead, where the user can handle it.
    /// </summary>
    public class SafeDeflateStream : DeflateStream
    {
        private static readonly FieldInfo CallbackFieldInfo;
        private static readonly MethodInfo ReadCallbackMethodInfo;
        private static readonly MethodInfo InvokeCallbackMethodInfo;

        #region Constructors

        /// <summary>
        /// Cache the necessary reflection evilness once
        /// </summary>
        /// <remarks>
        /// Since we use nasty reflection,
        /// there is no guarantee this will work
        /// on future version of the .NET framework.
        /// Then again, hopefully MS will fix the problem by then.
        /// </remarks>
        static SafeDeflateStream()
        {
            Type t = typeof(DeflateStream);
            MemberInfo[] mis = t.GetMember("*", 
              BindingFlags.NonPublic | BindingFlags.Instance);
            foreach (MemberInfo mi in mis)
            {
                if (mi.Name == "m_CallBack")
                {
                    CallbackFieldInfo = (FieldInfo)mi;
                }
                else if (mi.Name == "ReadCallback")
                {
                    ReadCallbackMethodInfo = (MethodInfo)mi;
                }
            }
            t = t.Assembly.GetType("System.IO." + 
                "Compression.DeflateStreamAsyncResult");
            mis = t.GetMember("*", BindingFlags.NonPublic | 
                                   BindingFlags.Instance);
            foreach (MemberInfo mi in mis)
            {
                if (mi.Name == "Complete")
                {
                    MethodInfo mei = (MethodInfo)mi;
                    if (mei.GetParameters().Length == 1)
                    {
                        InvokeCallbackMethodInfo = mei;
                        break;
                    }
                }
            }
            if (InvokeCallbackMethodInfo == null || 
                CallbackFieldInfo == null || 
                ReadCallbackMethodInfo == null)
            {
                //If we could not find the underlying 
                //members, throw an exception and crash.
                //This should only happen if MS updates 
                //their code with a new version of the framework
                throw new ApplicationException("Could" + 
                          " not find the hidden stuff");
            }
        }

        /// <remarks>We want to keep the constructors
        /// the same as in DeflateStream so callers
        /// can simply drop in this class</remarks>
        public SafeDeflateStream(Stream st, CompressionMode mode)
                                 : this(st, mode, false)
        { }

        /// <remarks>We want to keep the constructors
        /// the same as in DeflateStream so callers
        /// can simply drop in this class</remarks>
        public SafeDeflateStream(Stream st, 
               CompressionMode mode, bool leaveOpen)
               : base(st, mode, leaveOpen)
        {
            if (mode == CompressionMode.Decompress)
            {
                //If this is a decompression stream
                //(the only one which causes the problem),
                //we overwrite the private 
                //readonly callback with ours
                CallbackFieldInfo.SetValue(this, 
                  new AsyncCallback(SafeReadCallback));
            }
        }

        #endregion

#if DEBUG
        private static int TotalCallbackCount;
        private static int BadCallbackCount;
#endif

        private void SafeReadCallback(IAsyncResult ar)
        {
#if DEBUG
            ++TotalCallbackCount;
#endif
            try
            {
                ReadCallbackMethodInfo.Invoke(this, 
                              new object[] { ar });
            }
            catch (Exception ex)
            {
#if DEBUG
                ++BadCallbackCount;
#endif
                Debug.WriteLine(ex);
                try
                {
                    InvokeCallbackMethodInfo.Invoke(ar.AsyncState, 
                                    new object[] { (object) ex });
                }
                catch (Exception ex2)
                {
                    Trace.WriteLine(ex2);
                    throw;
                }
            }
        }
    }
}

Now... to test it. If you compile and run the following code, then carefully read the output, you'll see why Microsoft's implementation is... ungraceful.

using System;
using System.Diagnostics;
using System.IO;
using System.IO.Compression;
using System.Threading;
using RevolutionaryStuff.JBT.Compression;

namespace ConsoleApplication1
{
    class Program
    {
        private static readonly AutoResetEvent 
                TestCompleteEvent = new AutoResetEvent(false);

        private static void Echo(string format, params object[] args)
        {
            string s = string.Format(format, args);
            Console.WriteLine(s);
            Trace.WriteLine(s);
        }

        private static void 
                CurrentDomain_UnhandledException(object 
                sender, UnhandledExceptionEventArgs e)
        {
            Echo("Exception in CurrentDomain_UnhandledException" + 
                 " (bad because we can't properly recover)" + 
                 "\nIsTerminating={0}\nExceptionObject={1}", 
                 e.IsTerminating, e.ExceptionObject);
            TestCompleteEvent.Set();
        }

        private static void ReadComplete(IAsyncResult ar)
        {
            Echo("ReadComplete Starting vvvvvvv");
            try
            {
                Stream st = (Stream)ar.AsyncState;
                int bytesRead = st.EndRead(ar);
                Echo("Read {0} bytes", bytesRead);
                if (bytesRead > 0)
                {
                    TestRead(st);
                }
                else
                {
                    TestCompleteEvent.Set();
                }
            }
            catch (Exception ex)
            {
                Echo("Exception in ReadComplete" + 
                     " (where it should be)\n{0}", ex);
                TestCompleteEvent.Set();
            }
            Echo("ReadComplete Ending ^^^^^^^");
        }

        private static void TestRead(Stream st)
        {
            byte[] readBuf = new byte[2048];
            Echo("TestRead.{0} Starting vvvvvvvv", st.GetType());
            try
            {
                st.BeginRead(readBuf, 0, readBuf.Length, 
                   new AsyncCallback(ReadComplete), st);
            }
            catch (Exception ex)
            {
                Echo("Exception in TestRead\n{0}", ex);
            }
            Echo("TestRead.{0} Ending ^^^^^^^^", st.GetType());
        }

        [STAThread]
        static void Main(string[] args)
        {
            //Create a stream that we know to not be properly "deflated"
            MemoryStream corruptCompressedStream = new MemoryStream();
            using (DeflateStream compressedStream = new 
              DeflateStream(corruptCompressedStream, 
              CompressionMode.Compress, true))
            {
                byte[] buf = new byte[128];
                Random r = new Random(04091974);
                r.NextBytes(buf);
                compressedStream.Write(buf, 0, buf.Length);
                compressedStream.Flush();
            }
            //Corrupt the stream by overriting key header info
            corruptCompressedStream.Position = 5;
            corruptCompressedStream.Write(new byte[32], 0, 32);

            //Setup global exception handlers
            AppDomain.CurrentDomain.UnhandledException += 
                        CurrentDomain_UnhandledException;

            //Test the SafeDeflateStream
            corruptCompressedStream.Position = 0;
            TestRead(new SafeDeflateStream(corruptCompressedStream, 
                                CompressionMode.Decompress, true));
            TestCompleteEvent.WaitOne();

            //Test the DeflateStream
            //Must be run 2nd since it crashes the console app!
            corruptCompressedStream.Position = 0;
            TestRead(new DeflateStream(corruptCompressedStream, 
                            CompressionMode.Decompress, true));
            TestCompleteEvent.WaitOne();
        }
    }
}

Happy coding :)

History

• 1/18/2006 - First submission.
• 1/24/2006 - Includes code to prove the error condition.