Fix OBU total length error

src/media/video/encode/aom/aom_av1_encoder.cpp

@@ -20,8 +20,6 @@ extern "C" {
 };
 #endif
 
-#include "obu_parser.h"
-
 #define SAVE_NV12_STREAM 0
 #define SAVE_H264_STREAM 1
 
@@ -342,13 +340,6 @@ int AomAv1Encoder::Encode(
       SET_ENCODER_PARAM_OR_RETURN_ERROR(AOME_GET_LAST_QUANTIZER, &qp);
       LOG_INFO("Encoded frame qp = {}", qp);
 
-      std::vector<Obu> obus = ParseObus(encoded_frame_, encoded_frame_size_);
-      for (int i = 0; i < obus.size(); i++) {
-        LOG_ERROR("[{}] Obu size = [{}], Obu type [{}]", i,
-                  obus[i].payload_size_,
-                  ObuTypeToString((OBU_TYPE)ObuType(obus[i].header_)));
-      }
-
       if (on_encoded_image) {
         on_encoded_image((char *)encoded_frame_, encoded_frame_size_);
         if (SAVE_H264_STREAM) {

src/rtp/byte_buffer.cpp

@@ -29,7 +29,7 @@ bool ByteBufferReader::ReadUInt8(uint8_t* val) {
   return ReadBytes(reinterpret_cast<char*>(val), 1);
 }
 
-bool ByteBufferReader::ReadUVarint(uint64_t* val) {
+bool ByteBufferReader::ReadUVarint(uint64_t* val, size_t* len) {
   if (!val) {
     return false;
   }
@@ -46,6 +46,9 @@ bool ByteBufferReader::ReadUVarint(uint64_t* val) {
     // True if the msb is not a continuation byte.
     if (static_cast<uint64_t>(byte) < 0x80) {
       *val = v;
+      if (len) {
+        *len = i / 8 + (i % 8 ? 1 : 0) + 1;
+      }
       return true;
     }
   }

src/rtp/byte_buffer.h

@@ -24,7 +24,7 @@ class ByteBufferReader {
 
   bool ReadBytes(char* val, size_t len);
   bool ReadUInt8(uint8_t* val);
-  bool ReadUVarint(uint64_t* val);
+  bool ReadUVarint(uint64_t* val, size_t* len);
 
   bool Consume(size_t size);
 

src/rtp/obu.cpp

@@ -15,9 +15,11 @@ Obu::Obu(const Obu &obu) {
       memcpy(payload_, obu.payload_, obu.payload_size_);
     }
     payload_size_ = obu.payload_size_;
+    size_ = obu.size_;
     header_ = obu.header_;
     extension_header_ = obu.extension_header_;
   } else {
+    size_ = obu.size_;
     header_ = obu.header_;
     extension_header_ = obu.extension_header_;
   }
@@ -26,6 +28,7 @@ Obu::Obu(const Obu &obu) {
 Obu::Obu(Obu &&obu)
     : payload_((uint8_t *)std::move(obu.payload_)),
       payload_size_(obu.payload_size_),
+      size_(obu.size_),
       header_(obu.header_),
       extension_header_(obu.extension_header_) {
   obu.payload_ = nullptr;
@@ -38,6 +41,7 @@ Obu &Obu::operator=(const Obu &obu) {
       payload_ = (uint8_t *)realloc(payload_, obu.payload_size_);
       memcpy(payload_, obu.payload_, obu.payload_size_);
       payload_size_ = obu.payload_size_;
+      size_ = obu.size_;
       header_ = obu.header_;
       extension_header_ = obu.extension_header_;
     } else {
@@ -54,6 +58,8 @@ Obu &Obu::operator=(Obu &&obu) {
     obu.payload_ = nullptr;
     payload_size_ = obu.payload_size_;
     obu.payload_size_ = 0;
+    size_ = obu.size_;
+    obu.size_ = 0;
     header_ = obu.header_;
     obu.header_ = 0;
     extension_header_ = obu.extension_header_;
@@ -68,6 +74,7 @@ Obu::~Obu() {
     payload_ = nullptr;
   }
   payload_size_ = 0;
+  size_ = 0;
   header_ = 0;
   extension_header_ = 0;
 }

src/rtp/obu_parser.cpp

@@ -70,18 +70,21 @@ std::vector<Obu> ParseObus(uint8_t* payload, int payload_size) {
                      payload_reader.Length());
       payload_reader.Consume(payload_reader.Length());
     } else {
-      uint64_t size = 0;
+      uint64_t payload_size = 0;
-      if (!payload_reader.ReadUVarint(&size) ||
+      size_t len = 0;
-          size > payload_reader.Length()) {
+      if (!payload_reader.ReadUVarint(&payload_size, &len) ||
+          payload_size > payload_reader.Length()) {
         LOG_ERROR(
-            "Malformed AV1 input: declared size {} is larger than remaining "
+            "Malformed AV1 input: declared payload_size {} is larger than "
+            "remaining "
             "buffer size {}",
-            size, payload_reader.Length());
+            payload_size, payload_reader.Length());
         return {};
       }
       obu.SetPayload(reinterpret_cast<const uint8_t*>(payload_reader.Data()),
-                     size);
+                     payload_size);
-      payload_reader.Consume(size);
+      obu.size_ += len;
+      payload_reader.Consume(payload_size);
     }
     obu.size_ += obu.payload_size_;
     // Skip obus that shouldn't be transfered over rtp.

src/rtp/rtp_codec.cpp

@@ -3,22 +3,13 @@
 #include <chrono>
 
 #include "log.h"
+#include "obu_parser.h"
 
 #define RTP_VERSION 2
 #define NALU 1
 #define FU_A 28
 #define FU_B 29
 
-constexpr int kObuTypeSequenceHeader = 1;
-constexpr int kObuTypeTemporalDelimiter = 2;
-constexpr int kObuTypeFrameHeader = 3;
-constexpr int kObuTypeTileGroup = 4;
-constexpr int kObuTypeMetadata = 5;
-constexpr int kObuTypeFrame = 6;
-constexpr int kObuTypeRedundantFrameHeader = 7;
-constexpr int kObuTypeTileList = 8;
-constexpr int kObuTypePadding = 15;
-
 RtpCodec ::RtpCodec(RtpPacket::PAYLOAD_TYPE payload_type)
     : version_(RTP_VERSION),
       has_padding_(false),
@@ -227,6 +218,13 @@ void RtpCodec::Encode(uint8_t* buffer, size_t size,
       }
     }
   } else if (RtpPacket::PAYLOAD_TYPE::AV1 == payload_type_) {
+    std::vector<Obu> obus = ParseObus(buffer, size);
+    LOG_ERROR("Total size = [{}]", size);
+    for (int i = 0; i < obus.size(); i++) {
+      LOG_ERROR("[{}] Obu size = [{}], Obu type [{}]", i, obus[i].size_,
+                ObuTypeToString((OBU_TYPE)ObuType(obus[i].header_)));
+    }
+
     if (size <= MAX_NALU_LEN) {
       RtpPacket rtp_packet;
       rtp_packet.SetVerion(version_);