Remove client-side decryption key ID verification in Python.

This isn't necessary since it duplicates the server side check done when the
KeyID is included in the decryption request.

This also enables support for using key aliases.

PiperOrigin-RevId: 538511498

3 files changed, 22 insertions, 5 deletions

diff --git a/python/tink/integration/awskms/_aws_kms_client.py b/python/tink/integration/awskms/_aws_kms_client.py
index be368b078..d17d70b7c 100644
--- a/python/tink/integration/awskms/_aws_kms_client.py
+++ b/python/tink/integration/awskms/_aws_kms_client.py
@@ -63,11 +63,6 @@ class _AwsKmsAead(aead.Aead):
           CiphertextBlob=ciphertext,
           EncryptionContext=_encryption_context(associated_data),
       )
-      if response['KeyId'] != self.key_arn:
-        raise tink.TinkError(
-            'invalid key id: got %s, want %s'
-            % (self.key_arn, response['KeyId'])
-        )
       return response['Plaintext']
     except exceptions.ClientError as e:
       raise tink.TinkError(e)
diff --git a/python/tink/integration/awskms/_aws_kms_client_test.py b/python/tink/integration/awskms/_aws_kms_client_test.py
index 129f2de9e..1968c7bac 100644
--- a/python/tink/integration/awskms/_aws_kms_client_test.py
+++ b/python/tink/integration/awskms/_aws_kms_client_test.py
@@ -27,10 +27,17 @@ from tink.testing import helper
 
 CREDENTIAL_PATH = os.path.join(helper.tink_py_testdata_path(),
                                'aws/credentials.ini')
+
 KEY_URI = ('aws-kms://arn:aws:kms:us-east-2:235739564943:key/'
            '3ee50705-5a82-4f5b-9753-05c4f473922f')
+
+# An alias for KEY_URI.
+KEY_ALIAS_URI = ('aws-kms://arn:aws:kms:us-east-2:235739564943:alias/'
+                 'unit-and-integration-testing')
+
 KEY_URI_2 = ('aws-kms://arn:aws:kms:us-east-2:235739564943:key/'
              'b3ca2efd-a8fb-47f2-b541-7e20f8c5cd11')
+
 GCP_KEY_URI = ('gcp-kms://projects/tink-test-infrastructure/locations/global/'
                'keyRings/unit-and-integration-testing/cryptoKeys/aead-key')
 
@@ -41,6 +48,7 @@ class AwsKmsClientTest(absltest.TestCase):
     aws_client = awskms.AwsKmsClient(KEY_URI, CREDENTIAL_PATH)
 
     self.assertEqual(aws_client.does_support(KEY_URI), True)
+    self.assertEqual(aws_client.does_support(KEY_ALIAS_URI), False)
     self.assertEqual(aws_client.does_support(KEY_URI_2), False)
     self.assertEqual(aws_client.does_support(GCP_KEY_URI), False)
 
@@ -48,6 +56,7 @@ class AwsKmsClientTest(absltest.TestCase):
     aws_client = awskms.AwsKmsClient('', CREDENTIAL_PATH)
 
     self.assertEqual(aws_client.does_support(KEY_URI), True)
+    self.assertEqual(aws_client.does_support(KEY_ALIAS_URI), True)
     self.assertEqual(aws_client.does_support(KEY_URI_2), True)
     self.assertEqual(aws_client.does_support(GCP_KEY_URI), False)
 
diff --git a/python/tink/integration/awskms/_aws_kms_integration_test.py b/python/tink/integration/awskms/_aws_kms_integration_test.py
index b670a4e7e..eeb1e210c 100644
--- a/python/tink/integration/awskms/_aws_kms_integration_test.py
+++ b/python/tink/integration/awskms/_aws_kms_integration_test.py
@@ -69,6 +69,19 @@ class AwsKmsAeadTest(absltest.TestCase):
     ciphertext = aws_aead.encrypt(plaintext, b'')
     self.assertEqual(plaintext, aws_aead.decrypt(ciphertext, b''))
 
+  def test_encrypt_decrypt_with_key_alias(self):
+    aws_client = awskms.AwsKmsClient(KEY_ALIAS_URI, CREDENTIAL_PATH)
+    aws_aead = aws_client.get_aead(KEY_ALIAS_URI)
+
+    plaintext = b'hello'
+    associated_data = b'world'
+    ciphertext = aws_aead.encrypt(plaintext, associated_data)
+    self.assertEqual(plaintext, aws_aead.decrypt(ciphertext, associated_data))
+
+    plaintext = b'hello'
+    ciphertext = aws_aead.encrypt(plaintext, b'')
+    self.assertEqual(plaintext, aws_aead.decrypt(ciphertext, b''))
+
   def test_corrupted_ciphertext(self):
     aws_client = awskms.AwsKmsClient(KEY_URI, CREDENTIAL_PATH)
     aws_aead = aws_client.get_aead(KEY_URI)