Commit 03b82b3a authored by Kieran Kunhya's avatar Kieran Kunhya Committed by Your Name

h2645_parse: Allocate a single buffer per packet

Drastically reduces memory usage on pathological streams.
Fixes ticket #6789
parent 3357b68b
......@@ -31,11 +31,10 @@
#include "h2645_parse.h"
int ff_h2645_extract_rbsp(const uint8_t *src, int length,
H2645NAL *nal, int small_padding)
H2645RBSP *rbsp, H2645NAL *nal, int small_padding)
{
int i, si, di;
uint8_t *dst;
int64_t padding = small_padding ? 0 : MAX_MBPAIR_SIZE;
nal->skipped_bytes = 0;
#define STARTCODE_TEST \
......@@ -92,11 +91,7 @@ int ff_h2645_extract_rbsp(const uint8_t *src, int length,
} else if (i > length)
i = length;
av_fast_padded_malloc(&nal->rbsp_buffer, &nal->rbsp_buffer_size,
length + padding);
if (!nal->rbsp_buffer)
return AVERROR(ENOMEM);
nal->rbsp_buffer = &rbsp->rbsp_buffer[rbsp->rbsp_buffer_size];
dst = nal->rbsp_buffer;
memcpy(dst, src, i);
......@@ -145,6 +140,8 @@ nsc:
nal->size = di;
nal->raw_data = src;
nal->raw_size = si;
rbsp->rbsp_buffer_size += si;
return si;
}
......@@ -270,9 +267,14 @@ int ff_h2645_packet_split(H2645Packet *pkt, const uint8_t *buf, int length,
GetByteContext bc;
int consumed, ret = 0;
int next_avc = is_nalff ? 0 : length;
int64_t padding = small_padding ? 0 : MAX_MBPAIR_SIZE;
bytestream2_init(&bc, buf, length);
av_fast_padded_malloc(&pkt->rbsp.rbsp_buffer, &pkt->rbsp.rbsp_buffer_alloc_size, length + padding);
if (!pkt->rbsp.rbsp_buffer)
return AVERROR(ENOMEM);
pkt->rbsp.rbsp_buffer_size = 0;
pkt->nb_nals = 0;
while (bytestream2_get_bytes_left(&bc) >= 4) {
H2645NAL *nal;
......@@ -341,7 +343,7 @@ int ff_h2645_packet_split(H2645Packet *pkt, const uint8_t *buf, int length,
}
nal = &pkt->nals[pkt->nb_nals];
consumed = ff_h2645_extract_rbsp(bc.buffer, extract_length, nal, small_padding);
consumed = ff_h2645_extract_rbsp(bc.buffer, extract_length, &pkt->rbsp, nal, small_padding);
if (consumed < 0)
return consumed;
......@@ -385,9 +387,10 @@ void ff_h2645_packet_uninit(H2645Packet *pkt)
{
int i;
for (i = 0; i < pkt->nals_allocated; i++) {
av_freep(&pkt->nals[i].rbsp_buffer);
av_freep(&pkt->nals[i].skipped_bytes_pos);
}
av_freep(&pkt->nals);
pkt->nals_allocated = 0;
av_freep(&pkt->rbsp.rbsp_buffer);
pkt->rbsp.rbsp_buffer_alloc_size = pkt->rbsp.rbsp_buffer_size = 0;
}
......@@ -30,7 +30,6 @@
typedef struct H2645NAL {
uint8_t *rbsp_buffer;
int rbsp_buffer_size;
int size;
const uint8_t *data;
......@@ -65,9 +64,16 @@ typedef struct H2645NAL {
int ref_idc;
} H2645NAL;
typedef struct H2645RBSP {
uint8_t *rbsp_buffer;
int rbsp_buffer_alloc_size;
int rbsp_buffer_size;
} H2645RBSP;
/* an input packet split into unescaped NAL units */
typedef struct H2645Packet {
H2645NAL *nals;
H2645RBSP rbsp;
int nb_nals;
int nals_allocated;
} H2645Packet;
......@@ -75,7 +81,7 @@ typedef struct H2645Packet {
/**
* Extract the raw (unescaped) bitstream.
*/
int ff_h2645_extract_rbsp(const uint8_t *src, int length,
int ff_h2645_extract_rbsp(const uint8_t *src, int length, H2645RBSP *rbsp,
H2645NAL *nal, int small_padding);
/**
......
......@@ -243,6 +243,7 @@ static inline int parse_nal_units(AVCodecParserContext *s,
const uint8_t * const buf, int buf_size)
{
H264ParseContext *p = s->priv_data;
H2645RBSP rbsp = { NULL };
H2645NAL nal = { NULL };
int buf_index, next_avc;
unsigned int pps_id;
......@@ -263,6 +264,10 @@ static inline int parse_nal_units(AVCodecParserContext *s,
if (!buf_size)
return 0;
av_fast_padded_malloc(&rbsp.rbsp_buffer, &rbsp.rbsp_buffer_alloc_size, buf_size);
if (!rbsp.rbsp_buffer)
return AVERROR(ENOMEM);
buf_index = 0;
next_avc = p->is_avc ? 0 : buf_size;
for (;;) {
......@@ -300,7 +305,7 @@ static inline int parse_nal_units(AVCodecParserContext *s,
}
break;
}
consumed = ff_h2645_extract_rbsp(buf + buf_index, src_length, &nal, 1);
consumed = ff_h2645_extract_rbsp(buf + buf_index, src_length, &rbsp, &nal, 1);
if (consumed < 0)
break;
......@@ -544,18 +549,18 @@ static inline int parse_nal_units(AVCodecParserContext *s,
p->last_frame_num = p->poc.frame_num;
}
av_freep(&nal.rbsp_buffer);
av_freep(&rbsp.rbsp_buffer);
return 0; /* no need to evaluate the rest */
}
}
if (q264) {
av_freep(&nal.rbsp_buffer);
av_freep(&rbsp.rbsp_buffer);
return 0;
}
/* didn't find a picture! */
av_log(avctx, AV_LOG_ERROR, "missing picture in access unit with size %d\n", buf_size);
fail:
av_freep(&nal.rbsp_buffer);
av_freep(&rbsp.rbsp_buffer);
return -1;
}
......
......@@ -56,6 +56,7 @@ static int generate_fake_vps(QSVEncContext *q, AVCodecContext *avctx)
PutByteContext pbc;
GetBitContext gb;
H2645RBSP sps_rbsp = { NULL };
H2645NAL sps_nal = { NULL };
HEVCSPS sps = { 0 };
HEVCVPS vps = { 0 };
......@@ -69,8 +70,12 @@ static int generate_fake_vps(QSVEncContext *q, AVCodecContext *avctx)
return AVERROR_UNKNOWN;
}
av_fast_padded_malloc(&sps_rbsp.rbsp_buffer, &sps_rbsp.rbsp_buffer_alloc_size, avctx->extradata_size);
if (!sps_rbsp.rbsp_buffer)
return AVERROR(ENOMEM);
/* parse the SPS */
ret = ff_h2645_extract_rbsp(avctx->extradata + 4, avctx->extradata_size - 4, &sps_nal, 1);
ret = ff_h2645_extract_rbsp(avctx->extradata + 4, avctx->extradata_size - 4, &sps_rbsp, &sps_nal, 1);
if (ret < 0) {
av_log(avctx, AV_LOG_ERROR, "Error unescaping the SPS buffer\n");
return ret;
......@@ -78,7 +83,7 @@ static int generate_fake_vps(QSVEncContext *q, AVCodecContext *avctx)
ret = init_get_bits8(&gb, sps_nal.data, sps_nal.size);
if (ret < 0) {
av_freep(&sps_nal.rbsp_buffer);
av_freep(&sps_rbsp.rbsp_buffer);
return ret;
}
......@@ -87,13 +92,13 @@ static int generate_fake_vps(QSVEncContext *q, AVCodecContext *avctx)
if (type != HEVC_NAL_SPS) {
av_log(avctx, AV_LOG_ERROR, "Unexpected NAL type in the extradata: %d\n",
type);
av_freep(&sps_nal.rbsp_buffer);
av_freep(&sps_rbsp.rbsp_buffer);
return AVERROR_INVALIDDATA;
}
get_bits(&gb, 9);
ret = ff_hevc_parse_sps(&sps, &gb, &sps_id, 0, NULL, avctx);
av_freep(&sps_nal.rbsp_buffer);
av_freep(&sps_rbsp.rbsp_buffer);
if (ret < 0) {
av_log(avctx, AV_LOG_ERROR, "Error parsing the SPS\n");
return ret;
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment