Bug 1059797 - Pre-allocate zlib inflate buffers in faulty.lib. r=froydnj

Original patch from James Willcox <snorp@snorp.net>

mozglue/linker/Mappable.cpp

@@ -104,7 +104,7 @@ MappableExtractFile::Create(const char *name, Zip *zip, Zip::Stream *stream)
       return nullptr;
     }
 
-    z_stream zStream = stream->GetZStream(buffer);
+    zxx_stream zStream = stream->GetZStream(buffer);
 
     /* Decompress */
     if (inflateInit2(&zStream, -MAX_WBITS) != Z_OK) {

mozglue/linker/Mappable.h

@@ -180,7 +180,7 @@ private:
   mozilla::UniquePtr<_MappableBuffer> buffer;
 
   /* Zlib data */
-  z_stream zStream;
+  zxx_stream zStream;
 };
 
 /**

mozglue/linker/SeekableZStream.cpp

@@ -72,8 +72,7 @@ SeekableZStream::DecompressChunk(void *where, size_t chunk, size_t length)
 
   DEBUG_LOG("DecompressChunk #%" PRIdSize " @%p (%" PRIdSize "/% " PRIdSize ")",
         chunk, where, length, chunkLen);
-  z_stream zStream;
-  memset(&zStream, 0, sizeof(zStream));
+  zxx_stream zStream;
   zStream.avail_in = (isLastChunk ? totalSize : uint32_t(offsetTable[chunk + 1]))
                      - uint32_t(offsetTable[chunk]);
   zStream.next_in = const_cast<Bytef *>(buffer + uint32_t(offsetTable[chunk]));

mozglue/linker/Zip.h

@@ -10,7 +10,98 @@
 #include <vector>
 #include <zlib.h>
 #include "Utils.h"
+#include "mozilla/Assertions.h"
 #include "mozilla/RefPtr.h"
+#include "mozilla/UniquePtr.h"
+
+/**
+ * Helper class wrapping z_stream to avoid malloc() calls during
+ * inflate. Do not use for deflate.
+ * inflateInit allocates two buffers:
+ * - one for its internal state, which is "approximately 10K bytes" according
+ *   to inflate.h from zlib.
+ * - one for the compression window, which depends on the window size passed
+ *   to inflateInit2, but is never greater than 32K (1 << MAX_WBITS).
+ * Those buffers are created at instantiation time instead of when calling
+ * inflateInit2. When inflateInit2 is called, it will call zxx_stream::Alloc
+ * to get both these buffers. zxx_stream::Alloc will choose one of the
+ * pre-allocated buffers depending on the requested size.
+ */
+class zxx_stream: public z_stream
+{
+public:
+  zxx_stream() {
+    memset(this, 0, sizeof(z_stream));
+    zalloc = Alloc;
+    zfree = Free;
+    opaque = this;
+  }
+
+private:
+  static void *Alloc(void *data, uInt items, uInt size)
+  {
+    size_t buf_size = items * size;
+    zxx_stream *zStream = reinterpret_cast<zxx_stream *>(data);
+
+    if (items == 1 && buf_size <= zStream->stateBuf.size) {
+      return zStream->stateBuf.get();
+    } else if (buf_size == zStream->windowBuf.size) {
+      return zStream->windowBuf.get();
+    } else {
+      MOZ_CRASH("No ZStreamBuf for allocation");
+    }
+  }
+
+  static void Free(void *data, void *ptr)
+  {
+    zxx_stream *zStream = reinterpret_cast<zxx_stream *>(data);
+
+    if (zStream->stateBuf.Equals(ptr)) {
+      zStream->stateBuf.Release();
+    } else if (zStream->windowBuf.Equals(ptr)) {
+      zStream->windowBuf.Release();
+    } else {
+      MOZ_CRASH("Pointer doesn't match a ZStreamBuf");
+    }
+  }
+
+  /**
+   * Helper class for each buffer.
+   */
+  template <size_t Size>
+  class ZStreamBuf
+  {
+  public:
+    ZStreamBuf() : buf(new char[Size]), inUse(false) { }
+
+    char *get()
+    {
+      if (!inUse) {
+        inUse = true;
+        return buf.get();
+      } else {
+        MOZ_CRASH("ZStreamBuf already in use");
+      }
+    }
+
+    void Release()
+    {
+      memset(buf.get(), 0, Size);
+      inUse = false;
+    }
+
+    bool Equals(const void *other) { return other == buf.get(); }
+
+    static const size_t size = Size;
+
+  private:
+    mozilla::UniquePtr<char[]> buf;
+    bool inUse;
+  };
+
+  ZStreamBuf<0x3000> stateBuf; // 0x3000 is an arbitrary size above 10K.
+  ZStreamBuf<1 << MAX_WBITS> windowBuf;
+};
 
 /**
  * Forward declaration
@@ -87,14 +178,13 @@ public:
     Type GetType() { return type; }
 
     /**
-     * Returns a z_stream for use with inflate functions using the given
+     * Returns a zxx_stream for use with inflate functions using the given
      * buffer as inflate output. The caller is expected to allocate enough
      * memory for the Stream uncompressed size.
      */
-    z_stream GetZStream(void *buf)
+    zxx_stream GetZStream(void *buf)
     {
-      z_stream zStream;
-      memset(&zStream, 0, sizeof(zStream));
+      zxx_stream zStream;
       zStream.avail_in = compressedSize;
       zStream.next_in = reinterpret_cast<Bytef *>(
                         const_cast<void *>(compressedBuf));