Implement WiFi Display media packetizer.

extensions/renderer/api/display_source/wifi_display/wifi_display_media_packetizer_unittest.cc

Index: extensions/renderer/api/display_source/wifi_display/wifi_display_media_packetizer_unittest.cc
diff --git a/extensions/renderer/api/display_source/wifi_display/wifi_display_media_packetizer_unittest.cc b/extensions/renderer/api/display_source/wifi_display/wifi_display_media_packetizer_unittest.cc
index a8b769a4421ace00b07782a61c9e5344a614f764..19569bb58ef3a79e3ed319299be54e6b5b87d0d8 100644
--- a/extensions/renderer/api/display_source/wifi_display/wifi_display_media_packetizer_unittest.cc
+++ b/extensions/renderer/api/display_source/wifi_display/wifi_display_media_packetizer_unittest.cc
@@ -2,12 +2,40 @@
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
+#include <algorithm>
+#include <array>
+#include <list>
+
 #include "base/big_endian.h"
+#include "base/bind.h"
+#include "extensions/renderer/api/display_source/wifi_display/wifi_display_elementary_stream_info.h"
 #include "extensions/renderer/api/display_source/wifi_display/wifi_display_elementary_stream_packetizer.h"
+#include "extensions/renderer/api/display_source/wifi_display/wifi_display_media_packetizer.h"
+#include "extensions/renderer/api/display_source/wifi_display/wifi_display_transport_stream_packetizer.h"
 #include "testing/gtest/include/gtest/gtest.h"
 
+using PacketPart = extensions::WiFiDisplayStreamPacketPart;
+
 namespace extensions {
 
+std::ostream& operator<<(std::ostream& os, const PacketPart& part) {
+  const auto flags = os.flags();
+  os << "{" << std::hex << std::noshowbase;
+  for (const auto& item : part) {
+    if (&item != &*part.begin())
+      os << ", ";
+    os << "0x" << static_cast<unsigned>(item);
+  }
+  os.setf(flags, std::ios::basefield | std::ios::showbase);
+  return os << "}";
+}
+
+bool operator==(const PacketPart& a, const PacketPart& b) {
+  if (a.size() != b.size())
+    return false;
+  return std::equal(a.begin(), a.end(), b.begin());
+}
+
 namespace {
 
 namespace pes {
@@ -17,6 +45,129 @@ const unsigned kPtsFlag = 0x0080u;
 const size_t kUnitDataAlignment = sizeof(uint32_t);
 }
 
+namespace rtp {
+const unsigned kVersionMask = 0xC000u;
+const unsigned kVersion2 = 0x8000u;
+const unsigned kPaddingFlag = 0x2000u;
+const unsigned kExtensionFlag = 0x1000u;
+const unsigned kContributingSourceCountMask = 0x0F00u;
+const unsigned kMarkerFlag = 0x0010u;
+const unsigned kPayloadTypeMask = 0x007Fu;
+const unsigned kPayloadTypeMP2T = 0x0021u;
+}  // namespace rtp
+
+namespace ts {
+const uint64_t kTimeStampSecond = 90000u;  // 90 kHz
+
+// Packet header:
+const size_t kPacketHeaderSize = 4u;
+}  // namespace ts
+
+namespace widi {
+const unsigned kProgramClockReferencePacketId = 0x1000u;
+const size_t kMaxTransportStreamPacketCountPerDatagramPacket = 7u;
+}  // namespace widi
+
+template <typename PacketContainer>
+class PacketCollector {
+ public:
+  PacketContainer FetchPackets() {
+    PacketContainer container;
+    container.swap(packets_);
+    return container;
+  }
+
+ protected:
+  PacketContainer packets_;
+};
+
+class FakeMediaPacketizer
+    : public WiFiDisplayMediaPacketizer,
+      public PacketCollector<std::vector<std::vector<uint8_t>>> {
+ public:
+  FakeMediaPacketizer(const base::TimeDelta& delay_for_unit_time_stamps,
+                      std::vector<WiFiDisplayElementaryStreamInfo> stream_infos)
+      : WiFiDisplayMediaPacketizer(
+            delay_for_unit_time_stamps,
+            std::move(stream_infos),
+            base::Bind(&FakeMediaPacketizer::OnPacketizedMediaDatagramPacket,
+                       base::Unretained(this))) {}
+
+  // Extend the interface in order to allow to bypass packetization of units to
+  // Packetized Elementary Stream (PES) packets and further to Transport Stream
+  // (TS) packets and to test only packetization of TS packets to media
+  // datagram packets.
+  bool EncodeTransportStreamPacket(
+      const WiFiDisplayTransportStreamPacket& transport_stream_packet,
+      bool flush) {
+    return OnPacketizedTransportStreamPacket(transport_stream_packet, flush);
+  }
+
+ private:
+  bool OnPacketizedMediaDatagramPacket(
+      WiFiDisplayMediaDatagramPacket media_datagram_packet) {
+    packets_.emplace_back(std::move(media_datagram_packet));
+    return true;
+  }
+};
+
+class FakeTransportStreamPacketizer
+    : public WiFiDisplayTransportStreamPacketizer,
+      public PacketCollector<std::list<WiFiDisplayTransportStreamPacket>> {
+ public:
+  FakeTransportStreamPacketizer(
+      const base::TimeDelta& delay_for_unit_time_stamps,
+      std::vector<WiFiDisplayElementaryStreamInfo> stream_infos)
+      : WiFiDisplayTransportStreamPacketizer(delay_for_unit_time_stamps,
+                                             std::move(stream_infos)) {}
+
+ protected:
+  bool OnPacketizedTransportStreamPacket(
+      const WiFiDisplayTransportStreamPacket& transport_stream_packet,
+      bool flush) override {
+    // Make a copy of header bytes as they are in stack.
+    headers_.emplace_back(transport_stream_packet.header().begin(),
+                          transport_stream_packet.header().end());
+    const auto& header = headers_.back();
+    if (transport_stream_packet.payload().empty()) {
+      packets_.emplace_back(header.data(), header.size());
+    } else {
+      packets_.emplace_back(header.data(), header.size(),
+                            transport_stream_packet.payload().begin());
+    }
+    EXPECT_EQ(transport_stream_packet.header().size(),
+              packets_.back().header().size());
+    EXPECT_EQ(transport_stream_packet.payload().size(),
+              packets_.back().payload().size());
+    EXPECT_EQ(transport_stream_packet.filler().size(),
+              packets_.back().filler().size());
+    return true;
+  }
+
+ private:
+  std::vector<std::vector<uint8_t>> headers_;
+};
+
+struct ProgramClockReference {
+  enum { kInvalidBase = ~static_cast<uint64_t>(0u) };
+  uint64_t base;
+  uint16_t extension;
+};
+
+ProgramClockReference ParseProgramClockReference(const uint8_t pcr_bytes[6]) {
+  const uint8_t reserved_pcr_bits = pcr_bytes[4] & 0x7Eu;
+  EXPECT_EQ(0x7Eu, reserved_pcr_bits);
+  ProgramClockReference pcr;
+  pcr.base = pcr_bytes[0];
+  pcr.base = (pcr.base << 8) | pcr_bytes[1];
+  pcr.base = (pcr.base << 8) | pcr_bytes[2];
+  pcr.base = (pcr.base << 8) | pcr_bytes[3];
+  pcr.base = (pcr.base << 1) | ((pcr_bytes[4] & 0x80u) >> 7);
+  pcr.extension = pcr_bytes[4] & 0x01u;
+  pcr.extension = (pcr.extension << 8) | pcr_bytes[5];
+  return pcr;
+}
+
 uint64_t ParseTimeStamp(const uint8_t ts_bytes[5], uint8_t pts_dts_indicator) {
   EXPECT_EQ(pts_dts_indicator, (ts_bytes[0] & 0xF0u) >> 4);
   EXPECT_EQ(0x01u, ts_bytes[0] & 0x01u);
@@ -31,6 +182,14 @@ uint64_t ParseTimeStamp(const uint8_t ts_bytes[5], uint8_t pts_dts_indicator) {
   return ts;
 }
 
+unsigned ParseTransportStreamPacketId(
+    const WiFiDisplayTransportStreamPacket& packet) {
+  if (packet.header().size() < ts::kPacketHeaderSize)
+    return ~0u;
+  return (((packet.header().begin()[1] & 0x001Fu) << 8) |
+          packet.header().begin()[2]);
+}
+
 class WiFiDisplayElementaryStreamUnitPacketizationTest
     : public testing::TestWithParam<
           testing::tuple<unsigned, base::TimeDelta, base::TimeDelta>> {
@@ -162,5 +321,150 @@ INSTANTIATE_TEST_CASE_P(
                                      base::TimeDelta::FromMicroseconds(
                                          1000 * INT64_C(0x123456789) / 90))));
 
+TEST(WiFiDisplayTransportStreamPacketizationTest, EncodeToMediaDatagramPacket) {
+  const size_t kPacketHeaderSize = 12u;
+
+  // Create fake units.
+  const size_t kUnitCount = 12u;
+  const size_t kUnitSize =
+      WiFiDisplayTransportStreamPacket::kPacketSize - 4u - 12u;
+  std::vector<std::array<uint8_t, kUnitSize>> units(kUnitCount);
+  for (auto& unit : units)
+    unit.fill(static_cast<uint8_t>(&unit - units.data()));
+
+  // Create transport stream packets.
+  std::vector<WiFiDisplayElementaryStreamInfo> stream_infos;
+  stream_infos.emplace_back(WiFiDisplayElementaryStreamInfo::VIDEO_H264);
+  FakeTransportStreamPacketizer transport_stream_packetizer(
+      base::TimeDelta::FromMilliseconds(0), std::move(stream_infos));
+  for (const auto& unit : units) {
+    EXPECT_TRUE(transport_stream_packetizer.EncodeElementaryStreamUnit(
+        0u, unit.data(), unit.size(), false, base::TimeTicks(),
+        base::TimeTicks(), &unit == &units.back()));
+  }
+  auto transport_stream_packets = transport_stream_packetizer.FetchPackets();
+  // There should be exactly one transport stream payload packet for each unit.
+  // There should also be some but not too many transport stream meta
+  // information packets.
+  EXPECT_EQ(1u, transport_stream_packets.size() / kUnitCount);
+
+  // Encode transport stream packets to datagram packets.
+  FakeMediaPacketizer packetizer(
+      base::TimeDelta::FromMilliseconds(0),
+      std::vector<WiFiDisplayElementaryStreamInfo>());
+  for (const auto& transport_stream_packet : transport_stream_packets) {
+    EXPECT_TRUE(packetizer.EncodeTransportStreamPacket(
+        transport_stream_packet,
+        &transport_stream_packet == &transport_stream_packets.back()));
+  }
+  auto packets = packetizer.FetchPackets();
+
+  // Check datagram packets.
+  ProgramClockReference pcr = {ProgramClockReference::kInvalidBase, 0u};
+  uint16_t sequence_number;
+  uint32_t synchronization_source_identifier;
+  auto transport_stream_packet_it = transport_stream_packets.cbegin();
+  for (const auto& packet : packets) {
+    base::BigEndianReader header_reader(
+        reinterpret_cast<const char*>(packet.data()),
+        std::min(kPacketHeaderSize, packet.size()));
+
+    // Packet flags.
+    uint16_t parsed_u16;
+    EXPECT_TRUE(header_reader.ReadU16(&parsed_u16));
+    EXPECT_EQ(rtp::kVersion2, parsed_u16 & rtp::kVersionMask);
+    EXPECT_FALSE(parsed_u16 & rtp::kPaddingFlag);
+    EXPECT_FALSE(parsed_u16 & rtp::kExtensionFlag);
+    EXPECT_EQ(0u, parsed_u16 & rtp::kContributingSourceCountMask);
+    EXPECT_FALSE(parsed_u16 & rtp::kMarkerFlag);
+    EXPECT_EQ(rtp::kPayloadTypeMP2T, parsed_u16 & rtp::kPayloadTypeMask);
+
+    // Packet sequence number.
+    uint16_t parsed_sequence_number;
+    EXPECT_TRUE(header_reader.ReadU16(&parsed_sequence_number));
+    if (&packet == &packets.front())
+      sequence_number = parsed_sequence_number;
+    EXPECT_EQ(sequence_number++, parsed_sequence_number);
+
+    // Packet time stamp.
+    uint32_t parsed_time_stamp;
+    EXPECT_TRUE(header_reader.ReadU32(&parsed_time_stamp));
+    if (pcr.base == ProgramClockReference::kInvalidBase) {
+      // This happens only for the first datagram packet.
+      EXPECT_TRUE(&packet == &packets.front());
+      // Ensure that the next datagram packet reaches the else branch.
+      EXPECT_FALSE(&packet == &packets.back());
+    } else {
+      // Check that
+      //   0 <= parsed_time_stamp - pcr.base <= kTimeStampSecond
+      // but allow pcr.base and parsed_time_stamp to wrap around in 32 bits.
+      EXPECT_LE((parsed_time_stamp - pcr.base) & 0xFFFFFFFFu,
+                ts::kTimeStampSecond)
+          << " Time stamp must not be smaller than PCR!";
+    }
+
+    // Packet synchronization source identifier.
+    uint32_t parsed_synchronization_source_identifier;
+    EXPECT_TRUE(
+        header_reader.ReadU32(&parsed_synchronization_source_identifier));
+    if (&packet == &packets.front()) {
+      synchronization_source_identifier =
+          parsed_synchronization_source_identifier;
+    }
+    EXPECT_EQ(synchronization_source_identifier,
+              parsed_synchronization_source_identifier);
+
+    EXPECT_EQ(0, header_reader.remaining());
+
+    // Packet payload.
+    size_t offset = kPacketHeaderSize;
+    while (offset + WiFiDisplayTransportStreamPacket::kPacketSize <=
+               packet.size() &&
+           transport_stream_packet_it != transport_stream_packets.end()) {
+      const auto& transport_stream_packet = *transport_stream_packet_it++;
+      const PacketPart parsed_transport_stream_packet_header(
+          packet.data() + offset, transport_stream_packet.header().size());
+      const PacketPart parsed_transport_stream_packet_payload(
+          parsed_transport_stream_packet_header.end(),
+          transport_stream_packet.payload().size());
+      const PacketPart parsed_transport_stream_packet_filler(
+          parsed_transport_stream_packet_payload.end(),
+          transport_stream_packet.filler().size());
+      offset += WiFiDisplayTransportStreamPacket::kPacketSize;
+
+      // Check bytes.
+      EXPECT_EQ(transport_stream_packet.header(),
+                parsed_transport_stream_packet_header);
+      EXPECT_EQ(transport_stream_packet.payload(),
+                parsed_transport_stream_packet_payload);
+      EXPECT_EQ(transport_stream_packet.filler().size(),
+                std::count(parsed_transport_stream_packet_filler.begin(),
+                           parsed_transport_stream_packet_filler.end(),
+                           transport_stream_packet.filler().value()));
+
+      if (ParseTransportStreamPacketId(transport_stream_packet) ==
+          widi::kProgramClockReferencePacketId) {
+        pcr = ParseProgramClockReference(
+            &transport_stream_packet.header().begin()[6]);
+      }
+    }
+    EXPECT_EQ(offset, packet.size()) << " Extra packet payload bytes.";
+
+    // Check that the payload contains a correct number of transport stream
+    // packets.
+    const size_t transport_stream_packet_count_in_datagram_packet =
+        packet.size() / WiFiDisplayTransportStreamPacket::kPacketSize;
+    if (&packet == &packets.back()) {
+      EXPECT_GE(transport_stream_packet_count_in_datagram_packet, 1u);
+      EXPECT_LE(transport_stream_packet_count_in_datagram_packet,
+                widi::kMaxTransportStreamPacketCountPerDatagramPacket);
+    } else {
+      EXPECT_EQ(widi::kMaxTransportStreamPacketCountPerDatagramPacket,
+                transport_stream_packet_count_in_datagram_packet);
+    }
+  }
+  EXPECT_EQ(transport_stream_packets.end(), transport_stream_packet_it);
+}
+
 }  // namespace
 }  // namespace extensions