1 /*
2  * AVC helper functions for muxers
3  * Copyright (c) 2006 Baptiste Coudurier <baptiste.coudurier@smartjog.com>
4  *
5  * This file is part of FFmpeg.
6  *
7  * FFmpeg is free software; you can redistribute it and/or
8  * modify it under the terms of the GNU Lesser General Public
9  * License as published by the Free Software Foundation; either
10  * version 2.1 of the License, or (at your option) any later version.
11  *
12  * FFmpeg is distributed in the hope that it will be useful,
13  * but WITHOUT ANY WARRANTY; without even the implied warranty of
14  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
15  * Lesser General Public License for more details.
16  *
17  * You should have received a copy of the GNU Lesser General Public
18  * License along with FFmpeg; if not, write to the Free Software
19  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
20  */
21 
22 #include "libavutil/intreadwrite.h"
23 #include "libavcodec/h264.h"
24 #include "avformat.h"
25 #include "avio.h"
26 #include "avc.h"
27 
ff_avc_find_startcode_internal(const uint8_t * p,const uint8_t * end)28 static const uint8_t *ff_avc_find_startcode_internal(const uint8_t *p, const uint8_t *end)
29 {
30     const uint8_t *a = p + 4 - ((intptr_t)p & 3);
31 
32     for (end -= 3; p < a && p < end; p++) {
33         if (p[0] == 0 && p[1] == 0 && p[2] == 1)
34             return p;
35     }
36 
37     for (end -= 3; p < end; p += 4) {
38         uint32_t x = *(const uint32_t*)p;
39 //      if ((x - 0x01000100) & (~x) & 0x80008000) // little endian
40 //      if ((x - 0x00010001) & (~x) & 0x00800080) // big endian
41         if ((x - 0x01010101) & (~x) & 0x80808080) { // generic
42             if (p[1] == 0) {
43                 if (p[0] == 0 && p[2] == 1)
44                     return p;
45                 if (p[2] == 0 && p[3] == 1)
46                     return p+1;
47             }
48             if (p[3] == 0) {
49                 if (p[2] == 0 && p[4] == 1)
50                     return p+2;
51                 if (p[4] == 0 && p[5] == 1)
52                     return p+3;
53             }
54         }
55     }
56 
57     for (end += 3; p < end; p++) {
58         if (p[0] == 0 && p[1] == 0 && p[2] == 1)
59             return p;
60     }
61 
62     return end + 3;
63 }
64 
ff_avc_find_startcode(const uint8_t * p,const uint8_t * end)65 const uint8_t *ff_avc_find_startcode(const uint8_t *p, const uint8_t *end){
66     const uint8_t *out= ff_avc_find_startcode_internal(p, end);
67     if(p<out && out<end && !out[-1]) out--;
68     return out;
69 }
70 
ff_avc_parse_nal_units(AVIOContext * pb,const uint8_t * buf_in,int size)71 int ff_avc_parse_nal_units(AVIOContext *pb, const uint8_t *buf_in, int size)
72 {
73     const uint8_t *p = buf_in;
74     const uint8_t *end = p + size;
75     const uint8_t *nal_start, *nal_end;
76 
77     size = 0;
78     nal_start = ff_avc_find_startcode(p, end);
79     for (;;) {
80         while (nal_start < end && !*(nal_start++));
81         if (nal_start == end)
82             break;
83 
84         nal_end = ff_avc_find_startcode(nal_start, end);
85         avio_wb32(pb, nal_end - nal_start);
86         avio_write(pb, nal_start, nal_end - nal_start);
87         size += 4 + nal_end - nal_start;
88         nal_start = nal_end;
89     }
90     return size;
91 }
92 
ff_avc_parse_nal_units_buf(const uint8_t * buf_in,uint8_t ** buf,int * size)93 int ff_avc_parse_nal_units_buf(const uint8_t *buf_in, uint8_t **buf, int *size)
94 {
95     AVIOContext *pb;
96     int ret = avio_open_dyn_buf(&pb);
97     if(ret < 0)
98         return ret;
99 
100     ff_avc_parse_nal_units(pb, buf_in, *size);
101 
102     av_freep(buf);
103     *size = avio_close_dyn_buf(pb, buf);
104     return 0;
105 }
106 
ff_isom_write_avcc(AVIOContext * pb,const uint8_t * data,int len)107 int ff_isom_write_avcc(AVIOContext *pb, const uint8_t *data, int len)
108 {
109     AVIOContext *sps_pb = NULL, *pps_pb = NULL;
110     uint8_t *buf = NULL, *end, *start = NULL;
111     uint8_t *sps = NULL, *pps = NULL;
112     uint32_t sps_size = 0, pps_size = 0;
113     int ret, nb_sps = 0, nb_pps = 0;
114 
115     if (len <= 6)
116         return AVERROR_INVALIDDATA;
117 
118     /* check for H.264 start code */
119     if (AV_RB32(data) != 0x00000001 &&
120         AV_RB24(data) != 0x000001) {
121         avio_write(pb, data, len);
122         return 0;
123     }
124 
125     ret = ff_avc_parse_nal_units_buf(data, &buf, &len);
126     if (ret < 0)
127         return ret;
128     start = buf;
129     end = buf + len;
130 
131     ret = avio_open_dyn_buf(&sps_pb);
132     if (ret < 0)
133         goto fail;
134     ret = avio_open_dyn_buf(&pps_pb);
135     if (ret < 0)
136         goto fail;
137 
138     /* look for sps and pps */
139     while (end - buf > 4) {
140         uint32_t size;
141         uint8_t nal_type;
142         size = FFMIN(AV_RB32(buf), end - buf - 4);
143         buf += 4;
144         nal_type = buf[0] & 0x1f;
145 
146         if (nal_type == 7) { /* SPS */
147             nb_sps++;
148             if (size > UINT16_MAX || nb_sps >= H264_MAX_SPS_COUNT) {
149                 ret = AVERROR_INVALIDDATA;
150                 goto fail;
151             }
152             avio_wb16(sps_pb, size);
153             avio_write(sps_pb, buf, size);
154         } else if (nal_type == 8) { /* PPS */
155             nb_pps++;
156             if (size > UINT16_MAX || nb_pps >= H264_MAX_PPS_COUNT) {
157                 ret = AVERROR_INVALIDDATA;
158                 goto fail;
159             }
160             avio_wb16(pps_pb, size);
161             avio_write(pps_pb, buf, size);
162         }
163 
164         buf += size;
165     }
166     sps_size = avio_close_dyn_buf(sps_pb, &sps);
167     pps_size = avio_close_dyn_buf(pps_pb, &pps);
168 
169     if (sps_size < 6 || !pps_size) {
170         ret = AVERROR_INVALIDDATA;
171         goto fail;
172     }
173 
174     avio_w8(pb, 1); /* version */
175     avio_w8(pb, sps[3]); /* profile */
176     avio_w8(pb, sps[4]); /* profile compat */
177     avio_w8(pb, sps[5]); /* level */
178     avio_w8(pb, 0xff); /* 6 bits reserved (111111) + 2 bits nal size length - 1 (11) */
179     avio_w8(pb, 0xe0 | nb_sps); /* 3 bits reserved (111) + 5 bits number of sps */
180 
181     avio_write(pb, sps, sps_size);
182     avio_w8(pb, nb_pps); /* number of pps */
183     avio_write(pb, pps, pps_size);
184 
185 fail:
186     if (!sps)
187         avio_close_dyn_buf(sps_pb, &sps);
188     if (!pps)
189         avio_close_dyn_buf(pps_pb, &pps);
190     av_free(sps);
191     av_free(pps);
192     av_free(start);
193 
194     return ret;
195 }
196 
ff_avc_write_annexb_extradata(const uint8_t * in,uint8_t ** buf,int * size)197 int ff_avc_write_annexb_extradata(const uint8_t *in, uint8_t **buf, int *size)
198 {
199     uint16_t sps_size, pps_size;
200     uint8_t *out;
201     int out_size;
202 
203     *buf = NULL;
204     if (*size >= 4 && (AV_RB32(in) == 0x00000001 || AV_RB24(in) == 0x000001))
205         return 0;
206     if (*size < 11 || in[0] != 1)
207         return AVERROR_INVALIDDATA;
208 
209     sps_size = AV_RB16(&in[6]);
210     if (11 + sps_size > *size)
211         return AVERROR_INVALIDDATA;
212     pps_size = AV_RB16(&in[9 + sps_size]);
213     if (11 + sps_size + pps_size > *size)
214         return AVERROR_INVALIDDATA;
215     out_size = 8 + sps_size + pps_size;
216     out = av_mallocz(out_size + AV_INPUT_BUFFER_PADDING_SIZE);
217     if (!out)
218         return AVERROR(ENOMEM);
219     AV_WB32(&out[0], 0x00000001);
220     memcpy(out + 4, &in[8], sps_size);
221     AV_WB32(&out[4 + sps_size], 0x00000001);
222     memcpy(out + 8 + sps_size, &in[11 + sps_size], pps_size);
223     *buf = out;
224     *size = out_size;
225     return 0;
226 }
227 
ff_avc_mp4_find_startcode(const uint8_t * start,const uint8_t * end,int nal_length_size)228 const uint8_t *ff_avc_mp4_find_startcode(const uint8_t *start,
229                                          const uint8_t *end,
230                                          int nal_length_size)
231 {
232     unsigned int res = 0;
233 
234     if (end - start < nal_length_size)
235         return NULL;
236     while (nal_length_size--)
237         res = (res << 8) | *start++;
238 
239     if (res > end - start)
240         return NULL;
241 
242     return start + res;
243 }
244