apedec: check output buffer size after calculating actual output size

diff --git a/libavcodec/apedec.c b/libavcodec/apedec.c
index 0619358dad6a0410cb6a842410f035e6c61292bf..1567025183e128788389e9cd078e397e2e12a4ce 100644 (file)
--- a/libavcodec/apedec.c
+++ b/libavcodec/apedec.c
@@ -816,15 +816,9 @@ static int ape_decode_frame(AVCodecContext *avctx,
     int16_t *samples = data;
     uint32_t nblocks;
     int i;
-    int blockstodecode;
+    int blockstodecode, out_size;
     int bytes_used;
 
-    /* should not happen but who knows */
-    if (BLOCKS_PER_LOOP * 2 * avctx->channels > *data_size) {
-        av_log (avctx, AV_LOG_ERROR, "Output buffer is too small.\n");
-        return AVERROR(EINVAL);
-    }
-
     /* this should never be negative, but bad things will happen if it is, so
        check it just to make sure. */
     av_assert0(s->samples >= 0);
@@ -883,6 +877,13 @@ static int ape_decode_frame(AVCodecContext *avctx,
     nblocks = s->samples;
     blockstodecode = FFMIN(BLOCKS_PER_LOOP, nblocks);
 
+    out_size = blockstodecode * avctx->channels *
+               av_get_bytes_per_sample(avctx->sample_fmt);
+    if (*data_size < out_size) {
+        av_log(avctx, AV_LOG_ERROR, "Output buffer is too small.\n");
+        return AVERROR(EINVAL);
+    }
+
     s->error=0;
 
     if ((s->channels == 1) || (s->frameflags & APE_FRAMECODE_PSEUDO_STEREO))
@@ -905,9 +906,10 @@ static int ape_decode_frame(AVCodecContext *avctx,
 
     s->samples -= blockstodecode;
 
-    *data_size = blockstodecode * 2 * s->channels;
     bytes_used = s->samples ? s->ptr - s->last_ptr : buf_size;
     s->last_ptr = s->ptr;
+
+    *data_size = out_size;
     return bytes_used;
 }