Merge pull request #66 from opentensor/feat/roman/extend-encrypt_mlkem768-logic

Extend `encrypt_mlkem768` logic

Author: basfroman

Cargo.lock

@@ -1,6 +1,6 @@
 [package]
 name = "bittensor-drand"
-version = "1.2.0"
+version = "1.3.0"
 edition = "2021"
 
 [lib]
@@ -23,6 +23,7 @@ libc = "0.2.172"
 # ML-KEM dependencies
 ml-kem = "0.2.1"
 chacha20poly1305 = "0.10"
+twox-hash = { version = "2.1", default-features = false, features = ["xxhash64"] }
 
 [features]
 default = ["extension-module"]

Cargo.toml

@@ -179,26 +179,28 @@ def get_latest_round() -> int:
     return _get_latest_round()
 
 
-def encrypt_mlkem768(pk_bytes: bytes, plaintext: bytes) -> bytes:
+def encrypt_mlkem768(pk_bytes: bytes, plaintext: bytes, include_key_hash: bool = False) -> bytes:
     """Encrypts data using ML-KEM-768 + XChaCha20Poly1305.
 
     This function encrypts plaintext using ML-KEM-768 key encapsulation followed by XChaCha20Poly1305 authenticated
     encryption. The public key is rotated every block and can be queried from the NextKey storage item.
 
-    Blob format: [u16 kem_len LE][kem_ct][nonce24][aead_ct]
+    Blob format (include_key_hash=False): [u16 kem_len LE][kem_ct][nonce24][aead_ct]
+    Blob format (include_key_hash=True):  [key_hash(16)][u16 kem_len LE][kem_ct][nonce24][aead_ct]
 
     Arguments:
         pk_bytes: ML-KEM-768 public key bytes (from NextKey storage, 1184 bytes)
-        plaintext: Data to encrypt. For MEV Shield, this should be: payload_core + b"\\x01" + signature where
-            payload_core = signer_bytes (32B) + key_hash_bytes (32B) + SCALE(call)
+        plaintext: Data to encrypt.
+        include_key_hash: If True, prepends the twox_128 hash of pk_bytes (16 bytes) to the output.
+            Required for the MEV Shield wire format (pallet-shield v2).
 
     Returns:
         bytes: Encrypted blob
 
     Raises:
         ValueError: If encryption fails (invalid public key, buffer too small, etc.)
     """
-    return _encrypt_mlkem768(pk_bytes, plaintext)
+    return _encrypt_mlkem768(pk_bytes, plaintext, include_key_hash)
 
 
 def mlkem_kdf_id() -> bytes:

bittensor_drand/__init__.py

@@ -1,6 +1,6 @@
 [project]
 name = "bittensor-drand"
-version = "1.2.0"
+version = "1.3.0"
 description = ""
 readme = "README.md"
 license = {file = "LICENSE"}