Add parser for OBU packets

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

@@ -20,6 +20,8 @@ extern "C" {
 };
 #endif
 
+#include "obu_parser.h"
+
 #define SAVE_NV12_STREAM 0
 #define SAVE_H264_STREAM 1
 
@@ -340,6 +342,13 @@ 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_);
+      // LOG_ERROR("Obu size = [{}]", obus.size());
+      // for (int i = 0; i < obus.size(); i++) {
+      //   LOG_ERROR("Obu size = {} [{} {}]", i, obus[i].size_,
+      //             obus[i].payload_size_);
+      // }
+
       // int consumed_size = 0;
       // int offset = 0;
       // int unit = 0;

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

@@ -0,0 +1,59 @@
+#include "byte_buffer.h"
+
+#include <string.h>
+
+ByteBufferReader::ByteBufferReader(const char* bytes, size_t len) {
+  Construct(bytes, len);
+}
+
+void ByteBufferReader::Construct(const char* bytes, size_t len) {
+  bytes_ = bytes;
+  size_ = len;
+  start_ = 0;
+  end_ = len;
+}
+
+bool ByteBufferReader::ReadBytes(char* val, size_t len) {
+  if (len > Length()) {
+    return false;
+  } else {
+    memcpy(val, bytes_ + start_, len);
+    start_ += len;
+    return true;
+  }
+}
+
+bool ByteBufferReader::ReadUInt8(uint8_t* val) {
+  if (!val) return false;
+
+  return ReadBytes(reinterpret_cast<char*>(val), 1);
+}
+
+bool ByteBufferReader::ReadUVarint(uint64_t* val) {
+  if (!val) {
+    return false;
+  }
+  // Integers are deserialized 7 bits at a time, with each byte having a
+  // continuation byte (msb=1) if there are more bytes to be read.
+  uint64_t v = 0;
+  for (int i = 0; i < 64; i += 7) {
+    char byte;
+    if (!ReadBytes(&byte, 1)) {
+      return false;
+    }
+    // Read the first 7 bits of the byte, then offset by bits read so far.
+    v |= (static_cast<uint64_t>(byte) & 0x7F) << i;
+    // True if the msb is not a continuation byte.
+    if (static_cast<uint64_t>(byte) < 0x80) {
+      *val = v;
+      return true;
+    }
+  }
+  return false;
+}
+
+bool ByteBufferReader::Consume(size_t size) {
+  if (size > Length()) return false;
+  start_ += size;
+  return true;
+}

src/media/video/encode/aom/byte_buffer.h

@@ -0,0 +1,40 @@
+/*
+ * @Author: DI JUNKUN
+ * @Date: 2024-04-22
+ * Copyright (c) 2024 by DI JUNKUN, All Rights Reserved.
+ */
+
+#ifndef _BYTE_BUFFER_H_
+#define _BYTE_BUFFER_H_
+
+#include <stddef.h>
+#include <stdint.h>
+
+class ByteBufferReader {
+ public:
+  ByteBufferReader(const char* bytes, size_t len);
+
+  ByteBufferReader(const ByteBufferReader&) = delete;
+  ByteBufferReader& operator=(const ByteBufferReader&) = delete;
+
+  // Returns start of unprocessed data.
+  const char* Data() const { return bytes_ + start_; }
+  // Returns number of unprocessed bytes.
+  size_t Length() const { return end_ - start_; }
+
+  bool ReadBytes(char* val, size_t len);
+  bool ReadUInt8(uint8_t* val);
+  bool ReadUVarint(uint64_t* val);
+
+  bool Consume(size_t size);
+
+ protected:
+  void Construct(const char* bytes, size_t size);
+
+  const char* bytes_;
+  size_t size_;
+  size_t start_;
+  size_t end_;
+};
+
+#endif

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

@@ -0,0 +1,80 @@
+#include "obu.h"
+
+#include <string>
+
+#include "log.h"
+
+Obu::Obu() {}
+
+Obu::Obu(const Obu &obu) {
+  if (obu.payload_size_ > 0) {
+    payload_ = (uint8_t *)malloc(obu.payload_size_);
+    if (NULL == payload_) {
+      LOG_ERROR("Malloc failed");
+    } else {
+      memcpy(payload_, obu.payload_, obu.payload_size_);
+    }
+    payload_size_ = obu.payload_size_;
+    header_ = obu.header_;
+    extension_header_ = obu.extension_header_;
+  }
+}
+
+Obu::Obu(Obu &&obu)
+    : payload_((uint8_t *)std::move(obu.payload_)),
+      payload_size_(obu.payload_size_),
+      header_(obu.header_),
+      extension_header_(obu.extension_header_) {
+  obu.payload_ = nullptr;
+  obu.payload_size_ = 0;
+}
+
+Obu &Obu::operator=(const Obu &obu) {
+  if (&obu != this) {
+    payload_ = (uint8_t *)realloc(payload_, obu.payload_size_);
+    memcpy(payload_, obu.payload_, obu.payload_size_);
+    payload_size_ = obu.payload_size_;
+    header_ = obu.header_;
+    extension_header_ = obu.extension_header_;
+  }
+  return *this;
+}
+
+Obu &Obu::operator=(Obu &&obu) {
+  if (&obu != this) {
+    payload_ = std::move(obu.payload_);
+    obu.payload_ = nullptr;
+    payload_size_ = obu.payload_size_;
+    obu.payload_size_ = 0;
+    header_ = obu.header_;
+    obu.header_ = 0;
+    extension_header_ = obu.extension_header_;
+    obu.extension_header_ = 0;
+  }
+  return *this;
+}
+
+Obu::~Obu() {
+  if (payload_) {
+    free(payload_);
+    payload_ = nullptr;
+  }
+  payload_size_ = 0;
+  header_ = 0;
+  extension_header_ = 0;
+}
+
+bool Obu::SetPayload(const uint8_t *payload, int size) {
+  if (payload_) {
+    free(payload_);
+    payload_ = nullptr;
+  }
+  payload_ = (uint8_t *)malloc(size);
+  memcpy(payload_, payload, size);
+  payload_size_ = size;
+
+  if (payload_)
+    return true;
+  else
+    return false;
+}

src/media/video/encode/aom/obu.h

@@ -0,0 +1,34 @@
+/*
+ * @Author: DI JUNKUN
+ * @Date: 2024-04-22
+ * Copyright (c) 2024 by DI JUNKUN, All Rights Reserved.
+ */
+
+#ifndef _OBU_H_
+#define _OBU_H_
+
+#include <stddef.h>
+#include <stdint.h>
+
+#include "aom/aom_codec.h"
+
+class Obu {
+ public:
+  Obu();
+  Obu(const Obu &obu);
+  Obu(Obu &&obu);
+  Obu &operator=(const Obu &obu);
+  Obu &operator=(Obu &&obu);
+
+  ~Obu();
+
+  bool SetPayload(const uint8_t *payload, int size);
+
+  uint8_t header_ = 0;
+  uint8_t extension_header_ = 0;  // undefined if (header & kXbit) == 0
+  uint8_t *payload_ = nullptr;
+  int size_ = 0;  // size of the header and payload combined.
+  int payload_size_ = 0;
+};
+
+#endif

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

@@ -0,0 +1,112 @@
+#include "obu_parser.h"
+
+#include "byte_buffer.h"
+#include "log.h"
+
+constexpr int kAggregationHeaderSize = 1;
+// when there are 3 or less OBU (fragments) in a packet, size of the last one
+// can be omited.
+constexpr int kMaxNumObusToOmitSize = 3;
+constexpr uint8_t kObuSizePresentBit = 0b0'0000'010;
+constexpr int kObuTypeSequenceHeader = 1;
+constexpr int kObuTypeTemporalDelimiter = 2;
+constexpr int kObuTypeTileList = 8;
+constexpr int kObuTypePadding = 15;
+
+const char* ObuTypeToString(OBU_TYPE type) {
+  switch (type) {
+    case OBU_SEQUENCE_HEADER:
+      return "OBU_SEQUENCE_HEADER";
+    case OBU_TEMPORAL_DELIMITER:
+      return "OBU_TEMPORAL_DELIMITER";
+    case OBU_FRAME_HEADER:
+      return "OBU_FRAME_HEADER";
+    case OBU_REDUNDANT_FRAME_HEADER:
+      return "OBU_REDUNDANT_FRAME_HEADER";
+    case OBU_FRAME:
+      return "OBU_FRAME";
+    case OBU_TILE_GROUP:
+      return "OBU_TILE_GROUP";
+    case OBU_METADATA:
+      return "OBU_METADATA";
+    case OBU_TILE_LIST:
+      return "OBU_TILE_LIST";
+    case OBU_PADDING:
+      return "OBU_PADDING";
+    default:
+      break;
+  }
+  return "<Invalid OBU Type>";
+}
+
+bool ObuHasExtension(uint8_t obu_header) { return obu_header & 0b0'0000'100; }
+
+bool ObuHasSize(uint8_t obu_header) { return obu_header & kObuSizePresentBit; }
+
+int ObuType(uint8_t obu_header) { return (obu_header & 0b0'1111'000) >> 3; }
+
+std::vector<Obu> ParseObus(uint8_t* payload, int payload_size) {
+  std::vector<Obu> result;
+  ByteBufferReader payload_reader(reinterpret_cast<const char*>(payload),
+                                  payload_size);
+  while (payload_reader.Length() > 0) {
+    Obu obu;
+    payload_reader.ReadUInt8(&obu.header_);
+    LOG_ERROR("Get obu type = [{}]",
+              ObuTypeToString((OBU_TYPE)ObuType(obu.header_)))
+    obu.size_ = 1;
+    if (ObuHasExtension(obu.header_)) {
+      if (payload_reader.Length() == 0) {
+        LOG_ERROR(
+            "Malformed AV1 input: expected extension_header, no more bytes in "
+            "the buffer. Offset: {}",
+            (payload_size - payload_reader.Length()));
+        return {};
+      }
+      payload_reader.ReadUInt8(&obu.extension_header_);
+      ++obu.size_;
+    }
+    if (!ObuHasSize(obu.header_)) {
+      obu.SetPayload(reinterpret_cast<const uint8_t*>(payload_reader.Data()),
+                     payload_reader.Length());
+      payload_reader.Consume(payload_reader.Length());
+      LOG_ERROR("payload_reader.Length() = {}", payload_reader.Length());
+    } else {
+      uint64_t size = 0;
+      if (!payload_reader.ReadUVarint(&size) ||
+          size > payload_reader.Length()) {
+        LOG_ERROR(
+            "Malformed AV1 input: declared size {} is larger than remaining "
+            "buffer size {}",
+            size, payload_reader.Length());
+        return {};
+      }
+      obu.SetPayload(reinterpret_cast<const uint8_t*>(payload_reader.Data()),
+                     size);
+      payload_reader.Consume(size);
+      LOG_ERROR("size = {}", size);
+    }
+    obu.size_ += obu.payload_size_;
+    // Skip obus that shouldn't be transfered over rtp.
+    int obu_type = ObuType(obu.header_);
+    // if (obu_type != kObuTypeTemporalDelimiter &&  //
+    //     obu_type != kObuTypeTileList &&           //
+    //     obu_type != kObuTypePadding) {
+    //   result.push_back(obu);
+    // }
+
+    LOG_ERROR("Obu size = [{}], Obu type [{}]", obu.size_,
+              ObuTypeToString((OBU_TYPE)ObuType(obu.header_)));
+    if (1) {
+      result.push_back(obu);
+    }
+  }
+
+  LOG_ERROR("Obu size = [{}]", result.size());
+  for (int i = 0; i < result.size(); i++) {
+    LOG_ERROR("[{}] Obu size = [{}], Obu type [{}]", i, result[i].payload_size_,
+              ObuTypeToString((OBU_TYPE)ObuType(result[i].header_)));
+  }
+
+  return result;
+}

src/media/video/encode/aom/obu_parser.h

@@ -0,0 +1,18 @@
+/*
+ * @Author: DI JUNKUN
+ * @Date: 2024-04-22
+ * Copyright (c) 2024 by DI JUNKUN, All Rights Reserved.
+ */
+
+#ifndef _OBU_PARSER_H_
+#define _OBU_PARSER_H_
+
+#include <stddef.h>
+#include <stdint.h>
+
+#include <vector>
+
+#include "obu.h"
+
+std::vector<Obu> ParseObus(uint8_t* payload, int payload_size);
+#endif