Add functions to get/set rtp header extension by id.

webrtc/modules/rtp_rtcp/source/rtp_packet.cc

Index: webrtc/modules/rtp_rtcp/source/rtp_packet.cc
diff --git a/webrtc/modules/rtp_rtcp/source/rtp_packet.cc b/webrtc/modules/rtp_rtcp/source/rtp_packet.cc
index ec240a86359f2a970bf957887b268916493ffd1e..b66f001bd8b31411a667ed8e397e6843ea83c26b 100644
--- a/webrtc/modules/rtp_rtcp/source/rtp_packet.cc
+++ b/webrtc/modules/rtp_rtcp/source/rtp_packet.cc
@@ -30,6 +30,11 @@ constexpr uint16_t kOneByteExtensionId = 0xBEDE;
 constexpr size_t kOneByteHeaderSize = 1;
 constexpr size_t kDefaultPacketSize = 1500;
 }  // namespace
+
+constexpr size_t Packet::kMaxExtensionHeaders;
+constexpr int Packet::kMinExtensionId;
+constexpr int Packet::kMaxExtensionId;
+
 //  0                   1                   2                   3
 //  0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
 // +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
@@ -273,6 +278,97 @@ void Packet::SetCsrcs(const std::vector<uint32_t>& csrcs) {
   buffer_.SetSize(payload_offset_);
 }
 
+bool Packet::HasRawExtension(int id) const {
+  RTC_DCHECK_GE(id, kMinExtensionId);
+  RTC_DCHECK_LE(id, kMaxExtensionId);
+  return extension_entries_[id - 1].offset != 0;
+}
+
+rtc::ArrayView<const uint8_t> Packet::GetRawExtension(int id) const {
+  RTC_DCHECK_GE(id, kMinExtensionId);
+  RTC_DCHECK_LE(id, kMaxExtensionId);
+  const ExtensionInfo& extension = extension_entries_[id - 1];
+  if (extension.offset == 0)
+    return nullptr;
+  return rtc::MakeArrayView(data() + extension.offset, extension.length);
+}
+
+bool Packet::SetRawExtension(int id, rtc::ArrayView<const uint8_t> data) {
+  auto buffer = AllocateRawExtension(id, data.size());
+  if (buffer.empty())
+    return false;
+  RTC_DCHECK_EQ(buffer.size(), data.size());
+  memcpy(buffer.data(), data.data(), data.size());
+  return true;
+}
+
+rtc::ArrayView<uint8_t> Packet::AllocateRawExtension(int id, size_t length) {
+  RTC_DCHECK_GE(id, kMinExtensionId);
+  RTC_DCHECK_LE(id, kMaxExtensionId);
+  RTC_DCHECK_GE(length, 1);
+  RTC_DCHECK_LE(length, 16);
+
+  ExtensionInfo* extension_entry = &extension_entries_[id - 1];
+  if (extension_entry->offset != 0) {
+    // Extension already reserved. Check if same length is used.
+    if (extension_entry->length == length)
+      return rtc::MakeArrayView(WriteAt(extension_entry->offset), length);
+
+    LOG(LS_ERROR) << "Length mismatch for extension id " << id << " type "
+                  << static_cast<int>(extension_entry->type) << ": expected "
+                  << static_cast<int>(extension_entry->length) << ". received "
+                  << length;
+    return nullptr;
+  }
+  if (payload_size_ > 0) {
+    LOG(LS_ERROR) << "Can't add new extension id " << id
+                  << " after payload was set.";
+    return nullptr;
+  }
+  if (padding_size_ > 0) {
+    LOG(LS_ERROR) << "Can't add new extension id " << id
+                  << " after padding was set.";
+    return nullptr;
+  }
+
+  size_t num_csrc = data()[0] & 0x0F;
+  size_t extensions_offset = kFixedHeaderSize + (num_csrc * 4) + 4;
+  size_t new_extensions_size = extensions_size_ + kOneByteHeaderSize + length;
+  if (extensions_offset + new_extensions_size > capacity()) {
+    LOG(LS_ERROR)
+        << "Extension cannot be registered: Not enough space left in buffer.";
+    return nullptr;
+  }
+
+  // All checks passed, write down the extension headers.
+  if (extensions_size_ == 0) {
+    RTC_DCHECK_EQ(payload_offset_, kFixedHeaderSize + (num_csrc * 4));
+    WriteAt(0, data()[0] | 0x10);  // Set extension bit.
+    // Profile specific ID always set to OneByteExtensionHeader.
+    ByteWriter<uint16_t>::WriteBigEndian(WriteAt(extensions_offset - 4),
+                                         kOneByteExtensionId);
+  }
+
+  WriteAt(extensions_offset + extensions_size_, (id << 4) | (length - 1));
+
+  extension_entry->offset =
+      extensions_offset + extensions_size_ + kOneByteHeaderSize;
+  extension_entry->length = length;
+  extensions_size_ = new_extensions_size;
+
+  // Update header length field.
+  uint16_t extensions_words = (extensions_size_ + 3) / 4;  // Wrap up to 32bit.
+  ByteWriter<uint16_t>::WriteBigEndian(WriteAt(extensions_offset - 2),
+                                       extensions_words);
+  // Fill extension padding place with zeroes.
+  size_t extension_padding_size = 4 * extensions_words - extensions_size_;
+  memset(WriteAt(extensions_offset + extensions_size_), 0,
+         extension_padding_size);
+  payload_offset_ = extensions_offset + 4 * extensions_words;
+  buffer_.SetSize(payload_offset_);
+  return rtc::MakeArrayView(WriteAt(extension_entry->offset), length);
+}
+
 uint8_t* Packet::AllocatePayload(size_t size_bytes) {
   // Reset payload size to 0. If CopyOnWrite buffer_ was shared, this will cause
   // reallocation and memcpy. Keeping just header reduces memcpy size.
@@ -460,86 +556,16 @@ bool Packet::FindExtension(ExtensionType type,
   return false;
 }
 
-bool Packet::AllocateExtension(ExtensionType type,
-                               uint8_t length,
-                               uint16_t* offset) {
-  uint8_t extension_id = ExtensionManager::kInvalidId;
-  ExtensionInfo* extension_entry = nullptr;
+rtc::ArrayView<uint8_t> Packet::AllocateExtension(ExtensionType type,
+                                                  size_t length) {
   for (size_t i = 0; i < kMaxExtensionHeaders; ++i) {
     if (extension_entries_[i].type == type) {
-      extension_id = i + 1;
-      extension_entry = &extension_entries_[i];
-      break;
-    }
-  }
-
-  if (!extension_entry)  // Extension not registered.
-    return false;
-
-  if (extension_entry->length != 0) {  // Already allocated.
-    if (length != extension_entry->length) {
-      LOG(LS_WARNING) << "Length mismatch for extension '" << type
-                      << "': expected " << static_cast<int>(length)
-                      << ", received "
-                      << static_cast<int>(extension_entry->length);
-      return false;
+      int extension_id = i + 1;
+      return AllocateRawExtension(extension_id, length);
     }
-    *offset = extension_entry->offset;
-    return true;
-  }
-
-  // Can't add new extension after payload/padding was set.
-  if (payload_size_ > 0) {
-    return false;
-  }
-  if (padding_size_ > 0) {
-    return false;
-  }
-
-  RTC_DCHECK_GT(length, 0);
-  RTC_DCHECK_LE(length, 16);
-
-  size_t num_csrc = data()[0] & 0x0F;
-  size_t extensions_offset = kFixedHeaderSize + (num_csrc * 4) + 4;
-  if (extensions_offset + extensions_size_ + kOneByteHeaderSize + length >
-      capacity()) {
-    LOG(LS_WARNING) << "Extension cannot be registered: "
-                       "Not enough space left in buffer.";
-    return false;
   }
-
-  uint16_t new_extensions_size =
-      extensions_size_ + kOneByteHeaderSize + length;
-  uint16_t extensions_words =
-      (new_extensions_size + 3) / 4;  // Wrap up to 32bit.
-
-  // All checks passed, write down the extension.
-  if (extensions_size_ == 0) {
-    RTC_DCHECK_EQ(payload_offset_, kFixedHeaderSize + (num_csrc * 4));
-    WriteAt(0, data()[0] | 0x10);  // Set extension bit.
-    // Profile specific ID always set to OneByteExtensionHeader.
-    ByteWriter<uint16_t>::WriteBigEndian(WriteAt(extensions_offset - 4),
-                                         kOneByteExtensionId);
-  }
-
-  WriteAt(extensions_offset + extensions_size_,
-          (extension_id << 4) | (length - 1));
-
-  extension_entry->length = length;
-  *offset = extensions_offset + kOneByteHeaderSize + extensions_size_;
-  extension_entry->offset = *offset;
-  extensions_size_ = new_extensions_size;
-
-  // Update header length field.
-  ByteWriter<uint16_t>::WriteBigEndian(WriteAt(extensions_offset - 2),
-                                       extensions_words);
-  // Fill extension padding place with zeroes.
-  size_t extension_padding_size = 4 * extensions_words - extensions_size_;
-  memset(WriteAt(extensions_offset + extensions_size_), 0,
-         extension_padding_size);
-  payload_offset_ = extensions_offset + 4 * extensions_words;
-  buffer_.SetSize(payload_offset_);
-  return true;
+  // Extension not registered.
+  return nullptr;
 }
 
 uint8_t* Packet::WriteAt(size_t offset) {