lavc/videotoolboxenc: fix allocated buffer size

Allocates a buffer with padding, and with enough room for start codes when
length codes are smaller.
Signed-off-by: Rick Kern <kernrj@gmail.com>

libavcodec/videotoolboxenc.c

@@ -199,6 +199,45 @@ static void vtenc_q_push(VTEncContext *vtctx, CMSampleBufferRef buffer)
     pthread_mutex_unlock(&vtctx->lock);
 }
 
+static int count_nalus(size_t length_code_size,
+                       CMSampleBufferRef sample_buffer,
+                       int *count)
+{
+    size_t offset = 0;
+    int status;
+    int nalu_ct = 0;
+    uint8_t size_buf[4];
+    size_t src_size = CMSampleBufferGetTotalSampleSize(sample_buffer);
+    CMBlockBufferRef block = CMSampleBufferGetDataBuffer(sample_buffer);
+
+    if (length_code_size > 4)
+        return AVERROR_INVALIDDATA;
+
+    while (offset < src_size) {
+        size_t curr_src_len;
+        size_t box_len = 0;
+        size_t i;
+
+        status = CMBlockBufferCopyDataBytes(block,
+                                            offset,
+                                            length_code_size,
+                                            size_buf);
+
+        for (i = 0; i < length_code_size; i++) {
+            box_len <<= 8;
+            box_len |= size_buf[i];
+        }
+
+        curr_src_len = box_len + length_code_size;
+        offset += curr_src_len;
+
+        nalu_ct++;
+    }
+
+    *count = nalu_ct;
+    return 0;
+}
+
 static CMVideoCodecType get_cm_codec_type(enum AVCodecID id)
 {
     switch (id) {
@@ -207,12 +246,6 @@ static CMVideoCodecType get_cm_codec_type(enum AVCodecID id)
     }
 }
 
-static void vtenc_free_block(void *opaque, uint8_t *data)
-{
-    CMBlockBufferRef block = opaque;
-    CFRelease(block);
-}
-
 /**
  * Get the parameter sets from a CMSampleBufferRef.
  * @param dst If *dst isn't NULL, the parameters are copied into existing
@@ -1080,54 +1113,6 @@ static void vtenc_get_frame_info(CMSampleBufferRef buffer, bool *is_key_frame)
     }
 }
 
-/**
- * Replaces length codes with H.264 Annex B start codes.
- * length_code_size must equal sizeof(start_code).
- * On failure, the contents of data may have been modified.
- *
- * @param length_code_size Byte length of each length code
- * @param data Call with NAL units prefixed with length codes.
- *             On success, the length codes are replace with
- *             start codes.
- * @param size Length of data, excluding any padding.
- * @return 0 on success
- *         AVERROR_BUFFER_TOO_SMALL if length code size is smaller
- *         than a start code or if a length_code in data specifies
- *         data beyond the end of its buffer.
- */
-static int replace_length_codes(size_t  length_code_size,
-                                uint8_t *data,
-                                size_t  size)
-{
-    size_t remaining_size = size;
-
-    if (length_code_size != sizeof(start_code)) {
-        av_log(NULL, AV_LOG_ERROR, "Start code size and length code size not equal.\n");
-        return AVERROR_BUFFER_TOO_SMALL;
-    }
-
-    while (remaining_size > 0) {
-        size_t box_len = 0;
-        size_t i;
-
-        for (i = 0; i < length_code_size; i++) {
-            box_len <<= 8;
-            box_len |= data[i];
-        }
-
-        if (remaining_size < box_len + sizeof(start_code)) {
-            av_log(NULL, AV_LOG_ERROR, "Length is out of range.\n");
-            AVERROR_BUFFER_TOO_SMALL;
-        }
-
-        memcpy(data, start_code, sizeof(start_code));
-        data += box_len + sizeof(start_code);
-        remaining_size -= box_len + sizeof(start_code);
-    }
-
-    return 0;
-}
-
 /**
  * Copies NAL units and replaces length codes with
  * H.264 Annex B start codes. On failure, the contents of
@@ -1148,14 +1133,19 @@ static int replace_length_codes(size_t  length_code_size,
  *         the end of its buffer.
  */
 static int copy_replace_length_codes(
+    AVCodecContext *avctx,
     size_t        length_code_size,
-    const uint8_t *src_data,
-    size_t        src_size,
+    CMSampleBufferRef sample_buffer,
     uint8_t       *dst_data,
     size_t        dst_size)
 {
+    size_t src_size = CMSampleBufferGetTotalSampleSize(sample_buffer);
     size_t remaining_src_size = src_size;
     size_t remaining_dst_size = dst_size;
+    size_t src_offset = 0;
+    int status;
+    uint8_t size_buf[4];
+    CMBlockBufferRef block = CMSampleBufferGetDataBuffer(sample_buffer);
 
     if (length_code_size > 4) {
         return AVERROR_INVALIDDATA;
@@ -1168,11 +1158,19 @@ static int copy_replace_length_codes(
         size_t i;
 
         uint8_t       *dst_box;
-        const uint8_t *src_box;
+
+        status = CMBlockBufferCopyDataBytes(block,
+                                            src_offset,
+                                            length_code_size,
+                                            size_buf);
+        if (status) {
+            av_log(avctx, AV_LOG_ERROR, "Cannot copy length: %d\n", status);
+            return AVERROR_EXTERNAL;
+        }
 
         for (i = 0; i < length_code_size; i++) {
             box_len <<= 8;
-            box_len |= src_data[i];
+            box_len |= size_buf[i];
         }
 
         curr_src_len = box_len + length_code_size;
@@ -1187,12 +1185,19 @@ static int copy_replace_length_codes(
         }
 
         dst_box = dst_data + sizeof(start_code);
-        src_box = src_data + length_code_size;
 
         memcpy(dst_data, start_code, sizeof(start_code));
-        memcpy(dst_box,  src_box,    box_len);
+        status = CMBlockBufferCopyDataBytes(block,
+                                            src_offset + length_code_size,
+                                            box_len,
+                                            dst_box);
+
+        if (status) {
+            av_log(avctx, AV_LOG_ERROR, "Cannot copy data: %d\n", status);
+            return AVERROR_EXTERNAL;
+        }
 
-        src_data += curr_src_len;
+        src_offset += curr_src_len;
         dst_data += curr_dst_len;
 
         remaining_src_size -= curr_src_len;
@@ -1212,16 +1217,15 @@ static int vtenc_cm_to_avpacket(
     int     status;
     bool    is_key_frame;
     bool    add_header;
-    char    *buf_data;
     size_t  length_code_size;
     size_t  header_size = 0;
     size_t  in_buf_size;
+    size_t  out_buf_size;
     int64_t dts_delta;
     int64_t time_base_num;
+    int nalu_count;
     CMTime  pts;
     CMTime  dts;
-
-    CMBlockBufferRef            block;
     CMVideoFormatDescriptionRef vid_fmt;
 
 
@@ -1235,82 +1239,42 @@ static int vtenc_cm_to_avpacket(
         vid_fmt = CMSampleBufferGetFormatDescription(sample_buffer);
         if (!vid_fmt) {
             av_log(avctx, AV_LOG_ERROR, "Cannot get format description.\n");
+            return AVERROR_EXTERNAL;
         }
 
         int status = get_params_size(avctx, vid_fmt, &header_size);
         if (status) return status;
     }
 
-    block = CMSampleBufferGetDataBuffer(sample_buffer);
-    if (!block) {
-        av_log(avctx, AV_LOG_ERROR, "Could not get block buffer from sample buffer.\n");
-        return AVERROR_EXTERNAL;
-    }
-
-
-    status = CMBlockBufferGetDataPointer(block, 0, &in_buf_size, NULL, &buf_data);
-    if (status) {
-        av_log(avctx, AV_LOG_ERROR, "Error: cannot get data pointer: %d\n", status);
-        return AVERROR_EXTERNAL;
-    }
-
-    size_t out_buf_size = header_size + in_buf_size;
-    bool can_reuse_cmbuffer = !add_header &&
-                              !pkt->data  &&
-                              length_code_size == sizeof(start_code);
-
-    av_init_packet(pkt);
-
-    if (can_reuse_cmbuffer) {
-        AVBufferRef* buf_ref = av_buffer_create(
-            buf_data,
-            out_buf_size,
-            vtenc_free_block,
-            block,
-            0
-        );
-
-        if (!buf_ref) return AVERROR(ENOMEM);
-
-        CFRetain(block);
-
-        pkt->buf  = buf_ref;
-        pkt->data = buf_data;
-        pkt->size = in_buf_size;
+    status = count_nalus(length_code_size, sample_buffer, &nalu_count);
+    if(status)
+        return status;
 
-        status = replace_length_codes(length_code_size, pkt->data, pkt->size);
-        if (status) {
-            av_log(avctx, AV_LOG_ERROR, "Error replacing length codes: %d\n", status);
-            return status;
-        }
-    } else {
-        if (!pkt->data) {
-            status = av_new_packet(pkt, out_buf_size);
-            if(status) return status;
-        }
+    in_buf_size = CMSampleBufferGetTotalSampleSize(sample_buffer);
+    out_buf_size = header_size +
+                   in_buf_size +
+                   nalu_count * ((int)sizeof(start_code) - (int)length_code_size);
 
-        if (pkt->size < out_buf_size) {
-            av_log(avctx, AV_LOG_ERROR, "Error: packet's buffer is too small.\n");
-            return AVERROR_BUFFER_TOO_SMALL;
-        }
+    status = ff_alloc_packet2(avctx, pkt, out_buf_size, out_buf_size);
+    if (status < 0)
+        return status;
 
-        if (add_header) {
-            status = copy_param_sets(avctx, vid_fmt, pkt->data, out_buf_size);
-            if(status) return status;
-        }
+    if (add_header) {
+        status = copy_param_sets(avctx, vid_fmt, pkt->data, out_buf_size);
+        if(status) return status;
+    }
 
-        status = copy_replace_length_codes(
-            length_code_size,
-            buf_data,
-            in_buf_size,
-            pkt->data + header_size,
-            pkt->size - header_size
-        );
+    status = copy_replace_length_codes(
+        avctx,
+        length_code_size,
+        sample_buffer,
+        pkt->data + header_size,
+        pkt->size - header_size
+    );
 
-        if (status) {
-            av_log(avctx, AV_LOG_ERROR, "Error copying packet data: %d", status);
-            return status;
-        }
+    if (status) {
+        av_log(avctx, AV_LOG_ERROR, "Error copying packet data: %d", status);
+        return status;
     }
 
     if (is_key_frame) {
@@ -1333,6 +1297,7 @@ static int vtenc_cm_to_avpacket(
     time_base_num = avctx->time_base.num;
     pkt->pts = pts.value / time_base_num;
     pkt->dts = dts.value / time_base_num - dts_delta;
+    pkt->size = out_buf_size;
 
     return 0;
 }