feat: use two buffer pointers for init and media segments separately (#9

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

Hide file tree

Showing 3 changed files with 53 additions and 16 deletions.
diff --git a/include/packager/live_packager.h b/include/packager/live_packager.h
@@ -100,10 +100,13 @@ class LivePackager {
   Status PackageInit(const Segment& init_segment, FullSegmentBuffer& output);
 
   /// Performs packaging of segment data.
-  /// @param full_segment contains the full segment data (init + media).
+  /// @param init_segment contains the init segment data.
+  /// @param media_segment contains the media segment data.
   /// @param output contains the packaged segment data (init + media).
   /// @return OK on success, an appropriate error code on failure.
-  Status Package(const Segment& full_segment, FullSegmentBuffer& output);
+  Status Package(const Segment& init_segment,
+                 const Segment& media_segment,
+                 FullSegmentBuffer& output);
 
   LivePackager(const LivePackager&) = delete;
   LivePackager& operator=(const LivePackager&) = delete;

diff --git a/packager/live_packager.cc b/packager/live_packager.cc
@@ -80,6 +80,38 @@ class SegmentDataReader {
   uint64_t position_ = 0;
 };
 
+class MultiSegmentDataReader {
+ public:
+  MultiSegmentDataReader(const Segment& init_segment,
+                         const Segment& media_segment)
+      : init_segment_(init_segment), media_segment_(media_segment) {}
+
+  uint64_t Read(void* buffer, uint64_t size) {
+    if (position_ < init_segment_.Size()) {
+      const uint64_t first_chunk_size =
+          std::min(size, init_segment_.Size() - position_);
+      memcpy(buffer, init_segment_.Data() + position_, first_chunk_size);
+
+      position_ += first_chunk_size;
+      return first_chunk_size;
+    }
+    auto segment_position = position_ - init_segment_.Size();
+    if (segment_position >= media_segment_.Size()) {
+      return 0;
+    }
+    const uint64_t second_chunk_size =
+        std::min(size, media_segment_.Size() - segment_position);
+    memcpy(buffer, media_segment_.Data() + segment_position, second_chunk_size);
+    position_ += second_chunk_size;
+    return second_chunk_size;
+  }
+
+ private:
+  const Segment& init_segment_;
+  const Segment& media_segment_;
+  uint64_t position_ = 0;
+};
+
 class SegmentManager {
  public:
   explicit SegmentManager();
@@ -242,8 +274,10 @@ Status LivePackager::PackageInit(const Segment& init_segment,
   return packager.Run();
 }
 
-Status LivePackager::Package(const Segment& in, FullSegmentBuffer& out) {
-  SegmentDataReader reader(in);
+Status LivePackager::Package(const Segment& init_segment,
+                             const Segment& media_segment,
+                             FullSegmentBuffer& out) {
+  MultiSegmentDataReader reader(init_segment, media_segment);
   shaka::BufferCallbackParams callback_params;
   callback_params.read_func = [&reader](const std::string& name, void* buffer,
                                         uint64_t size) {

diff --git a/packager/live_packager_test.cc b/packager/live_packager_test.cc
@@ -127,9 +127,9 @@ TEST_F(LivePackagerBaseTest, VerifyAes128WithDecryption) {
     std::vector<uint8_t> segment_buffer = ReadTestDataFile(segment_num);
     ASSERT_FALSE(segment_buffer.empty());
 
-    FullSegmentBuffer in;
-    in.SetInitSegment(init_segment_buffer.data(), init_segment_buffer.size());
-    in.AppendData(segment_buffer.data(), segment_buffer.size());
+    SegmentData init_seg(init_segment_buffer.data(),
+                         init_segment_buffer.size());
+    SegmentData media_seg(segment_buffer.data(), segment_buffer.size());
 
     FullSegmentBuffer out;
 
@@ -139,7 +139,7 @@ TEST_F(LivePackagerBaseTest, VerifyAes128WithDecryption) {
     live_config.protection_scheme = LiveConfig::EncryptionScheme::AES_128;
 
     SetupLivePackagerConfig(live_config);
-    ASSERT_EQ(Status::OK, live_packager_->Package(in, out));
+    ASSERT_EQ(Status::OK, live_packager_->Package(init_seg, media_seg, out));
     ASSERT_GT(out.SegmentSize(), 0);
 
     std::string exp_segment_num = absl::StrFormat("expected/ts/%04d.ts", i + 1);
@@ -168,9 +168,9 @@ TEST_F(LivePackagerBaseTest, EncryptionFailure) {
     std::vector<uint8_t> segment_buffer = ReadTestDataFile(segment_num);
     ASSERT_FALSE(segment_buffer.empty());
 
-    FullSegmentBuffer in;
-    in.SetInitSegment(init_segment_buffer.data(), init_segment_buffer.size());
-    in.AppendData(segment_buffer.data(), segment_buffer.size());
+    SegmentData init_seg(init_segment_buffer.data(),
+                         init_segment_buffer.size());
+    SegmentData media_seg(segment_buffer.data(), segment_buffer.size());
 
     FullSegmentBuffer out;
 
@@ -182,7 +182,7 @@ TEST_F(LivePackagerBaseTest, EncryptionFailure) {
     SetupLivePackagerConfig(live_config);
     ASSERT_EQ(Status(error::INVALID_ARGUMENT,
                      "invalid key and IV supplied to encryptor"),
-              live_packager_->Package(in, out));
+              live_packager_->Package(init_seg, media_seg, out));
   }
 }
 
@@ -226,13 +226,13 @@ TEST_P(LivePackagerEncryptionTest, VerifyWithEncryption) {
     std::vector<uint8_t> segment_buffer = ReadTestDataFile(format_output);
     ASSERT_FALSE(segment_buffer.empty());
 
-    FullSegmentBuffer in;
-    in.SetInitSegment(init_segment_buffer.data(), init_segment_buffer.size());
-    in.AppendData(segment_buffer.data(), segment_buffer.size());
+    SegmentData init_seg(init_segment_buffer.data(),
+                         init_segment_buffer.size());
+    SegmentData media_seg(segment_buffer.data(), segment_buffer.size());
 
     FullSegmentBuffer out;
 
-    ASSERT_EQ(Status::OK, live_packager_->Package(in, out));
+    ASSERT_EQ(Status::OK, live_packager_->Package(init_seg, media_seg, out));
     ASSERT_GT(out.SegmentSize(), 0);
   }
 }