fix(idempotency): validate before saving to cache (#3822)

* fix(parameters): make cache aware of single vs multiple calls

Signed-off-by: heitorlessa <lessa@amazon.co.uk>

* chore: cleanup, add test for single and nested

Signed-off-by: heitorlessa <lessa@amazon.co.uk>

* refactor: validate before saving to cache

* fix: use multiprocess LIB to allow local functions by using dill over pickle

* chore: add test to prevent regression

Signed-off-by: heitorlessa <lessa@amazon.co.uk>

---------

Signed-off-by: heitorlessa <lessa@amazon.co.uk>

Author: heitorlessa

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
 

poetry.lock

@@ -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"]

pyproject.toml

@@ -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/persistence/test_redis_layer.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