avcodec/av1_parse: return size of the parsed OBU in parse_obu_header()

Signed-off-by: James Almer <jamrial@gmail.com>

3 files changed, 26 insertions, 29 deletions

diff --git a/libavcodec/av1_parse.c b/libavcodec/av1_parse.c
index 48feb9fb8a..50dd940f03 100644
--- a/libavcodec/av1_parse.c
+++ b/libavcodec/av1_parse.c
@@ -29,11 +29,12 @@ int ff_av1_extract_obu(AV1OBU *obu, const uint8_t *buf, int length, void *logctx
 {
     int64_t obu_size;
     int start_pos, type, temporal_id, spatial_id;
+    int len, ret;
 
-    int ret = parse_obu_header(buf, length, &obu_size, &start_pos,
-                               &type, &temporal_id, &spatial_id);
-    if (ret < 0)
-        return ret;
+    len = parse_obu_header(buf, length, &obu_size, &start_pos,
+                           &type, &temporal_id, &spatial_id);
+    if (len < 0)
+        return len;
 
     if (obu_size > INT_MAX / 8 || obu_size < 0)
         return AVERROR(ERANGE);
@@ -42,12 +43,10 @@ int ff_av1_extract_obu(AV1OBU *obu, const uint8_t *buf, int length, void *logctx
     obu->temporal_id = temporal_id;
     obu->spatial_id  = spatial_id;
 
-    length = obu_size + start_pos;
-
     obu->data     = buf + start_pos;
     obu->size     = obu_size;
     obu->raw_data = buf;
-    obu->raw_size = length;
+    obu->raw_size = len;
 
     ret = init_get_bits(&obu->gb, obu->data, obu->size * 8);
     if (ret < 0)
@@ -57,7 +56,7 @@ int ff_av1_extract_obu(AV1OBU *obu, const uint8_t *buf, int length, void *logctx
            "obu_type: %d, temporal_id: %d, spatial_id: %d, payload size: %d\n",
            obu->type, obu->temporal_id, obu->spatial_id, obu->size);
 
-    return length;
+    return len;
 }
 
 int ff_av1_packet_split(AV1Packet *pkt, const uint8_t *buf, int length, void *logctx)
diff --git a/libavcodec/av1_parse.h b/libavcodec/av1_parse.h
index 3a4151491a..9a6e6835ab 100644
--- a/libavcodec/av1_parse.h
+++ b/libavcodec/av1_parse.h
@@ -95,6 +95,7 @@ static inline int parse_obu_header(const uint8_t *buf, int buf_size,
 {
     GetBitContext gb;
     int ret, extension_flag, has_size_flag;
+    int64_t size;
 
     ret = init_get_bits8(&gb, buf, FFMIN(buf_size, 2 + 8)); // OBU header fields + max leb128 length
     if (ret < 0)
@@ -124,7 +125,12 @@ static inline int parse_obu_header(const uint8_t *buf, int buf_size,
 
     *start_pos = get_bits_count(&gb) / 8;
 
-    return 0;
+    size = *obu_size + *start_pos;
+
+    if (size > INT_MAX)
+        return AVERROR(ERANGE);
+
+    return size;
 }
 
 #endif /* AVCODEC_AV1_PARSE_H */
diff --git a/libavformat/av1.c b/libavformat/av1.c
index 7c55a100bf..fa34820e12 100644
--- a/libavformat/av1.c
+++ b/libavformat/av1.c
@@ -33,14 +33,10 @@ int ff_av1_filter_obus(AVIOContext *pb, const uint8_t *buf, int size)
 
     size = 0;
     while (buf < end) {
-        int ret = parse_obu_header(buf, end - buf, &obu_size, &start_pos,
+        int len = parse_obu_header(buf, end - buf, &obu_size, &start_pos,
                                    &type, &temporal_id, &spatial_id);
-        if (ret < 0)
-            return ret;
-
-        obu_size += start_pos;
-        if (obu_size > INT_MAX)
-            return AVERROR_INVALIDDATA;
+        if (len < 0)
+            return len;
 
         switch (type) {
         case AV1_OBU_TEMPORAL_DELIMITER:
@@ -48,11 +44,11 @@ int ff_av1_filter_obus(AVIOContext *pb, const uint8_t *buf, int size)
         case AV1_OBU_PADDING:
             break;
         default:
-            avio_write(pb, buf, obu_size);
-            size += obu_size;
+            avio_write(pb, buf, len);
+            size += len;
             break;
         }
-        buf += obu_size;
+        buf += len;
     }
 
     return size;
@@ -86,25 +82,21 @@ int ff_isom_write_av1c(AVIOContext *pb, const uint8_t *buf, int size)
         return AVERROR_INVALIDDATA;
 
     while (size > 0) {
-        int ret = parse_obu_header(buf, size, &obu_size, &start_pos,
+        int len = parse_obu_header(buf, size, &obu_size, &start_pos,
                                    &type, &temporal_id, &spatial_id);
-        if (ret < 0)
-            return ret;
-
-        obu_size += start_pos;
-        if (obu_size > INT_MAX)
-            return AVERROR_INVALIDDATA;
+        if (len < 0)
+            return len;
 
         switch (type) {
         case AV1_OBU_SEQUENCE_HEADER:
         case AV1_OBU_METADATA:
-            avio_write(pb, buf, obu_size);
+            avio_write(pb, buf, len);
             break;
         default:
             break;
         }
-        size -= obu_size;
-        buf  += obu_size;
+        size -= len;
+        buf  += len;
     }
 
     return 0;