Switch MacWrapper to using full primitives, among other things reworking the internal structure and logging the input data length instead of the processed data length in the LEGACY case, and registering LegacyFullMac to handle potential old/custom primitives using the old APIs.

PiperOrigin-RevId: 555869169

6 files changed, 127 insertions, 95 deletions

diff --git a/java_src/src/main/java/com/google/crypto/tink/mac/BUILD.bazel b/java_src/src/main/java/com/google/crypto/tink/mac/BUILD.bazel
index 94ef3db7f..bcfc32fc4 100644
--- a/java_src/src/main/java/com/google/crypto/tink/mac/BUILD.bazel
+++ b/java_src/src/main/java/com/google/crypto/tink/mac/BUILD.bazel
@@ -109,17 +109,19 @@ java_library(
     srcs = ["MacWrapper.java"],
     deps = [
         ":mac_key",
-        "//proto:tink_java_proto",
         "//src/main/java/com/google/crypto/tink:crypto_format",
         "//src/main/java/com/google/crypto/tink:mac",
         "//src/main/java/com/google/crypto/tink:primitive_set",
         "//src/main/java/com/google/crypto/tink:primitive_wrapper",
         "//src/main/java/com/google/crypto/tink:registry",
+        "//src/main/java/com/google/crypto/tink/internal:legacy_proto_key",
         "//src/main/java/com/google/crypto/tink/internal:monitoring_util",
         "//src/main/java/com/google/crypto/tink/internal:mutable_monitoring_registry",
+        "//src/main/java/com/google/crypto/tink/internal:mutable_primitive_registry",
+        "//src/main/java/com/google/crypto/tink/internal:primitive_constructor",
+        "//src/main/java/com/google/crypto/tink/mac/internal:legacy_full_mac",
         "//src/main/java/com/google/crypto/tink/monitoring:monitoring_client",
         "//src/main/java/com/google/crypto/tink/monitoring:monitoring_keyset_info",
-        "//src/main/java/com/google/crypto/tink/subtle:bytes",
         "//src/main/java/com/google/crypto/tink/util:bytes",
     ],
 )
@@ -129,17 +131,19 @@ android_library(
     srcs = ["MacWrapper.java"],
     deps = [
         ":mac_key-android",
-        "//proto:tink_java_proto_lite",
         "//src/main/java/com/google/crypto/tink:crypto_format-android",
         "//src/main/java/com/google/crypto/tink:mac-android",
         "//src/main/java/com/google/crypto/tink:primitive_set-android",
         "//src/main/java/com/google/crypto/tink:primitive_wrapper-android",
         "//src/main/java/com/google/crypto/tink:registry-android",
+        "//src/main/java/com/google/crypto/tink/internal:legacy_proto_key-android",
         "//src/main/java/com/google/crypto/tink/internal:monitoring_util-android",
         "//src/main/java/com/google/crypto/tink/internal:mutable_monitoring_registry-android",
+        "//src/main/java/com/google/crypto/tink/internal:mutable_primitive_registry-android",
+        "//src/main/java/com/google/crypto/tink/internal:primitive_constructor-android",
+        "//src/main/java/com/google/crypto/tink/mac/internal:legacy_full_mac-android",
         "//src/main/java/com/google/crypto/tink/monitoring:monitoring_client-android",
         "//src/main/java/com/google/crypto/tink/monitoring:monitoring_keyset_info-android",
-        "//src/main/java/com/google/crypto/tink/subtle:bytes-android",
         "//src/main/java/com/google/crypto/tink/util:bytes-android",
     ],
 )
diff --git a/java_src/src/main/java/com/google/crypto/tink/mac/MacWrapper.java b/java_src/src/main/java/com/google/crypto/tink/mac/MacWrapper.java
index f0ab5bb17..5b8d2aa9a 100644
--- a/java_src/src/main/java/com/google/crypto/tink/mac/MacWrapper.java
+++ b/java_src/src/main/java/com/google/crypto/tink/mac/MacWrapper.java
@@ -21,11 +21,14 @@ import com.google.crypto.tink.Mac;
 import com.google.crypto.tink.PrimitiveSet;
 import com.google.crypto.tink.PrimitiveWrapper;
 import com.google.crypto.tink.Registry;
+import com.google.crypto.tink.internal.LegacyProtoKey;
 import com.google.crypto.tink.internal.MonitoringUtil;
 import com.google.crypto.tink.internal.MutableMonitoringRegistry;
+import com.google.crypto.tink.internal.MutablePrimitiveRegistry;
+import com.google.crypto.tink.internal.PrimitiveConstructor;
+import com.google.crypto.tink.mac.internal.LegacyFullMac;
 import com.google.crypto.tink.monitoring.MonitoringClient;
 import com.google.crypto.tink.monitoring.MonitoringKeysetInfo;
-import com.google.crypto.tink.proto.OutputPrefixType;
 import com.google.crypto.tink.util.Bytes;
 import java.security.GeneralSecurityException;
 import java.util.Arrays;
@@ -44,8 +47,10 @@ import java.util.logging.Logger;
 class MacWrapper implements PrimitiveWrapper<Mac, Mac> {
   private static final Logger logger = Logger.getLogger(MacWrapper.class.getName());
 
-  private static final byte[] FORMAT_VERSION = new byte[] {0};
   private static final MacWrapper WRAPPER = new MacWrapper();
+  private static final PrimitiveConstructor<LegacyProtoKey, Mac>
+      LEGACY_FULL_MAC_PRIMITIVE_CONSTRUCTOR =
+          PrimitiveConstructor.create(LegacyFullMac::create, LegacyProtoKey.class, Mac.class);
 
   private static class WrappedMac implements Mac {
     private final PrimitiveSet<Mac> primitives;
@@ -67,16 +72,9 @@ class MacWrapper implements PrimitiveWrapper<Mac, Mac> {
 
     @Override
     public byte[] computeMac(final byte[] data) throws GeneralSecurityException {
-      byte[] data2 = data;
-      if (primitives.getPrimary().getOutputPrefixType().equals(OutputPrefixType.LEGACY)) {
-        data2 = com.google.crypto.tink.subtle.Bytes.concat(data, FORMAT_VERSION);
-      }
       try {
-        byte[] output =
-            com.google.crypto.tink.subtle.Bytes.concat(
-                primitives.getPrimary().getIdentifier(),
-                primitives.getPrimary().getPrimitive().computeMac(data2));
-        computeLogger.log(primitives.getPrimary().getKeyId(), data2.length);
+        byte[] output = primitives.getPrimary().getFullPrimitive().computeMac(data);
+        computeLogger.log(primitives.getPrimary().getKeyId(), data.length);
         return output;
       } catch (GeneralSecurityException e) {
         computeLogger.logFailure();
@@ -93,16 +91,11 @@ class MacWrapper implements PrimitiveWrapper<Mac, Mac> {
         throw new GeneralSecurityException("tag too short");
       }
       byte[] prefix = Arrays.copyOf(mac, CryptoFormat.NON_RAW_PREFIX_SIZE);
-      byte[] macNoPrefix = Arrays.copyOfRange(mac, CryptoFormat.NON_RAW_PREFIX_SIZE, mac.length);
       List<PrimitiveSet.Entry<Mac>> entries = primitives.getPrimitive(prefix);
       for (PrimitiveSet.Entry<Mac> entry : entries) {
-        byte[] data2 = data;
-        if (entry.getOutputPrefixType().equals(OutputPrefixType.LEGACY)) {
-          data2 = com.google.crypto.tink.subtle.Bytes.concat(data, FORMAT_VERSION);
-        }
         try {
-          entry.getPrimitive().verifyMac(macNoPrefix, data2);
-          verifyLogger.log(entry.getKeyId(), data2.length);
+          entry.getFullPrimitive().verifyMac(mac, data);
+          verifyLogger.log(entry.getKeyId(), data.length);
           // If there is no exception, the MAC is valid and we can return.
           return;
         } catch (GeneralSecurityException e) {
@@ -115,7 +108,7 @@ class MacWrapper implements PrimitiveWrapper<Mac, Mac> {
       entries = primitives.getRawPrimitives();
       for (PrimitiveSet.Entry<Mac> entry : entries) {
         try {
-          entry.getPrimitive().verifyMac(mac, data);
+          entry.getFullPrimitive().verifyMac(mac, data);
           verifyLogger.log(entry.getKeyId(), data.length);
           // If there is no exception, the MAC is valid and we can return.
           return;
@@ -169,7 +162,9 @@ class MacWrapper implements PrimitiveWrapper<Mac, Mac> {
     return Mac.class;
   }
 
- public static void register() throws GeneralSecurityException {
+  public static void register() throws GeneralSecurityException {
     Registry.registerPrimitiveWrapper(WRAPPER);
+    MutablePrimitiveRegistry.globalInstance()
+        .registerPrimitiveConstructor(LEGACY_FULL_MAC_PRIMITIVE_CONSTRUCTOR);
   }
 }
diff --git a/java_src/src/test/java/com/google/crypto/tink/internal/BUILD.bazel b/java_src/src/test/java/com/google/crypto/tink/internal/BUILD.bazel
index 5f6b30d6c..b441ac2fd 100644
--- a/java_src/src/test/java/com/google/crypto/tink/internal/BUILD.bazel
+++ b/java_src/src/test/java/com/google/crypto/tink/internal/BUILD.bazel
@@ -476,13 +476,17 @@ java_test(
         "//src/main/java/com/google/crypto/tink:registry_cluster",
         "//src/main/java/com/google/crypto/tink/aead:aes_eax_key",
         "//src/main/java/com/google/crypto/tink/aead:aes_eax_parameters",
+        "//src/main/java/com/google/crypto/tink/internal:legacy_proto_key",
         "//src/main/java/com/google/crypto/tink/internal:mutable_primitive_registry",
+        "//src/main/java/com/google/crypto/tink/internal:mutable_serialization_registry",
+        "//src/main/java/com/google/crypto/tink/internal:proto_key_serialization",
         "//src/main/java/com/google/crypto/tink/internal:registry_configuration",
         "//src/main/java/com/google/crypto/tink/mac:chunked_mac",
         "//src/main/java/com/google/crypto/tink/mac:chunked_mac_computation",
         "//src/main/java/com/google/crypto/tink/mac:hmac_key",
         "//src/main/java/com/google/crypto/tink/mac:hmac_parameters",
         "//src/main/java/com/google/crypto/tink/mac:mac_config",
+        "//src/main/java/com/google/crypto/tink/mac/internal:legacy_full_mac",
         "//src/main/java/com/google/crypto/tink/util:secret_bytes",
         "@maven//:com_google_protobuf_protobuf_java",
         "@maven//:com_google_truth_truth",
diff --git a/java_src/src/test/java/com/google/crypto/tink/internal/RegistryConfigurationTest.java b/java_src/src/test/java/com/google/crypto/tink/internal/RegistryConfigurationTest.java
index 5a7ef7596..4b24e1aeb 100644
--- a/java_src/src/test/java/com/google/crypto/tink/internal/RegistryConfigurationTest.java
+++ b/java_src/src/test/java/com/google/crypto/tink/internal/RegistryConfigurationTest.java
@@ -35,6 +35,7 @@ import com.google.crypto.tink.mac.HmacKey;
 import com.google.crypto.tink.mac.HmacParameters;
 import com.google.crypto.tink.mac.HmacParameters.HashType;
 import com.google.crypto.tink.mac.MacConfig;
+import com.google.crypto.tink.mac.internal.LegacyFullMac;
 import com.google.crypto.tink.proto.HmacParams;
 import com.google.crypto.tink.proto.KeyData;
 import com.google.crypto.tink.proto.KeyData.KeyMaterialType;
@@ -59,6 +60,7 @@ public class RegistryConfigurationTest {
   private static HmacKey rawKey;
   private static KeyData rawKeyData;
   private static Keyset.Key rawKeysetKey;
+  private static LegacyProtoKey legacyProtoRawKey;
 
   @Before
   public void setUp() throws GeneralSecurityException {
@@ -110,6 +112,11 @@ public class RegistryConfigurationTest {
               .setKeyId(keysetHandle.getKeysetInfo().getPrimaryKeyId())
               .setOutputPrefixType(OutputPrefixType.RAW)
               .build();
+      legacyProtoRawKey =
+          new LegacyProtoKey(
+              MutableSerializationRegistry.globalInstance()
+                  .serializeKey(rawKey, ProtoKeySerialization.class, InsecureSecretKeyAccess.get()),
+              InsecureSecretKeyAccess.get());
     } catch (GeneralSecurityException e) {
       throw new IllegalStateException(e);
     }
@@ -148,13 +155,13 @@ public class RegistryConfigurationTest {
     byte[] plaintext = "plaintext".getBytes(UTF_8);
 
     Mac registryMac = Registry.getPrimitive(rawKeyData, Mac.class);
-    Mac configurationMac = RegistryConfiguration.get().getLegacyPrimitive(rawKeyData, Mac.class);
+    // The following relies on the fact that internally LegacyFullMac uses RegistryConfiguration.
     Mac wrappedConfigurationMac =
         RegistryConfiguration.get()
             .wrap(
                 PrimitiveSet.newBuilder(Mac.class)
                     .addPrimaryFullPrimitiveAndOptionalPrimitive(
-                        null, configurationMac, rawKeysetKey)
+                        LegacyFullMac.create(legacyProtoRawKey), null, rawKeysetKey)
                     .build(),
                 Mac.class);
 
diff --git a/java_src/src/test/java/com/google/crypto/tink/mac/BUILD.bazel b/java_src/src/test/java/com/google/crypto/tink/mac/BUILD.bazel
index 1a46ce179..fe7d2013c 100644
--- a/java_src/src/test/java/com/google/crypto/tink/mac/BUILD.bazel
+++ b/java_src/src/test/java/com/google/crypto/tink/mac/BUILD.bazel
@@ -118,25 +118,25 @@ java_test(
     size = "small",
     srcs = ["MacWrapperTest.java"],
     deps = [
-        "//proto:tink_java_proto",
         "//src/main/java/com/google/crypto/tink:insecure_secret_key_access",
         "//src/main/java/com/google/crypto/tink:mac",
-        "//src/main/java/com/google/crypto/tink:primitive_set",
+        "//src/main/java/com/google/crypto/tink:registry",
         "//src/main/java/com/google/crypto/tink:registry_cluster",
         "//src/main/java/com/google/crypto/tink/internal:mutable_monitoring_registry",
+        "//src/main/java/com/google/crypto/tink/internal:mutable_primitive_registry",
+        "//src/main/java/com/google/crypto/tink/internal:primitive_constructor",
         "//src/main/java/com/google/crypto/tink/internal/testing:fake_monitoring_client",
         "//src/main/java/com/google/crypto/tink/mac:aes_cmac_key",
         "//src/main/java/com/google/crypto/tink/mac:aes_cmac_parameters",
         "//src/main/java/com/google/crypto/tink/mac:aes_cmac_proto_serialization",
         "//src/main/java/com/google/crypto/tink/mac:hmac_key",
+        "//src/main/java/com/google/crypto/tink/mac:hmac_key_manager",
         "//src/main/java/com/google/crypto/tink/mac:hmac_parameters",
         "//src/main/java/com/google/crypto/tink/mac:mac_config",
         "//src/main/java/com/google/crypto/tink/mac:mac_wrapper",
         "//src/main/java/com/google/crypto/tink/mac/internal:hmac_proto_serialization",
         "//src/main/java/com/google/crypto/tink/monitoring:monitoring_annotations",
         "//src/main/java/com/google/crypto/tink/subtle:hex",
-        "//src/main/java/com/google/crypto/tink/subtle:random",
-        "//src/main/java/com/google/crypto/tink/testing:test_util",
         "//src/main/java/com/google/crypto/tink/util:secret_bytes",
         "@maven//:com_google_truth_truth",
         "@maven//:junit_junit",
diff --git a/java_src/src/test/java/com/google/crypto/tink/mac/MacWrapperTest.java b/java_src/src/test/java/com/google/crypto/tink/mac/MacWrapperTest.java
index a60dce6cb..443f3870f 100644
--- a/java_src/src/test/java/com/google/crypto/tink/mac/MacWrapperTest.java
+++ b/java_src/src/test/java/com/google/crypto/tink/mac/MacWrapperTest.java
@@ -23,18 +23,15 @@ import static org.junit.Assert.assertThrows;
 import com.google.crypto.tink.InsecureSecretKeyAccess;
 import com.google.crypto.tink.KeysetHandle;
 import com.google.crypto.tink.Mac;
-import com.google.crypto.tink.PrimitiveSet;
+import com.google.crypto.tink.Registry;
 import com.google.crypto.tink.internal.MutableMonitoringRegistry;
+import com.google.crypto.tink.internal.MutablePrimitiveRegistry;
+import com.google.crypto.tink.internal.PrimitiveConstructor;
 import com.google.crypto.tink.internal.testing.FakeMonitoringClient;
 import com.google.crypto.tink.mac.HmacParameters.HashType;
 import com.google.crypto.tink.mac.internal.HmacProtoSerialization;
 import com.google.crypto.tink.monitoring.MonitoringAnnotations;
-import com.google.crypto.tink.proto.KeyStatusType;
-import com.google.crypto.tink.proto.Keyset.Key;
-import com.google.crypto.tink.proto.OutputPrefixType;
 import com.google.crypto.tink.subtle.Hex;
-import com.google.crypto.tink.subtle.Random;
-import com.google.crypto.tink.testing.TestUtil;
 import com.google.crypto.tink.util.SecretBytes;
 import java.security.GeneralSecurityException;
 import java.util.List;
@@ -197,6 +194,9 @@ public class MacWrapperTest {
 
   @Test
   public void testComputeVerifyMac_throwsOnWrongKey() throws Exception {
+    MutablePrimitiveRegistry.resetGlobalInstanceTestOnly();
+    MacConfig.register();
+
     byte[] plaintext = "plaintext".getBytes(UTF_8);
     KeysetHandle computeKeysetHandle =
         KeysetHandle.newBuilder()
@@ -216,6 +216,9 @@ public class MacWrapperTest {
 
   @Test
   public void testVerifyMac_checksAllNecessaryRawKeys() throws Exception {
+    MutablePrimitiveRegistry.resetGlobalInstanceTestOnly();
+    MacConfig.register();
+
     byte[] plaintext = "plaintext".getBytes(UTF_8);
     KeysetHandle computeKeysetHandle =
         KeysetHandle.newBuilder()
@@ -236,6 +239,9 @@ public class MacWrapperTest {
 
   @Test
   public void testVerifyMac_checksRawKeysWhenTagHasTinkKeyPrefix() throws Exception {
+    MutablePrimitiveRegistry.resetGlobalInstanceTestOnly();
+    MacConfig.register();
+
     byte[] plaintext = "plaintext".getBytes(UTF_8);
     byte[] tag = Hex.decode("0152af9740d2fab0cf3f");
     HmacKey rawKey5 =
@@ -267,6 +273,9 @@ public class MacWrapperTest {
 
   @Test
   public void computeMac_usesPrimaryKey() throws Exception {
+    MutablePrimitiveRegistry.resetGlobalInstanceTestOnly();
+    MacConfig.register();
+
     byte[] plaintext = "plaintext".getBytes(UTF_8);
     KeysetHandle keysetHandle =
         KeysetHandle.newBuilder()
@@ -288,6 +297,9 @@ public class MacWrapperTest {
 
   @Test
   public void testComputeVerifyMac_manyKeysWork() throws Exception {
+    MutablePrimitiveRegistry.resetGlobalInstanceTestOnly();
+    MacConfig.register();
+
     byte[] plaintext = "plaintext".getBytes(UTF_8);
     KeysetHandle assortedKeysetHandle =
         KeysetHandle.newBuilder()
@@ -311,6 +323,9 @@ public class MacWrapperTest {
 
   @Test
   public void testVerifyMac_shiftedPrimaryWithManyKeysWorks() throws Exception {
+    MutablePrimitiveRegistry.resetGlobalInstanceTestOnly();
+    MacConfig.register();
+
     byte[] plaintext = "plaintext".getBytes(UTF_8);
     KeysetHandle assortedKeysetHandle0 =
         KeysetHandle.newBuilder()
@@ -339,6 +354,9 @@ public class MacWrapperTest {
 
   @Test
   public void testMultipleKeysWithoutAnnotation() throws Exception {
+    MutablePrimitiveRegistry.resetGlobalInstanceTestOnly();
+    MacConfig.register();
+
     FakeMonitoringClient fakeMonitoringClient = new FakeMonitoringClient();
     MutableMonitoringRegistry.globalInstance().clear();
     MutableMonitoringRegistry.globalInstance().registerMonitoringClient(fakeMonitoringClient);
@@ -388,45 +406,36 @@ public class MacWrapperTest {
 
   @Test
   public void testWithAnnotation_hasMonitoring() throws Exception {
+    MutablePrimitiveRegistry.resetGlobalInstanceTestOnly();
+    MacConfig.register();
+
     FakeMonitoringClient fakeMonitoringClient = new FakeMonitoringClient();
     MutableMonitoringRegistry.globalInstance().clear();
     MutableMonitoringRegistry.globalInstance().registerMonitoringClient(fakeMonitoringClient);
 
-    Key tinkKey =
-        TestUtil.createKey(
-            TestUtil.createHmacKeyData(Random.randBytes(HMAC_KEY_SIZE), 16),
-            42,
-            KeyStatusType.ENABLED,
-            OutputPrefixType.TINK);
-    Key rawKey =
-        TestUtil.createKey(
-            TestUtil.createHmacKeyData(Random.randBytes(HMAC_KEY_SIZE), 16),
-            43,
-            KeyStatusType.ENABLED,
-            OutputPrefixType.RAW);
-    Key legacyKey =
-        TestUtil.createKey(
-            TestUtil.createHmacKeyData(Random.randBytes(HMAC_KEY_SIZE), 16),
-            44,
-            KeyStatusType.ENABLED,
-            OutputPrefixType.LEGACY);
     MonitoringAnnotations annotations =
         MonitoringAnnotations.newBuilder().add("annotation_name", "annotation_value").build();
-    Mac rawMac =
-        new MacWrapper()
-            .wrap(
-                TestUtil.createPrimitiveSetWithAnnotations(
-                    TestUtil.createKeyset(rawKey), annotations, Mac.class));
-    Mac legacyMac =
-        new MacWrapper()
-            .wrap(
-                TestUtil.createPrimitiveSetWithAnnotations(
-                    TestUtil.createKeyset(legacyKey), annotations, Mac.class));
-    Mac mac =
-        new MacWrapper()
-            .wrap(
-                TestUtil.createPrimitiveSetWithAnnotations(
-                    TestUtil.createKeyset(tinkKey, rawKey, legacyKey), annotations, Mac.class));
+    KeysetHandle rawKeysetHandle =
+        KeysetHandle.newBuilder()
+            .addEntry(KeysetHandle.importKey(rawKey0).withFixedId(43).makePrimary())
+            .setMonitoringAnnotations(annotations)
+            .build();
+    KeysetHandle legacyKeysetHandle =
+        KeysetHandle.newBuilder()
+            .addEntry(KeysetHandle.importKey(legacyKey0).makePrimary())
+            .setMonitoringAnnotations(annotations)
+            .build();
+    KeysetHandle mixedKeysetHandle =
+        KeysetHandle.newBuilder()
+            .addEntry(KeysetHandle.importKey(tinkKey1).makePrimary())
+            .addEntry(KeysetHandle.importKey(rawKey0).withFixedId(43))
+            .addEntry(KeysetHandle.importKey(legacyKey0))
+            .setMonitoringAnnotations(annotations)
+            .build();
+    Mac rawMac = rawKeysetHandle.getPrimitive(Mac.class);
+    Mac legacyMac = legacyKeysetHandle.getPrimitive(Mac.class);
+    Mac mac = mixedKeysetHandle.getPrimitive(Mac.class);
+
     byte[] plaintext = "plaintext".getBytes(UTF_8);
     byte[] tinkTag = mac.computeMac(plaintext);
     byte[] rawTag = rawMac.computeMac(plaintext);
@@ -440,7 +449,8 @@ public class MacWrapperTest {
     assertThat(logEntries).hasSize(6);
 
     FakeMonitoringClient.LogEntry tinkComputeEntry = logEntries.get(0);
-    assertThat(tinkComputeEntry.getKeyId()).isEqualTo(42);
+    // 5 is tinkKey1's id.
+    assertThat(tinkComputeEntry.getKeyId()).isEqualTo(5);
     assertThat(tinkComputeEntry.getPrimitive()).isEqualTo("mac");
     assertThat(tinkComputeEntry.getApi()).isEqualTo("compute");
     assertThat(tinkComputeEntry.getNumBytesAsInput()).isEqualTo(plaintext.length);
@@ -454,15 +464,16 @@ public class MacWrapperTest {
     assertThat(rawComputeEntry.getKeysetInfo().getAnnotations()).isEqualTo(annotations);
 
     FakeMonitoringClient.LogEntry legacyComputeEntry = logEntries.get(2);
-    assertThat(legacyComputeEntry.getKeyId()).isEqualTo(44);
+    // 8 is legacyKey0's id.
+    assertThat(legacyComputeEntry.getKeyId()).isEqualTo(8);
     assertThat(legacyComputeEntry.getPrimitive()).isEqualTo("mac");
     assertThat(legacyComputeEntry.getApi()).isEqualTo("compute");
-    // legacy mac appends one byte to the input data, therefore the input length is one longer.
-    assertThat(legacyComputeEntry.getNumBytesAsInput()).isEqualTo(plaintext.length + 1);
+    assertThat(legacyComputeEntry.getNumBytesAsInput()).isEqualTo(plaintext.length);
     assertThat(legacyComputeEntry.getKeysetInfo().getAnnotations()).isEqualTo(annotations);
 
     FakeMonitoringClient.LogEntry tinkVerifyEntry = logEntries.get(3);
-    assertThat(tinkVerifyEntry.getKeyId()).isEqualTo(42);
+    // 5 is tinkKey1's id.
+    assertThat(tinkVerifyEntry.getKeyId()).isEqualTo(5);
     assertThat(tinkVerifyEntry.getPrimitive()).isEqualTo("mac");
     assertThat(tinkVerifyEntry.getApi()).isEqualTo("verify");
     assertThat(tinkVerifyEntry.getNumBytesAsInput()).isEqualTo(plaintext.length);
@@ -476,11 +487,11 @@ public class MacWrapperTest {
     assertThat(rawVerifyEntry.getKeysetInfo().getAnnotations()).isEqualTo(annotations);
 
     FakeMonitoringClient.LogEntry legacyVerifyEntry = logEntries.get(5);
-    assertThat(legacyVerifyEntry.getKeyId()).isEqualTo(44);
+    // 8 is legacyKey0's id.
+    assertThat(legacyVerifyEntry.getKeyId()).isEqualTo(8);
     assertThat(legacyVerifyEntry.getPrimitive()).isEqualTo("mac");
     assertThat(legacyVerifyEntry.getApi()).isEqualTo("verify");
-    // legacy mac appends one byte to the input data, therefore the input length is one longer.
-    assertThat(legacyVerifyEntry.getNumBytesAsInput()).isEqualTo(plaintext.length + 1);
+    assertThat(legacyVerifyEntry.getNumBytesAsInput()).isEqualTo(plaintext.length);
     assertThat(legacyVerifyEntry.getKeysetInfo().getAnnotations()).isEqualTo(annotations);
 
     List<FakeMonitoringClient.LogFailureEntry> failures =
@@ -489,11 +500,15 @@ public class MacWrapperTest {
     FakeMonitoringClient.LogFailureEntry verifyFailure = failures.get(0);
     assertThat(verifyFailure.getPrimitive()).isEqualTo("mac");
     assertThat(verifyFailure.getApi()).isEqualTo("verify");
-    assertThat(verifyFailure.getKeysetInfo().getPrimaryKeyId()).isEqualTo(42);
+    // 5 is tinkKey1's id.
+    assertThat(verifyFailure.getKeysetInfo().getPrimaryKeyId()).isEqualTo(5);
     assertThat(verifyFailure.getKeysetInfo().getAnnotations()).isEqualTo(annotations);
   }
 
   public static class AlwaysFailingMac implements Mac {
+
+    AlwaysFailingMac(HmacKey key) {}
+
     @Override
     public byte[] computeMac(final byte[] data) throws GeneralSecurityException {
       throw new GeneralSecurityException("fail");
@@ -507,35 +522,39 @@ public class MacWrapperTest {
 
   @Test
   public void testAlwaysFailingWithAnnotation_hasMonitoring() throws Exception {
+    // Test setup.
+    MutablePrimitiveRegistry.resetGlobalInstanceTestOnly();
+    MutablePrimitiveRegistry.globalInstance()
+        .registerPrimitiveConstructor(
+            PrimitiveConstructor.create(AlwaysFailingMac::new, HmacKey.class, Mac.class));
+    MacWrapper.register();
+    HmacProtoSerialization.register();
+    Registry.registerKeyManager(new HmacKeyManager(), true);
+
     FakeMonitoringClient fakeMonitoringClient = new FakeMonitoringClient();
     MutableMonitoringRegistry.globalInstance().clear();
     MutableMonitoringRegistry.globalInstance().registerMonitoringClient(fakeMonitoringClient);
 
     MonitoringAnnotations annotations =
         MonitoringAnnotations.newBuilder().add("annotation_name", "annotation_value").build();
-    PrimitiveSet<Mac> primitives =
-        PrimitiveSet.newBuilder(Mac.class)
-            .setAnnotations(annotations)
-            .addPrimaryPrimitive(
-                new AlwaysFailingMac(),
-                TestUtil.createKey(
-                    TestUtil.createHmacKeyData(Random.randBytes(HMAC_KEY_SIZE), 16),
-                    42,
-                    KeyStatusType.ENABLED,
-                    OutputPrefixType.TINK))
+    KeysetHandle keysetHandle =
+        KeysetHandle.newBuilder()
+            .addEntry(KeysetHandle.importKey(tinkKey0).makePrimary())
+            .setMonitoringAnnotations(annotations)
             .build();
-    Mac mac = new MacWrapper().wrap(primitives);
+    Mac mac = keysetHandle.getPrimitive(Mac.class);
 
     byte[] data = "some data".getBytes(UTF_8);
     byte[] invalidTag = "an invalid tag".getBytes(UTF_8);
+    byte[] shortInvalidTag = "t".getBytes(UTF_8);
 
+    // Test active work, including a test with a short tag, because there is a different code path
+    // for this.
     assertThrows(GeneralSecurityException.class, () -> mac.computeMac(data));
     assertThrows(GeneralSecurityException.class, () -> mac.verifyMac(invalidTag, data));
-
-    // Test short tags, because there is a different code path for this.
-    byte[] shortInvalidTag = "t".getBytes(UTF_8);
     assertThrows(GeneralSecurityException.class, () -> mac.verifyMac(shortInvalidTag, data));
 
+    // Assert correctness.
     assertThat(fakeMonitoringClient.getLogEntries()).isEmpty();
 
     List<FakeMonitoringClient.LogFailureEntry> failures =
@@ -544,19 +563,22 @@ public class MacWrapperTest {
     FakeMonitoringClient.LogFailureEntry compFailure = failures.get(0);
     assertThat(compFailure.getPrimitive()).isEqualTo("mac");
     assertThat(compFailure.getApi()).isEqualTo("compute");
-    assertThat(compFailure.getKeysetInfo().getPrimaryKeyId()).isEqualTo(42);
+    // 4 is tinkKey0's id.
+    assertThat(compFailure.getKeysetInfo().getPrimaryKeyId()).isEqualTo(4);
     assertThat(compFailure.getKeysetInfo().getAnnotations()).isEqualTo(annotations);
 
     FakeMonitoringClient.LogFailureEntry verifyFailure = failures.get(1);
     assertThat(verifyFailure.getPrimitive()).isEqualTo("mac");
     assertThat(verifyFailure.getApi()).isEqualTo("verify");
-    assertThat(verifyFailure.getKeysetInfo().getPrimaryKeyId()).isEqualTo(42);
+    // 4 is tinkKey0's id.
+    assertThat(verifyFailure.getKeysetInfo().getPrimaryKeyId()).isEqualTo(4);
     assertThat(verifyFailure.getKeysetInfo().getAnnotations()).isEqualTo(annotations);
 
     FakeMonitoringClient.LogFailureEntry verifyFailure2 = failures.get(2);
     assertThat(verifyFailure2.getPrimitive()).isEqualTo("mac");
     assertThat(verifyFailure2.getApi()).isEqualTo("verify");
-    assertThat(verifyFailure2.getKeysetInfo().getPrimaryKeyId()).isEqualTo(42);
+    // 4 is tinkKey0's id.
+    assertThat(verifyFailure2.getKeysetInfo().getPrimaryKeyId()).isEqualTo(4);
     assertThat(verifyFailure2.getKeysetInfo().getAnnotations()).isEqualTo(annotations);
   }
 }