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

[Michael Niedermayer]

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);
+
+/**
  * @}
  * @}
  */
-- 
2.10.2