Let KMS envelope key manager reject invalid DEK templates.

PiperOrigin-RevId: 543424290

3 files changed, 20 insertions, 12 deletions

diff --git a/python/tink/aead/_aead_key_manager_test.py b/python/tink/aead/_aead_key_manager_test.py
index 25332fc13..ec8671d0a 100644
--- a/python/tink/aead/_aead_key_manager_test.py
+++ b/python/tink/aead/_aead_key_manager_test.py
@@ -188,9 +188,8 @@ class AeadKeyManagerTest(parameterized.TestCase):
         kek_uri=FAKE_KMS_URI,
         dek_template=mac.mac_key_templates.HMAC_SHA256_128BITTAG,
     )
-    handle = tink.new_keyset_handle(template)
     with self.assertRaises(tink.TinkError):
-      handle.primitive(aead.Aead)
+      _ = tink.new_keyset_handle(template)
 
   def test_kms_envelope_aead_with_envelope_template_as_dek_template_fails(self):
     env_template = (
@@ -203,9 +202,8 @@ class AeadKeyManagerTest(parameterized.TestCase):
         kek_uri=FAKE_KMS_URI,
         dek_template=env_template,
     )
-    handle = tink.new_keyset_handle(template)
     with self.assertRaises(tink.TinkError):
-      _ = handle.primitive(aead.Aead)
+      _ = tink.new_keyset_handle(template)
 
   def test_kms_envelope_aead_with_kms_template_as_dek_template_fails(self):
     kms_template = aead.aead_key_templates.create_kms_aead_key_template(
@@ -215,9 +213,8 @@ class AeadKeyManagerTest(parameterized.TestCase):
         kek_uri=FAKE_KMS_URI,
         dek_template=kms_template,
     )
-    handle = tink.new_keyset_handle(template)
     with self.assertRaises(tink.TinkError):
-      _ = handle.primitive(aead.Aead)
+      _ = tink.new_keyset_handle(template)
 
   def test_kms_envelope_aead_decrypt_fixed_ciphertext_success(self):
     # This keyset contains a single KmsEnvelopeAeadKey with
diff --git a/python/tink/aead/_kms_aead_key_manager.py b/python/tink/aead/_kms_aead_key_manager.py
index 10a5720ca..cffa65f68 100644
--- a/python/tink/aead/_kms_aead_key_manager.py
+++ b/python/tink/aead/_kms_aead_key_manager.py
@@ -124,11 +124,18 @@ class KmsEnvelopeAeadKeyManager(core.KeyManager[_aead.Aead]):
   ) -> tink_pb2.KeyData:
     if key_template.type_url != _KMS_ENVELOPE_AEAD_KEY_TYPE_URL:
       raise core.TinkError('wrong key type: ' + key_template.type_url)
+    params = kms_envelope_pb2.KmsEnvelopeAeadKeyFormat.FromString(
+        key_template.value
+    )
+    if not _kms_envelope_aead.is_supported_dek_key_type(
+        params.dek_template.type_url
+    ):
+      raise core.TinkError(
+          'Unsupported DEK key type: %s' % key_template.type_url
+      )
     env_key = kms_envelope_pb2.KmsEnvelopeAeadKey(
         version=0,
-        params=kms_envelope_pb2.KmsEnvelopeAeadKeyFormat.FromString(
-            key_template.value
-        ),
+        params=params,
     )
     return tink_pb2.KeyData(
         type_url=_KMS_ENVELOPE_AEAD_KEY_TYPE_URL,
diff --git a/python/tink/aead/_kms_envelope_aead.py b/python/tink/aead/_kms_envelope_aead.py
index b4a0ab3dc..823c40e17 100644
--- a/python/tink/aead/_kms_envelope_aead.py
+++ b/python/tink/aead/_kms_envelope_aead.py
@@ -19,13 +19,17 @@ from tink.proto import tink_pb2
 from tink import core
 from tink.aead import _aead
 
-_SUPPORTED_DEK_KEY_TYPES = {
+_SUPPORTED_DEK_KEY_TYPES = frozenset({
     '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',
-}
+})
+
+
+def is_supported_dek_key_type(type_url: str) -> bool:
+  return type_url in _SUPPORTED_DEK_KEY_TYPES
 
 
 class KmsEnvelopeAead(_aead.Aead):
@@ -56,7 +60,7 @@ class KmsEnvelopeAead(_aead.Aead):
   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:
+    if not is_supported_dek_key_type(key_template.type_url):
       raise core.TinkError(
           'Unsupported DEK key type: %s' % key_template.type_url
       )