Switch MacWrapper to using full primitives. That includes reworking the internal structure and logging the input data length instead of the processed data length in the LEGACY case.

PiperOrigin-RevId: 555189951

6 files changed, 110 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..f809860e0 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,7 +109,6 @@ 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",
@@ -119,7 +118,6 @@ 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",
     ],
 )
@@ -129,7 +127,6 @@ 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",
@@ -139,7 +136,6 @@ 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 f0ab5bb17..010317e81 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,7 +25,6 @@ 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;
@@ -44,7 +43,6 @@ 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 {
@@ -67,16 +65,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 +84,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 +101,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 +155,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 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);
   }
 }