fix(idempotency): validate before saving to cache

aws_lambda_powertools/utilities/idempotency/persistence/base.py

@@ -1,6 +1,7 @@
 """
 Persistence layers supporting idempotency
 """
+
 import datetime
 import hashlib
 import json
@@ -383,9 +384,9 @@ def get_record(self, data: Dict[str, Any]) -> Optional[DataRecord]:
 
  record = self._get_record(idempotency_key=idempotency_key)
 
+ self._validate_payload(data_payload=data, stored_data_record=record)
  self._save_to_cache(data_record=record)
 
- self._validate_payload(data_payload=data, stored_data_record=record)
  return record
 
  @abstractmethod

aws_lambda_powertools/utilities/idempotency/persistence/dynamodb.py

@@ -259,10 +259,10 @@ def _put_record(self, data_record: DataRecord) -> None:
  f"expiry_timestamp: {old_data_record.expiry_timestamp}, "
  f"and in_progress_expiry_timestamp: {old_data_record.in_progress_expiry_timestamp}",
  )
- self._save_to_cache(data_record=old_data_record)
 
  try:
  self._validate_payload(data_payload=data_record, stored_data_record=old_data_record)
+ self._save_to_cache(data_record=old_data_record)
  except IdempotencyValidationError as idempotency_validation_error:
  raise idempotency_validation_error from exc
 

pyproject.toml

@@ -120,6 +120,7 @@ retry2 = "^0.9.5"
 pytest-socket = ">=0.6,<0.8"
 types-redis = "^4.6.0.7"
 testcontainers = { extras = ["redis"], version = "^3.7.1" }
+multiprocess = "^0.70.16"
 
 [tool.coverage.run]
 source = ["aws_lambda_powertools"]

tests/functional/idempotency/persistence/test_redis_layer.py

@@ -1,36 +1,33 @@
 # ruff: noqa
 import copy
+import datetime
 import json
 import time as t
+from unittest import mock
 
 import pytest
-from unittest.mock import patch
-
-from aws_lambda_powertools.utilities.idempotency.persistence.redis import (
- RedisCachePersistenceLayer,
-)
-import datetime
+from multiprocess import Lock, Manager, Process
 
-from aws_lambda_powertools.utilities.idempotency.persistence.base import (
- STATUS_CONSTANTS,
- DataRecord,
-)
-
-from unittest import mock
-from multiprocessing import Process, Manager, Lock
 from aws_lambda_powertools.utilities.idempotency.exceptions import (
  IdempotencyAlreadyInProgressError,
  IdempotencyItemAlreadyExistsError,
  IdempotencyItemNotFoundError,
- IdempotencyPersistenceConnectionError,
  IdempotencyPersistenceConfigError,
+ IdempotencyPersistenceConnectionError,
  IdempotencyPersistenceConsistencyError,
  IdempotencyValidationError,
 )
 from aws_lambda_powertools.utilities.idempotency.idempotency import (
+ IdempotencyConfig,
  idempotent,
  idempotent_function,
- IdempotencyConfig,
+)
+from aws_lambda_powertools.utilities.idempotency.persistence.base import (
+ STATUS_CONSTANTS,
+ DataRecord,
+)
+from aws_lambda_powertools.utilities.idempotency.persistence.redis import (
+ RedisCachePersistenceLayer,
 )
 
 redis_badhost = "badhost"
@@ -557,6 +554,7 @@ def test_redis_orphan_record_race_condition(lambda_context):
  port="63005",
  mock_latency_ms=50,
  )
+
  manager = Manager()
  # use a thread safe dict
  redis_client.expire_dict = manager.dict()
@@ -576,11 +574,13 @@ def lambda_handler(event, context):
 
  # run handler for the first time to create a valid record in cache
  lambda_handler(mock_event, lambda_context)
+
  # modify the cache expiration to create the orphan record
  for key, item in redis_client.cache.items():
  json_dict = json.loads(item)
  json_dict["expiration"] = int(t.time()) - 4000
  redis_client.cache[key] = json.dumps(json_dict).encode()
+
  # Given orphan idempotency record with same payload already in Redis
  # When running two lambda handler at the same time
  redis_client.cache["exec_count"] = 0
@@ -590,6 +590,7 @@ def lambda_handler(event, context):
  p2.start()
  p1.join()
  p2.join()
+
  # Then only one handler will actually run
  assert redis_client.cache["exec_count"] == 1
 

tests/functional/idempotency/test_idempotency.py

@@ -9,6 +9,7 @@
 from botocore.config import Config
 from pydantic import BaseModel
 from pytest import FixtureRequest
+from pytest_mock import MockerFixture
 
 from aws_lambda_powertools.utilities.data_classes import (
  APIGatewayProxyEventV2,
@@ -1928,3 +1929,60 @@ def lambda_handler(event, context):
 
  stubber.assert_no_pending_responses()
  stubber.deactivate()
+
+
+def test_idempotency_cache_with_payload_tampering(
+ persistence_store: DynamoDBPersistenceLayer,
+ timestamp_future,
+ lambda_context,
+ request: FixtureRequest,
+ mocker: MockerFixture,
+):
+ # GIVEN an idempotency config with a compound idempotency key (refund, customer_id)
+ # AND with payload validation key to prevent tampering
+
+ cache_spy = mocker.spy(persistence_store, "_save_to_cache")
+
+ validation_key = "amount"
+ idempotency_config = IdempotencyConfig(
+ event_key_jmespath='["refund_id", "customer_id"]',
+ payload_validation_jmespath=validation_key,
+ use_local_cache=True,
+ )
+
+ # AND a previous transaction already processed in the persistent store
+ transaction = {
+ "refund_id": "ffd11882-d476-4598-bbf1-643f2be5addf",
+ "customer_id": "9e9fc440-9e65-49b5-9e71-1382ea1b1658",
+ "amount": 100,
+ }
+
+ stubber = stub.Stubber(persistence_store.client)
+ ddb_response = build_idempotency_put_item_response_stub(
+ data=transaction,
+ expiration=timestamp_future,
+ status="COMPLETED",
+ request=request,
+ validation_data=transaction[validation_key],
+ )
+
+ stubber.add_client_error("put_item", "ConditionalCheckFailedException", modeled_fields=ddb_response)
+ stubber.activate()
+
+ # AND an upcoming tampered transaction
+ tampered_transaction = copy.deepcopy(transaction)
+ tampered_transaction["amount"] = 10_000
+
+ @idempotent(config=idempotency_config, persistence_store=persistence_store)
+ def lambda_handler(event, context):
+ return event
+
+ # WHEN the tampered request is made
+ with pytest.raises(IdempotencyValidationError):
+ lambda_handler(tampered_transaction, lambda_context)
+
+ stubber.assert_no_pending_responses()
+ stubber.deactivate()
+
+ # THEN we should not cache a transaction that failed validation
+ assert cache_spy.call_count == 0