feat: support user-specified sequence number in output moof header (#11)

Pluto-tv · Jan 9, 2024 · df8aae4 · df8aae4
1 parent 2e17fd0
commit df8aae4
Show file tree

Hide file tree

Showing 5 changed files with 222 additions and 1 deletion.
diff --git a/include/packager/live_packager.h b/include/packager/live_packager.h
@@ -86,6 +86,13 @@ struct LiveConfig {
   std::vector<uint8_t> key;
   std::vector<uint8_t> key_id;
   EncryptionScheme protection_scheme;
+
+  /// User-specified segment number.
+  /// For FMP4 output:
+  ///   It can be used to set the moof header sequence number if > 0.
+  /// For M2TS output:
+  ///   It is be used to set the continuity counter (TODO: UNIMPLEMENTED).
+  uint32_t segment_number = 0;
 };
 
 class LivePackager {

diff --git a/include/packager/mp4_output_params.h b/include/packager/mp4_output_params.h
@@ -26,6 +26,11 @@ struct Mp4OutputParams {
   /// and mdat atom. Each chunk is uploaded immediately upon creation,
   /// decoupling latency from segment duration.
   bool low_latency_dash_mode = false;
+
+  /// User-specified sequence number to be set in the moof header.
+  /// The moof header sequence number starts at 1 so values less than 1 will be
+  /// set to 1.
+  uint32_t sequence_number = 1;
 };
 
 }  // namespace shaka

diff --git a/packager/live_packager.cc b/packager/live_packager.cc
@@ -308,6 +308,8 @@ Status LivePackager::Package(const Segment& init_segment,
   packaging_params.chunking_params.segment_duration_in_seconds =
       config_.segment_duration_sec;
 
+  packaging_params.mp4_output_params.sequence_number = config_.segment_number;
+
   EncryptionParams& encryption_params = packaging_params.encryption_params;
   // As a side effect of InitializeEncryption, encryption_params will be
   // modified.

diff --git a/packager/live_packager_test.cc b/packager/live_packager_test.cc
@@ -16,6 +16,8 @@
 #include <packager/file.h>
 #include <packager/live_packager.h>
 #include <packager/media/base/aes_decryptor.h>
+#include <packager/media/formats/mp4/box_definitions.h>
+#include <packager/media/formats/mp4/box_reader.h>
 
 namespace shaka {
 namespace {
@@ -59,6 +61,178 @@ std::vector<uint8_t> ReadTestDataFile(const std::string& name) {
 
   return data;
 }
+
+bool ParseAndCheckType(media::mp4::Box& box, media::mp4::BoxReader* reader) {
+  box.Parse(reader);
+  return box.BoxType() == reader->type();
+}
+
+struct SegmentIndexBoxChecker {
+  SegmentIndexBoxChecker(media::mp4::SegmentIndex box)
+      : sidx_(std::move(box)) {}
+
+  void Check(media::mp4::BoxReader* reader) {
+    media::mp4::SegmentIndex box;
+    CHECK(ParseAndCheckType(box, reader));
+    EXPECT_EQ(sidx_.timescale, box.timescale);
+  }
+
+ private:
+  media::mp4::SegmentIndex sidx_;
+};
+
+struct MovieFragmentBoxChecker {
+  MovieFragmentBoxChecker(media::mp4::MovieFragment box)
+      : moof_(std::move(box)) {}
+
+  void Check(media::mp4::BoxReader* reader) {
+    media::mp4::MovieFragment box;
+    CHECK(ParseAndCheckType(box, reader));
+    EXPECT_EQ(moof_.header.sequence_number, box.header.sequence_number);
+  }
+
+ private:
+  media::mp4::MovieFragment moof_;
+};
+
+struct SegmentTypeBoxChecker {
+  void Check(media::mp4::BoxReader* reader) {
+    media::mp4::SegmentType box;
+    CHECK(ParseAndCheckType(box, reader));
+    EXPECT_EQ(media::FourCC::FOURCC_mp41, box.major_brand);
+  }
+};
+
+struct FileTypeBoxChecker {
+  void Check(media::mp4::BoxReader* reader) {
+    media::mp4::FileType box;
+    CHECK(ParseAndCheckType(box, reader));
+    EXPECT_EQ(media::FourCC::FOURCC_mp41, box.major_brand);
+  }
+};
+
+struct MovieBoxChecker {
+  MovieBoxChecker(media::mp4::Movie movie) : moov_(std::move(movie)) {}
+
+  void Check(media::mp4::BoxReader* reader) {
+    media::mp4::Movie moov;
+    CHECK(ParseAndCheckType(moov, reader));
+
+    EXPECT_EQ(moov_.tracks.size(), moov.tracks.size());
+
+    for (unsigned i(0); i < moov_.tracks.size(); ++i) {
+      const auto& exp_track = moov_.tracks[i];
+      const auto& act_track = moov.tracks[i];
+
+      EXPECT_EQ(exp_track.media.handler.handler_type,
+                act_track.media.handler.handler_type);
+
+      const auto& exp_video_entries =
+          exp_track.media.information.sample_table.description.video_entries;
+      const auto& act_video_entries =
+          act_track.media.information.sample_table.description.video_entries;
+
+      EXPECT_EQ(exp_video_entries.size(), act_video_entries.size());
+
+      for (unsigned j(0); j < exp_video_entries.size(); ++j) {
+        const auto& exp_entry = exp_video_entries[j];
+        const auto& act_entry = act_video_entries[j];
+
+        EXPECT_EQ(exp_entry.BoxType(), act_entry.BoxType());
+        EXPECT_EQ(exp_entry.width, act_entry.width);
+        EXPECT_EQ(exp_entry.height, act_entry.height);
+      }
+    }
+  }
+
+ private:
+  media::mp4::Movie moov_;
+};
+
+void CheckVideoInitSegment(const FullSegmentBuffer& buffer) {
+  bool err(true);
+  size_t bytes_to_read(buffer.InitSegmentSize());
+  const uint8_t* data(buffer.InitSegmentData());
+
+  {
+    std::unique_ptr<media::mp4::BoxReader> reader(
+        media::mp4::BoxReader::ReadBox(data, bytes_to_read, &err));
+    EXPECT_FALSE(err);
+
+    FileTypeBoxChecker checker;
+    checker.Check(reader.get());
+
+    data += reader->size();
+    bytes_to_read -= reader->size();
+  }
+
+  {
+    std::unique_ptr<media::mp4::BoxReader> reader(
+        media::mp4::BoxReader::ReadBox(data, bytes_to_read, &err));
+    EXPECT_FALSE(err);
+
+    media::mp4::VideoSampleEntry entry;
+    entry.format = media::FOURCC_avc1;
+    entry.width = 1024;
+    entry.height = 576;
+
+    media::mp4::Track track;
+    track.media.handler.handler_type = media::FourCC::FOURCC_vide;
+    track.media.information.sample_table.description.video_entries.push_back(
+        entry);
+
+    media::mp4::Movie expected;
+    expected.tracks.push_back(track);
+
+    MovieBoxChecker checker(expected);
+    checker.Check(reader.get());
+  }
+}
+
+void CheckSegment(const LiveConfig& config, const FullSegmentBuffer& buffer) {
+  bool err(true);
+  size_t bytes_to_read(buffer.SegmentSize());
+  const uint8_t* data(buffer.SegmentData());
+
+  {
+    std::unique_ptr<media::mp4::BoxReader> reader(
+        media::mp4::BoxReader::ReadBox(data, bytes_to_read, &err));
+    EXPECT_FALSE(err);
+
+    SegmentTypeBoxChecker checker;
+    checker.Check(reader.get());
+
+    data += reader->size();
+    bytes_to_read -= reader->size();
+  }
+
+  {
+    std::unique_ptr<media::mp4::BoxReader> reader(
+        media::mp4::BoxReader::ReadBox(data, bytes_to_read, &err));
+    EXPECT_FALSE(err);
+
+    media::mp4::SegmentIndex expected;
+    expected.timescale = 10000000;
+    SegmentIndexBoxChecker checker(expected);
+    checker.Check(reader.get());
+
+    data += reader->size();
+    bytes_to_read -= reader->size();
+  }
+
+  {
+    std::unique_ptr<media::mp4::BoxReader> reader(
+        media::mp4::BoxReader::ReadBox(data, bytes_to_read, &err));
+    EXPECT_FALSE(err);
+
+    media::mp4::MovieFragment expected;
+    expected.header.sequence_number = config.segment_number;
+
+    MovieFragmentBoxChecker checker(expected);
+    checker.Check(reader.get());
+  }
+}
+
 }  // namespace
 
 class LivePackagerBaseTest : public ::testing::Test {
@@ -111,6 +285,8 @@ TEST_F(LivePackagerBaseTest, InitSegmentOnly) {
   ASSERT_EQ(Status::OK, live_packager_->PackageInit(in, out));
   ASSERT_GT(out.InitSegmentSize(), 0);
   ASSERT_EQ(out.SegmentSize(), 0);
+
+  CheckVideoInitSegment(out);
 }
 
 TEST_F(LivePackagerBaseTest, VerifyAes128WithDecryption) {
@@ -186,6 +362,35 @@ TEST_F(LivePackagerBaseTest, EncryptionFailure) {
   }
 }
 
+TEST_F(LivePackagerBaseTest, CustomMoofSequenceNumber) {
+  std::vector<uint8_t> init_segment_buffer = ReadTestDataFile("input/init.mp4");
+  ASSERT_FALSE(init_segment_buffer.empty());
+  LiveConfig live_config;
+  live_config.format = LiveConfig::OutputFormat::FMP4;
+  live_config.track_type = LiveConfig::TrackType::VIDEO;
+  live_config.protection_scheme = LiveConfig::EncryptionScheme::NONE;
+  live_config.segment_duration_sec = kSegmentDurationInSeconds;
+
+  for (unsigned int i = 0; i < kNumSegments; i++) {
+    live_config.segment_number = i + 1;
+    std::string segment_num = absl::StrFormat("input/%04d.m4s", i);
+    std::vector<uint8_t> segment_buffer = ReadTestDataFile(segment_num);
+    ASSERT_FALSE(segment_buffer.empty());
+
+    SegmentData init_seg(init_segment_buffer.data(),
+                         init_segment_buffer.size());
+    SegmentData media_seg(segment_buffer.data(), segment_buffer.size());
+
+    FullSegmentBuffer out;
+    LivePackager packager(live_config);
+
+    ASSERT_EQ(Status::OK, packager.Package(init_seg, media_seg, out));
+    ASSERT_GT(out.SegmentSize(), 0);
+
+    CheckSegment(live_config, out);
+  }
+}
+
 struct LivePackagerTestCase {
   unsigned int num_segments;
   std::string init_segment_name;

diff --git a/packager/media/formats/mp4/segmenter.cc b/packager/media/formats/mp4/segmenter.cc
@@ -92,7 +92,9 @@ Status Segmenter::Initialize(
 
   // Use the reference stream's time scale as movie time scale.
   moov_->header.timescale = sidx_->timescale;
-  moof_->header.sequence_number = 1;
+  moof_->header.sequence_number = options_.mp4_params.sequence_number > 0
+                                      ? options_.mp4_params.sequence_number
+                                      : 1;
 
   // Fill in version information.
   const std::string version = GetPackagerVersion();