From 2e17fd0e6bc0f5d325163497ccd78e662294fd5e Mon Sep 17 00:00:00 2001 From: Lee Fordyce Date: Tue, 9 Jan 2024 12:45:08 -0700 Subject: [PATCH] feat: use two buffer pointers for init and media segments separately (#9) --- include/packager/live_packager.h | 7 ++++-- packager/live_packager.cc | 38 ++++++++++++++++++++++++++++++-- packager/live_packager_test.cc | 24 ++++++++++---------- 3 files changed, 53 insertions(+), 16 deletions(-) diff --git a/include/packager/live_packager.h b/include/packager/live_packager.h index 16d03004ea7..43d00f73182 100644 --- 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 index 44b1251c23c..82533b0a5cc 100644 --- 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 index 40404c4467a..8e6dd68a1a8 100644 --- a/packager/live_packager_test.cc +++ b/packager/live_packager_test.cc @@ -127,9 +127,9 @@ TEST_F(LivePackagerBaseTest, VerifyAes128WithDecryption) { std::vector 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 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 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); } }