From ea85067bebd5fcaaa779096e864fc636f154cb14 Mon Sep 17 00:00:00 2001
From: Khwahish Patel <khwahish.p1@ahduni.edu.in>
Date: Sun, 16 Mar 2025 15:32:25 +0530
Subject: [PATCH] check other cases and add test_rsa.py

---
 libp2p/crypto/rsa.py          | 30 ++++++++++++++++++++----------
 tests/core/crypto/test_rsa.py | 30 ++++++++++++++++++++++++++++++
 2 files changed, 50 insertions(+), 10 deletions(-)
 create mode 100644 tests/core/crypto/test_rsa.py

diff --git a/libp2p/crypto/rsa.py b/libp2p/crypto/rsa.py
index 6673919a..68963989 100644
--- a/libp2p/crypto/rsa.py
+++ b/libp2p/crypto/rsa.py
@@ -22,18 +22,31 @@ from libp2p.crypto.keys import (
 MAX_RSA_KEY_SIZE = 4096
 
 
+def validate_rsa_key_length(key_length: int) -> None:
+    """
+    Validate that the RSA key length is positive and within the allowed maximum.
+
+    :param key_length: RSA key size in bits.
+    :raises CryptographyError:
+        If the key size is not positive or exceeds MAX_RSA_KEY_SIZE.
+    """
+    if key_length <= 0:
+        raise CryptographyError("RSA key size must be positive")
+    if key_length > MAX_RSA_KEY_SIZE:
+        raise CryptographyError(
+            f"RSA key size {key_length} exceeds maximum allowed size {MAX_RSA_KEY_SIZE}"
+        )
+
+
 def validate_rsa_key_size(key: RsaKey) -> None:
     """
     Validate that an RSA key's size is within acceptable bounds.
 
-    :param key: The RSA key to validate
-    :raises CryptographyError: If the key size exceeds the maximum allowed size
+    :param key: The RSA key to validate.
+    :raises CryptographyError: If the key size is invalid.
     """
     key_size = key.size_in_bits()
-    if key_size > MAX_RSA_KEY_SIZE:
-        msg = f"RSA key size {key_size} "
-        msg += f"exceeds maximum allowed size {MAX_RSA_KEY_SIZE}"
-        raise CryptographyError(msg)
+    validate_rsa_key_length(key_size)
 
 
 class RSAPublicKey(PublicKey):
@@ -69,10 +82,7 @@ class RSAPrivateKey(PrivateKey):
 
     @classmethod
     def new(cls, bits: int = 2048, e: int = 65537) -> "RSAPrivateKey":
-        if bits > MAX_RSA_KEY_SIZE:
-            msg = f"Requested RSA key size {bits} "
-            msg += f"exceeds maximum allowed size {MAX_RSA_KEY_SIZE}"
-            raise CryptographyError(msg)
+        validate_rsa_key_length(bits)
         private_key_impl = RSA.generate(bits, e=e)
         return cls(private_key_impl)
 
diff --git a/tests/core/crypto/test_rsa.py b/tests/core/crypto/test_rsa.py
new file mode 100644
index 00000000..21b1af7d
--- /dev/null
+++ b/tests/core/crypto/test_rsa.py
@@ -0,0 +1,30 @@
+import pytest
+
+from libp2p.crypto.exceptions import (
+    CryptographyError,
+)
+from libp2p.crypto.rsa import (
+    MAX_RSA_KEY_SIZE,
+    RSAPrivateKey,
+    validate_rsa_key_size,
+)
+
+
+def test_validate_rsa_key_size():
+    # Test valid key size
+    key = RSAPrivateKey.new(2048)
+    validate_rsa_key_size(key.impl)
+
+    # Test key size too large
+    with pytest.raises(
+        CryptographyError, match=f".*exceeds maximum allowed size {MAX_RSA_KEY_SIZE}"
+    ):
+        RSAPrivateKey.new(MAX_RSA_KEY_SIZE + 1)
+
+    # Test negative key size (this would be caught when creating the key)
+    with pytest.raises(CryptographyError, match="RSA key size must be positive"):
+        RSAPrivateKey.new(-1)
+
+    # Test zero key size
+    with pytest.raises(CryptographyError, match="RSA key size must be positive"):
+        RSAPrivateKey.new(0)