From a39e323a4db80a57feecf2ae212c08070234050c Mon Sep 17 00:00:00 2001
From: Nathan Moinvaziri <nathan@nathanm.com>
Date: Tue, 4 Jan 2022 10:30:54 -0800
Subject: [PATCH] Added memory alignment compensation functions for users who
 may be using custom allocators that don't align on the same boundary zlib-ng
 expects.

---
 zutil.c | 39 +++++++++++++++++++++++++++++++++++++++
 zutil.h | 11 +++++++++--
 2 files changed, 48 insertions(+), 2 deletions(-)

diff --git a/zutil.c b/zutil.c
index 18db2fe0..20f5ac12 100644
--- a/zutil.c
+++ b/zutil.c
@@ -109,3 +109,42 @@ void Z_INTERNAL zng_cfree(void *opaque, void *ptr) {
     Z_UNUSED(opaque);
     zng_free(ptr);
 }
+
+/* Since we support custom memory allocators, some which might not align memory as we expect,
+ * we have to ask for extra memory and return an aligned pointer. */
+void Z_INTERNAL *zng_calloc_aligned(zng_calloc_func zalloc, void *opaque, unsigned items, unsigned size, unsigned align) {
+    uintptr_t return_ptr, original_ptr;
+    int32_t alloc_size, align_diff;
+    void *ptr;
+
+    /* Allocate enough memory for proper alignment and to store the original memory pointer */
+    alloc_size = sizeof(void *) + (items * size) + align;
+    ptr = zalloc(opaque, 1, alloc_size);
+    if (!ptr)
+        return NULL;
+
+    /* Calculate return pointer address with space enough to store original pointer */
+    align_diff = align - ((uintptr_t)ptr % align);
+    return_ptr = (uintptr_t)ptr + align_diff;
+    if (align_diff < sizeof(void *))
+        return_ptr += align;
+
+    /* Store the original pointer for free() */
+    original_ptr = return_ptr - sizeof(void *);
+    memcpy((void *)original_ptr, &ptr, sizeof(void *));
+
+    /* Return properly aligned pointer in allocation */
+    return (void *)return_ptr;
+}
+
+void Z_INTERNAL zng_cfree_aligned(zng_cfree_func zfree, void *opaque, void *ptr) {
+    if (!ptr)
+        return;
+
+    /* Calculate offset to original memory allocation pointer */
+    void *original_ptr = (void *)((uintptr_t)ptr - sizeof(void *));
+    void *free_ptr = *(void **)original_ptr;
+
+    /* Free original memory allocation */
+    zfree(opaque, free_ptr);
+}
diff --git a/zutil.h b/zutil.h
index 2a30c6d5..5022b36f 100644
--- a/zutil.h
+++ b/zutil.h
@@ -125,8 +125,15 @@ extern z_const char * const PREFIX(z_errmsg)[10]; /* indexed by 2-zlib_error */
 void Z_INTERNAL *zng_calloc(void *opaque, unsigned items, unsigned size);
 void Z_INTERNAL   zng_cfree(void *opaque, void *ptr);
 
-#define ZALLOC(strm, items, size) (*((strm)->zalloc))((strm)->opaque, (items), (size))
-#define ZFREE(strm, addr)         (*((strm)->zfree))((strm)->opaque, (void *)(addr))
+typedef void *zng_calloc_func(void *opaque, unsigned items, unsigned size);
+typedef void  zng_cfree_func(void *opaque, void *ptr);
+
+void Z_INTERNAL *zng_calloc_aligned(zng_calloc_func zalloc, void *opaque, unsigned items, unsigned size, unsigned align);
+void Z_INTERNAL  zng_cfree_aligned(zng_cfree_func zfree, void *opaque, void *ptr);
+
+#define ZALLOC(strm, items, size) zng_calloc_aligned((strm)->zalloc, (strm)->opaque, (items), (size), 64)
+#define ZFREE(strm, addr)         zng_cfree_aligned((strm)->zfree, (strm)->opaque, (void *)(addr))
+
 #define TRY_FREE(s, p)            {if (p) ZFREE(s, p);}
 
 #endif /* ZUTIL_H_ */