feat: PSSH boxes generator for live packager

include/packager/live_packager.h

@@ -8,8 +8,6 @@
 #define PACKAGER_LIVE_PACKAGER_H_
 
 #include <packager/packager.h>
-#include <memory>
-#include <string>
 
 namespace shaka {
 
@@ -125,6 +123,34 @@ class LivePackager {
   LiveConfig config_;
 };
 
+struct PSSHData {
+  std::vector<uint8_t> cenc_box;
+  std::vector<uint8_t> mspr_box;
+  std::vector<uint8_t> mspr_pro;
+  std::vector<uint8_t> wv_box;
+};
+
+struct PSSHGeneratorInput {
+  enum struct MP4ProtectionSchemeFourCC : uint32_t {
+    CBCS = 0x63626373,
+    CENC = 0x63656e63,
+  };
+
+  MP4ProtectionSchemeFourCC protection_scheme;
+
+  // key of a single adaption set for DRM systems that don't support
+  // multile keys (i.e PlayReady)
+  std::vector<uint8_t> key;
+  // key id of the key for DRM systems that don't support
+  // multile keys (i.e PlayReady)
+  std::vector<uint8_t> key_id;
+  // key ids of all adaptation sets for DRM systems that support
+  // multiple keys (i.e Widevine, Common Encryption)
+  std::vector<std::vector<uint8_t>> key_ids;
+};
+
+Status GeneratePSSHData(const PSSHGeneratorInput& in, PSSHData* out);
+
 }  // namespace shaka
 
 #endif  // PACKAGER_LIVE_PACKAGER_H_

packager/live_packager.cc

@@ -15,10 +15,18 @@
 #include <absl/log/log.h>
 #include <packager/chunking_params.h>
 #include <packager/file.h>
+
 #include <packager/live_packager.h>
+#include <packager/macros/status.h>
 #include <packager/media/base/aes_encryptor.h>
 #include <packager/packager.h>
 
+#include "media/base/common_pssh_generator.h"
+#include "media/base/playready_pssh_generator.h"
+#include "media/base/protection_system_ids.h"
+#include "media/base/pssh_generator.h"
+#include "media/base/widevine_pssh_generator.h"
+
 namespace shaka {
 
 namespace {
@@ -418,4 +426,108 @@ Status Aes128EncryptedSegmentManager::InitializeEncryption(
   }
   return Status::OK;
 }
+
+void FillPSSHBoxByDRM(const media::ProtectionSystemSpecificInfo& pssh_info,
+                      PSSHData* data) {
+  if (std::equal(std::begin(media::kCommonSystemId),
+                 std::end(media::kCommonSystemId),
+                 pssh_info.system_id.begin())) {
+    data->cenc_box = pssh_info.psshs;
+    return;
+  }
+
+  if (std::equal(std::begin(media::kWidevineSystemId),
+                 std::end(media::kWidevineSystemId),
+                 pssh_info.system_id.begin())) {
+    data->wv_box = pssh_info.psshs;
+    return;
+  }
+
+  if (std::equal(std::begin(media::kPlayReadySystemId),
+                 std::end(media::kPlayReadySystemId),
+                 pssh_info.system_id.begin())) {
+    data->mspr_box = pssh_info.psshs;
+
+    std::unique_ptr<media::PsshBoxBuilder> box_builder =
+        media::PsshBoxBuilder::ParseFromBox(pssh_info.psshs.data(),
+                                            pssh_info.psshs.size());
+    data->mspr_pro = box_builder->pssh_data();
+  }
+}
+
+Status ValidatePSSHGeneratorInput(const PSSHGeneratorInput& input) {
+  constexpr int kKeySize = 16;
+
+  if (input.protection_scheme !=
+          PSSHGeneratorInput::MP4ProtectionSchemeFourCC::CBCS &&
+      input.protection_scheme !=
+          PSSHGeneratorInput::MP4ProtectionSchemeFourCC::CENC) {
+    LOG(WARNING) << "invalid encryption scheme in PSSH generator input";
+    return Status(error::INVALID_ARGUMENT,
+                  "invalid encryption scheme in PSSH generator input");
+  }
+
+  if (input.key.size() != kKeySize) {
+    LOG(WARNING) << "invalid key length in PSSH generator input";
+    return Status(error::INVALID_ARGUMENT,
+                  "invalid key length in PSSH generator input");
+  }
+
+  if (input.key_id.size() != kKeySize) {
+    LOG(WARNING) << "invalid key id length in PSSH generator input";
+    return Status(error::INVALID_ARGUMENT,
+                  "invalid key id length in PSSH generator input");
+  }
+
+  if (input.key_ids.empty()) {
+    LOG(WARNING) << "key ids cannot be empty in PSSH generator input";
+    return Status(error::INVALID_ARGUMENT,
+                  "key ids cannot be empty in PSSH generator input");
+  }
+
+  for (size_t i = 0; i < input.key_ids.size(); ++i) {
+    if (input.key_ids[i].size() != kKeySize) {
+      LOG(WARNING) << "invalid key id length in key ids array in PSSH "
+                      "generator input, index " +
+                          std::to_string(i);
+      return Status(error::INVALID_ARGUMENT,
+                    "invalid key id length in key ids array in PSSH generator "
+                    "input, index " +
+                        std::to_string(i));
+    }
+  }
+
+  return Status::OK;
+}
+
+Status GeneratePSSHData(const PSSHGeneratorInput& in, PSSHData* out) {
+  const char* kNoExtraHeadersForPlayReady = "";
+
+  RETURN_IF_ERROR(ValidatePSSHGeneratorInput(in));
+  if (!out) {
+    return Status(error::INVALID_ARGUMENT, "output data cannot be null");
+  }
+
+  std::vector<std::unique_ptr<media::PsshGenerator>> pssh_generators;
+  pssh_generators.emplace_back(std::make_unique<media::CommonPsshGenerator>());
+  pssh_generators.emplace_back(std::make_unique<media::PlayReadyPsshGenerator>(
+      kNoExtraHeadersForPlayReady,
+      static_cast<media::FourCC>(in.protection_scheme)));
+  pssh_generators.emplace_back(std::make_unique<media::WidevinePsshGenerator>(
+      static_cast<media::FourCC>(in.protection_scheme)));
+
+  for (const auto& pssh_generator : pssh_generators) {
+    media::ProtectionSystemSpecificInfo info;
+    if (pssh_generator->SupportMultipleKeys()) {
+      RETURN_IF_ERROR(
+          pssh_generator->GeneratePsshFromKeyIds(in.key_ids, &info));
+    } else {
+      RETURN_IF_ERROR(pssh_generator->GeneratePsshFromKeyIdAndKey(
+          in.key_id, in.key, &info));
+    }
+    FillPSSHBoxByDRM(info, out);
+  }
+
+  return Status::OK;
+}
 }  // namespace shaka

packager/live_packager_test.cc

@@ -19,6 +19,8 @@
 #include <packager/media/formats/mp4/box_definitions.h>
 #include <packager/media/formats/mp4/box_reader.h>
 
+#include "absl/strings/escaping.h"
+
 namespace shaka {
 namespace {
 
@@ -62,6 +64,22 @@ std::vector<uint8_t> ReadTestDataFile(const std::string& name) {
   return data;
 }
 
+std::vector<uint8_t> unhex(const std::string& in) {
+  auto converted = absl::HexStringToBytes(in);
+  return {converted.begin(), converted.end()};
+}
+
+std::vector<uint8_t> unbase64(const std::string& base64_string) {
+  std::string str;
+  std::vector<uint8_t> bytes;
+  if (!absl::Base64Unescape(base64_string, &str)) {
+    return {};
+  }
+
+  bytes.assign(str.begin(), str.end());
+  return bytes;
+}
+
 bool ParseAndCheckType(media::mp4::Box& box, media::mp4::BoxReader* reader) {
   box.Parse(reader);
   return box.BoxType() == reader->type();
@@ -235,6 +253,95 @@ void CheckSegment(const LiveConfig& config, const FullSegmentBuffer& buffer) {
 
 }  // namespace
 
+TEST(GeneratePSSHData, GeneratesPSSHBoxesAndMSPRObject) {
+  PSSHGeneratorInput in{.protection_scheme = PSSHGeneratorInput::MP4ProtectionSchemeFourCC::CENC,
+                        .key_id = unhex("00000000621f2afe7ab2c868d5fd2e2e"),
+                        .key = unhex("1af987fa084ff3c0f4ad35a6bdab98e2"),
+                        .key_ids = {unhex("00000000621f2afe7ab2c868d5fd2e2e"),
+                                    unhex("00000000621f2afe7ab2c868d5fd2e2f")}};
+
+  PSSHData expected{
+      .cenc_box = unbase64("AAAARHBzc2gBAAAAEHfv7MCyTQKs4zweUuL7SwAAAAIAAAAAYh8"
+                           "q/nqyyGjV/S4uAAAAAGIfKv56ssho1f0uLwAAAAA="),
+      .mspr_box = unbase64(
+          "AAACJnBzc2gAAAAAmgTweZhAQoarkuZb4IhflQAAAgYGAgAAAQABAPwBPABXAFIATQBI"
+          "AEUAQQBEAEUAUgAgAHgAbQBsAG4AcwA9ACIAaAB0AHQAcAA6AC8ALwBzAGMAaABlAG0A"
+          "YQBzAC4AbQBpAGMAcgBvAHMAbwBmAHQALgBjAG8AbQAvAEQAUgBNAC8AMgAwADAANwAv"
+          "ADAAMwAvAFAAbABhAHkAUgBlAGEAZAB5AEgAZQBhAGQAZQByACIAIAB2AGUAcgBzAGkA"
+          "bwBuAD0AIgA0AC4AMAAuADAALgAwACIAPgA8AEQAQQBUAEEAPgA8AFAAUgBPAFQARQBD"
+          "AFQASQBOAEYATwA+"
+          "ADwASwBFAFkATABFAE4APgAxADYAPAAvAEsARQBZAEwARQBOAD4APABBAEwARwBJAEQA"
+          "PgBBAEUAUwBDAFQAUgA8AC8AQQBMAEcASQBEAD4APAAvAFAAUgBPAFQARQBDAFQASQBO"
+          "AEYATwA+"
+          "ADwASwBJAEQAPgBBAEEAQQBBAEEAQgA5AGkALwBpAHAANgBzAHMAaABvADEAZgAwAHUA"
+          "TABnAD0APQA8AC8ASwBJAEQAPgA8AEMASABFAEMASwBTAFUATQA+"
+          "ADQAZgB1AEIAdABEAFUAKwBLAGsARQA9ADwALwBDAEgARQBDAEsAUwBVAE0APgA8AC8A"
+          "RABBAFQAQQA+ADwALwBXAFIATQBIAEUAQQBEAEUAUgA+AA=="),
+      .mspr_pro = unbase64(
+          "BgIAAAEAAQD8ATwAVwBSAE0ASABFAEEARABFAFIAIAB4AG0AbABuAHMAPQAiAGgAdAB0"
+          "AHAAOgAvAC8AcwBjAGgAZQBtAGEAcwAuAG0AaQBjAHIAbwBzAG8AZgB0AC4AYwBvAG0A"
+          "LwBEAFIATQAvADIAMAAwADcALwAwADMALwBQAGwAYQB5AFIAZQBhAGQAeQBIAGUAYQBk"
+          "AGUAcgAiACAAdgBlAHIAcwBpAG8AbgA9ACIANAAuADAALgAwAC4AMAAiAD4APABEAEEA"
+          "VABBAD4APABQAFIATwBUAEUAQwBUAEkATgBGAE8APgA8AEsARQBZAEwARQBOAD4AMQA2"
+          "ADwALwBLAEUAWQBMAEUATgA+"
+          "ADwAQQBMAEcASQBEAD4AQQBFAFMAQwBUAFIAPAAvAEEATABHAEkARAA+"
+          "ADwALwBQAFIATwBUAEUAQwBUAEkATgBGAE8APgA8AEsASQBEAD4AQQBBAEEAQQBBAEIA"
+          "OQBpAC8AaQBwADYAcwBzAGgAbwAxAGYAMAB1AEwAZwA9AD0APAAvAEsASQBEAD4APABD"
+          "AEgARQBDAEsAUwBVAE0APgA0AGYAdQBCAHQARABVACsASwBrAEUAPQA8AC8AQwBIAEUA"
+          "QwBLAFMAVQBNAD4APAAvAEQAQQBUAEEAPgA8AC8AVwBSAE0ASABFAEEARABFAFIAPgA"
+          "="),
+      .wv_box =
+          unbase64("AAAASnBzc2gAAAAA7e+LqXnWSs6jyCfc1R0h7QAAACoSEAAAAABiHyr+"
+                   "erLIaNX9Li4SEAAAAABiHyr+erLIaNX9Li9I49yVmwY=")};
+  PSSHData actual{};
+
+  ASSERT_EQ(Status::OK, GeneratePSSHData(in, &actual));
+  ASSERT_EQ(expected.cenc_box, actual.cenc_box);
+  ASSERT_EQ(expected.mspr_box, actual.mspr_box);
+  ASSERT_EQ(expected.mspr_pro, actual.mspr_pro);
+  ASSERT_EQ(expected.wv_box, actual.wv_box);
+}
+
+TEST(GeneratePSSHData, FailsOnInvalidInput) {
+  const PSSHGeneratorInput valid_input{
+      .protection_scheme = PSSHGeneratorInput::MP4ProtectionSchemeFourCC::CENC,
+      .key_id = unhex("00000000621f2afe7ab2c868d5fd2e2e"),
+      .key = unhex("1af987fa084ff3c0f4ad35a6bdab98e2"),
+      .key_ids = {unhex("00000000621f2afe7ab2c868d5fd2e2e"),
+                  unhex("00000000621f2afe7ab2c868d5fd2e2f")}};
+
+  PSSHGeneratorInput in;
+  ASSERT_EQ(Status(error::INVALID_ARGUMENT,
+                   "invalid encryption scheme in PSSH generator input"),
+            GeneratePSSHData(in, nullptr));
+
+  in.protection_scheme = valid_input.protection_scheme;
+  ASSERT_EQ(Status(error::INVALID_ARGUMENT,
+                   "invalid key length in PSSH generator input"),
+            GeneratePSSHData(in, nullptr));
+
+  in.key = valid_input.key;
+  ASSERT_EQ(Status(error::INVALID_ARGUMENT,
+                   "invalid key id length in PSSH generator input"),
+            GeneratePSSHData(in, nullptr));
+
+  in.key_id = valid_input.key_id;
+  ASSERT_EQ(Status(error::INVALID_ARGUMENT,
+                   "key ids cannot be empty in PSSH generator input"),
+            GeneratePSSHData(in, nullptr));
+
+  in.key_ids = valid_input.key_ids;
+  in.key_ids[1] = {};
+  ASSERT_EQ(Status(error::INVALID_ARGUMENT,
+                   "invalid key id length in key ids array in PSSH generator "
+                   "input, index 1"),
+            GeneratePSSHData(in, nullptr));
+
+  in.key_ids = valid_input.key_ids;
+  ASSERT_EQ(Status(error::INVALID_ARGUMENT, "output data cannot be null"),
+            GeneratePSSHData(in, nullptr));
+}
+
 class LivePackagerBaseTest : public ::testing::Test {
  public:
   void SetUp() override {