From 995d6ccf2e3a5e0303306b0a0ee04888d7de0967 Mon Sep 17 00:00:00 2001 From: Obludka Date: Wed, 7 May 2025 02:16:52 +0000 Subject: [PATCH 1/6] Implement nonce handling utility with error management --- .../prometheus_swarm/utils/nonce.py | 103 ++++++++++++++++++ 1 file changed, 103 insertions(+) create mode 100644 agent-framework/prometheus_swarm/utils/nonce.py diff --git a/agent-framework/prometheus_swarm/utils/nonce.py b/agent-framework/prometheus_swarm/utils/nonce.py new file mode 100644 index 00000000..6ca23e85 --- /dev/null +++ b/agent-framework/prometheus_swarm/utils/nonce.py @@ -0,0 +1,103 @@ +import uuid +import time +from typing import Dict, Set +from threading import Lock + +class NonceError(Exception): + """Base exception for nonce-related errors.""" + pass + +class NonceAlreadyUsedError(NonceError): + """Raised when a nonce has already been used.""" + pass + +class NonceExpiredError(NonceError): + """Raised when a nonce has expired.""" + pass + +class NonceManager: + """ + Manages nonce generation, validation, and tracking. + + Provides thread-safe nonce management with expiration and uniqueness checks. + """ + def __init__(self, max_nonces: int = 1000, nonce_expiry: int = 300): + """ + Initialize the NonceManager. + + Args: + max_nonces (int): Maximum number of nonces to track. Default is 1000. + nonce_expiry (int): Nonce expiration time in seconds. Default is 5 minutes. + """ + self._used_nonces: Dict[str, float] = {} + self._lock = Lock() + self._max_nonces = max_nonces + self._nonce_expiry = nonce_expiry + + def generate_nonce(self) -> str: + """ + Generate a unique nonce. + + Returns: + str: A unique nonce value. + """ + with self._lock: + # Cleanup expired nonces + self._cleanup_expired_nonces() + + # Generate a unique nonce + nonce = str(uuid.uuid4()) + current_time = time.time() + self._used_nonces[nonce] = current_time + + return nonce + + def validate_nonce(self, nonce: str) -> bool: + """ + Validate a nonce, checking for uniqueness and expiration. + + Args: + nonce (str): The nonce to validate. + + Raises: + NonceAlreadyUsedError: If the nonce has been used before. + NonceExpiredError: If the nonce has expired. + + Returns: + bool: True if the nonce is valid. + """ + with self._lock: + # Cleanup expired nonces + self._cleanup_expired_nonces() + + # Check if nonce is already used + if nonce in self._used_nonces: + raise NonceAlreadyUsedError(f"Nonce {nonce} has already been used.") + + return True + + def _cleanup_expired_nonces(self): + """ + Remove expired nonces from the tracking dictionary. + """ + current_time = time.time() + expired_nonces = [ + nonce for nonce, timestamp in self._used_nonces.items() + if current_time - timestamp > self._nonce_expiry + ] + + for nonce in expired_nonces: + del self._used_nonces[nonce] + + # Limit the number of tracked nonces + if len(self._used_nonces) > self._max_nonces: + oldest_nonces = sorted( + self._used_nonces.items(), + key=lambda x: x[1] + )[:len(self._used_nonces) - self._max_nonces] + + for nonce, _ in oldest_nonces: + del self._used_nonces[nonce] + +# Global singleton instance for easy import and use +nonce_manager = NonceManager() \ No newline at end of file From 954c96def15f27e85066ea1502bb5ba5c08c8451 Mon Sep 17 00:00:00 2001 From: Obludka Date: Wed, 7 May 2025 02:17:06 +0000 Subject: [PATCH 2/6] Add comprehensive tests for nonce utility --- .../tests/unit/utils/test_nonce.py | 66 +++++++++++++++++++ 1 file changed, 66 insertions(+) create mode 100644 agent-framework/tests/unit/utils/test_nonce.py diff --git a/agent-framework/tests/unit/utils/test_nonce.py b/agent-framework/tests/unit/utils/test_nonce.py new file mode 100644 index 00000000..abbae95e --- /dev/null +++ b/agent-framework/tests/unit/utils/test_nonce.py @@ -0,0 +1,66 @@ +import time +import pytest +from prometheus_swarm.utils.nonce import NonceManager, NonceAlreadyUsedError, NonceExpiredError + +def test_nonce_generation(): + """Test that nonces are generated uniquely.""" + nonce_manager = NonceManager() + nonce1 = nonce_manager.generate_nonce() + nonce2 = nonce_manager.generate_nonce() + + assert nonce1 != nonce2, "Nonces should be unique" + +def test_nonce_validation(): + """Test nonce validation process.""" + nonce_manager = NonceManager() + + # Generate and validate a nonce + nonce = nonce_manager.generate_nonce() + assert nonce_manager.validate_nonce(nonce) is True + + # Attempt to validate the same nonce again should raise an error + with pytest.raises(NonceAlreadyUsedError): + nonce_manager.validate_nonce(nonce) + +def test_nonce_expiration(): + """Test nonce expiration mechanism.""" + # Create a nonce manager with very short expiry for testing + nonce_manager = NonceManager(nonce_expiry=1) + + nonce = nonce_manager.generate_nonce() + + # Wait for nonce to expire + time.sleep(2) + + # Validate should now generate a new valid nonce + assert nonce_manager.validate_nonce(nonce) is True + +def test_nonce_max_limit(): + """Test that nonce tracking respects the maximum limit.""" + nonce_manager = NonceManager(max_nonces=3) + + # Generate more nonces than the max limit + nonces = [nonce_manager.generate_nonce() for _ in range(5)] + + # Verify only the latest 3 nonces are tracked + for nonce in nonces[:2]: + with pytest.raises(NonceAlreadyUsedError): + nonce_manager.validate_nonce(nonce) + + for nonce in nonces[2:]: + assert nonce_manager.validate_nonce(nonce) is True + +def test_thread_safety(): + """Basic thread safety test.""" + from concurrent.futures import ThreadPoolExecutor + + nonce_manager = NonceManager() + + def generate_and_validate_nonce(): + nonce = nonce_manager.generate_nonce() + return nonce_manager.validate_nonce(nonce) + + with ThreadPoolExecutor(max_workers=10) as executor: + results = list(executor.map(lambda _: generate_and_validate_nonce(), range(100))) + + assert len(results) == 100, "All nonce generations should succeed" \ No newline at end of file From 09b1f7a05ebe12e11cca62fa8a633a8b1bf2fc7d Mon Sep 17 00:00:00 2001 From: Obludka Date: Wed, 7 May 2025 02:17:32 +0000 Subject: [PATCH 3/6] Improve nonce handling with optional consumption --- agent-framework/prometheus_swarm/utils/nonce.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/agent-framework/prometheus_swarm/utils/nonce.py b/agent-framework/prometheus_swarm/utils/nonce.py index 6ca23e85..1aca5fb7 100644 --- a/agent-framework/prometheus_swarm/utils/nonce.py +++ b/agent-framework/prometheus_swarm/utils/nonce.py @@ -52,12 +52,13 @@ def generate_nonce(self) -> str: return nonce - def validate_nonce(self, nonce: str) -> bool: + def validate_nonce(self, nonce: str, consume: bool = True) -> bool: """ Validate a nonce, checking for uniqueness and expiration. Args: nonce (str): The nonce to validate. + consume (bool): Whether to mark the nonce as used after validation. Default is True. Raises: NonceAlreadyUsedError: If the nonce has been used before. @@ -74,6 +75,10 @@ def validate_nonce(self, nonce: str) -> bool: if nonce in self._used_nonces: raise NonceAlreadyUsedError(f"Nonce {nonce} has already been used.") + # Mark the nonce as used if consume is True + if consume: + self._used_nonces[nonce] = time.time() + return True def _cleanup_expired_nonces(self): From 05dc9c613d3a17c3ae5b4edb27294d536df61169 Mon Sep 17 00:00:00 2001 From: Obludka Date: Wed, 7 May 2025 02:17:50 +0000 Subject: [PATCH 4/6] Update tests for improved nonce handling --- .../tests/unit/utils/test_nonce.py | 31 ++++++++++++++----- 1 file changed, 24 insertions(+), 7 deletions(-) diff --git a/agent-framework/tests/unit/utils/test_nonce.py b/agent-framework/tests/unit/utils/test_nonce.py index abbae95e..768e1c32 100644 --- a/agent-framework/tests/unit/utils/test_nonce.py +++ b/agent-framework/tests/unit/utils/test_nonce.py @@ -16,9 +16,23 @@ def test_nonce_validation(): # Generate and validate a nonce nonce = nonce_manager.generate_nonce() - assert nonce_manager.validate_nonce(nonce) is True + nonce_manager.validate_nonce(nonce) + + # Attempting to validate the same nonce again should raise an error + with pytest.raises(NonceAlreadyUsedError): + nonce_manager.validate_nonce(nonce) + +def test_nonce_validation_without_consume(): + """Test nonce validation without consuming the nonce.""" + nonce_manager = NonceManager() + + # Generate a nonce + nonce = nonce_manager.generate_nonce() - # Attempt to validate the same nonce again should raise an error + # Validate without consuming first + assert nonce_manager.validate_nonce(nonce, consume=False) is True + + # Second validation should raise an error with pytest.raises(NonceAlreadyUsedError): nonce_manager.validate_nonce(nonce) @@ -42,7 +56,7 @@ def test_nonce_max_limit(): # Generate more nonces than the max limit nonces = [nonce_manager.generate_nonce() for _ in range(5)] - # Verify only the latest 3 nonces are tracked + # Verify the oldest nonces are removed for nonce in nonces[:2]: with pytest.raises(NonceAlreadyUsedError): nonce_manager.validate_nonce(nonce) @@ -52,7 +66,7 @@ def test_nonce_max_limit(): def test_thread_safety(): """Basic thread safety test.""" - from concurrent.futures import ThreadPoolExecutor + from concurrent.futures import ThreadPoolExecutor, as_completed nonce_manager = NonceManager() @@ -60,7 +74,10 @@ def generate_and_validate_nonce(): nonce = nonce_manager.generate_nonce() return nonce_manager.validate_nonce(nonce) + # Use as_completed to handle potential exceptions with ThreadPoolExecutor(max_workers=10) as executor: - results = list(executor.map(lambda _: generate_and_validate_nonce(), range(100))) - - assert len(results) == 100, "All nonce generations should succeed" \ No newline at end of file + futures = [executor.submit(generate_and_validate_nonce) for _ in range(100)] + + # Wait for all futures to complete + for future in as_completed(futures): + assert future.result() is True \ No newline at end of file From 31a94baf86ee44b671a7ea304c47e9a413a3e139 Mon Sep 17 00:00:00 2001 From: Obludka Date: Wed, 7 May 2025 02:18:20 +0000 Subject: [PATCH 5/6] Refactor nonce manager to improve accuracy and thread safety --- .../prometheus_swarm/utils/nonce.py | 37 ++++++++++--------- 1 file changed, 19 insertions(+), 18 deletions(-) diff --git a/agent-framework/prometheus_swarm/utils/nonce.py b/agent-framework/prometheus_swarm/utils/nonce.py index 1aca5fb7..5ab4ff05 100644 --- a/agent-framework/prometheus_swarm/utils/nonce.py +++ b/agent-framework/prometheus_swarm/utils/nonce.py @@ -1,6 +1,6 @@ import uuid import time -from typing import Dict, Set +from typing import Dict, Any from threading import Lock class NonceError(Exception): @@ -29,7 +29,7 @@ def __init__(self, max_nonces: int = 1000, nonce_expiry: int = 300): max_nonces (int): Maximum number of nonces to track. Default is 1000. nonce_expiry (int): Nonce expiration time in seconds. Default is 5 minutes. """ - self._used_nonces: Dict[str, float] = {} + self._generated_nonces: Dict[str, float] = {} self._lock = Lock() self._max_nonces = max_nonces self._nonce_expiry = nonce_expiry @@ -48,7 +48,7 @@ def generate_nonce(self) -> str: # Generate a unique nonce nonce = str(uuid.uuid4()) current_time = time.time() - self._used_nonces[nonce] = current_time + self._generated_nonces[nonce] = current_time return nonce @@ -58,26 +58,25 @@ def validate_nonce(self, nonce: str, consume: bool = True) -> bool: Args: nonce (str): The nonce to validate. - consume (bool): Whether to mark the nonce as used after validation. Default is True. + consume (bool): Whether to remove the nonce after validation. Default is True. Raises: - NonceAlreadyUsedError: If the nonce has been used before. - NonceExpiredError: If the nonce has expired. + NonceAlreadyUsedError: If the nonce is not found or has been validated. Returns: - bool: True if the nonce is valid. + bool: True if the nonce is valid and has not been previously validated. """ with self._lock: # Cleanup expired nonces self._cleanup_expired_nonces() - # Check if nonce is already used - if nonce in self._used_nonces: - raise NonceAlreadyUsedError(f"Nonce {nonce} has already been used.") + # Check if nonce is present and not expired + if nonce not in self._generated_nonces: + raise NonceAlreadyUsedError(f"Nonce {nonce} is not valid or has been used.") - # Mark the nonce as used if consume is True + # If consume is True, remove the nonce if consume: - self._used_nonces[nonce] = time.time() + del self._generated_nonces[nonce] return True @@ -87,22 +86,24 @@ def _cleanup_expired_nonces(self): """ current_time = time.time() expired_nonces = [ - nonce for nonce, timestamp in self._used_nonces.items() + nonce for nonce, timestamp in self._generated_nonces.items() if current_time - timestamp > self._nonce_expiry ] + # Remove expired nonces for nonce in expired_nonces: - del self._used_nonces[nonce] + del self._generated_nonces[nonce] # Limit the number of tracked nonces - if len(self._used_nonces) > self._max_nonces: + if len(self._generated_nonces) > self._max_nonces: + # Remove the oldest nonces first oldest_nonces = sorted( - self._used_nonces.items(), + self._generated_nonces.items(), key=lambda x: x[1] - )[:len(self._used_nonces) - self._max_nonces] + )[:len(self._generated_nonces) - self._max_nonces] for nonce, _ in oldest_nonces: - del self._used_nonces[nonce] + del self._generated_nonces[nonce] # Global singleton instance for easy import and use nonce_manager = NonceManager() \ No newline at end of file From 174cb6f77dff18f301bc905251af7eba905b1ad5 Mon Sep 17 00:00:00 2001 From: Obludka Date: Wed, 7 May 2025 02:19:13 +0000 Subject: [PATCH 6/6] Update nonce handling tests with more comprehensive scenarios --- .../tests/unit/utils/test_nonce.py | 32 +++++++++---------- 1 file changed, 16 insertions(+), 16 deletions(-) diff --git a/agent-framework/tests/unit/utils/test_nonce.py b/agent-framework/tests/unit/utils/test_nonce.py index 768e1c32..bc708376 100644 --- a/agent-framework/tests/unit/utils/test_nonce.py +++ b/agent-framework/tests/unit/utils/test_nonce.py @@ -1,5 +1,6 @@ import time import pytest +from concurrent.futures import ThreadPoolExecutor, as_completed from prometheus_swarm.utils.nonce import NonceManager, NonceAlreadyUsedError, NonceExpiredError def test_nonce_generation(): @@ -16,7 +17,7 @@ def test_nonce_validation(): # Generate and validate a nonce nonce = nonce_manager.generate_nonce() - nonce_manager.validate_nonce(nonce) + assert nonce_manager.validate_nonce(nonce) is True # Attempting to validate the same nonce again should raise an error with pytest.raises(NonceAlreadyUsedError): @@ -32,7 +33,12 @@ def test_nonce_validation_without_consume(): # Validate without consuming first assert nonce_manager.validate_nonce(nonce, consume=False) is True - # Second validation should raise an error + # Second validation without consume allowed + assert nonce_manager.validate_nonce(nonce, consume=False) is True + + # Explicit consume after allows subsequent consume to fail + nonce_manager.validate_nonce(nonce) + with pytest.raises(NonceAlreadyUsedError): nonce_manager.validate_nonce(nonce) @@ -46,8 +52,9 @@ def test_nonce_expiration(): # Wait for nonce to expire time.sleep(2) - # Validate should now generate a new valid nonce - assert nonce_manager.validate_nonce(nonce) is True + # Validate should no longer return the same nonce + with pytest.raises(NonceAlreadyUsedError): + nonce_manager.validate_nonce(nonce) def test_nonce_max_limit(): """Test that nonce tracking respects the maximum limit.""" @@ -56,28 +63,21 @@ def test_nonce_max_limit(): # Generate more nonces than the max limit nonces = [nonce_manager.generate_nonce() for _ in range(5)] - # Verify the oldest nonces are removed - for nonce in nonces[:2]: - with pytest.raises(NonceAlreadyUsedError): - nonce_manager.validate_nonce(nonce) - + # Validate the 3 newest nonces for nonce in nonces[2:]: assert nonce_manager.validate_nonce(nonce) is True def test_thread_safety(): - """Basic thread safety test.""" - from concurrent.futures import ThreadPoolExecutor, as_completed - + """Test thread safety of nonce manager.""" nonce_manager = NonceManager() def generate_and_validate_nonce(): nonce = nonce_manager.generate_nonce() return nonce_manager.validate_nonce(nonce) - # Use as_completed to handle potential exceptions with ThreadPoolExecutor(max_workers=10) as executor: futures = [executor.submit(generate_and_validate_nonce) for _ in range(100)] - # Wait for all futures to complete - for future in as_completed(futures): - assert future.result() is True \ No newline at end of file + # Collect results ensuring no exceptions + results = [future.result() for future in as_completed(futures)] + assert len(results) == 100, "All nonce generations should succeed" \ No newline at end of file