[FFmpeg-devel] [PATCH 3/3] avutil/mem: Support limiting the heap usage

[wm4]

On Mon, 21 Nov 2016 16:18:58 +0100
Michael Niedermayer <michael at niedermayer.cc> wrote:

> With this it is possible to prevent OOM with untrusted media
> 
> As this fundamentally requires keeping track of allocated memory sizes
> and the memalign hack code does nearly exactly that already, this feature
> depends on it being enabled.
> 
> Signed-off-by: Michael Niedermayer <michael at niedermayer.cc>
> ---
>  libavutil/mem.c | 45 ++++++++++++++++++++++++++++++++++++++++++++-
>  libavutil/mem.h | 17 +++++++++++++++++
>  2 files changed, 61 insertions(+), 1 deletion(-)
> 
> diff --git a/libavutil/mem.c b/libavutil/mem.c
> index 6273d89..6ee07b7 100644
> --- a/libavutil/mem.c
> +++ b/libavutil/mem.c
> @@ -37,6 +37,7 @@
>  #endif
>  
>  #include "avassert.h"
> +#include "atomic.h"
>  #include "avutil.h"
>  #include "common.h"
>  #include "dynarray.h"
> @@ -69,11 +70,39 @@ void  free(void *ptr);
>   * Note that this will cost performance. */
>  
>  static size_t max_alloc_size= INT_MAX;
> +static int64_t max_heap_size= 0;
> +static volatile int64_t heap_size = 0;
>  
>  void av_max_alloc(size_t max){
>      max_alloc_size = max;
>  }
>  
> +int av_max_heap(uint64_t max)
> +{
> +#if CONFIG_MEMALIGN_HACK
> +    max_heap_size = max;
> +    return 0;
> +#else
> +    return AVERROR(ENOSYS);
> +#endif
> +}
> +
> +static int heap_change(int64_t size, int force)
> +{
> +    int64_t oldval, newval;
> +    if (!max_heap_size)
> +        return 0;
> +
> +    do {
> +        oldval = heap_size;
> +        if (!force && max_heap_size - oldval < size)
> +            return AVERROR(ENOMEM);
> +        newval = oldval + size;
> +    } while (oldval != avpriv_atomic_int64_cas(&heap_size, oldval, newval));
> +
> +    return 0;
> +}
> +
>  void *av_malloc(size_t size)
>  {
>      void *ptr = NULL;
> @@ -86,9 +115,15 @@ void *av_malloc(size_t size)
>          return NULL;
>  
>  #if CONFIG_MEMALIGN_HACK
> +
> +    if (heap_change(size + ALIGN + 8, 0) < 0)
> +        return NULL;
> +
>      ptr = malloc(size + ALIGN + 8);
> -    if (!ptr)
> +    if (!ptr) {
> +        heap_change(-((int64_t)size + ALIGN + 8), 1);
>          return ptr;
> +    }
>  
>      diff = ((((uintptr_t)ptr) + 9 + ALIGN - 1) & ~(ALIGN - 1)) - (uintptr_t)ptr;
>      av_assert0(diff>0 && diff<=ALIGN + 8);  
> @@ -164,9 +199,15 @@ void *av_realloc(void *ptr, size_t size)
>  
>      // Quick path for apparently and likly aligned realloc()
>      if (diff == ALIGN) {
> +        int64_t size_delta = size - ((int64_t *)ptr)[-1];
> +        if (heap_change(size_delta, 0) < 0)
> +            return NULL;
> +
>          ptr = realloc((char *)ptr - diff, size + diff);
>          if (ptr)
>              ptr = (char *)ptr + diff;
> +        else
> +            heap_change(-size_delta, 1);
>          if (!(((uintptr_t)ptr) & (ALIGN - 1))) {
>              if (ptr)
>                  ((int64_t *)ptr)[-1] = size;
> @@ -250,8 +291,10 @@ void av_free(void *ptr)
>  #if CONFIG_MEMALIGN_HACK
>      if (ptr) {
>          int v= ((char *)ptr)[-9];
> +        int64_t size = ((int64_t *)ptr)[-1];
>          av_assert0(v>0 && v<=ALIGN + 8);  
>          free((char *)ptr - v);
> +        heap_change(-size, 1);
>      }
>  #elif HAVE_ALIGNED_MALLOC
>      _aligned_free(ptr);
> diff --git a/libavutil/mem.h b/libavutil/mem.h
> index f9d8884..9cce433 100644
> --- a/libavutil/mem.h
> +++ b/libavutil/mem.h
> @@ -689,6 +689,23 @@ static inline int av_size_mult(size_t a, size_t b, size_t *r)
>  void av_max_alloc(size_t max);
>  
>  /**
> + * Set the maximum heap size that may be allocated.
> + *
> + * The value specified with this function is effective for all libavutil's @ref
> + * lavu_mem_funcs "heap management functions."
> + *
> + * By default, there is no limit.
> + * If this is set lower than the currently allocated space then behavior is
> + * undefined.
> + * Allocations prior to this being set are potentially not kept track of.
> + *
> + * @param max Value to be set as the new maximum size
> + *
> + * @return >= 0 or a negative error code in case of failure
> + */
> +int av_max_heap(uint64_t max);
> +
> +/**
>   * @}
>   * @}
>   */

"No."

Not another questionable global state thing please.

This isn't any better than ulimit anyway.

Besides, using this in practice would only trigger tons of malloc
failure handling bugs, so practical uses would probably involve
tearing down the whole process.