Merge b7496a18a8 into 4fc2bba1f2

litellm/proxy/auth/litellm_license.py

@@ -176,7 +176,7 @@ class LicenseCheck:
     def verify_license_without_api_request(self, public_key, license_key):
         try:
             from cryptography.hazmat.primitives import hashes
-            from cryptography.hazmat.primitives.asymmetric import padding
+            from cryptography.hazmat.primitives.asymmetric import padding, rsa
 
             from litellm.proxy._types import EnterpriseLicenseData
 
@@ -187,7 +187,14 @@ class LicenseCheck:
                 license_key += "=" * (4 - padding_needed)
 
             decoded: Final = base64.b64decode(license_key)
-            message, signature = decoded.split(b".", 1)
+            if not isinstance(public_key, rsa.RSAPublicKey):
+                self.airgapped_license_data = None
+                return False
+            signature_size: Final = (public_key.key_size + 7) // 8
+            if decoded[-signature_size - 1 : -signature_size] != b".":
+                self.airgapped_license_data = None
+                return False
+            message, signature = decoded[: -signature_size - 1], decoded[-signature_size:]
 
             # Verify the signature
             public_key.verify(

tests/unit/proxy/auth/test_litellm_license.py

@@ -1,10 +1,13 @@
+import base64
 import json
 import logging
+from types import MappingProxyType
 from typing import Final
 from unittest.mock import MagicMock
 
 import httpx
 import pytest
+from cryptography.hazmat.primitives import serialization
 from cryptography.hazmat.primitives.asymmetric.rsa import RSAPublicKey
 
 from litellm.llms.custom_httpx.http_handler import HTTPHandler
@@ -70,18 +73,20 @@ def test_auto_router_capability_limit() -> None:
 
 
 def _signed_license(
-    expiration_date: str, allowed_features: tuple[str, ...] = ("auto_router",)
+    expiration_date: str,
+    allowed_features: tuple[str, ...] = ("auto_router",),
+    *,
+    user_id: str = "u",
+    key_size: int = 2048,
 ) -> tuple[RSAPublicKey, str]:
-    import base64
-
     from cryptography.hazmat.primitives import hashes
     from cryptography.hazmat.primitives.asymmetric import padding, rsa
 
-    private_key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
-    message = json.dumps(
-        {"expiration_date": expiration_date, "user_id": "u", "allowed_features": list(allowed_features)}
+    private_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=key_size)
+    message: Final = json.dumps(
+        {"expiration_date": expiration_date, "user_id": user_id, "allowed_features": list(allowed_features)}
     ).encode()
-    signature = private_key.sign(
+    signature: Final = private_key.sign(
         message,
         padding.PSS(mgf=padding.MGF1(hashes.SHA256()), salt_length=padding.PSS.MAX_LENGTH),
         hashes.SHA256(),
@@ -89,6 +94,83 @@ def _signed_license(
     return private_key.public_key(), base64.b64encode(message + b"." + signature).decode()
 
 
+@pytest.mark.parametrize("key_size", [2048, 2050, 3072])
+@pytest.mark.parametrize(
+    ("user_id", "omit_padding"),
+    [("u", False), ("qa.example.test", False), ("user@qa.example.test", True)],
+)
+def test_offline_license_preserves_dotted_user_ids(key_size: int, user_id: str, omit_padding: bool) -> None:
+    public_key, signed_key = _signed_license("2999-01-01", user_id=user_id, key_size=key_size)
+    license_key: Final = signed_key.rstrip("=") if omit_padding else signed_key
+    license_check: Final = LicenseCheck()
+
+    assert license_check.verify_license_without_api_request(public_key, license_key) is True
+    assert license_check.airgapped_license_data == {
+        "expiration_date": "2999-01-01",
+        "user_id": user_id,
+        "allowed_features": ["auto_router"],
+    }
+
+
+def test_offline_license_accepts_dot_bytes_inside_the_rsa_signature() -> None:
+    public_key: Final = serialization.load_pem_public_key(
+        b"""-----BEGIN PUBLIC KEY-----
+MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEAz8McedOGZ4LV1b7Pyqn0
+U2/LjBgKDiI+RPDC2fuP1ZGO5FWay/uxq2vb1foUde5W7+EXEpWsbRP18dWzJXXo
+/QWm0VOEHyfK9JPHtblShLK/6MpsEif9ck+iGVoirJz7DhyY4U8PzzHu7TUbivt7
+6LR+ntT1xBtMbZVuVPGclpF/qhwi1hcZaYodzJjdpLvnpl21w4CVtPRHGY6dKgs7
+BFRzfxIXqbPReJlUQsDxNfZryKMtTAwnK9veiQMFrBNaA0DmCFmGc34jq671RLqh
+P/pZAoERnJiRvqXEkSfioFjfahT6y+m/L/Feb4phpvlYHWF/dM9m7Lechagg7mWU
+2wIDAQAB
+-----END PUBLIC KEY-----
+"""
+    )
+    message: Final = b'{"expiration_date":"2999-01-01","user_id":"u"}'
+    signature: Final = base64.b64decode(
+        "QwLSziyXpGpf95yT9lp2/L5qq6r9laCKV/IIxbMuciWuLT32a2wqi1118ZxwC54WMVNU51aF9LUfEIx1tr+yriEb09qzQvHNM6bx41agflZ/"
+        "NPbQoIBTjtRbZVc3RqupJmMb0F2mkwATA3lHqLAJ6EEATprnvOuMN8T0ucqBiQqk7FYGQ4spAxPsw/9IEBECfEV0VQX/PujIqJP7gu9qZ2eqI"
+        "+tZLYjx3RfzebRpMfH1/3DpCimDIVs3p7dpSlAfqnv9NXfL20WYVOIVAOIDKto9Not3Ivwq/bgjY2QBFY5CkHvLCe5BIWzWejHIyCHaryqsphUkFx"
+        "+C7Ekmkgkvig=="
+    )
+    license_key: Final = base64.b64encode(message + b"." + signature).decode()
+    license_check: Final = LicenseCheck()
+
+    assert b"." in signature
+    assert license_check.verify_license_without_api_request(public_key, license_key) is True
+    assert license_check.airgapped_license_data == {"expiration_date": "2999-01-01", "user_id": "u"}
+
+
+@pytest.mark.parametrize("corruption", ["separator", "truncated", "appended", "payload", "signature"])
+def test_offline_license_rejects_corruption_and_clears_prior_grants(corruption: str) -> None:
+    public_key, license_key = _signed_license("2999-01-01")
+    decoded: Final = base64.b64decode(license_key)
+    corrupted: Final = MappingProxyType(
+        {
+            "separator": decoded.replace(b".", b"!", 1),
+            "truncated": decoded[:-1],
+            "appended": decoded + b"!",
+            "payload": decoded.replace(b'"u"', b'"v"', 1),
+            "signature": decoded[:-1] + bytes([decoded[-1] ^ 1]),
+        }
+    )[corruption]
+    invalid_key: Final = base64.b64encode(corrupted).decode()
+    license_check: Final = LicenseCheck()
+    assert license_check.verify_license_without_api_request(public_key, license_key) is True
+
+    assert license_check.verify_license_without_api_request(public_key, invalid_key) is False
+    assert license_check.airgapped_license_data is None
+    assert license_check.auto_router_capability_limit() == 1
+
+
+def test_offline_license_without_a_public_key_clears_prior_grants() -> None:
+    public_key, license_key = _signed_license("2999-01-01")
+    license_check: Final = LicenseCheck()
+    assert license_check.verify_license_without_api_request(public_key, license_key) is True
+
+    assert license_check.verify_license_without_api_request(None, license_key) is False
+    assert license_check.airgapped_license_data is None
+
+
 def test_expired_or_unreadable_license_grants_no_features() -> None:
     """The verifier stores the signed payload only after the expiry check passes and clears it when a
     later verify rejects the license, so a stale payload cannot keep lifting the heuristic_v2 limit."""
@@ -270,7 +352,7 @@ def test_local_license_exception_does_not_log_license(
     license_logs: pytest.LogCaptureFixture,
 ) -> None:
     _public_key, license_value = _signed_license("2999-01-01")
-    failing_public_key: Final = MagicMock()
+    failing_public_key: Final = MagicMock(spec=RSAPublicKey, key_size=_public_key.key_size)
     failing_public_key.verify.side_effect = ValueError(license_value)
     license_check: Final = LicenseCheck()
     assert license_check.verify_license_without_api_request(failing_public_key, license_value) is False