imprv: Add more validation for bundle

src/omemo_dr/exceptions.py

@@ -73,3 +73,7 @@ class InvalidMessage(Exception):
 
 class DuplicateMessage(Exception):
     pass
+
+
+class BundleValidationError(Exception):
+    pass

src/omemo_dr/state/prekeybundle.py

 from __future__ import annotations
 
+from ..const import MAX_INT
 from ..const import NS_OMEMO_2
 from ..const import NS_OMEMO_TMP
 from ..ecc.djbec import CurvePublicKey
 from ..ecc.djbec import EdPublicKey
-from ..exceptions import InvalidKeyException
+from ..exceptions import BundleValidationError
 from ..identitykey import IdentityKey
 from ..structs import OMEMOBundleProto
 
@@ -42,10 +43,19 @@ class PreKeyBundle:
         elif ns == NS_OMEMO_2:
             ik_pub = EdPublicKey.from_bytes(bundle.ik).to_curve()
         else:
-            raise InvalidKeyException("Unknown namespace on bundle: %s", ns)
+            raise BundleValidationError("Unknown namespace on bundle: %s", ns)
 
         ik = IdentityKey(ik_pub)
 
+        if not 1 <= bundle.device_id <= MAX_INT:
+            raise BundleValidationError("Device id out of range")
+
+        if not 1 <= prekey["id"] <= MAX_INT:
+            raise BundleValidationError("Prekey id out of range")
+
+        if not 1 <= bundle.spk["id"] <= MAX_INT:
+            raise BundleValidationError("Signed pre key id out of range")
+
         return cls(
             bundle.device_id,
             bundle.namespace,