dwt: check malloc calls

Signed-off-by: Diego Biurrun <diego@biurrun.de>

1 files changed, 15 insertions, 1 deletions

diff --git a/libavcodec/dwt.c b/libavcodec/dwt.c
index f9577fd1bd..675644dbe6 100644
--- a/libavcodec/dwt.c
+++ b/libavcodec/dwt.c
@@ -33,10 +33,24 @@ void ff_slice_buffer_init(slice_buffer *buf, int line_count,
     buf->line_width  = line_width;
     buf->data_count  = max_allocated_lines;
     buf->line        = av_mallocz(sizeof(IDWTELEM *) * line_count);
+    if (!buf->line)
+        return AVERROR(ENOMEM);
     buf->data_stack  = av_malloc(sizeof(IDWTELEM *) * max_allocated_lines);
+    if (!buf->data_stack) {
+        av_free(buf->line);
+        return AVERROR(ENOMEM);
+    }
 
-    for (i = 0; i < max_allocated_lines; i++)
+    for (i = 0; i < max_allocated_lines; i++) {
         buf->data_stack[i] = av_malloc(sizeof(IDWTELEM) * line_width);
+        if (!buf->data_stack[i]) {
+            for (i--; i >=0; i--)
+                av_free(buf->data_stack[i]);
+            av_free(buf->data_stack);
+            av_free(buf->line);
+            return AVERROR(ENOMEM);
+        }
+    }
 
     buf->data_stack_top = max_allocated_lines - 1;
 }