add proxy tests for the KMS service (#116)

Author: whummer

.github/workflows/aws-proxy.yml

@@ -39,6 +39,12 @@ jobs:
       - name: Set up Terraform CLI
         uses: hashicorp/setup-terraform@v2
 
+      - name: Run linter
+        run: |
+          cd aws-proxy
+          make install
+          make lint
+
       - name: Install LocalStack and extension
         env:
           LOCALSTACK_AUTH_TOKEN: ${{ secrets.LOCALSTACK_AUTH_TOKEN }}
@@ -58,8 +64,6 @@ jobs:
           # build and install extension
           localstack extensions init
           (
-            make install
-            . .venv/bin/activate
             make build
             make enable
           )
@@ -73,12 +77,6 @@ jobs:
           DEBUG=1 GATEWAY_SERVER=hypercorn localstack start -d
           localstack wait
 
-      - name: Run linter
-        run: |
-          cd aws-proxy
-          (. .venv/bin/activate; pip install --upgrade --pre localstack localstack-ext)
-          make lint
-
       - name: Run integration tests
         env:
           AWS_DEFAULT_REGION: us-east-1

.github/workflows/typedb.yml

@@ -1,7 +1,15 @@
 name: LocalStack TypeDB Extension Tests
 
 on:
+  push:
+    paths:
+      - typedb/**
+    branches:
+      - main
   pull_request:
+    paths:
+      - .github/workflows/typedb.yml
+      - typedb/**
   workflow_dispatch:
 
 env:

aws-proxy/tests/proxy/test_kms.py

@@ -0,0 +1,241 @@
+# Note/disclosure: This file has been (partially or fully) generated by an AI agent.
+import boto3
+import pytest
+from botocore.exceptions import ClientError
+from localstack.aws.connect import connect_to
+from localstack.utils.strings import short_uid
+
+
+def test_kms_key_operations(start_aws_proxy, cleanups):
+    """Test basic KMS key operations with proxy."""
+    key_description = f"test-key-{short_uid()}"
+
+    # start proxy - forwarding requests for KMS keys with specific alias pattern
+    config = {
+        "services": {
+            "kms": {"resources": [".*:key/.*"]},
+        }
+    }
+    start_aws_proxy(config)
+
+    # create clients
+    kms_client = connect_to().kms
+    kms_client_aws = boto3.client("kms")
+
+    # create key in AWS
+    create_response = kms_client_aws.create_key(Description=key_description)
+    key_id_aws = create_response["KeyMetadata"]["KeyId"]
+    key_arn_aws = create_response["KeyMetadata"]["Arn"]
+    cleanups.append(
+        lambda: kms_client_aws.schedule_key_deletion(
+            KeyId=key_id_aws, PendingWindowInDays=7
+        )
+    )
+
+    # assert that local call for this key is proxied
+    key_aws = kms_client_aws.describe_key(KeyId=key_id_aws)
+    key_local = kms_client.describe_key(KeyId=key_id_aws)
+    assert key_local["KeyMetadata"]["KeyId"] == key_aws["KeyMetadata"]["KeyId"]
+    assert key_local["KeyMetadata"]["Arn"] == key_arn_aws
+
+    # test encryption with AWS client
+    plaintext = b"test message"
+    encrypt_response_aws = kms_client_aws.encrypt(KeyId=key_id_aws, Plaintext=plaintext)
+    ciphertext = encrypt_response_aws["CiphertextBlob"]
+
+    # decrypt with local client (proxied to AWS)
+    decrypt_response_local = kms_client.decrypt(CiphertextBlob=ciphertext)
+    assert decrypt_response_local["Plaintext"] == plaintext
+    assert decrypt_response_local["KeyId"] == key_arn_aws
+
+    # encrypt with local client (proxied to AWS)
+    plaintext2 = b"another test message"
+    encrypt_response_local = kms_client.encrypt(KeyId=key_id_aws, Plaintext=plaintext2)
+    ciphertext2 = encrypt_response_local["CiphertextBlob"]
+
+    # decrypt with AWS client
+    decrypt_response_aws = kms_client_aws.decrypt(CiphertextBlob=ciphertext2)
+    assert decrypt_response_aws["Plaintext"] == plaintext2
+
+
+def test_kms_key_alias_operations(start_aws_proxy, cleanups):
+    """Test KMS key alias operations with proxy."""
+    alias_name = f"alias/test-alias-{short_uid()}"
+    key_description = f"test-key-with-alias-{short_uid()}"
+
+    # start proxy - forwarding requests for KMS
+    config = {
+        "services": {
+            "kms": {"resources": [".*:key/.*", f".*:{alias_name}"]},
+        }
+    }
+    start_aws_proxy(config)
+
+    # create clients
+    kms_client = connect_to().kms
+    kms_client_aws = boto3.client("kms")
+
+    # create key in AWS
+    create_response = kms_client_aws.create_key(Description=key_description)
+    key_id_aws = create_response["KeyMetadata"]["KeyId"]
+    cleanups.append(
+        lambda: kms_client_aws.schedule_key_deletion(
+            KeyId=key_id_aws, PendingWindowInDays=7
+        )
+    )
+
+    # create alias in AWS
+    kms_client_aws.create_alias(AliasName=alias_name, TargetKeyId=key_id_aws)
+    cleanups.append(lambda: kms_client_aws.delete_alias(AliasName=alias_name))
+
+    # assert that local call for alias operations is proxied
+    aliases_aws = kms_client_aws.list_aliases(KeyId=key_id_aws)["Aliases"]
+    aliases_local = kms_client.list_aliases(KeyId=key_id_aws)["Aliases"]
+
+    # filter for our specific alias
+    alias_aws = [a for a in aliases_aws if a["AliasName"] == alias_name][0]
+    alias_local = [a for a in aliases_local if a["AliasName"] == alias_name][0]
+
+    assert alias_local["AliasName"] == alias_aws["AliasName"]
+    assert alias_local["TargetKeyId"] == alias_aws["TargetKeyId"]
+
+    # test encryption with alias via local client
+    plaintext = b"test with alias"
+    encrypt_response_local = kms_client.encrypt(KeyId=alias_name, Plaintext=plaintext)
+    ciphertext = encrypt_response_local["CiphertextBlob"]
+
+    # decrypt with AWS client
+    decrypt_response_aws = kms_client_aws.decrypt(CiphertextBlob=ciphertext)
+    assert decrypt_response_aws["Plaintext"] == plaintext
+
+
+def test_kms_readonly_operations(start_aws_proxy, cleanups):
+    """Test KMS operations in read-only proxy mode."""
+    key_description = f"test-readonly-key-{short_uid()}"
+
+    # start proxy - forwarding requests for KMS in read-only mode
+    config = {
+        "services": {
+            "kms": {"resources": [".*:key/.*"], "read_only": True},
+        }
+    }
+    start_aws_proxy(config)
+
+    # create clients
+    kms_client = connect_to().kms
+    kms_client_aws = boto3.client("kms")
+
+    # create key in AWS (this should succeed as it's direct AWS client)
+    create_response = kms_client_aws.create_key(Description=key_description)
+    key_id_aws = create_response["KeyMetadata"]["KeyId"]
+    cleanups.append(
+        lambda: kms_client_aws.schedule_key_deletion(
+            KeyId=key_id_aws, PendingWindowInDays=7
+        )
+    )
+
+    # assert that local call for describe_key is proxied and results are consistent
+    key_aws = kms_client_aws.describe_key(KeyId=key_id_aws)
+    key_local = kms_client.describe_key(KeyId=key_id_aws)
+    assert key_local["KeyMetadata"]["KeyId"] == key_aws["KeyMetadata"]["KeyId"]
+
+    # assert that local call for list_keys is proxied
+    keys_local = kms_client.list_keys()["Keys"]
+    keys_aws = kms_client_aws.list_keys()["Keys"]
+
+    # filter for our specific key
+    key_local_filtered = [k for k in keys_local if k["KeyId"] == key_id_aws]
+    key_aws_filtered = [k for k in keys_aws if k["KeyId"] == key_id_aws]
+    assert key_local_filtered == key_aws_filtered
+
+    # Negative test: attempt write operations with proxied client
+    # Create a new key using the proxied client (should succeed in LocalStack)
+    new_key_description = f"no-proxy-key-{short_uid()}"
+    new_key_local = kms_client.create_key(Description=new_key_description)
+    new_key_id_local = new_key_local["KeyMetadata"]["KeyId"]
+    cleanups.append(
+        lambda: kms_client.schedule_key_deletion(
+            KeyId=new_key_id_local, PendingWindowInDays=7
+        )
+    )
+
+    # Verify that this new key does NOT exist in real AWS
+    keys_aws_after_create = kms_client_aws.list_keys()["Keys"]
+    assert not any(k for k in keys_aws_after_create if k["KeyId"] == new_key_id_local)
+
+    # Attempt to encrypt data with the AWS key using proxied client (should fail)
+    plaintext = b"this should not work"
+    with pytest.raises(ClientError) as excinfo:
+        kms_client.encrypt(KeyId=key_id_aws, Plaintext=plaintext)
+    # In read-only mode, the key exists in AWS but LocalStack doesn't have it locally
+    # so encrypt operation should fail
+    assert excinfo.value.response["Error"]["Code"] in [
+        "NotFoundException",
+        "InvalidKeyId.NotFound",
+    ]
+
+
+def test_kms_selective_resource_matching(start_aws_proxy, cleanups):
+    """Test that proxy forwards requests for specific KMS keys matching ARN pattern."""
+    key_description_1 = f"test-proxied-key-1-{short_uid()}"
+    key_description_2 = f"test-proxied-key-2-{short_uid()}"
+
+    # create clients
+    kms_client_aws = boto3.client("kms")
+
+    # create two keys in AWS
+    create_response_1 = kms_client_aws.create_key(Description=key_description_1)
+    key_id_1 = create_response_1["KeyMetadata"]["KeyId"]
+    key_arn_1 = create_response_1["KeyMetadata"]["Arn"]
+    cleanups.append(
+        lambda: kms_client_aws.schedule_key_deletion(
+            KeyId=key_id_1, PendingWindowInDays=7
+        )
+    )
+
+    create_response_2 = kms_client_aws.create_key(Description=key_description_2)
+    key_id_2 = create_response_2["KeyMetadata"]["KeyId"]
+    key_arn_2 = create_response_2["KeyMetadata"]["Arn"]
+    cleanups.append(
+        lambda: kms_client_aws.schedule_key_deletion(
+            KeyId=key_id_2, PendingWindowInDays=7
+        )
+    )
+
+    # start proxy - forwarding requests for both keys using wildcard pattern
+    config = {
+        "services": {
+            "kms": {"resources": [".*:key/.*"]},
+        }
+    }
+    start_aws_proxy(config)
+
+    # create LocalStack client after proxy is started
+    kms_client = connect_to().kms
+
+    # verify that both keys are accessible via local client (proxied)
+    key_1_local = kms_client.describe_key(KeyId=key_id_1)
+    assert key_1_local["KeyMetadata"]["KeyId"] == key_id_1
+    assert key_1_local["KeyMetadata"]["Arn"] == key_arn_1
+
+    key_2_local = kms_client.describe_key(KeyId=key_id_2)
+    assert key_2_local["KeyMetadata"]["KeyId"] == key_id_2
+    assert key_2_local["KeyMetadata"]["Arn"] == key_arn_2
+
+    # verify we can encrypt with both proxied keys
+    plaintext = b"test with first key"
+    encrypt_response_1 = kms_client.encrypt(KeyId=key_id_1, Plaintext=plaintext)
+    ciphertext_1 = encrypt_response_1["CiphertextBlob"]
+
+    # decrypt with AWS client to confirm it went through proxy
+    decrypt_response_1 = kms_client_aws.decrypt(CiphertextBlob=ciphertext_1)
+    assert decrypt_response_1["Plaintext"] == plaintext
+
+    # encrypt with second key
+    plaintext_2 = b"test with second key"
+    encrypt_response_2 = kms_client.encrypt(KeyId=key_id_2, Plaintext=plaintext_2)
+    ciphertext_2 = encrypt_response_2["CiphertextBlob"]
+
+    # decrypt with AWS client
+    decrypt_response_2 = kms_client_aws.decrypt(CiphertextBlob=ciphertext_2)
+    assert decrypt_response_2["Plaintext"] == plaintext_2