From bbd06e242ae4975b98390cefa0fdd1a66d01705f Mon Sep 17 00:00:00 2001 From: mohammadnaseri Date: Thu, 13 Jun 2024 12:18:14 +0100 Subject: [PATCH 1/2] Verify TTL when storing taskin and taskres --- src/py/flwr/server/utils/validator.py | 6 ++++++ src/py/flwr/server/utils/validator_test.py | 18 ++++++++++++++++++ 2 files changed, 24 insertions(+) diff --git a/src/py/flwr/server/utils/validator.py b/src/py/flwr/server/utils/validator.py index c0b0ec85761c..e05f0b48b69a 100644 --- a/src/py/flwr/server/utils/validator.py +++ b/src/py/flwr/server/utils/validator.py @@ -18,6 +18,7 @@ from typing import List, Union from flwr.proto.task_pb2 import TaskIns, TaskRes # pylint: disable=E0611 +import time # pylint: disable-next=too-many-branches,too-many-statements @@ -47,6 +48,11 @@ def validate_task_ins_or_res(tasks_ins_res: Union[TaskIns, TaskRes]) -> List[str # unix timestamp of 27 March 2024 00h:00m:00s UTC validation_errors.append("`pushed_at` is not a recent timestamp") + # Verify TTL and created_at time + current_time = time.time() + if tasks_ins_res.task.created_at + tasks_ins_res.task.ttl <= current_time: + validation_errors.append("Task TTL has expired") + # TaskIns specific if isinstance(tasks_ins_res, TaskIns): # Task producer diff --git a/src/py/flwr/server/utils/validator_test.py b/src/py/flwr/server/utils/validator_test.py index 61fe094c23d4..d9e5d67cf20d 100644 --- a/src/py/flwr/server/utils/validator_test.py +++ b/src/py/flwr/server/utils/validator_test.py @@ -77,6 +77,24 @@ def test_is_valid_task_res(self) -> None: val_errors = validate_task_ins_or_res(msg) self.assertTrue(val_errors, (producer_node_id, anonymous, ancestry)) + def test_task_ttl_expired(self) -> None: + """Test validation for expired Task TTL.""" + # Prepare an expired TaskIns + expired_task_ins = create_task_ins(0, True) + expired_task_ins.task.created_at = time.time() - 10 # 10 seconds ago + expired_task_ins.task.ttl = 6 # 6 seconds TTL + + expired_task_res = create_task_res(0, True, ["1"]) + expired_task_res.task.created_at = time.time() - 10 # 10 seconds ago + expired_task_res.task.ttl = 6 # 6 seconds TTL + + # Execute & Assert + val_errors_ins = validate_task_ins_or_res(expired_task_ins) + self.assertIn("Task TTL has expired", val_errors_ins) + + val_errors_res = validate_task_ins_or_res(expired_task_res) + self.assertIn("Task TTL has expired", val_errors_res) + def create_task_ins( consumer_node_id: int, From 1a4708021a15c63bf47a007833b8fc63d70d7827 Mon Sep 17 00:00:00 2001 From: mohammadnaseri Date: Thu, 13 Jun 2024 12:19:09 +0100 Subject: [PATCH 2/2] Format --- src/py/flwr/server/utils/validator.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/py/flwr/server/utils/validator.py b/src/py/flwr/server/utils/validator.py index e05f0b48b69a..6c782f0177cb 100644 --- a/src/py/flwr/server/utils/validator.py +++ b/src/py/flwr/server/utils/validator.py @@ -15,10 +15,10 @@ """Validators.""" +import time from typing import List, Union from flwr.proto.task_pb2 import TaskIns, TaskRes # pylint: disable=E0611 -import time # pylint: disable-next=too-many-branches,too-many-statements