[release/7.0-staging] Zlib: Add some protections to the allocator used by zlib (#89532)

* Add guards to zlib memory allocator

* Fix whitespace

* Remove heap LFH flag, simplify publishing logic

* Remove dead code from cmake file

* Normalize cmake files

* Slight zlib PAL & allocator restructuring - Move allocator files under Compression.Native dir - Update pal_zlib includes, use calloc instead of malloc - Remove custom typedefs from zlib unix allocator

* Add env var for Windows

* Add env var for Unix

* Fix mono compilation
- Mono defines HOST_WIN32, not CLR_CMAKE_HOST_WIN32
- Follow same pattern from src\native\eventpipe\CMakeLists.txt

---------

Co-authored-by: Levi Broderick <levib@microsoft.com>

Author: github-actions[bot]

src/native/external/zlib-intel.cmake

@@ -20,6 +20,10 @@ set(ZLIB_SOURCES_BASE
     trees.c
     x86.c
     zutil.c
+    ../../libs/System.IO.Compression.Native/zlib_allocator_win.c
 )
 
+# enable custom zlib allocator
+add_definitions(-DMY_ZCALLOC)
+
 addprefix(ZLIB_SOURCES "${CMAKE_CURRENT_LIST_DIR}/zlib-intel"  "${ZLIB_SOURCES_BASE}")

src/native/external/zlib.cmake

@@ -29,4 +29,12 @@ set(ZLIB_SOURCES_BASE
     zutil.h
 )
 
+# enable custom zlib allocator
+add_definitions(-DMY_ZCALLOC)
+if(HOST_WIN32 OR CLR_CMAKE_TARGET_WIN32)
+    set(ZLIB_SOURCES_BASE ${ZLIB_SOURCES_BASE} ../../libs/System.IO.Compression.Native/zlib_allocator_win.c)
+else()
+    set(ZLIB_SOURCES_BASE ${ZLIB_SOURCES_BASE} ../../libs/System.IO.Compression.Native/zlib_allocator_unix.c)
+endif()
+
 addprefix(ZLIB_SOURCES "${CMAKE_CURRENT_LIST_DIR}/zlib"  "${ZLIB_SOURCES_BASE}")

src/native/libs/System.IO.Compression.Native/CMakeLists.txt

@@ -101,6 +101,7 @@ else ()
 
     if (CLR_CMAKE_HOST_ARCH_I386 OR CLR_CMAKE_HOST_ARCH_AMD64)
         include(${CLR_SRC_NATIVE_DIR}/external/zlib-intel.cmake)
+        add_definitions(-DINTERNAL_ZLIB_INTEL)
     else ()
         include(${CLR_SRC_NATIVE_DIR}/external/zlib.cmake)
     endif ()

src/native/libs/System.IO.Compression.Native/pal_zlib.c

@@ -9,7 +9,11 @@
     #ifdef  _WIN32
         #define c_static_assert(e) static_assert((e),"")
     #endif
-    #include <external/zlib/zlib.h>
+    #ifdef INTERNAL_ZLIB_INTEL
+        #include <external/zlib-intel/zlib.h>
+    #else
+        #include <external/zlib/zlib.h>
+    #endif
 #else
     #include "pal_utilities.h"
     #include <zlib.h>
@@ -39,14 +43,11 @@ Initializes the PAL_ZStream by creating and setting its underlying z_stream.
 */
 static int32_t Init(PAL_ZStream* stream)
 {
-    z_stream* zStream = (z_stream*)malloc(sizeof(z_stream));
+    z_stream* zStream = (z_stream*)calloc(1, sizeof(z_stream));
     stream->internalState = zStream;
 
     if (zStream != NULL)
     {
-        zStream->zalloc = Z_NULL;
-        zStream->zfree = Z_NULL;
-        zStream->opaque = Z_NULL;
         return PAL_Z_OK;
     }
     else

src/native/libs/System.IO.Compression.Native/zlib_allocator_unix.c

@@ -0,0 +1,201 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+#include <stdbool.h>
+#include <stdint.h>
+#include <external/zlib/zutil.h>
+
+/* A custom allocator for zlib that provides some defense-in-depth over standard malloc / free.
+ * (non-Windows version)
+ *
+ * 1. When zlib allocates fixed-length data structures for containing stream metadata, we zero
+ *    the memory before using it, preventing use of uninitialized memory within these structures.
+ *    Ideally we would do this for dynamically-sized buffers as well, but there is a measurable
+ *    perf impact to doing this. Zeroing fixed structures seems like a good trade-off here, since
+ *    these data structures contain most of the metadata used for managing the variable-length
+ *    dynamically allocated buffers.
+ *
+ * 2. We put a cookie both before and after any allocated memory, which allows us to detect local
+ *    buffer overruns on the call to free(). The cookie values are tied to the addresses where
+ *    the data is located in memory.
+ *
+ * 3. We trash the aforementioned cookie on free(), which allows us to detect double-free.
+ *
+ * If any of these checks fails, the application raises SIGABRT.
+ */
+
+static bool IsMitigationDisabled()
+{
+    enum _MitigationEnablementTristate
+    {
+        MITIGATION_NOT_YET_QUERIED = 0,
+        MITIGATION_DISABLED = 1,
+        MITIGATION_ENABLED = 2 // really, anything other than 0 or 1
+    };
+    static int s_fMitigationEnablementState = MITIGATION_NOT_YET_QUERIED;
+
+    // If already initialized, return immediately.
+    // We don't need a volatile read here since the publish is performed with release semantics.
+    if (s_fMitigationEnablementState != MITIGATION_NOT_YET_QUERIED)
+    {
+        return (s_fMitigationEnablementState == MITIGATION_DISABLED);
+    }
+
+    // Initialize the tri-state now.
+    // It's ok for multiple threads to do this simultaneously. Only one thread will win.
+    // Valid env var values to disable mitigation: "true" and "1"
+    // All other env var values (or error) leaves mitigation enabled.
+
+    char* pchEnvVar = getenv("DOTNET_SYSTEM_IO_COMPRESSION_DISABLEZLIBMITIGATIONS");
+    bool fMitigationDisabled = (pchEnvVar && (strcmp(pchEnvVar, "1") == 0 || strcmp(pchEnvVar, "true") == 0));
+    
+    // We really don't care about the return value of the ICE operation. If another thread
+    // beat us to it, so be it. The recursive call will figure it out.
+    __sync_val_compare_and_swap(
+        /* destination: */ &s_fMitigationEnablementState,
+        /* comparand:   */ MITIGATION_NOT_YET_QUERIED,
+        /* exchange:    */ fMitigationDisabled ? MITIGATION_DISABLED : MITIGATION_ENABLED);
+    return IsMitigationDisabled();
+}
+
+#ifndef MEMORY_ALLOCATION_ALIGNMENT
+// malloc() returns an address suitably aligned for any built-in data type.
+// Historically, this has been twice the arch's natural word size.
+#ifdef HOST_64BIT
+#define MEMORY_ALLOCATION_ALIGNMENT 16
+#else
+#define MEMORY_ALLOCATION_ALIGNMENT 8
+#endif
+#endif
+
+typedef struct _DOTNET_ALLOC_COOKIE
+{
+    void* Address;
+    size_t Size;
+} DOTNET_ALLOC_COOKIE;
+
+static bool SafeAdd(size_t a, size_t b, size_t* sum)
+{
+    if (SIZE_MAX - a >= b) { *sum = a + b; return true; }
+    else { *sum = 0; return false; }
+}
+
+static bool SafeMult(size_t a, size_t b, size_t* product)
+{
+    if (SIZE_MAX / a >= b) { *product = a * b; return true; }
+    else { *product = 0; return false; }
+}
+
+static DOTNET_ALLOC_COOKIE ReadAllocCookieUnaligned(const void* pSrc)
+{
+    DOTNET_ALLOC_COOKIE vCookie;
+    memcpy(&vCookie, pSrc, sizeof(DOTNET_ALLOC_COOKIE));
+    return vCookie;
+}
+
+static void WriteAllocCookieUnaligned(void* pDest, DOTNET_ALLOC_COOKIE vCookie)
+{
+    memcpy(pDest, &vCookie, sizeof(DOTNET_ALLOC_COOKIE));
+}
+
+// Historically, the memory allocator always returns addresses aligned to some
+// particular boundary. We'll make that same guarantee here just in case somebody
+// depends on it.
+const size_t DOTNET_ALLOC_HEADER_COOKIE_SIZE_WITH_PADDING = (sizeof(DOTNET_ALLOC_COOKIE) + MEMORY_ALLOCATION_ALIGNMENT - 1) & ~((size_t)MEMORY_ALLOCATION_ALIGNMENT  - 1);
+const size_t DOTNET_ALLOC_TRAILER_COOKIE_SIZE = sizeof(DOTNET_ALLOC_COOKIE);
+
+voidpf ZLIB_INTERNAL zcalloc(opaque, items, size)
+    voidpf opaque;
+    unsigned items;
+    unsigned size;
+{
+    (void)opaque; // unreferenced formal parameter
+
+    if (IsMitigationDisabled())
+    {
+        // fallback logic copied from zutil.c
+        return sizeof(uInt) > 2 ? (voidpf)malloc(items * size) :
+                                  (voidpf)calloc(items, size);
+    }
+
+    // If initializing a fixed-size structure, zero the memory.
+    bool fZeroMemory = (items == 1);
+    
+    size_t cbRequested;
+    if (sizeof(items) + sizeof(size) <= sizeof(cbRequested))
+    {
+        // multiplication can't overflow; no need for safeint
+        cbRequested = (size_t)items * (size_t)size;
+    }
+    else
+    {
+        // multiplication can overflow; go through safeint
+        if (!SafeMult((size_t)items, (size_t)size, &cbRequested)) { return NULL; }
+    }
+
+    // Make sure the actual allocation has enough room for our frontside & backside cookies.
+    size_t cbActualAllocationSize;
+    if (!SafeAdd(cbRequested, DOTNET_ALLOC_HEADER_COOKIE_SIZE_WITH_PADDING + DOTNET_ALLOC_TRAILER_COOKIE_SIZE, &cbActualAllocationSize)) { return NULL; }
+
+    void* pAlloced = (fZeroMemory) ? calloc(1, cbActualAllocationSize) : malloc(cbActualAllocationSize);
+    if (pAlloced == NULL) { return NULL; } // OOM
+
+    DOTNET_ALLOC_COOKIE* pHeaderCookie = (DOTNET_ALLOC_COOKIE*)pAlloced;
+    uint8_t* pReturnToCaller = (uint8_t*)pAlloced + DOTNET_ALLOC_HEADER_COOKIE_SIZE_WITH_PADDING;
+    uint8_t* pTrailerCookie = pReturnToCaller + cbRequested;
+
+    // Write out the same cookie for the header & the trailer, then we're done.
+
+    DOTNET_ALLOC_COOKIE vCookie = { 0 };
+    vCookie.Address = pReturnToCaller;
+    vCookie.Size = cbRequested;
+    *pHeaderCookie = vCookie; // aligned
+    WriteAllocCookieUnaligned(pTrailerCookie, vCookie);
+
+    return pReturnToCaller;
+}
+
+static void zcfree_trash_cookie(void* pCookie)
+{
+    memset(pCookie, 0, sizeof(DOTNET_ALLOC_COOKIE));
+}
+
+void ZLIB_INTERNAL zcfree(opaque, ptr)
+    voidpf opaque;
+    voidpf ptr;
+{
+    (void)opaque; // unreferenced formal parameter
+
+    if (IsMitigationDisabled())
+    {
+        // fallback logic copied from zutil.c
+        free(ptr);
+        return;
+    }
+
+    if (ptr == NULL) { return; } // ok to free nullptr
+
+    // Check cookie at beginning
+
+    DOTNET_ALLOC_COOKIE* pHeaderCookie = (DOTNET_ALLOC_COOKIE*)((uint8_t*)ptr - DOTNET_ALLOC_HEADER_COOKIE_SIZE_WITH_PADDING);
+    if (pHeaderCookie->Address != ptr) { goto Fail; }
+    size_t cbRequested = pHeaderCookie->Size;
+
+    // Check cookie at end
+
+    uint8_t* pTrailerCookie = (uint8_t*)ptr + cbRequested;
+    DOTNET_ALLOC_COOKIE vTrailerCookie = ReadAllocCookieUnaligned(pTrailerCookie);
+    if (vTrailerCookie.Address != ptr) { goto Fail; }
+    if (vTrailerCookie.Size != cbRequested) { goto Fail; }
+
+    // Checks passed - now trash the cookies and free memory
+
+    zcfree_trash_cookie(pHeaderCookie);
+    zcfree_trash_cookie(pTrailerCookie);
+
+    free(pHeaderCookie);
+    return;
+
+Fail:
+    abort(); // cookie check failed
+}