Skip to content

Commit

Permalink
Rewrite inflate memory allocation.
Browse files Browse the repository at this point in the history 
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 <[email protected]>
  • Loading branch information
@Dead2 @iii-i
Dead2 and iii-i committed May 28, 2024
1 parent 130055e commit 63e1d46
Show file tree
Hide file tree
Showing 5 changed files with 137 additions and 59 deletions.
15 changes: 11 additions & 4 deletions infback.c
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;
}
149 changes: 106 additions & 43 deletions inflate.c
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
Expand Down Expand Up @@ -120,21 +120,98 @@ 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;
state->wbits = (unsigned)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;
Expand All @@ -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);
strm->state = NULL;
free_inflate(strm);
}
return ret;
}
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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)
ZFREE_WINDOW(strm, state->window);
ZFREE(strm, strm->state);
strm->state = NULL;

/* Free allocated buffers */
free_inflate(strm);

Tracev((stderr, "inflate: end\n"));
return Z_OK;
}
Expand Down Expand Up @@ -1332,30 +1397,28 @@ 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));
if (copy == NULL)
inflate_allocs *alloc_bufs = alloc_inflate(dest);
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) {
copy->lencode = copy->codes + (state->lencode - state->codes);
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;
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);
}
memcpy(copy->window, state->window, INFLATE_ADJUST_WINDOW_SIZE((size_t)state->wsize));

dest->state = (struct internal_state *)copy;
return Z_OK;
Expand Down
14 changes: 13 additions & 1 deletion inflate.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 */
Expand Down Expand Up @@ -136,12 +145,15 @@ 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
};

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_ */
11 changes: 7 additions & 4 deletions inflate_p.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 ZALLOC_WINDOW(strm, items, size) ZALLOC(strm, items, size)
# define ZCOPY_WINDOW(dest, src, n) memcpy(dest, src, n)
# define ZFREE_WINDOW(strm, addr) ZFREE(strm, addr)
# define PAD_WINDOW PAD_64
# define WINDOW_PAD_SIZE 64
# define HINT_ALIGNED_WINDOW HINT_ALIGNED_64
/* 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. */
Expand Down
7 changes: 0 additions & 7 deletions test/infcover.c
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down Expand Up @@ -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);
Expand Down

0 comments on commit 63e1d46

Please sign in to comment.