Fix entitlement tests

whitebox/BUILD

@@ -1,8 +1,9 @@
 # Copyright 2022 Google LLC. All Rights Reserved.
 
-load("//:defs.bzl", "build_type")
+load("//:defs.bzl", "build_type", "disable_features")
 
 build_type(name="build_type", build_setting_default="chrome")
+disable_features(name="disable_features", build_setting_default="")
 
 config_setting(
     name="is_chrome",
@@ -38,3 +39,11 @@ config_setting(
         "//:build_type": "old_vmpra",
     }
 )
+
+config_setting(
+    name="is_ce_disable_entitlement",
+    flag_values={
+        "//:build_type": "ce",
+        "//:disable_features": "entitlement",
+    }
+)

whitebox/api/BUILD

@@ -5,9 +5,12 @@ package(default_visibility = ["//visibility:private"])
 cc_library(
     name = "shared_settings",
     defines = select({
+        "//:is_ce_disable_entitlement": [],
+        "//:is_ce": ["HAS_ENTITLEMENT"],
+        "//conditions:default": [],
+    }) + select({
         "//:is_ce": [
             "ALWAYS_DECRYPT_TO_CLEAR",
-            "HAS_ENTITLEMENT",
             "HAS_SIGN_PST_REPORT",
         ],
         "//:is_old_api": [],

whitebox/api/golden_data.cc

@@ -97,6 +97,7 @@ GoldenData::GoldenData() {
   //         -iv 6fc04cd8423d5f660ca045769a200048 |
   //       xxd -i
   entitlement_.entitlement_key.id = GetFreeId();
+  entitlement_.entitlement_key.level = SecurityLevel::kSoftwareSecureCrypto;
   entitlement_.entitlement_key.key = {
       0x8e, 0x68, 0x24, 0x47, 0xb6, 0xb4, 0x66, 0x96,
       0xeb, 0x87, 0x4d, 0x1e, 0x38, 0x46, 0x77, 0x84,

whitebox/api/license_whitebox_entitlement_content_key_test.cc

@@ -21,6 +21,7 @@ class LicenseWhiteboxEntitlementContentKeyTest
     server_ = TestServer::CreateDualKey();
 
     TestLicenseBuilder builder;
+    builder.GetSettings().odk_version = TestLicenseBuilder::OdkVersion::k16_5;
     builder.AddSigningKey(TestLicenseBuilder::DefaultSigningKey());
     builder.AddEntitlementKey(
         golden_data_.EntitlementContent().entitlement_key);

whitebox/api/license_whitebox_process_license_response_test.cc

@@ -45,6 +45,7 @@ class LicenseWhiteboxProcessLicenseResponseTest
 
   void UseLicenseWithEntitlementKey() {
     TestLicenseBuilder builder;
+    builder.GetSettings().odk_version = TestLicenseBuilder::OdkVersion::k16_5;
     builder.AddSigningKey(TestLicenseBuilder::DefaultSigningKey());
     builder.AddEntitlementKey(
         golden_data_.EntitlementContent().entitlement_key);

whitebox/api/test_key_types.h

@@ -44,6 +44,8 @@ struct EntitlementKeyData {
   // The unique key id for this key. Any instance with this id should contain
   // the same level and key as this.
   KeyId id;
+
+  SecurityLevel level;
   Aes256Key key;
 };
 

whitebox/api/test_license_builder.cc

@@ -380,6 +380,12 @@ void AddEntitlementKeyToContainer(
   std::vector<uint8_t> key(key_data.key.begin(), key_data.key.end());
   auto encrypted_key = Encrypt(container_key, key_iv, key);
   container->set_key(encrypted_key);
+
+  auto* key_control = container->mutable_key_control();
+  const auto key_control_block =
+      CreateKeyControlBlock(SecurityLevelToProto(key_data.level), key_control);
+  key_control->set_key_control_block(key_control_block.data(),
+                                     key_control_block.size());
 }
 
 void AddSigningKeyToContainer(const TestLicenseBuilder::SigningKey& key_data,

whitebox/defs.bzl

@@ -12,3 +12,19 @@ def _impl(ctx):
   return BuildType(type=ctx.build_setting_value)
 
 build_type = rule(implementation=_impl, build_setting=config.string(flag=True))
+
+
+Features = provider(fields = ["type"])
+
+possible_features = ["entitlement"]
+
+def _feature_impl(ctx):
+  # Allow an empty value since that is the default value.
+  for v in ctx.build_setting_value:
+    if v and v not in possible_features:
+      fail("Invalid features: " + v)
+  return Features(type=ctx.build_setting_value)
+
+disable_features = rule(
+    implementation=_feature_impl,
+    build_setting=config.string(flag=True, allow_multiple=True))