Restrict KMS envelope AEAD to only use Tink AEAD key types as DEK.

KMS Envelope AEAD is a simple way to encrypt data. So using custom
implementations of AEAD is not needed for this.

This also prevents the user from using the envelope encryption key type as DEK by accident.

This change will help us simplify our implementation of KMS Envelope AEAD,
because we will be able to remove the dependency on Tink's Registry.

PiperOrigin-RevId: 542549805

2 files changed, 58 insertions, 11 deletions

diff --git a/python/tink/aead/_aead_key_manager_test.py b/python/tink/aead/_aead_key_manager_test.py
index 63db927aa..25332fc13 100644
--- a/python/tink/aead/_aead_key_manager_test.py
+++ b/python/tink/aead/_aead_key_manager_test.py
@@ -192,6 +192,33 @@ class AeadKeyManagerTest(parameterized.TestCase):
     with self.assertRaises(tink.TinkError):
       handle.primitive(aead.Aead)
 
+  def test_kms_envelope_aead_with_envelope_template_as_dek_template_fails(self):
+    env_template = (
+        aead.aead_key_templates.create_kms_envelope_aead_key_template(
+            kek_uri=FAKE_KMS_URI,
+            dek_template=aead.aead_key_templates.AES128_GCM,
+        )
+    )
+    template = aead.aead_key_templates.create_kms_envelope_aead_key_template(
+        kek_uri=FAKE_KMS_URI,
+        dek_template=env_template,
+    )
+    handle = tink.new_keyset_handle(template)
+    with self.assertRaises(tink.TinkError):
+      _ = handle.primitive(aead.Aead)
+
+  def test_kms_envelope_aead_with_kms_template_as_dek_template_fails(self):
+    kms_template = aead.aead_key_templates.create_kms_aead_key_template(
+        key_uri=FAKE_KMS_URI,
+    )
+    template = aead.aead_key_templates.create_kms_envelope_aead_key_template(
+        kek_uri=FAKE_KMS_URI,
+        dek_template=kms_template,
+    )
+    handle = tink.new_keyset_handle(template)
+    with self.assertRaises(tink.TinkError):
+      _ = handle.primitive(aead.Aead)
+
   def test_kms_envelope_aead_decrypt_fixed_ciphertext_success(self):
     # This keyset contains a single KmsEnvelopeAeadKey with
     # kek_uri = FAKE_KMS_URI and dek_template = AES128_GCM.
diff --git a/python/tink/aead/_kms_envelope_aead.py b/python/tink/aead/_kms_envelope_aead.py
index e135ae813..b4a0ab3dc 100644
--- a/python/tink/aead/_kms_envelope_aead.py
+++ b/python/tink/aead/_kms_envelope_aead.py
@@ -19,27 +19,47 @@ from tink.proto import tink_pb2
 from tink import core
 from tink.aead import _aead
 
+_SUPPORTED_DEK_KEY_TYPES = {
+    'type.googleapis.com/google.crypto.tink.AesGcmKey',
+    'type.googleapis.com/google.crypto.tink.XChaCha20Poly1305Key',
+    'type.googleapis.com/google.crypto.tink.AesCtrHmacAeadKey',
+    'type.googleapis.com/google.crypto.tink.AesEaxKey',
+    'type.googleapis.com/google.crypto.tink.AesGcmSivKey',
+}
+
 
 class KmsEnvelopeAead(_aead.Aead):
   """Implements envelope encryption.
 
-     Envelope encryption generates a data encryption key (DEK) which is used
-     to encrypt the payload. The DEK is then send to a KMS to be encrypted and
-     the encrypted DEK is attached to the ciphertext. In order to decrypt the
-     ciphertext, the DEK first has to be decrypted by the KMS, and then the DEK
-     can be used to decrypt the ciphertext. For further information see
-     https://cloud.google.com/kms/docs/envelope-encryption.
-
-     The ciphertext structure is as follows:
-     * Length of the encrypted DEK: 4 bytes (big endian)
-     * Encrypted DEK: variable length, specified by the previous 4 bytes
-     * AEAD payload: variable length
+  Envelope encryption generates a data encryption key (DEK) which is used
+  to encrypt the payload. The DEK is then send to a KMS to be encrypted and
+  the encrypted DEK is attached to the ciphertext. In order to decrypt the
+  ciphertext, the DEK first has to be decrypted by the KMS, and then the DEK
+  can be used to decrypt the ciphertext. For further information see
+  https://cloud.google.com/kms/docs/envelope-encryption.
+
+  DEK key template must be a KeyTemplate for any of these Tink AEAD key types
+  (any other key template will be rejected):
+  * AesGcmKey
+  * XChaCha20Poly1305
+  * AesCtrHmacAeadKey
+  * AesEaxKey
+  * AesGcmSivKey
+
+  The ciphertext structure is as follows:
+  * Length of the encrypted DEK: 4 bytes (big endian)
+  * Encrypted DEK: variable length, specified by the previous 4 bytes
+  * AEAD payload: variable length
   """
 
   # Defines in how many bytes the DEK length will be encoded.
   DEK_LEN_BYTES = 4
 
   def __init__(self, key_template: tink_pb2.KeyTemplate, remote: _aead.Aead):
+    if key_template.type_url not in _SUPPORTED_DEK_KEY_TYPES:
+      raise core.TinkError(
+          'Unsupported DEK key type: %s' % key_template.type_url
+      )
     # Create a dek to make sure that it works, so that KmsEnvelopeAead already
     # fails when it is created, and not just when it is used.
     # The C++ implementation does the same check, and we want this