Add a create function to KmsEnvelopeAead.

This function uses the new Parameters instead of the old proto KeyTemplates.

PiperOrigin-RevId: 550614620

4 files changed, 125 insertions, 18 deletions

diff --git a/java_src/src/main/java/com/google/crypto/tink/aead/BUILD.bazel b/java_src/src/main/java/com/google/crypto/tink/aead/BUILD.bazel
index e59dbe6e8..f42a44759 100644
--- a/java_src/src/main/java/com/google/crypto/tink/aead/BUILD.bazel
+++ b/java_src/src/main/java/com/google/crypto/tink/aead/BUILD.bazel
@@ -124,9 +124,12 @@ java_library(
     name = "kms_envelope_aead",
     srcs = ["KmsEnvelopeAead.java"],
     deps = [
+        ":aead_parameters",
         "//proto:tink_java_proto",
         "//src/main/java/com/google/crypto/tink:aead",
         "//src/main/java/com/google/crypto/tink:registry",
+        "//src/main/java/com/google/crypto/tink:tink_proto_parameters_format",
+        "@maven//:com_google_protobuf_protobuf_java",
     ],
 )
 
@@ -554,9 +557,12 @@ android_library(
     name = "kms_envelope_aead-android",
     srcs = ["KmsEnvelopeAead.java"],
     deps = [
+        ":aead_parameters-android",
         "//proto:tink_java_proto_lite",
         "//src/main/java/com/google/crypto/tink:aead-android",
         "//src/main/java/com/google/crypto/tink:registry-android",
+        "//src/main/java/com/google/crypto/tink:tink_proto_parameters_format-android",
+        "@maven//:com_google_protobuf_protobuf_javalite",
     ],
 )
 
diff --git a/java_src/src/main/java/com/google/crypto/tink/aead/KmsEnvelopeAead.java b/java_src/src/main/java/com/google/crypto/tink/aead/KmsEnvelopeAead.java
index de641f46b..1532add3e 100644
--- a/java_src/src/main/java/com/google/crypto/tink/aead/KmsEnvelopeAead.java
+++ b/java_src/src/main/java/com/google/crypto/tink/aead/KmsEnvelopeAead.java
@@ -18,7 +18,10 @@ package com.google.crypto.tink.aead; // instead of subtle, because it depends on
 
 import com.google.crypto.tink.Aead;
 import com.google.crypto.tink.Registry;
+import com.google.crypto.tink.TinkProtoParametersFormat;
 import com.google.crypto.tink.proto.KeyTemplate;
+import com.google.protobuf.ExtensionRegistryLite;
+import com.google.protobuf.InvalidProtocolBufferException;
 import java.nio.BufferUnderflowException;
 import java.nio.ByteBuffer;
 import java.security.GeneralSecurityException;
@@ -78,6 +81,28 @@ public final class KmsEnvelopeAead implements Aead {
     this.remote = remote;
   }
 
+  /**
+   * Creates a new instance of Tink's KMS Envelope AEAD.
+   *
+   * <p>{@code dekParameters} must be any of these Tink AEAD parameters (any other will be
+   * rejected): {@link AesGcmParameters}, {@link ChaCha20Poly1305Parameters}, {@link
+   * XChaCha20Poly1305Parameters}, {@link AesCtrHmacAeadParameters}, {@link AesGcmSivParameters}, or
+   * {@link AesEaxParameters}.
+   */
+  public static Aead create(AeadParameters dekParameters, Aead remote)
+      throws GeneralSecurityException {
+    KeyTemplate dekTemplate;
+    try {
+      dekTemplate =
+          KeyTemplate.parseFrom(
+              TinkProtoParametersFormat.serialize(dekParameters),
+              ExtensionRegistryLite.getEmptyRegistry());
+    } catch (InvalidProtocolBufferException e) {
+      throw new GeneralSecurityException(e);
+    }
+    return new KmsEnvelopeAead(dekTemplate, remote);
+  }
+
   @Override
   public byte[] encrypt(final byte[] plaintext, final byte[] associatedData)
       throws GeneralSecurityException {
diff --git a/java_src/src/test/java/com/google/crypto/tink/aead/BUILD.bazel b/java_src/src/test/java/com/google/crypto/tink/aead/BUILD.bazel
index 2ac2a9af0..ef643dbe8 100644
--- a/java_src/src/test/java/com/google/crypto/tink/aead/BUILD.bazel
+++ b/java_src/src/test/java/com/google/crypto/tink/aead/BUILD.bazel
@@ -54,11 +54,19 @@ java_test(
         "//src/main/java/com/google/crypto/tink:aead",
         "//src/main/java/com/google/crypto/tink:key_template",
         "//src/main/java/com/google/crypto/tink:key_templates",
+        "//src/main/java/com/google/crypto/tink:kms_client",
+        "//src/main/java/com/google/crypto/tink:kms_clients",
         "//src/main/java/com/google/crypto/tink:registry_cluster",
         "//src/main/java/com/google/crypto/tink/aead:aead_config",
+        "//src/main/java/com/google/crypto/tink/aead:aead_parameters",
+        "//src/main/java/com/google/crypto/tink/aead:aes_ctr_hmac_aead_key_manager",
         "//src/main/java/com/google/crypto/tink/aead:kms_envelope_aead",
+        "//src/main/java/com/google/crypto/tink/aead:kms_envelope_aead_key_manager",
+        "//src/main/java/com/google/crypto/tink/aead:predefined_aead_parameters",
         "//src/main/java/com/google/crypto/tink/internal:key_template_proto_converter",
         "//src/main/java/com/google/crypto/tink/mac:hmac_key_manager",
+        "//src/main/java/com/google/crypto/tink/subtle:random",
+        "//src/main/java/com/google/crypto/tink/testing:fake_kms_client",
         "@maven//:com_google_truth_truth",
         "@maven//:junit_junit",
     ],
diff --git a/java_src/src/test/java/com/google/crypto/tink/aead/KmsEnvelopeAeadTest.java b/java_src/src/test/java/com/google/crypto/tink/aead/KmsEnvelopeAeadTest.java
index 96634ebeb..51460acc4 100644
--- a/java_src/src/test/java/com/google/crypto/tink/aead/KmsEnvelopeAeadTest.java
+++ b/java_src/src/test/java/com/google/crypto/tink/aead/KmsEnvelopeAeadTest.java
@@ -24,8 +24,12 @@ import com.google.crypto.tink.Aead;
 import com.google.crypto.tink.KeyTemplate;
 import com.google.crypto.tink.KeyTemplates;
 import com.google.crypto.tink.KeysetHandle;
+import com.google.crypto.tink.KmsClient;
+import com.google.crypto.tink.KmsClients;
 import com.google.crypto.tink.internal.KeyTemplateProtoConverter;
 import com.google.crypto.tink.mac.HmacKeyManager;
+import com.google.crypto.tink.subtle.Random;
+import com.google.crypto.tink.testing.FakeKmsClient;
 import java.security.GeneralSecurityException;
 import org.junit.BeforeClass;
 import org.junit.Test;
@@ -50,6 +54,33 @@ public final class KmsEnvelopeAeadTest {
     return keysetHandle.getPrimitive(Aead.class);
   }
 
+  @DataPoints("dekParameters")
+  public static final AeadParameters[] DEK_PARAMETERS =
+      new AeadParameters[] {
+        PredefinedAeadParameters.AES128_GCM,
+        PredefinedAeadParameters.AES256_GCM,
+        PredefinedAeadParameters.AES128_EAX,
+        PredefinedAeadParameters.AES256_EAX,
+        PredefinedAeadParameters.AES128_CTR_HMAC_SHA256,
+        PredefinedAeadParameters.AES256_CTR_HMAC_SHA256,
+        PredefinedAeadParameters.CHACHA20_POLY1305,
+        PredefinedAeadParameters.XCHACHA20_POLY1305,
+      };
+
+  @Theory
+  public void createEncryptDecrypt_works(
+      @FromDataPoints("dekParameters") AeadParameters dekParameters) throws Exception {
+    Aead remoteAead = this.generateNewRemoteAead();
+    Aead envAead = KmsEnvelopeAead.create(dekParameters, remoteAead);
+    byte[] plaintext = "plaintext".getBytes(UTF_8);
+    byte[] associatedData = "associatedData".getBytes(UTF_8);
+    byte[] ciphertext = envAead.encrypt(plaintext, associatedData);
+    assertThat(envAead.decrypt(ciphertext, associatedData)).isEqualTo(plaintext);
+
+    assertThat(envAead.decrypt(envAead.encrypt(plaintext, EMPTY_ADD), EMPTY_ADD))
+        .isEqualTo(plaintext);
+  }
+
   @DataPoints("tinkDekTemplates")
   public static final String[] TINK_DEK_TEMPLATES =
       new String[] {
@@ -65,10 +96,10 @@ public final class KmsEnvelopeAeadTest {
       };
 
   @Theory
-  public void encryptDecrypt_works(@FromDataPoints("tinkDekTemplates") String dekTemplateName)
-      throws Exception {
+  public void legacyConstructorEncryptDecrypt_works(
+      @FromDataPoints("tinkDekTemplates") String dekTemplateName) throws Exception {
     Aead remoteAead = this.generateNewRemoteAead();
-    KmsEnvelopeAead envAead =
+    Aead envAead =
         new KmsEnvelopeAead(
             KeyTemplateProtoConverter.toProto(KeyTemplates.get(dekTemplateName)), remoteAead);
     byte[] plaintext = "plaintext".getBytes(UTF_8);
@@ -94,9 +125,7 @@ public final class KmsEnvelopeAeadTest {
   @Test
   public void decryptWithInvalidAssociatedData_fails() throws GeneralSecurityException {
     Aead remoteAead =  this.generateNewRemoteAead();
-    KmsEnvelopeAead envAead =
-        new KmsEnvelopeAead(
-            KeyTemplateProtoConverter.toProto(KeyTemplates.get("AES128_EAX")), remoteAead);
+    Aead envAead = KmsEnvelopeAead.create(PredefinedAeadParameters.AES128_EAX, remoteAead);
     byte[] plaintext = "plaintext".getBytes(UTF_8);
     byte[] associatedData = "associatedData".getBytes(UTF_8);
     byte[] ciphertext = envAead.encrypt(plaintext, associatedData);
@@ -109,9 +138,7 @@ public final class KmsEnvelopeAeadTest {
   @Test
   public void corruptedCiphertext_fails() throws GeneralSecurityException {
     Aead remoteAead =  this.generateNewRemoteAead();
-    KmsEnvelopeAead envAead =
-        new KmsEnvelopeAead(
-            KeyTemplateProtoConverter.toProto(KeyTemplates.get("AES128_EAX")), remoteAead);
+    Aead envAead = KmsEnvelopeAead.create(PredefinedAeadParameters.AES128_EAX, remoteAead);
     byte[] associatedData = "envelope_ad".getBytes(UTF_8);
     byte[] plaintext = "helloworld".getBytes(UTF_8);
     byte[] ciphertext = envAead.encrypt(plaintext, associatedData);
@@ -124,9 +151,7 @@ public final class KmsEnvelopeAeadTest {
   @Test
   public void corruptedDek_fails() throws GeneralSecurityException {
     Aead remoteAead =  this.generateNewRemoteAead();
-    KmsEnvelopeAead envAead =
-        new KmsEnvelopeAead(
-            KeyTemplateProtoConverter.toProto(KeyTemplates.get("AES128_EAX")), remoteAead);
+    Aead envAead = KmsEnvelopeAead.create(PredefinedAeadParameters.AES128_EAX, remoteAead);
     byte[] plaintext = "helloworld".getBytes(UTF_8);
     byte[] associatedData = "envelope_ad".getBytes(UTF_8);
     byte[] ciphertext = envAead.encrypt(plaintext, associatedData);
@@ -139,9 +164,7 @@ public final class KmsEnvelopeAeadTest {
   @Test
   public void ciphertextTooShort_fails() throws GeneralSecurityException {
     Aead remoteAead =  this.generateNewRemoteAead();
-    KmsEnvelopeAead envAead =
-        new KmsEnvelopeAead(
-            KeyTemplateProtoConverter.toProto(KeyTemplates.get("AES128_EAX")), remoteAead);
+    Aead envAead = KmsEnvelopeAead.create(PredefinedAeadParameters.AES128_EAX, remoteAead);
     assertThrows(
         GeneralSecurityException.class,
         () -> envAead.decrypt("foo".getBytes(UTF_8), "envelope_ad".getBytes(UTF_8)));
@@ -150,9 +173,7 @@ public final class KmsEnvelopeAeadTest {
   @Test
   public void malformedDekLength_fails() throws GeneralSecurityException {
     Aead remoteAead =  this.generateNewRemoteAead();
-    KmsEnvelopeAead envAead =
-        new KmsEnvelopeAead(
-            KeyTemplateProtoConverter.toProto(KeyTemplates.get("AES128_EAX")), remoteAead);
+    Aead envAead = KmsEnvelopeAead.create(PredefinedAeadParameters.AES128_EAX, remoteAead);
 
     byte[] plaintext = "helloworld".getBytes(UTF_8);
     byte[] associatedData = "envelope_ad".getBytes(UTF_8);
@@ -174,6 +195,53 @@ public final class KmsEnvelopeAeadTest {
         GeneralSecurityException.class,
         () -> envAead.decrypt(corruptedCiphertext2, associatedData));
   }
+
+  @Test
+  public void create_isCompatibleWithOldConstructor() throws Exception {
+    String kekUri = FakeKmsClient.createFakeKeyUri();
+    Aead remoteAead = new FakeKmsClient().getAead(kekUri);
+
+    Aead aead1 =
+        new KmsEnvelopeAead(
+            KeyTemplateProtoConverter.toProto(
+                AesCtrHmacAeadKeyManager.aes128CtrHmacSha256Template()),
+            remoteAead);
+    Aead aead2 =
+        KmsEnvelopeAead.create(PredefinedAeadParameters.AES128_CTR_HMAC_SHA256, remoteAead);
+
+    byte[] plaintext = Random.randBytes(20);
+    byte[] associatedData = Random.randBytes(20);
+    assertThat(aead1.decrypt(aead2.encrypt(plaintext, associatedData), associatedData))
+        .isEqualTo(plaintext);
+    assertThat(aead2.decrypt(aead1.encrypt(plaintext, associatedData), associatedData))
+        .isEqualTo(plaintext);
+  }
+
+  @Test
+  public void create_isCompatibleWithKmsEnvelopeAeadKey() throws Exception {
+    String kekUri = FakeKmsClient.createFakeKeyUri();
+    KeyTemplate dekTemplate = AesCtrHmacAeadKeyManager.aes128CtrHmacSha256Template();
+
+    // Register kmsClient and create a keyset with a KmsEnvelopeAeadKey key.
+    KmsClient kmsClient1 = new FakeKmsClient(kekUri);
+    KmsClients.add(kmsClient1);
+    KeysetHandle handle1 =
+        KeysetHandle.generateNew(KmsEnvelopeAeadKeyManager.createKeyTemplate(kekUri, dekTemplate));
+    Aead aead1 = handle1.getPrimitive(Aead.class);
+
+    // Get Aead object from the kmsClient, and create the envelope AEAD without the registry.
+    Aead remoteAead = new FakeKmsClient().getAead(kekUri);
+    Aead aead2 =
+        KmsEnvelopeAead.create(PredefinedAeadParameters.AES128_CTR_HMAC_SHA256, remoteAead);
+
+    // Check that aead1 and aead2 implement the same primitive
+    byte[] plaintext = Random.randBytes(20);
+    byte[] associatedData = Random.randBytes(20);
+    assertThat(aead1.decrypt(aead2.encrypt(plaintext, associatedData), associatedData))
+        .isEqualTo(plaintext);
+    assertThat(aead2.decrypt(aead1.encrypt(plaintext, associatedData), associatedData))
+        .isEqualTo(plaintext);
+  }
 }