Undoing the MacWrapper change.

PiperOrigin-RevId: 555434824

6 files changed, 95 insertions, 110 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 f809860e0..94ef3db7f 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,6 +109,7 @@ 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",
@@ -118,6 +119,7 @@ java_library(
         "//src/main/java/com/google/crypto/tink/internal:mutable_monitoring_registry",
         "//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",
     ],
 )
@@ -127,6 +129,7 @@ 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",
@@ -136,6 +139,7 @@ android_library(
         "//src/main/java/com/google/crypto/tink/internal:mutable_monitoring_registry-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 010317e81..f0ab5bb17 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
@@ -25,6 +25,7 @@ import com.google.crypto.tink.internal.MonitoringUtil;
 import com.google.crypto.tink.internal.MutableMonitoringRegistry;
 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;
@@ -43,6 +44,7 @@ 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 class WrappedMac implements Mac {
@@ -65,9 +67,16 @@ 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 = primitives.getPrimary().getFullPrimitive().computeMac(data);
-        computeLogger.log(primitives.getPrimary().getKeyId(), data.length);
+        byte[] output =
+            com.google.crypto.tink.subtle.Bytes.concat(
+                primitives.getPrimary().getIdentifier(),
+                primitives.getPrimary().getPrimitive().computeMac(data2));
+        computeLogger.log(primitives.getPrimary().getKeyId(), data2.length);
         return output;
       } catch (GeneralSecurityException e) {
         computeLogger.logFailure();
@@ -84,11 +93,16 @@ 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.getFullPrimitive().verifyMac(mac, data);
-          verifyLogger.log(entry.getKeyId(), data.length);
+          entry.getPrimitive().verifyMac(macNoPrefix, data2);
+          verifyLogger.log(entry.getKeyId(), data2.length);
           // If there is no exception, the MAC is valid and we can return.
           return;
         } catch (GeneralSecurityException e) {
@@ -101,7 +115,7 @@ class MacWrapper implements PrimitiveWrapper<Mac, Mac> {
       entries = primitives.getRawPrimitives();
       for (PrimitiveSet.Entry<Mac> entry : entries) {
         try {
-          entry.getFullPrimitive().verifyMac(mac, data);
+          entry.getPrimitive().verifyMac(mac, data);
           verifyLogger.log(entry.getKeyId(), data.length);
           // If there is no exception, the MAC is valid and we can return.
           return;
@@ -155,7 +169,7 @@ class MacWrapper implements PrimitiveWrapper<Mac, Mac> {
     return Mac.class;
   }
 
-  public static void register() throws GeneralSecurityException {
+ public static void register() throws GeneralSecurityException {
     Registry.registerPrimitiveWrapper(WRAPPER);
   }
 }
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 b441ac2fd..5f6b30d6c 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,17 +476,13 @@ 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 4b24e1aeb..5a7ef7596 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,7 +35,6 @@ 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;
@@ -60,7 +59,6 @@ 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 {
@@ -112,11 +110,6 @@ 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);
     }
@@ -155,13 +148,13 @@ public class RegistryConfigurationTest {
     byte[] plaintext = "plaintext".getBytes(UTF_8);
 
     Mac registryMac = Registry.getPrimitive(rawKeyData, Mac.class);
-    // The following relies on the fact that internally LegacyFullMac uses RegistryConfiguration.
+    Mac configurationMac = RegistryConfiguration.get().getLegacyPrimitive(rawKeyData, Mac.class);
     Mac wrappedConfigurationMac =
         RegistryConfiguration.get()
             .wrap(
                 PrimitiveSet.newBuilder(Mac.class)
                     .addPrimaryFullPrimitiveAndOptionalPrimitive(
-                        LegacyFullMac.create(legacyProtoRawKey), null, rawKeysetKey)
+                        null, configurationMac, 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 fe7d2013c..1a46ce179 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:registry",
+        "//src/main/java/com/google/crypto/tink:primitive_set",
         "//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 443f3870f..a60dce6cb 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,15 +23,18 @@ 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.Registry;
+import com.google.crypto.tink.PrimitiveSet;
 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;
@@ -194,9 +197,6 @@ 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,9 +216,6 @@ public class MacWrapperTest {
 
   @Test
   public void testVerifyMac_checksAllNecessaryRawKeys() throws Exception {
-    MutablePrimitiveRegistry.resetGlobalInstanceTestOnly();
-    MacConfig.register();
-
     byte[] plaintext = "plaintext".getBytes(UTF_8);
     KeysetHandle computeKeysetHandle =
         KeysetHandle.newBuilder()
@@ -239,9 +236,6 @@ 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 =
@@ -273,9 +267,6 @@ public class MacWrapperTest {
 
   @Test
   public void computeMac_usesPrimaryKey() throws Exception {
-    MutablePrimitiveRegistry.resetGlobalInstanceTestOnly();
-    MacConfig.register();
-
     byte[] plaintext = "plaintext".getBytes(UTF_8);
     KeysetHandle keysetHandle =
         KeysetHandle.newBuilder()
@@ -297,9 +288,6 @@ public class MacWrapperTest {
 
   @Test
   public void testComputeVerifyMac_manyKeysWork() throws Exception {
-    MutablePrimitiveRegistry.resetGlobalInstanceTestOnly();
-    MacConfig.register();
-
     byte[] plaintext = "plaintext".getBytes(UTF_8);
     KeysetHandle assortedKeysetHandle =
         KeysetHandle.newBuilder()
@@ -323,9 +311,6 @@ public class MacWrapperTest {
 
   @Test
   public void testVerifyMac_shiftedPrimaryWithManyKeysWorks() throws Exception {
-    MutablePrimitiveRegistry.resetGlobalInstanceTestOnly();
-    MacConfig.register();
-
     byte[] plaintext = "plaintext".getBytes(UTF_8);
     KeysetHandle assortedKeysetHandle0 =
         KeysetHandle.newBuilder()
@@ -354,9 +339,6 @@ public class MacWrapperTest {
 
   @Test
   public void testMultipleKeysWithoutAnnotation() throws Exception {
-    MutablePrimitiveRegistry.resetGlobalInstanceTestOnly();
-    MacConfig.register();
-
     FakeMonitoringClient fakeMonitoringClient = new FakeMonitoringClient();
     MutableMonitoringRegistry.globalInstance().clear();
     MutableMonitoringRegistry.globalInstance().registerMonitoringClient(fakeMonitoringClient);
@@ -406,36 +388,45 @@ 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();
-    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);
-
+    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));
     byte[] plaintext = "plaintext".getBytes(UTF_8);
     byte[] tinkTag = mac.computeMac(plaintext);
     byte[] rawTag = rawMac.computeMac(plaintext);
@@ -449,8 +440,7 @@ public class MacWrapperTest {
     assertThat(logEntries).hasSize(6);
 
     FakeMonitoringClient.LogEntry tinkComputeEntry = logEntries.get(0);
-    // 5 is tinkKey1's id.
-    assertThat(tinkComputeEntry.getKeyId()).isEqualTo(5);
+    assertThat(tinkComputeEntry.getKeyId()).isEqualTo(42);
     assertThat(tinkComputeEntry.getPrimitive()).isEqualTo("mac");
     assertThat(tinkComputeEntry.getApi()).isEqualTo("compute");
     assertThat(tinkComputeEntry.getNumBytesAsInput()).isEqualTo(plaintext.length);
@@ -464,16 +454,15 @@ public class MacWrapperTest {
     assertThat(rawComputeEntry.getKeysetInfo().getAnnotations()).isEqualTo(annotations);
 
     FakeMonitoringClient.LogEntry legacyComputeEntry = logEntries.get(2);
-    // 8 is legacyKey0's id.
-    assertThat(legacyComputeEntry.getKeyId()).isEqualTo(8);
+    assertThat(legacyComputeEntry.getKeyId()).isEqualTo(44);
     assertThat(legacyComputeEntry.getPrimitive()).isEqualTo("mac");
     assertThat(legacyComputeEntry.getApi()).isEqualTo("compute");
-    assertThat(legacyComputeEntry.getNumBytesAsInput()).isEqualTo(plaintext.length);
+    // 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.getKeysetInfo().getAnnotations()).isEqualTo(annotations);
 
     FakeMonitoringClient.LogEntry tinkVerifyEntry = logEntries.get(3);
-    // 5 is tinkKey1's id.
-    assertThat(tinkVerifyEntry.getKeyId()).isEqualTo(5);
+    assertThat(tinkVerifyEntry.getKeyId()).isEqualTo(42);
     assertThat(tinkVerifyEntry.getPrimitive()).isEqualTo("mac");
     assertThat(tinkVerifyEntry.getApi()).isEqualTo("verify");
     assertThat(tinkVerifyEntry.getNumBytesAsInput()).isEqualTo(plaintext.length);
@@ -487,11 +476,11 @@ public class MacWrapperTest {
     assertThat(rawVerifyEntry.getKeysetInfo().getAnnotations()).isEqualTo(annotations);
 
     FakeMonitoringClient.LogEntry legacyVerifyEntry = logEntries.get(5);
-    // 8 is legacyKey0's id.
-    assertThat(legacyVerifyEntry.getKeyId()).isEqualTo(8);
+    assertThat(legacyVerifyEntry.getKeyId()).isEqualTo(44);
     assertThat(legacyVerifyEntry.getPrimitive()).isEqualTo("mac");
     assertThat(legacyVerifyEntry.getApi()).isEqualTo("verify");
-    assertThat(legacyVerifyEntry.getNumBytesAsInput()).isEqualTo(plaintext.length);
+    // 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.getKeysetInfo().getAnnotations()).isEqualTo(annotations);
 
     List<FakeMonitoringClient.LogFailureEntry> failures =
@@ -500,15 +489,11 @@ public class MacWrapperTest {
     FakeMonitoringClient.LogFailureEntry verifyFailure = failures.get(0);
     assertThat(verifyFailure.getPrimitive()).isEqualTo("mac");
     assertThat(verifyFailure.getApi()).isEqualTo("verify");
-    // 5 is tinkKey1's id.
-    assertThat(verifyFailure.getKeysetInfo().getPrimaryKeyId()).isEqualTo(5);
+    assertThat(verifyFailure.getKeysetInfo().getPrimaryKeyId()).isEqualTo(42);
     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");
@@ -522,39 +507,35 @@ 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();
-    KeysetHandle keysetHandle =
-        KeysetHandle.newBuilder()
-            .addEntry(KeysetHandle.importKey(tinkKey0).makePrimary())
-            .setMonitoringAnnotations(annotations)
+    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))
             .build();
-    Mac mac = keysetHandle.getPrimitive(Mac.class);
+    Mac mac = new MacWrapper().wrap(primitives);
 
     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 =
@@ -563,22 +544,19 @@ public class MacWrapperTest {
     FakeMonitoringClient.LogFailureEntry compFailure = failures.get(0);
     assertThat(compFailure.getPrimitive()).isEqualTo("mac");
     assertThat(compFailure.getApi()).isEqualTo("compute");
-    // 4 is tinkKey0's id.
-    assertThat(compFailure.getKeysetInfo().getPrimaryKeyId()).isEqualTo(4);
+    assertThat(compFailure.getKeysetInfo().getPrimaryKeyId()).isEqualTo(42);
     assertThat(compFailure.getKeysetInfo().getAnnotations()).isEqualTo(annotations);
 
     FakeMonitoringClient.LogFailureEntry verifyFailure = failures.get(1);
     assertThat(verifyFailure.getPrimitive()).isEqualTo("mac");
     assertThat(verifyFailure.getApi()).isEqualTo("verify");
-    // 4 is tinkKey0's id.
-    assertThat(verifyFailure.getKeysetInfo().getPrimaryKeyId()).isEqualTo(4);
+    assertThat(verifyFailure.getKeysetInfo().getPrimaryKeyId()).isEqualTo(42);
     assertThat(verifyFailure.getKeysetInfo().getAnnotations()).isEqualTo(annotations);
 
     FakeMonitoringClient.LogFailureEntry verifyFailure2 = failures.get(2);
     assertThat(verifyFailure2.getPrimitive()).isEqualTo("mac");
     assertThat(verifyFailure2.getApi()).isEqualTo("verify");
-    // 4 is tinkKey0's id.
-    assertThat(verifyFailure2.getKeysetInfo().getPrimaryKeyId()).isEqualTo(4);
+    assertThat(verifyFailure2.getKeysetInfo().getPrimaryKeyId()).isEqualTo(42);
     assertThat(verifyFailure2.getKeysetInfo().getAnnotations()).isEqualTo(annotations);
   }
 }