Add functions to get/set rtp header extension by id.
This helps with dynamic size extensions.

BUG=webrtc:7433

Review-Url: https://codereview.webrtc.org/2789773004
Cr-Commit-Position: refs/heads/master@{#17505}
diff --git a/webrtc/modules/rtp_rtcp/source/rtp_packet.cc b/webrtc/modules/rtp_rtcp/source/rtp_packet.cc
index ec240a8..b66f001 100644
--- a/webrtc/modules/rtp_rtcp/source/rtp_packet.cc
+++ b/webrtc/modules/rtp_rtcp/source/rtp_packet.cc
@@ -30,6 +30,11 @@
 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 @@
   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 @@
   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;
+      int extension_id = i + 1;
+      return AllocateRawExtension(extension_id, length);
     }
   }
-
-  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;
-    }
-    *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) {
diff --git a/webrtc/modules/rtp_rtcp/source/rtp_packet.h b/webrtc/modules/rtp_rtcp/source/rtp_packet.h
index 75c73cf..072456d 100644
--- a/webrtc/modules/rtp_rtcp/source/rtp_packet.h
+++ b/webrtc/modules/rtp_rtcp/source/rtp_packet.h
@@ -28,6 +28,8 @@
   using ExtensionType = RTPExtensionType;
   using ExtensionManager = RtpHeaderExtensionMap;
   static constexpr size_t kMaxExtensionHeaders = 14;
+  static constexpr int kMinExtensionId = 1;
+  static constexpr int kMaxExtensionId = 14;
 
   // Parse and copy given buffer into Packet.
   bool Parse(const uint8_t* buffer, size_t size);
@@ -95,6 +97,21 @@
   template <typename Extension>
   bool ReserveExtension();
 
+  // Following 4 helpers identify rtp header extension by |id| negotiated with
+  // remote peer and written in an rtp packet.
+  bool HasRawExtension(int id) const;
+
+  // Returns place where extension with |id| is stored.
+  // Returns empty arrayview if extension is not present.
+  rtc::ArrayView<const uint8_t> GetRawExtension(int id) const;
+
+  // Allocates and store header extension. Returns true on success.
+  bool SetRawExtension(int id, rtc::ArrayView<const uint8_t> data);
+
+  // Allocates and returns place to store rtp header extension.
+  // Returns empty arrayview on failure.
+  rtc::ArrayView<uint8_t> AllocateRawExtension(int id, size_t length);
+
   // Reserve size_bytes for payload. Returns nullptr on failure.
   uint8_t* SetPayloadSize(size_t size_bytes);
   // Same as SetPayloadSize but doesn't guarantee to keep current payload.
@@ -133,14 +150,9 @@
                      uint8_t length,
                      uint16_t* offset) const;
 
-  // Find or allocate an extension, based on the type field of the parameter.
-  // If found, the length field be checked against what is already registered
-  // and the offset field will be set, then true is returned. If allocated, the
-  // length field will be used for allocation and the offset update to indicate
-  // position, the true is returned.
-  // If not found and allocations fails, false is returned and parameter remains
-  // unchanged.
-  bool AllocateExtension(ExtensionType type, uint8_t length, uint16_t* offset);
+  // Find or allocate an extension |type|. Returns view of size |length|
+  // to write raw extension to or an empty view on failure.
+  rtc::ArrayView<uint8_t> AllocateExtension(ExtensionType type, size_t length);
 
   uint8_t* WriteAt(size_t offset);
   void WriteAt(size_t offset, uint8_t byte);
@@ -176,18 +188,18 @@
 
 template <typename Extension, typename... Values>
 bool Packet::SetExtension(Values... values) {
-  uint16_t offset = 0;
-  if (!AllocateExtension(Extension::kId, Extension::kValueSizeBytes, &offset))
+  auto buffer = AllocateExtension(Extension::kId, Extension::kValueSizeBytes);
+  if (buffer.empty())
     return false;
-  return Extension::Write(WriteAt(offset), values...);
+  return Extension::Write(buffer.data(), values...);
 }
 
 template <typename Extension>
 bool Packet::ReserveExtension() {
-  uint16_t offset = 0;
-  if (!AllocateExtension(Extension::kId, Extension::kValueSizeBytes, &offset))
+  auto buffer = AllocateExtension(Extension::kId, Extension::kValueSizeBytes);
+  if (buffer.empty())
     return false;
-  memset(WriteAt(offset), 0, Extension::kValueSizeBytes);
+  memset(buffer.data(), 0, Extension::kValueSizeBytes);
   return true;
 }
 }  // namespace rtp
diff --git a/webrtc/modules/rtp_rtcp/source/rtp_packet_unittest.cc b/webrtc/modules/rtp_rtcp/source/rtp_packet_unittest.cc
index 2f1259e..d887065 100644
--- a/webrtc/modules/rtp_rtcp/source/rtp_packet_unittest.cc
+++ b/webrtc/modules/rtp_rtcp/source/rtp_packet_unittest.cc
@@ -116,6 +116,52 @@
               ElementsAreArray(packet.data(), packet.size()));
 }
 
+TEST(RtpPacketTest, CreateWithExtensionsWithoutManager) {
+  RtpPacketToSend packet(nullptr);
+  packet.SetPayloadType(kPayloadType);
+  packet.SetSequenceNumber(kSeqNum);
+  packet.SetTimestamp(kTimestamp);
+  packet.SetSsrc(kSsrc);
+
+  auto raw = packet.AllocateRawExtension(kTransmissionOffsetExtensionId,
+                                         TransmissionOffset::kValueSizeBytes);
+  EXPECT_EQ(raw.size(), TransmissionOffset::kValueSizeBytes);
+  TransmissionOffset::Write(raw.data(), kTimeOffset);
+
+  raw = packet.AllocateRawExtension(kAudioLevelExtensionId,
+                                    AudioLevel::kValueSizeBytes);
+  EXPECT_EQ(raw.size(), AudioLevel::kValueSizeBytes);
+  AudioLevel::Write(raw.data(), kVoiceActive, kAudioLevel);
+
+  EXPECT_THAT(kPacketWithTOAndAL,
+              ElementsAreArray(packet.data(), packet.size()));
+}
+
+TEST(RtpPacketTest, CreateWithMaxSizeHeaderExtension) {
+  const size_t kMaxExtensionSize = 16;
+  const int kId = 1;
+  const uint8_t kValue[16] = "123456789abcdef";
+
+  // Write packet with a custom extension.
+  RtpPacketToSend packet(nullptr);
+  packet.SetRawExtension(kId, kValue);
+  // Using different size for same id is not allowed.
+  EXPECT_TRUE(packet.AllocateRawExtension(kId, kMaxExtensionSize - 1).empty());
+
+  packet.SetPayloadSize(42);
+  // Rewriting allocated extension is allowed.
+  EXPECT_EQ(packet.AllocateRawExtension(kId, kMaxExtensionSize).size(),
+            kMaxExtensionSize);
+  // Adding another extension after payload is set is not allowed.
+  EXPECT_TRUE(packet.AllocateRawExtension(kId + 1, kMaxExtensionSize).empty());
+
+  // Read packet with the custom extension.
+  RtpPacketReceived parsed;
+  EXPECT_TRUE(parsed.Parse(packet.Buffer()));
+  auto read_raw = parsed.GetRawExtension(kId);
+  EXPECT_THAT(read_raw, ElementsAreArray(kValue, kMaxExtensionSize));
+}
+
 TEST(RtpPacketTest, SetReservedExtensionsAfterPayload) {
   const size_t kPayloadSize = 4;
   RtpPacketToSend::ExtensionManager extensions;
@@ -300,4 +346,21 @@
   EXPECT_EQ(0u, packet.padding_size());
 }
 
+TEST(RtpPacketTest, ParseWithoutExtensionManager) {
+  RtpPacketReceived packet;
+  EXPECT_TRUE(packet.Parse(kPacketWithTO, sizeof(kPacketWithTO)));
+
+  EXPECT_FALSE(packet.HasRawExtension(kAudioLevelExtensionId));
+  EXPECT_TRUE(packet.GetRawExtension(kAudioLevelExtensionId).empty());
+
+  EXPECT_TRUE(packet.HasRawExtension(kTransmissionOffsetExtensionId));
+
+  int32_t time_offset = 0;
+  auto raw_extension = packet.GetRawExtension(kTransmissionOffsetExtensionId);
+  EXPECT_EQ(raw_extension.size(), TransmissionOffset::kValueSizeBytes);
+  EXPECT_TRUE(TransmissionOffset::Parse(raw_extension.data(), &time_offset));
+
+  EXPECT_EQ(time_offset, kTimeOffset);
+}
+
 }  // namespace webrtc