diff mbox

[FFmpeg-devel,5/5,v2] avformat: add an AV1 Annex B demuxer

Message ID 20191112021611.2645-1-jamrial@gmail.com
State New
Headers show

Commit Message

James Almer Nov. 12, 2019, 2:16 a.m. UTC
Signed-off-by: James Almer <jamrial@gmail.com>
---
 configure                |   1 +
 libavformat/Makefile     |   1 +
 libavformat/allformats.c |   1 +
 libavformat/obu.c        | 277 +++++++++++++++++++++++++++++++++++++++
 4 files changed, 280 insertions(+)
 create mode 100644 libavformat/obu.c

Comments

Ronald S. Bultje Nov. 12, 2019, 5:52 a.m. UTC | #1
Hi,

On Mon, Nov 11, 2019 at 9:17 PM James Almer <jamrial@gmail.com> wrote:

> +static int leb(AVIOContext *pb, uint32_t *len) {
>

 this can overflow, should be uint64_t.


> +        unsigned bits;
>

Same.

Ronald
James Almer Nov. 12, 2019, 1:09 p.m. UTC | #2
On 11/12/2019 2:52 AM, Ronald S. Bultje wrote:
> Hi,
> 
> On Mon, Nov 11, 2019 at 9:17 PM James Almer <jamrial@gmail.com> wrote:
> 
>> +static int leb(AVIOContext *pb, uint32_t *len) {
>>
> 
>  this can overflow, should be uint64_t.
> 
> 
>> +        unsigned bits;
>>
> 
> Same.
> 
> Ronald

I used the same method as in dav1d, it will not overflow as it will not
try to load more than five bytes worth leb128 data, ensuring to only
accept up to 4 bits from the last, resulting in assembled values up to
UINT32_MAX.

> _______________________________________________
> ffmpeg-devel mailing list
> ffmpeg-devel@ffmpeg.org
> https://ffmpeg.org/mailman/listinfo/ffmpeg-devel
> 
> To unsubscribe, visit link above, or email
> ffmpeg-devel-request@ffmpeg.org with subject "unsubscribe".
>
Ronald S. Bultje Nov. 12, 2019, 9:19 p.m. UTC | #3
Hi,

On Tue, Nov 12, 2019 at 8:09 AM James Almer <jamrial@gmail.com> wrote:

> On 11/12/2019 2:52 AM, Ronald S. Bultje wrote:
> > Hi,
> >
> > On Mon, Nov 11, 2019 at 9:17 PM James Almer <jamrial@gmail.com> wrote:
> >
> >> +static int leb(AVIOContext *pb, uint32_t *len) {
> >>
> >
> >  this can overflow, should be uint64_t.
> >
> >
> >> +        unsigned bits;
> >>
> >
> > Same.
> >
> > Ronald
>
> I used the same method as in dav1d, it will not overflow as it will not
> try to load more than five bytes worth leb128 data, ensuring to only
> accept up to 4 bits from the last, resulting in assembled values up to
> UINT32_MAX.


Oh, you are right, I was looking at "if (++i == 8 && more)", but the line
above it ("if (i <= 3 || (i == 4 && bits < (1 << 4)))") demonstrates
quickly that this is different. It's OK here, although a bit confusing, but
I should probably fix that in dav1d first.

Rest of patchset lgtm, and thanks for adding this feature, this is
genuinely very useful for AV1 users in more complex AV1 processing
pipelines. Section5 would be nice also ;-).

Ronald
diff mbox

Patch

diff --git a/configure b/configure
index 70f60997c1..a8dbba879d 100755
--- a/configure
+++ b/configure
@@ -3293,6 +3293,7 @@  mxf_d10_muxer_select="mxf_muxer"
 mxf_opatom_muxer_select="mxf_muxer"
 nut_muxer_select="riffenc"
 nuv_demuxer_select="riffdec"
+obu_demuxer_select="av1_frame_merge_bsf"
 oga_muxer_select="ogg_muxer"
 ogg_demuxer_select="dirac_parse"
 ogv_muxer_select="ogg_muxer"
diff --git a/libavformat/Makefile b/libavformat/Makefile
index 8251f8f657..9057d0358a 100644
--- a/libavformat/Makefile
+++ b/libavformat/Makefile
@@ -350,6 +350,7 @@  OBJS-$(CONFIG_NULL_MUXER)                += nullenc.o
 OBJS-$(CONFIG_NUT_DEMUXER)               += nutdec.o nut.o isom.o
 OBJS-$(CONFIG_NUT_MUXER)                 += nutenc.o nut.o
 OBJS-$(CONFIG_NUV_DEMUXER)               += nuv.o
+OBJS-$(CONFIG_OBU_DEMUXER)               += obu.o
 OBJS-$(CONFIG_OGG_DEMUXER)               += oggdec.o         \
                                             oggparsecelt.o   \
                                             oggparsedaala.o  \
diff --git a/libavformat/allformats.c b/libavformat/allformats.c
index f7fea32b45..152644e9f9 100644
--- a/libavformat/allformats.c
+++ b/libavformat/allformats.c
@@ -282,6 +282,7 @@  extern AVInputFormat  ff_nut_demuxer;
 extern AVOutputFormat ff_nut_muxer;
 extern AVInputFormat  ff_nuv_demuxer;
 extern AVOutputFormat ff_oga_muxer;
+extern AVInputFormat  ff_obu_demuxer;
 extern AVInputFormat  ff_ogg_demuxer;
 extern AVOutputFormat ff_ogg_muxer;
 extern AVOutputFormat ff_ogv_muxer;
diff --git a/libavformat/obu.c b/libavformat/obu.c
new file mode 100644
index 0000000000..a8ba8a26f4
--- /dev/null
+++ b/libavformat/obu.c
@@ -0,0 +1,277 @@ 
+/*
+ * AV1 Annex B demuxer
+ * Copyright (c) 2019 James Almer <jamrial@gmail.com>
+ *
+ * This file is part of FFmpeg.
+ *
+ * FFmpeg is free software; you can redistribute it and/or
+ * modify it under the terms of the GNU Lesser General Public
+ * License as published by the Free Software Foundation; either
+ * version 2.1 of the License, or (at your option) any later version.
+ *
+ * FFmpeg is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+ * Lesser General Public License for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public
+ * License along with FFmpeg; if not, write to the Free Software
+ * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
+ */
+
+#include "config.h"
+
+#include "libavutil/common.h"
+#include "libavutil/opt.h"
+#include "libavcodec/av1_parse.h"
+#include "avformat.h"
+#include "avio_internal.h"
+#include "internal.h"
+
+typedef struct AnnexBContext {
+    const AVClass *class;
+    AVBSFContext *bsf;
+    uint32_t temporal_unit_size;
+    uint32_t frame_unit_size;
+    AVRational framerate;
+} AnnexBContext;
+
+static int leb(AVIOContext *pb, uint32_t *len) {
+    int more, i = 0;
+    uint8_t byte;
+    *len = 0;
+    do {
+        unsigned bits;
+        byte = avio_r8(pb);
+        more = byte & 0x80;
+        bits = byte & 0x7f;
+        if (i <= 3 || (i == 4 && bits < (1 << 4)))
+            *len |= bits << (i * 7);
+        else if (bits)
+            return AVERROR_INVALIDDATA;
+        if (++i == 8 && more)
+            return AVERROR_INVALIDDATA;
+        if (pb->eof_reached || pb->error)
+            return pb->error ? pb->error : AVERROR(EIO);
+    } while (more);
+    return i;
+}
+
+static int read_obu(const uint8_t *buf, int size, int64_t *obu_size, int *type)
+{
+    int start_pos, temporal_id, spatial_id;
+    int len;
+
+    len = parse_obu_header(buf, size, obu_size, &start_pos,
+                           type, &temporal_id, &spatial_id);
+    if (len < 0)
+        return len;
+
+    return 0;
+}
+
+static int annexb_probe(const AVProbeData *p)
+{
+    AVIOContext pb;
+    int64_t obu_size;
+    uint32_t temporal_unit_size, frame_unit_size, obu_unit_size;
+    int seq = 0, frame_header = 0;
+    int ret, type, cnt = 0;
+
+    ffio_init_context(&pb, p->buf, p->buf_size, 0,
+                      NULL, NULL, NULL, NULL);
+
+    ret = leb(&pb, &temporal_unit_size);
+    if (ret < 0)
+        return 0;
+    cnt += ret;
+    ret = leb(&pb, &frame_unit_size);
+    if (ret < 0 || ((int64_t)frame_unit_size + ret) > temporal_unit_size)
+        return 0;
+    cnt += ret;
+    temporal_unit_size -= ret;
+    ret = leb(&pb, &obu_unit_size);
+    if (ret < 0 || ((int64_t)obu_unit_size + ret) >= frame_unit_size)
+        return 0;
+    cnt += ret;
+
+    temporal_unit_size -= obu_unit_size + ret;
+    frame_unit_size -= obu_unit_size + ret;
+
+    avio_skip(&pb, obu_unit_size);
+    if (pb.eof_reached || pb.error)
+        return 0;
+
+    // Check that the first OBU is a Temporal Delimiter.
+    ret = read_obu(p->buf + cnt, FFMIN(p->buf_size - cnt, obu_unit_size), &obu_size, &type);
+    if (ret < 0 || type != AV1_OBU_TEMPORAL_DELIMITER || obu_size > 0)
+        return 0;
+    cnt += obu_unit_size;
+
+    do {
+        ret = leb(&pb, &obu_unit_size);
+        if (ret < 0 || ((int64_t)obu_unit_size + ret) > frame_unit_size)
+            return 0;
+        cnt += ret;
+
+        avio_skip(&pb, obu_unit_size);
+        if (pb.eof_reached || pb.error)
+            return 0;
+
+        ret = read_obu(p->buf + cnt, FFMIN(p->buf_size - cnt, obu_unit_size), &obu_size, &type);
+        if (ret < 0)
+            return 0;
+        cnt += obu_unit_size;
+
+        if (type == AV1_OBU_SEQUENCE_HEADER)
+            seq = 1;
+        if (type == AV1_OBU_FRAME || type == AV1_OBU_FRAME_HEADER) {
+            if (frame_header || !seq)
+                return 0;
+            frame_header = 1;
+            break;
+        }
+        if (type == AV1_OBU_TILE_GROUP && !frame_header)
+            return 0;
+
+        temporal_unit_size -= obu_unit_size + ret;
+        frame_unit_size -= obu_unit_size + ret;
+    } while (!frame_header && frame_unit_size);
+
+    return frame_header ? AVPROBE_SCORE_EXTENSION + 1 : 0;
+}
+
+static int annexb_read_header(AVFormatContext *s)
+{
+    AnnexBContext *c = s->priv_data;
+    const AVBitStreamFilter *filter = av_bsf_get_by_name("av1_frame_merge");
+    AVStream *st;
+    int ret;
+
+    if (!filter) {
+        av_log(c, AV_LOG_ERROR, "av1_frame_merge bitstream filter "
+               "not found. This is a bug, please report it.\n");
+        return AVERROR_BUG;
+    }
+
+    st = avformat_new_stream(s, NULL);
+    if (!st)
+        return AVERROR(ENOMEM);
+
+    st->codecpar->codec_type = AVMEDIA_TYPE_VIDEO;
+    st->codecpar->codec_id = AV_CODEC_ID_AV1;
+    st->need_parsing = AVSTREAM_PARSE_HEADERS;
+
+    st->internal->avctx->framerate = c->framerate;
+    // taken from rawvideo demuxers
+    avpriv_set_pts_info(st, 64, 1, 1200000);
+
+    ret = av_bsf_alloc(filter, &c->bsf);
+    if (ret < 0)
+        return ret;
+
+    ret = avcodec_parameters_copy(c->bsf->par_in, st->codecpar);
+    if (ret < 0) {
+        av_bsf_free(&c->bsf);
+        return ret;
+    }
+
+    ret = av_bsf_init(c->bsf);
+    if (ret < 0)
+        av_bsf_free(&c->bsf);
+
+    return ret;
+}
+
+static int annexb_read_packet(AVFormatContext *s, AVPacket *pkt)
+{
+    AnnexBContext *c = s->priv_data;
+    uint32_t obu_unit_size;
+    int ret, len;
+
+retry:
+    if (avio_feof(s->pb)) {
+        if (c->temporal_unit_size || c->frame_unit_size)
+            return AVERROR(EIO);
+        av_bsf_send_packet(c->bsf, NULL);
+        goto end;
+    }
+
+    if (!c->temporal_unit_size) {
+        len = leb(s->pb, &c->temporal_unit_size);
+        if (len < 0) return AVERROR_INVALIDDATA;
+    }
+
+    if (!c->frame_unit_size) {
+        len = leb(s->pb, &c->frame_unit_size);
+        if (len < 0 || ((int64_t)c->frame_unit_size + len) > c->temporal_unit_size)
+            return AVERROR_INVALIDDATA;
+        c->temporal_unit_size -= len;
+    }
+
+    len = leb(s->pb, &obu_unit_size);
+    if (len < 0 || ((int64_t)obu_unit_size + len) > c->frame_unit_size)
+        return AVERROR_INVALIDDATA;
+
+    ret = av_get_packet(s->pb, pkt, obu_unit_size);
+    if (ret < 0)
+        return ret;
+    if (ret != obu_unit_size)
+        return AVERROR(EIO);
+
+    c->temporal_unit_size -= obu_unit_size + len;
+    c->frame_unit_size -= obu_unit_size + len;
+
+    ret = av_bsf_send_packet(c->bsf, pkt);
+    if (ret < 0) {
+        av_log(s, AV_LOG_ERROR, "Failed to send packet to "
+                                "av1_frame_merge filter\n");
+        return ret;
+    }
+
+end:
+    ret = av_bsf_receive_packet(c->bsf, pkt);
+    if (ret < 0 && ret != AVERROR(EAGAIN) && ret != AVERROR_EOF)
+        av_log(s, AV_LOG_ERROR, "av1_frame_merge filter failed to "
+                                "send output packet\n");
+
+    if (ret == AVERROR(EAGAIN))
+        goto retry;
+
+    return ret;
+}
+
+static int annexb_read_close(AVFormatContext *s)
+{
+    AnnexBContext *c = s->priv_data;
+
+    av_bsf_free(&c->bsf);
+    return 0;
+}
+
+#define OFFSET(x) offsetof(AnnexBContext, x)
+#define DEC AV_OPT_FLAG_DECODING_PARAM
+static const AVOption annexb_options[] = {
+    { "framerate", "", OFFSET(framerate), AV_OPT_TYPE_VIDEO_RATE, {.str = "25"}, 0, INT_MAX, DEC},
+    { NULL },
+};
+
+static const AVClass annexb_demuxer_class = {
+    .class_name = "AV1 Annex-B demuxer",
+    .item_name  = av_default_item_name,
+    .option     = annexb_options,
+    .version    = LIBAVUTIL_VERSION_INT,
+};
+
+AVInputFormat ff_obu_demuxer = {
+    .name           = "obu",
+    .long_name      = NULL_IF_CONFIG_SMALL("OBU"),
+    .priv_data_size = sizeof(AnnexBContext),
+    .read_probe     = annexb_probe,
+    .read_header    = annexb_read_header,
+    .read_packet    = annexb_read_packet,
+    .read_close     = annexb_read_close,
+    .extensions     = "obu",
+    .flags          = AVFMT_GENERIC_INDEX,
+    .priv_class     = &annexb_demuxer_class,
+};