ts_muxer.cc

// Copyright 2016 Google LLC. All rights reserved.
//
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file or at
// https://developers.google.com/open-source/licenses/bsd
 
#include <packager/media/formats/mp2t/ts_muxer.h>
 
#include <vector>
 
#include <absl/log/check.h>
 
#include <packager/macros/status.h>
#include <packager/media/base/aes_encryptor.h>
#include <packager/media/base/fourccs.h>
#include <packager/media/base/muxer_util.h>
 
namespace shaka {
namespace media {
namespace mp2t {
 
namespace {
const int32_t kTsTimescale = 90000;
}  // namespace
 
TsMuxer::TsMuxer(const MuxerOptions& muxer_options) : Muxer(muxer_options) {}
TsMuxer::~TsMuxer() {}
 
Status TsMuxer::InitializeMuxer() {
  if (streams().size() > 1u)
    return Status(error::MUXER_FAILURE, "Cannot handle more than one streams.");
 
  if (options().segment_template.empty()) {
    const std::string& file_name = options().output_file_name;
    DCHECK(!file_name.empty());
    output_file_.reset(File::Open(file_name.c_str(), "w"));
    if (!output_file_) {
      return Status(error::FILE_FAILURE,
                    "Cannot open file for write " + file_name);
    }
  }
 
  segmenter_.reset(new TsSegmenter(options(), muxer_listener()));
  Status status = segmenter_->Initialize(*streams()[0]);
 
  // Capture AES-128 encryption config once so FinalizeSegment can encrypt.
  const auto& stream = *streams()[0];
  if (stream.is_encrypted() &&
      stream.encryption_config().protection_scheme == kAes128ProtectionScheme) {
    aes128_encryption_config_ = stream.encryption_config();
  }
 
  FireOnMediaStartEvent();
  return status;
}
 
Status TsMuxer::Finalize() {
  FireOnMediaEndEvent();
  return segmenter_->Finalize();
}
 
Status TsMuxer::AddMediaSample(size_t stream_id, const MediaSample& sample) {
  DCHECK_EQ(stream_id, 0u);
 
  // The duration of the first sample may have been adjusted, so use
  // the duration of the second sample instead.
  if (num_samples_ < 2) {
    sample_durations_[num_samples_] =
        sample.duration() * kTsTimescale / streams().front()->time_scale();
    if (num_samples_ == 1 && muxer_listener())
      muxer_listener()->OnSampleDurationReady(sample_durations_[num_samples_]);
    num_samples_++;
  }
  return segmenter_->AddSample(sample);
}
 
Status TsMuxer::FinalizeSegment(size_t stream_id,
                                const SegmentInfo& segment_info) {
  DCHECK_EQ(stream_id, 0u);
 
  if (segment_info.is_subsegment)
    return Status::OK;
 
  Status s = segmenter_->FinalizeSegment(segment_info.start_timestamp,
                                         segment_info.duration);
  if (!s.ok())
    return s;
  if (!segmenter_->segment_started())
    return Status::OK;
 
  int64_t segment_start_timestamp = segmenter_->segment_start_timestamp();
 
  std::string segment_path =
      options().segment_template.empty()
          ? options().output_file_name
          : GetSegmentName(options().segment_template, segment_start_timestamp,
                           segment_info.segment_number, options().bandwidth);
 
  const int64_t file_size = segmenter_->segment_buffer()->Size();
 
  EncryptSegmentIfNeeded(segmenter_->segment_buffer());
 
  RETURN_IF_ERROR(WriteSegment(segment_path, segmenter_->segment_buffer()));
 
  total_duration_ += segment_info.duration;
 
  if (muxer_listener()) {
    muxer_listener()->OnNewSegment(
        segment_path,
        segment_info.start_timestamp * segmenter_->timescale() +
            segmenter_->transport_stream_timestamp_offset(),
        segment_info.duration * segmenter_->timescale(), file_size,
        segment_info.segment_number);
  }
 
  segmenter_->set_segment_started(false);
 
  return Status::OK;
}
 
void TsMuxer::EncryptSegmentIfNeeded(BufferWriter* segment_buffer) {
  if (aes128_encryption_config_.protection_scheme != kAes128ProtectionScheme)
    return;
  if (segment_buffer->Size() == 0)
    return;
 
  // AES-128 HLS: encrypt the entire TS segment as one CBC stream.
  // Per RFC 8216 §5.2, PKCS7 padding is required; TS packet size (188 bytes)
  // is not a multiple of 16, so segments are not guaranteed to be
  // block-aligned.
  AesCbcEncryptor encryptor(kPkcs5Padding, AesCryptor::kUseConstantIv);
  if (!encryptor.InitializeWithIv(aes128_encryption_config_.key,
                                  aes128_encryption_config_.constant_iv)) {
    LOG(ERROR) << "AES-128: failed to initialize encryptor for segment.";
    return;
  }
 
  std::vector<uint8_t> plaintext(
      segment_buffer->Buffer(),
      segment_buffer->Buffer() + segment_buffer->Size());
  std::vector<uint8_t> ciphertext;
  if (!encryptor.Crypt(plaintext, &ciphertext)) {
    LOG(ERROR) << "AES-128: encryption failed for segment.";
    return;
  }
  segment_buffer->Clear();
  segment_buffer->AppendVector(ciphertext);
}
 
Status TsMuxer::WriteSegment(const std::string& segment_path,
                             BufferWriter* segment_buffer) {
  std::unique_ptr<File, FileCloser> file;
 
  if (output_file_) {
    // This is in single segment mode.
    Range range;
    range.start = media_ranges_.subsegment_ranges.empty()
                      ? 0
                      : (media_ranges_.subsegment_ranges.back().end + 1);
    range.end = range.start + segment_buffer->Size() - 1;
    media_ranges_.subsegment_ranges.push_back(range);
  } else {
    file.reset(File::Open(segment_path.c_str(), "w"));
    if (!file) {
      return Status(error::FILE_FAILURE,
                    "Cannot open file for write " + segment_path);
    }
  }
 
  RETURN_IF_ERROR(segment_buffer->WriteToFile(output_file_ ? output_file_.get()
                                                           : file.get()));
 
  if (file)
    RETURN_IF_ERROR(CloseFile(std::move(file)));
  return Status::OK;
}
 
Status TsMuxer::CloseFile(std::unique_ptr<File, FileCloser> file) {
  std::string file_name = file->file_name();
  if (!file.release()->Close()) {
    return Status(
        error::FILE_FAILURE,
        "Cannot close file " + file_name +
            ", possibly file permission issue or running out of disk space.");
  }
  return Status::OK;
}
 
void TsMuxer::FireOnMediaStartEvent() {
  if (!muxer_listener())
    return;
  muxer_listener()->OnMediaStart(options(), *streams().front(), kTsTimescale,
                                 MuxerListener::kContainerMpeg2ts);
}
 
void TsMuxer::FireOnMediaEndEvent() {
  if (!muxer_listener())
    return;
 
  muxer_listener()->OnMediaEnd(media_ranges_, total_duration_);
}
 
}  // namespace mp2t
}  // namespace media
}  // namespace shaka