Addresses ticket #6864
Reviewed-by: Michael Niedermayer <michael@niedermayer.cc>
Signed-off-by: James Almer <jamrial@gmail.com>
... | ... |
@@ -20,6 +20,7 @@ |
20 | 20 |
*/ |
21 | 21 |
|
22 | 22 |
#include "libavutil/intreadwrite.h" |
23 |
+#include "libavcodec/h264.h" |
|
23 | 24 |
#include "avformat.h" |
24 | 25 |
#include "avio.h" |
25 | 26 |
#include "avc.h" |
... | ... |
@@ -105,10 +106,11 @@ int ff_avc_parse_nal_units_buf(const uint8_t *buf_in, uint8_t **buf, int *size) |
105 | 105 |
|
106 | 106 |
int ff_isom_write_avcc(AVIOContext *pb, const uint8_t *data, int len) |
107 | 107 |
{ |
108 |
+ AVIOContext *sps_pb = NULL, *pps_pb = NULL; |
|
108 | 109 |
uint8_t *buf = NULL, *end, *start = NULL; |
109 | 110 |
uint8_t *sps = NULL, *pps = NULL; |
110 | 111 |
uint32_t sps_size = 0, pps_size = 0; |
111 |
- int ret; |
|
112 |
+ int ret, nb_sps = 0, nb_pps = 0; |
|
112 | 113 |
|
113 | 114 |
if (len <= 6) |
114 | 115 |
return AVERROR_INVALIDDATA; |
... | ... |
@@ -126,6 +128,13 @@ int ff_isom_write_avcc(AVIOContext *pb, const uint8_t *data, int len) |
126 | 126 |
start = buf; |
127 | 127 |
end = buf + len; |
128 | 128 |
|
129 |
+ ret = avio_open_dyn_buf(&sps_pb); |
|
130 |
+ if (ret < 0) |
|
131 |
+ goto fail; |
|
132 |
+ ret = avio_open_dyn_buf(&pps_pb); |
|
133 |
+ if (ret < 0) |
|
134 |
+ goto fail; |
|
135 |
+ |
|
129 | 136 |
/* look for sps and pps */ |
130 | 137 |
while (end - buf > 4) { |
131 | 138 |
uint32_t size; |
... | ... |
@@ -135,35 +144,51 @@ int ff_isom_write_avcc(AVIOContext *pb, const uint8_t *data, int len) |
135 | 135 |
nal_type = buf[0] & 0x1f; |
136 | 136 |
|
137 | 137 |
if (nal_type == 7) { /* SPS */ |
138 |
- sps = buf; |
|
139 |
- sps_size = size; |
|
138 |
+ nb_sps++; |
|
139 |
+ if (size > UINT16_MAX || nb_sps >= H264_MAX_SPS_COUNT) { |
|
140 |
+ ret = AVERROR_INVALIDDATA; |
|
141 |
+ goto fail; |
|
142 |
+ } |
|
143 |
+ avio_wb16(sps_pb, size); |
|
144 |
+ avio_write(sps_pb, buf, size); |
|
140 | 145 |
} else if (nal_type == 8) { /* PPS */ |
141 |
- pps = buf; |
|
142 |
- pps_size = size; |
|
146 |
+ nb_pps++; |
|
147 |
+ if (size > UINT16_MAX || nb_pps >= H264_MAX_PPS_COUNT) { |
|
148 |
+ ret = AVERROR_INVALIDDATA; |
|
149 |
+ goto fail; |
|
150 |
+ } |
|
151 |
+ avio_wb16(pps_pb, size); |
|
152 |
+ avio_write(pps_pb, buf, size); |
|
143 | 153 |
} |
144 | 154 |
|
145 | 155 |
buf += size; |
146 | 156 |
} |
157 |
+ sps_size = avio_close_dyn_buf(sps_pb, &sps); |
|
158 |
+ pps_size = avio_close_dyn_buf(pps_pb, &pps); |
|
147 | 159 |
|
148 |
- if (!sps || !pps || sps_size < 4 || sps_size > UINT16_MAX || pps_size > UINT16_MAX) { |
|
160 |
+ if (sps_size < 6 || !pps_size) { |
|
149 | 161 |
ret = AVERROR_INVALIDDATA; |
150 | 162 |
goto fail; |
151 | 163 |
} |
152 | 164 |
|
153 | 165 |
avio_w8(pb, 1); /* version */ |
154 |
- avio_w8(pb, sps[1]); /* profile */ |
|
155 |
- avio_w8(pb, sps[2]); /* profile compat */ |
|
156 |
- avio_w8(pb, sps[3]); /* level */ |
|
166 |
+ avio_w8(pb, sps[3]); /* profile */ |
|
167 |
+ avio_w8(pb, sps[4]); /* profile compat */ |
|
168 |
+ avio_w8(pb, sps[5]); /* level */ |
|
157 | 169 |
avio_w8(pb, 0xff); /* 6 bits reserved (111111) + 2 bits nal size length - 1 (11) */ |
158 |
- avio_w8(pb, 0xe1); /* 3 bits reserved (111) + 5 bits number of sps (00001) */ |
|
170 |
+ avio_w8(pb, 0xe0 | nb_sps); /* 3 bits reserved (111) + 5 bits number of sps */ |
|
159 | 171 |
|
160 |
- avio_wb16(pb, sps_size); |
|
161 | 172 |
avio_write(pb, sps, sps_size); |
162 |
- avio_w8(pb, 1); /* number of pps */ |
|
163 |
- avio_wb16(pb, pps_size); |
|
173 |
+ avio_w8(pb, nb_pps); /* number of pps */ |
|
164 | 174 |
avio_write(pb, pps, pps_size); |
165 | 175 |
|
166 | 176 |
fail: |
177 |
+ if (!sps) |
|
178 |
+ avio_close_dyn_buf(sps_pb, &sps); |
|
179 |
+ if (!pps) |
|
180 |
+ avio_close_dyn_buf(pps_pb, &pps); |
|
181 |
+ av_free(sps); |
|
182 |
+ av_free(pps); |
|
167 | 183 |
av_free(start); |
168 | 184 |
|
169 | 185 |
return ret; |