libavformat/nal.c

   1 /*
   2  * NAL helper functions for muxers
   3  *
   4  * This file is part of FFmpeg.
   5  *
   6  * FFmpeg is free software; you can redistribute it and/or
   7  * modify it under the terms of the GNU Lesser General Public
   8  * License as published by the Free Software Foundation; either
   9  * version 2.1 of the License, or (at your option) any later version.
  10  *
  11  * FFmpeg is distributed in the hope that it will be useful,
  12  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  13  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  14  * Lesser General Public License for more details.
  15  *
  16  * You should have received a copy of the GNU Lesser General Public
  17  * License along with FFmpeg; if not, write to the Free Software
  18  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
  19  */
  20
  21 #include <stdint.h>
  22 #include <string.h>
  23
  24 #include "libavutil/mem.h"
  25 #include "libavutil/error.h"
  26 #include "libavcodec/defs.h"
  27 #include "avio.h"
  28 #include "avio_internal.h"
  29 #include "nal.h"
  30
  31 static const uint8_t *nal_find_startcode_internal(const uint8_t *p, const uint8_t *end)
  32 {
  33     const uint8_t *a = p + 4 - ((intptr_t)p & 3);
  34
  35     for (end -= 3; p < a && p < end; p++) {
  36         if (p[0] == 0 && p[1] == 0 && p[2] == 1)
  37             return p;
  38     }
  39
  40     for (end -= 3; p < end; p += 4) {
  41         uint32_t x = *(const uint32_t*)p;
  42 //      if ((x - 0x01000100) & (~x) & 0x80008000) // little endian
  43 //      if ((x - 0x00010001) & (~x) & 0x00800080) // big endian
  44         if ((x - 0x01010101) & (~x) & 0x80808080) { // generic
  45             if (p[1] == 0) {
  46                 if (p[0] == 0 && p[2] == 1)
  47                     return p;
  48                 if (p[2] == 0 && p[3] == 1)
  49                     return p+1;
  50             }
  51             if (p[3] == 0) {
  52                 if (p[2] == 0 && p[4] == 1)
  53                     return p+2;
  54                 if (p[4] == 0 && p[5] == 1)
  55                     return p+3;
  56             }
  57         }
  58     }
  59
  60     for (end += 3; p < end; p++) {
  61         if (p[0] == 0 && p[1] == 0 && p[2] == 1)
  62             return p;
  63     }
  64
  65     return end + 3;
  66 }
  67
  68 const uint8_t *ff_nal_find_startcode(const uint8_t *p, const uint8_t *end){
  69     const uint8_t *out = nal_find_startcode_internal(p, end);
  70     if(p<out && out<end && !out[-1]) out--;
  71     return out;
  72 }
  73
  74 static int nal_parse_units(AVIOContext *pb, NALUList *list,
  75                            const uint8_t *buf_in, int size)
  76 {
  77     const uint8_t *p = buf_in;
  78     const uint8_t *end = p + size;
  79     const uint8_t *nal_start, *nal_end;
  80
  81     size = 0;
  82     nal_start = ff_nal_find_startcode(p, end);
  83     for (;;) {
  84         const size_t nalu_limit = SIZE_MAX / sizeof(*list->nalus);
  85         while (nal_start < end && !*(nal_start++));
  86         if (nal_start == end)
  87             break;
  88
  89         nal_end = ff_nal_find_startcode(nal_start, end);
  90         if (pb) {
  91             avio_wb32(pb, nal_end - nal_start);
  92             avio_write(pb, nal_start, nal_end - nal_start);
  93         } else if (list->nb_nalus >= nalu_limit) {
  94             return AVERROR(ERANGE);
  95         } else {
  96             NALU *tmp = av_fast_realloc(list->nalus, &list->nalus_array_size,
  97                                         (list->nb_nalus + 1) * sizeof(*list->nalus));
  98             if (!tmp)
  99                 return AVERROR(ENOMEM);
 100             list->nalus = tmp;
 101             tmp[list->nb_nalus++] = (NALU){ .offset = nal_start - p,
 102                                             .size   = nal_end - nal_start };
 103         }
 104         size += 4 + nal_end - nal_start;
 105         nal_start = nal_end;
 106     }
 107     return size;
 108 }
 109
 110 int ff_nal_parse_units(AVIOContext *pb, const uint8_t *buf_in, int size)
 111 {
 112     return nal_parse_units(pb, NULL, buf_in, size);
 113 }
 114
 115 int ff_nal_units_create_list(NALUList *list, const uint8_t *buf, int size)
 116 {
 117     list->nb_nalus = 0;
 118     return nal_parse_units(NULL, list, buf, size);
 119 }
 120
 121 void ff_nal_units_write_list(const NALUList *list, AVIOContext *pb,
 122                              const uint8_t *buf)
 123 {
 124     for (unsigned i = 0; i < list->nb_nalus; i++) {
 125         avio_wb32(pb, list->nalus[i].size);
 126         avio_write(pb, buf + list->nalus[i].offset, list->nalus[i].size);
 127     }
 128 }
 129
 130 int ff_nal_parse_units_buf(const uint8_t *buf_in, uint8_t **buf, int *size)
 131 {
 132     AVIOContext *pb;
 133     int ret = avio_open_dyn_buf(&pb);
 134     if(ret < 0)
 135         return ret;
 136
 137     ff_nal_parse_units(pb, buf_in, *size);
 138
 139     *size = avio_close_dyn_buf(pb, buf);
 140     return 0;
 141 }
 142
 143 const uint8_t *ff_nal_mp4_find_startcode(const uint8_t *start,
 144                                          const uint8_t *end,
 145                                          int nal_length_size)
 146 {
 147     unsigned int res = 0;
 148
 149     if (end - start < nal_length_size)
 150         return NULL;
 151     while (nal_length_size--)
 152         res = (res << 8) | *start++;
 153
 154     if (res > end - start)
 155         return NULL;
 156
 157     return start + res;
 158 }
 159
 160 uint8_t *ff_nal_unit_extract_rbsp(const uint8_t *src, uint32_t src_len,
 161                                   uint32_t *dst_len, int header_len)
 162 {
 163     uint8_t *dst;
 164     uint32_t i, len;
 165
 166     dst = av_malloc(src_len + AV_INPUT_BUFFER_PADDING_SIZE);
 167     if (!dst)
 168         return NULL;
 169
 170     /* NAL unit header */
 171     i = len = 0;
 172     while (i < header_len && i < src_len)
 173         dst[len++] = src[i++];
 174
 175     while (i + 2 < src_len)
 176         if (!src[i] && !src[i + 1] && src[i + 2] == 3) {
 177             dst[len++] = src[i++];
 178             dst[len++] = src[i++];
 179             i++; // remove emulation_prevention_three_byte
 180         } else
 181             dst[len++] = src[i++];
 182
 183     while (i < src_len)
 184         dst[len++] = src[i++];
 185
 186     memset(dst + len, 0, AV_INPUT_BUFFER_PADDING_SIZE);
 187
 188     *dst_len = len;
 189     return dst;
 190 }