Rewrite inflate memory allocation.

Inflate used to allocate state during init, but window would be allocated
when/if needed and could be resized and that required a new free/alloc round.

- Now, we allocate state and a 32K window during init, allowing the latency cost
  of allocs to be done during init instead of at one or more times later.
- Total memory allocation is about the same when requesting a 32K window, but
  if now window or a smaller window was requested, then it is an increase.
- While doing alloc(), we now store pointer to corresponding free(), avoiding crashes
  with applications that incorrectly set alloc/free pointers after running init function.
- After init has succeeded, inflate will no longer possibly fail due to a failing malloc.

Co-authored-by: Ilya Leoshkevich <iii@linux.ibm.com>

infback.c

@@ -43,10 +43,15 @@ int32_t ZNG_CONDEXPORT PREFIX(inflateBackInit)(PREFIX3(stream) *strm, int32_t wi
     }
     if (strm->zfree == NULL)
         strm->zfree = PREFIX(zcfree);
-    state = ZALLOC(strm, 1, sizeof(struct inflate_state));
+
-    if (state == NULL)
+    inflate_allocs *alloc_bufs = alloc_inflate(strm);
+    if (alloc_bufs == NULL)
         return Z_MEM_ERROR;
+
+    state = alloc_bufs->state;
+    state->alloc_bufs = alloc_bufs;
     Tracev((stderr, "inflate: allocated\n"));
+
     strm->state = (struct internal_state *)state;
     state->dmax = 32768U;
     state->wbits = (unsigned int)windowBits;
@@ -504,8 +509,10 @@ int32_t Z_EXPORT PREFIX(inflateBack)(PREFIX3(stream) *strm, in_func in, void *in
 int32_t Z_EXPORT PREFIX(inflateBackEnd)(PREFIX3(stream) *strm) {
     if (strm == NULL || strm->state == NULL || strm->zfree == NULL)
         return Z_STREAM_ERROR;
-    ZFREE(strm, strm->state);
+
-    strm->state = NULL;
+    /* Free allocated buffers */
+    free_inflate(strm);
+
     Tracev((stderr, "inflate: end\n"));
     return Z_OK;
 }

inflate.c

@@ -53,7 +53,7 @@ static int inflateStateCheck(PREFIX3(stream) *strm) {
     if (strm == NULL || strm->zalloc == NULL || strm->zfree == NULL)
         return 1;
     state = (struct inflate_state *)strm->state;
-    if (state == NULL || state->strm != strm || state->mode < HEAD || state->mode > SYNC)
+    if (state == NULL || state->alloc_bufs == NULL || state->strm != strm || state->mode < HEAD || state->mode > SYNC)
         return 1;
     return 0;
 }
@@ -120,13 +120,9 @@ int32_t Z_EXPORT PREFIX(inflateReset2)(PREFIX3(stream) *strm, int32_t windowBits
 #endif
     }
 
-    /* set number of window bits, free window if different */
+    /* set number of window bits */
     if (windowBits && (windowBits < MIN_WBITS || windowBits > MAX_WBITS))
         return Z_STREAM_ERROR;
-    if (state->window != NULL && state->wbits != (unsigned)windowBits) {
-        ZFREE_WINDOW(strm, state->window);
-        state->window = NULL;
-    }
 
     /* update state and reset the rest of it */
     state->wrap = wrap;
@@ -134,7 +130,88 @@ int32_t Z_EXPORT PREFIX(inflateReset2)(PREFIX3(stream) *strm, int32_t windowBits
     return PREFIX(inflateReset)(strm);
 }
 
-/* This function is hidden in ZLIB_COMPAT builds. */
+#ifdef INF_ALLOC_DEBUG
+#  include <stdio.h>
+#  define LOGSZ(name,size)           fprintf(stderr, "%s is %d bytes\n", name, size)
+#  define LOGSZP(name,size,loc,pad)  fprintf(stderr, "%s is %d bytes, offset %d, padded %d\n", name, size, loc, pad)
+#  define LOGSZPL(name,size,loc,pad) fprintf(stderr, "%s is %d bytes, offset %ld, padded %d\n", name, size, loc, pad)
+#else
+#  define LOGSZ(name,size)
+#  define LOGSZP(name,size,loc,pad)
+#  define LOGSZPL(name,size,loc,pad)
+#endif
+
+/* ===========================================================================
+ * Allocate a big buffer and divide it up into the various buffers inflate needs.
+ * Handles alignment of allocated buffer and alignment of individual buffers.
+ */
+Z_INTERNAL inflate_allocs* alloc_inflate(PREFIX3(stream) *strm) {
+    int curr_size = 0;
+
+    /* Define sizes */
+    int window_size = INFLATE_ADJUST_WINDOW_SIZE((1 << MAX_WBITS) + 64); /* 64B padding for chunksize */
+    int state_size = sizeof(inflate_state);
+    int alloc_size = sizeof(inflate_allocs);
+
+    /* Calculate relative buffer positions and paddings */
+    LOGSZP("window", window_size, PAD_WINDOW(curr_size), PADSZ(curr_size,WINDOW_PAD_SIZE));
+    int window_pos = PAD_WINDOW(curr_size);
+    curr_size = window_pos + window_size;
+
+    LOGSZP("state", state_size, PAD_64(curr_size), PADSZ(curr_size,64));
+    int state_pos = PAD_64(curr_size);
+    curr_size = state_pos + state_size;
+
+    LOGSZP("alloc", alloc_size, PAD_16(curr_size), PADSZ(curr_size,16));
+    int alloc_pos = PAD_16(curr_size);
+    curr_size = alloc_pos + alloc_size;
+
+    /* Add 64-1 or 4096-1 to allow window alignment, and round size of buffer up to multiple of 64 */
+    int total_size = PAD_64(curr_size + (WINDOW_PAD_SIZE - 1));
+
+    /* Allocate buffer, align to 64-byte cacheline, and zerofill the resulting buffer */
+    char *original_buf = strm->zalloc(strm->opaque, 1, total_size);
+    if (original_buf == NULL)
+        return NULL;
+
+    char *buff = (char *)HINT_ALIGNED_WINDOW((char *)PAD_WINDOW(original_buf));
+    LOGSZPL("Buffer alloc", total_size, PADSZ((uintptr_t)original_buf,WINDOW_PAD_SIZE), PADSZ(curr_size,WINDOW_PAD_SIZE));
+
+    /* Initialize alloc_bufs */
+    inflate_allocs *alloc_bufs  = (struct inflate_allocs_s *)(buff + alloc_pos);
+    alloc_bufs->buf_start = (char *)original_buf;
+    alloc_bufs->zfree = strm->zfree;
+
+    alloc_bufs->window =  (unsigned char *)HINT_ALIGNED_WINDOW((buff + window_pos));
+    alloc_bufs->state = (inflate_state *)HINT_ALIGNED_64((buff + state_pos));
+
+#ifdef Z_MEMORY_SANITIZER
+    /* This is _not_ to subvert the memory sanitizer but to instead unposion some
+       data we willingly and purposefully load uninitialized into vector registers
+       in order to safely read the last < chunksize bytes of the window. */
+    __msan_unpoison(alloc_bufs->window + window_size, 64);
+#endif
+
+    return alloc_bufs;
+}
+
+/* ===========================================================================
+ * Free all allocated inflate buffers
+ */
+Z_INTERNAL void free_inflate(PREFIX3(stream) *strm) {
+    struct inflate_state *state = (struct inflate_state *)strm->state;
+
+    if (state->alloc_bufs != NULL) {
+        inflate_allocs *alloc_bufs = state->alloc_bufs;
+        alloc_bufs->zfree(strm->opaque, alloc_bufs->buf_start);
+        strm->state = NULL;
+    }
+}
+
+/* ===========================================================================
+ * Initialize inflate state and buffers.
+ * This function is hidden in ZLIB_COMPAT builds.
+ */
 int32_t ZNG_CONDEXPORT PREFIX(inflateInit2)(PREFIX3(stream) *strm, int32_t windowBits) {
     int32_t ret;
     struct inflate_state *state;
@@ -151,19 +228,23 @@ int32_t ZNG_CONDEXPORT PREFIX(inflateInit2)(PREFIX3(stream) *strm, int32_t windo
     }
     if (strm->zfree == NULL)
         strm->zfree = PREFIX(zcfree);
-    state = ZALLOC(strm, 1, sizeof(struct inflate_state));
+
-    if (state == NULL)
+    inflate_allocs *alloc_bufs = alloc_inflate(strm);
+    if (alloc_bufs == NULL)
         return Z_MEM_ERROR;
+
+    state = alloc_bufs->state;
+    state->window = alloc_bufs->window;
+    state->alloc_bufs = alloc_bufs;
     Tracev((stderr, "inflate: allocated\n"));
+
     strm->state = (struct internal_state *)state;
     state->strm = strm;
-    state->window = NULL;
     state->mode = HEAD;     /* to pass state test in inflateReset2() */
     state->chunksize = FUNCTABLE_CALL(chunksize)();
     ret = PREFIX(inflateReset2)(strm, windowBits);
     if (ret != Z_OK) {
-        ZFREE(strm, state);
+        free_inflate(strm);
-        strm->state = NULL;
     }
     return ret;
 }
@@ -223,20 +304,6 @@ void Z_INTERNAL PREFIX(fixedtables)(struct inflate_state *state) {
 }
 
 int Z_INTERNAL PREFIX(inflate_ensure_window)(struct inflate_state *state) {
-    /* if it hasn't been done already, allocate space for the window */
-    if (state->window == NULL) {
-        unsigned wsize = 1U << state->wbits;
-        state->window = (unsigned char *)ZALLOC_WINDOW(state->strm, wsize + state->chunksize, sizeof(unsigned char));
-        if (state->window == NULL)
-            return Z_MEM_ERROR;
-#ifdef Z_MEMORY_SANITIZER
-        /* This is _not_ to subvert the memory sanitizer but to instead unposion some
-           data we willingly and purposefully load uninitialized into vector registers
-           in order to safely read the last < chunksize bytes of the window. */
-        __msan_unpoison(state->window + wsize, state->chunksize);
-#endif
-    }
-
     /* if window not in use yet, initialize */
     if (state->wsize == 0) {
         state->wsize = 1U << state->wbits;
@@ -1142,14 +1209,12 @@ int32_t Z_EXPORT PREFIX(inflate)(PREFIX3(stream) *strm, int32_t flush) {
 }
 
 int32_t Z_EXPORT PREFIX(inflateEnd)(PREFIX3(stream) *strm) {
-    struct inflate_state *state;
     if (inflateStateCheck(strm))
         return Z_STREAM_ERROR;
-    state = (struct inflate_state *)strm->state;
+
-    if (state->window != NULL)
+    /* Free allocated buffers */
-        ZFREE_WINDOW(strm, state->window);
+    free_inflate(strm);
-    ZFREE(strm, strm->state);
+
-    strm->state = NULL;
     Tracev((stderr, "inflate: end\n"));
     return Z_OK;
 }
@@ -1332,13 +1397,16 @@ int32_t Z_EXPORT PREFIX(inflateCopy)(PREFIX3(stream) *dest, PREFIX3(stream) *sou
         return Z_STREAM_ERROR;
     state = (struct inflate_state *)source->state;
 
+    /* copy stream */
+    memcpy((void *)dest, (void *)source, sizeof(PREFIX3(stream)));
+
     /* allocate space */
-    copy = ZALLOC(source, 1, sizeof(struct inflate_state));
+    inflate_allocs *alloc_bufs = alloc_inflate(dest);
-    if (copy == NULL)
+    if (alloc_bufs == NULL)
         return Z_MEM_ERROR;
+    copy = alloc_bufs->state;
 
     /* copy state */
-    memcpy((void *)dest, (void *)source, sizeof(PREFIX3(stream)));
     memcpy(copy, state, sizeof(struct inflate_state));
     copy->strm = dest;
     if (state->lencode >= state->codes && state->lencode <= state->codes + ENOUGH - 1) {
@@ -1346,16 +1414,11 @@ int32_t Z_EXPORT PREFIX(inflateCopy)(PREFIX3(stream) *dest, PREFIX3(stream) *sou
         copy->distcode = copy->codes + (state->distcode - state->codes);
     }
     copy->next = copy->codes + (state->next - state->codes);
+    copy->window = alloc_bufs->window;
+    copy->alloc_bufs = alloc_bufs;
 
     /* window */
-    copy->window = NULL;
+    memcpy(copy->window, state->window, INFLATE_ADJUST_WINDOW_SIZE((size_t)state->wsize));
-    if (state->window != NULL) {
-        if (PREFIX(inflate_ensure_window)(copy)) {
-            ZFREE(source, copy);
-            return Z_MEM_ERROR;
-        }
-        ZCOPY_WINDOW(copy->window, state->window, (size_t)state->wsize);
-    }
 
     dest->state = (struct internal_state *)copy;
     return Z_OK;

inflate.h

@@ -85,10 +85,19 @@ typedef enum {
     Process trailer:
         CHECK -> LENGTH -> DONE
  */
+typedef struct inflate_state inflate_state;
+
+/* Struct for memory allocation handling */
+typedef struct inflate_allocs_s {
+    char            *buf_start;
+    free_func        zfree;
+    inflate_state   *state;
+    unsigned char   *window;
+} inflate_allocs;
 
 /* State maintained between inflate() calls -- approximately 7K bytes, not
    including the allocated sliding window, which is up to 32K bytes. */
-struct inflate_state {
+struct ALIGNED_(64) inflate_state {
     PREFIX3(stream) *strm;             /* pointer back to this zlib stream */
     inflate_mode mode;          /* current inflate mode */
     int last;                   /* true if processing last block */
@@ -136,6 +145,7 @@ struct inflate_state {
     int back;                   /* bits back of last unprocessed length/lit */
     unsigned was;               /* initial length of match */
     uint32_t chunksize;         /* size of memory copying chunk */
+    inflate_allocs *alloc_bufs; /* struct for handling memory allocations */
 #ifdef HAVE_ARCH_INFLATE_STATE
     arch_inflate_state arch;    /* architecture-specific extensions */
 #endif
@@ -143,5 +153,7 @@ struct inflate_state {
 
 int Z_INTERNAL PREFIX(inflate_ensure_window)(struct inflate_state *state);
 void Z_INTERNAL PREFIX(fixedtables)(struct inflate_state *state);
+Z_INTERNAL inflate_allocs* alloc_inflate(PREFIX3(stream) *strm);
+Z_INTERNAL void free_inflate(PREFIX3(stream) *strm);
 
 #endif /* INFLATE_H_ */

inflate_p.h

@@ -10,11 +10,14 @@
 /* Architecture-specific hooks. */
 #ifdef S390_DFLTCC_INFLATE
 #  include "arch/s390/dfltcc_inflate.h"
+/* DFLTCC instructions require window to be page-aligned */
+#  define PAD_WINDOW            PAD_4096
+#  define WINDOW_PAD_SIZE       4096
+#  define HINT_ALIGNED_WINDOW   HINT_ALIGNED_4096
 #else
-/* Memory management for the window. Useful for allocation the aligned window. */
+#  define PAD_WINDOW            PAD_64
-#  define ZALLOC_WINDOW(strm, items, size) ZALLOC(strm, items, size)
+#  define WINDOW_PAD_SIZE       64
-#  define ZCOPY_WINDOW(dest, src, n) memcpy(dest, src, n)
+#  define HINT_ALIGNED_WINDOW   HINT_ALIGNED_64
-#  define ZFREE_WINDOW(strm, addr) ZFREE(strm, addr)
 /* Adjust the window size for the arch-specific inflate code. */
 #  define INFLATE_ADJUST_WINDOW_SIZE(n) (n)
 /* Invoked at the end of inflateResetKeep(). Useful for initializing arch-specific extension blocks. */

test/infcover.c

@@ -319,9 +319,6 @@ static void inf(char *hex, char *what, unsigned step, int win, unsigned len, int
         if (ret == Z_NEED_DICT) {
             ret = PREFIX(inflateSetDictionary)(&strm, in, 1);
                                                 assert(ret == Z_DATA_ERROR);
-            mem_limit(&strm, 1);
-            ret = PREFIX(inflateSetDictionary)(&strm, out, 0);
-                                                assert(ret == Z_MEM_ERROR);
             mem_limit(&strm, 0);
             ((struct inflate_state *)strm.state)->mode = DICT;
             ret = PREFIX(inflateSetDictionary)(&strm, out, 0);
@@ -418,10 +415,6 @@ static void cover_wrap(void) {
     strm.next_in = (void *)"\x63";
     strm.avail_out = 1;
     strm.next_out = (void *)&ret;
-    mem_limit(&strm, 1);
-    ret = PREFIX(inflate)(&strm, Z_NO_FLUSH);   assert(ret == Z_MEM_ERROR);
-    ret = PREFIX(inflate)(&strm, Z_NO_FLUSH);   assert(ret == Z_MEM_ERROR);
-    mem_limit(&strm, 0);
     memset(dict, 0, 257);
     ret = PREFIX(inflateSetDictionary)(&strm, dict, 257);
                                                 assert(ret == Z_OK);