lavc: check that extended_data is properly set in avcodec_encode_audio2().

[ffmpeg] / libavcodec / utils.c

diff --git a/libavcodec/utils.c b/libavcodec/utils.c
index bfd470b24e93abc769e40e553185680ae7efbed2..9ec4a9e8484568cb2fb001c3b2fb7d0d395a9769 100644 (file)
--- a/libavcodec/utils.c
+++ b/libavcodec/utils.c
@@ -857,11 +857,59 @@ int ff_alloc_packet(AVPacket *avpkt, int size)
     }
 }
 
+/**
+ * Pad last frame with silence.
+ */
+static int pad_last_frame(AVCodecContext *s, AVFrame **dst, const AVFrame *src)
+{
+    AVFrame *frame = NULL;
+    uint8_t *buf   = NULL;
+    int ret;
+
+    if (!(frame = avcodec_alloc_frame()))
+        return AVERROR(ENOMEM);
+    *frame = *src;
+
+    if ((ret = av_samples_get_buffer_size(&frame->linesize[0], s->channels,
+                                          s->frame_size, s->sample_fmt, 0)) < 0)
+        goto fail;
+
+    if (!(buf = av_malloc(ret))) {
+        ret = AVERROR(ENOMEM);
+        goto fail;
+    }
+
+    frame->nb_samples = s->frame_size;
+    if ((ret = avcodec_fill_audio_frame(frame, s->channels, s->sample_fmt,
+                                        buf, ret, 0)) < 0)
+        goto fail;
+    if ((ret = av_samples_copy(frame->extended_data, src->extended_data, 0, 0,
+                               src->nb_samples, s->channels, s->sample_fmt)) < 0)
+        goto fail;
+    if ((ret = av_samples_set_silence(frame->extended_data, src->nb_samples,
+                                      frame->nb_samples - src->nb_samples,
+                                      s->channels, s->sample_fmt)) < 0)
+        goto fail;
+
+    *dst = frame;
+
+    return 0;
+
+fail:
+    if (frame->extended_data != frame->data)
+        av_freep(&frame->extended_data);
+    av_freep(&buf);
+    av_freep(&frame);
+    return ret;
+}
+
 int attribute_align_arg avcodec_encode_audio2(AVCodecContext *avctx,
                                               AVPacket *avpkt,
                                               const AVFrame *frame,
                                               int *got_packet_ptr)
 {
+    AVFrame tmp;
+    AVFrame *padded_frame = NULL;
     int ret;
     int user_packet = !!avpkt->data;
 
@@ -870,23 +918,49 @@ int attribute_align_arg avcodec_encode_audio2(AVCodecContext *avctx,
     if (!(avctx->codec->capabilities & CODEC_CAP_DELAY) && !frame) {
         av_free_packet(avpkt);
         av_init_packet(avpkt);
-        avpkt->size = 0;
         return 0;
     }
 
+    /* ensure that extended_data is properly set */
+    if (frame && !frame->extended_data) {
+        if (av_sample_fmt_is_planar(avctx->sample_fmt) &&
+            avctx->channels > AV_NUM_DATA_POINTERS) {
+            av_log(avctx, AV_LOG_ERROR, "Encoding to a planar sample format, "
+                   "with more than %d channels, but extended_data is not set.\n",
+                   AV_NUM_DATA_POINTERS);
+            return AVERROR(EINVAL);
+        }
+        av_log(avctx, AV_LOG_WARNING, "extended_data is not set.\n");
+
+        tmp = *frame;
+        tmp.extended_data = tmp.data;
+        frame = &tmp;
+    }
+
     /* check for valid frame size */
     if (frame) {
         if (avctx->codec->capabilities & CODEC_CAP_SMALL_LAST_FRAME) {
             if (frame->nb_samples > avctx->frame_size)
                 return AVERROR(EINVAL);
         } else if (!(avctx->codec->capabilities & CODEC_CAP_VARIABLE_FRAME_SIZE)) {
+            if (frame->nb_samples < avctx->frame_size &&
+                !avctx->internal->last_audio_frame) {
+                ret = pad_last_frame(avctx, &padded_frame, frame);
+                if (ret < 0)
+                    return ret;
+
+                frame = padded_frame;
+                avctx->internal->last_audio_frame = 1;
+            }
+
             if (frame->nb_samples != avctx->frame_size)
                 return AVERROR(EINVAL);
         }
     }
 
-        ret = avctx->codec->encode2(avctx, avpkt, frame, got_packet_ptr);
-        if (!ret && *got_packet_ptr) {
+    ret = avctx->codec->encode2(avctx, avpkt, frame, got_packet_ptr);
+    if (!ret) {
+        if (*got_packet_ptr) {
             if (!(avctx->codec->capabilities & CODEC_CAP_DELAY)) {
                 if (avpkt->pts == AV_NOPTS_VALUE)
                     avpkt->pts = frame->pts;
@@ -899,7 +973,6 @@ int attribute_align_arg avcodec_encode_audio2(AVCodecContext *avctx,
             avpkt->size = 0;
         }
 
-    if (!ret) {
         if (!user_packet && avpkt->size) {
             uint8_t *new_data = av_realloc(avpkt->data, avpkt->size);
             if (new_data)
@@ -909,14 +982,24 @@ int attribute_align_arg avcodec_encode_audio2(AVCodecContext *avctx,
         avctx->frame_number++;
     }
 
-    if (ret < 0 || !*got_packet_ptr)
+    if (ret < 0 || !*got_packet_ptr) {
         av_free_packet(avpkt);
+        av_init_packet(avpkt);
+        return ret;
+    }
 
     /* NOTE: if we add any audio encoders which output non-keyframe packets,
              this needs to be moved to the encoders, but for now we can do it
              here to simplify things */
     avpkt->flags |= AV_PKT_FLAG_KEY;
 
+    if (padded_frame) {
+        av_freep(&padded_frame->data[0]);
+        if (padded_frame->extended_data != padded_frame->data)
+            av_freep(&padded_frame->extended_data);
+        av_freep(&padded_frame);
+    }
+
     return ret;
 }