8000 KMS: allow multiple keys to have the same grant name (#7911) · codeperl/localstack@80469cc · GitHub
[go: up one dir, main page]

Skip to content

Commit 80469cc

Browse files
authored
KMS: allow multiple keys to have the same grant name (localstack#7911)
1 parent f141482 commit 80469cc

File tree

3 files changed

+40
-8
lines changed

3 files changed

+40
-8
lines changed

localstack/services/kms/models.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@
88
import uuid
99
from collections import namedtuple
1010
from dataclasses import dataclass
11-
from typing import Dict, List
11+
from typing import Dict, List, Tuple
1212

1313
from cryptography.exceptions import InvalidSignature
1414
from cryptography.hazmat.primitives import hashes
@@ -517,8 +517,8 @@ class KmsStore(BaseStore):
517517
#
518518
# maps grant ids to grants
519519
grants: Dict[str, KmsGrant] = LocalAttribute(default=dict)
520-
# maps from grant names (used for idempotency) to grant ids
521-
grant_names: Dict[str, str] = LocalAttribute(default=dict)
520+
# maps from (grant names (used for idempotency), key id) to grant ids
521+
grant_names: Dict[Tuple[str, str], str] = LocalAttribute(default=dict)
522522
# maps grant tokens to grant ids
523523
grant_tokens: Dict[str, str] = LocalAttribute(default=dict)
524524

localstack/services/kms/provider.py

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -274,17 +274,18 @@ def create_grant(
274274
# KeyId can potentially hold one of multiple different types of key identifiers. Here we find a key no
275275
# matter which type of id is used.
276276
key = store.get_key(request.get("KeyId"))
277-
request["KeyId"] = key.metadata.get("KeyId")
277+
key_id = key.metadata.get("KeyId")
278+
request["KeyId"] = key_id
278279
self._validate_grant_request(request, store)
279280
grant_name = request.get("Name")
280-
if grant_name and grant_name in store.grant_names:
281-
grant = store.grants[store.grant_names[grant_name]]
281+
if grant_name and (grant_name, key_id) in store.grant_names:
282+
grant = store.grants[store.grant_names[(grant_name, key_id)]]
282283
else:
283284
grant = KmsGrant(request)
284285
grant_id = grant.metadata["GrantId"]
285286
store.grants[grant_id] = grant
286287
if grant_name:
287-
store.grant_names[grant_name] = grant_id
288+
store.grant_names[(grant_name, key_id)] = grant_id
288289
store.grant_tokens[grant.token] = grant_id
289290

290291
# At the moment we do not support multiple GrantTokens for grant creation request. Instead, we always use
@@ -379,7 +380,7 @@ def _delete_grant(
379380
# In AWS grants have one or more tokens. But we have a simplified modeling of grants, where they have exactly
380381
# one token.
381382
store.grant_tokens.pop(grant.token)
382-
store.grant_names.pop(grant.metadata.get("Name"), None)
383+
store.grant_names.pop((grant.metadata.get("Name"), key_id), None)
383384
store.grants.pop(grant_id)
384385

385386
def revoke_grant(

tests/integration/test_kms.py

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -256,6 +256,37 @@ def test_create_grant_with_valid_key(self, kms_client, kms_key, user_arn):
256256
grants_after = kms_client.list_grants(KeyId=key_id)["Grants"]
257257
assert len(grants_after) == len(grants_before) + 1
258258

259+
@pytest.mark.aws_validated
260+
def test_create_grant_with_same_name_two_keys(self, kms_client, kms_create_key, user_arn):
261+
first_key_id = kms_create_key()["KeyId"]
262+
second_key_id = kms_create_key()["KeyId"]
263+
264+
grant_name = "TestGrantName"
265+
266+
first_grant = kms_client.create_grant(
267+
KeyId=first_key_id,
268+
GranteePrincipal=user_arn,
269+
Name=grant_name,
270+
Operations=["Decrypt", "DescribeKey"],
271+
)
272+
assert "GrantId" in first_grant
273+
assert "GrantToken" in first_grant
274+
275+
second_grant = kms_client.create_grant(
276+
KeyId=second_key_id,
277+
GranteePrincipal=user_arn,
278+
Name=grant_name,
279+
Operations=["Decrypt", "DescribeKey"],
280+
)
281+
assert "GrantId" in second_grant
282+
assert "GrantToken" in second_grant
283+
284+
first_grants_after = kms_client.list_grants(KeyId=first_key_id)["Grants"]
285+
assert len(first_grants_after) == 1
286+
287+
second_grants_after = kms_client.list_grants(KeyId=second_key_id)["Grants"]
288+
assert len(second_grants_after) == 1
289+
259290
@pytest.mark.aws_validated
260291
def test_revoke_grant(self, kms_client, kms_grant_and_key):
261292
grant = kms_grant_and_key[0]

0 commit comments

Comments
 (0)
0