From 3426f347767b5fc8f69f07eca019d8191ed8fd97 Mon Sep 17 00:00:00 2001
From: Balaji Arun <balaji@aptoslabs.com>
Date: Tue, 12 Aug 2025 12:26:28 -0700
Subject: [PATCH] MultiEd25519: Check signature valid before verifying (#17316)

---
 crates/aptos-crypto/src/multi_ed25519.rs      | 17 ++++++-
 .../src/unit_tests/multi_ed25519_test.rs      | 46 ++++++++++++++++++-
 2 files changed, 60 insertions(+), 3 deletions(-)

diff --git a/crates/aptos-crypto/src/multi_ed25519.rs b/crates/aptos-crypto/src/multi_ed25519.rs
index 5d988c44d6d02..cdbfb25bdee51 100644
--- a/crates/aptos-crypto/src/multi_ed25519.rs
+++ b/crates/aptos-crypto/src/multi_ed25519.rs
@@ -525,7 +525,8 @@ impl Signature for MultiEd25519Signature {
                 ))
             },
         };
-        if bitmap_count_ones(self.bitmap) < public_key.threshold as u32 {
+        let num_ones_in_bitmap = bitmap_count_ones(self.bitmap);
+        if num_ones_in_bitmap < public_key.threshold as u32 {
             return Err(anyhow!(
                 "{}",
                 CryptoMaterialError::BitVecError(
@@ -533,13 +534,25 @@ impl Signature for MultiEd25519Signature {
                 )
             ));
         }
+        if num_ones_in_bitmap != self.signatures.len() as u32 {
+            return Err(anyhow!(
+                "{}",
+                CryptoMaterialError::BitVecError(
+                    "Bitmap ones and signatures count are not equal".to_string()
+                )
+            ));
+        }
         let mut bitmap_index = 0;
         // TODO: Eventually switch to deterministic batch verification
         for sig in &self.signatures {
             while !bitmap_get_bit(self.bitmap, bitmap_index) {
                 bitmap_index += 1;
             }
-            sig.verify_arbitrary_msg(message, &public_key.public_keys[bitmap_index])?;
+            let pk = public_key
+                .public_keys
+                .get(bitmap_index)
+                .ok_or_else(|| anyhow::anyhow!("Public key index {bitmap_index} out of bounds"))?;
+            sig.verify_arbitrary_msg(message, pk)?;
             bitmap_index += 1;
         }
         Ok(())
diff --git a/crates/aptos-crypto/src/unit_tests/multi_ed25519_test.rs b/crates/aptos-crypto/src/unit_tests/multi_ed25519_test.rs
index 8622a32734968..addf4d264a201 100644
--- a/crates/aptos-crypto/src/unit_tests/multi_ed25519_test.rs
+++ b/crates/aptos-crypto/src/unit_tests/multi_ed25519_test.rs
@@ -472,7 +472,7 @@ fn test_multi_ed25519_signature_verification() {
 fn test_invalid_multi_ed25519_signature_bitmap() {
     let priv_keys_3 = generate_keys(3);
 
-    let multi_private_key_2of3 = MultiEd25519PrivateKey::new(priv_keys_3, 2).unwrap();
+    let multi_private_key_2of3 = MultiEd25519PrivateKey::new(priv_keys_3.clone(), 2).unwrap();
     let multi_public_key_2of3 = MultiEd25519PublicKey::from(&multi_private_key_2of3);
 
     let multi_signature_2of3 = multi_private_key_2of3.sign(message()).unwrap();
@@ -502,6 +502,50 @@ fn test_invalid_multi_ed25519_signature_bitmap() {
     assert!(multi_signature_1of3
         .verify(message(), &multi_public_key_2of3)
         .is_err());
+
+    let multi_signature_3of3_invalid_bitmap = MultiEd25519Signature::new_with_signatures_and_bitmap(
+        priv_keys_3
+            .iter()
+            .map(|x| x.sign(message()).unwrap())
+            .collect(),
+        [0b1100_0000, 0u8, 0u8, 0u8],
+    );
+
+    // Bitmap is invalid
+    assert_eq!(
+        multi_signature_3of3_invalid_bitmap
+            .verify(message(), &multi_public_key_2of3)
+            .unwrap_err()
+            .to_string(),
+        "BitVecError(\"Bitmap ones and signatures count are not equal\")".to_string()
+    );
+
+    let multi_signature_3of3 = MultiEd25519Signature::new_with_signatures_and_bitmap(
+        priv_keys_3
+            .iter()
+            .map(|x| x.sign(message()).unwrap())
+            .collect(),
+        [0b1110_0000, 0u8, 0u8, 0u8],
+    );
+    let multi_public_key_2of3_2only = MultiEd25519PublicKey::new(
+        multi_public_key_2of3
+            .public_keys()
+            .iter()
+            .take(2)
+            .cloned()
+            .collect(),
+        2,
+    )
+    .unwrap();
+
+    // Bitmap is valid but fewer public keys provided
+    assert_eq!(
+        multi_signature_3of3
+            .verify(message(), &multi_public_key_2of3_2only)
+            .unwrap_err()
+            .to_string(),
+        "BitVecError(\"Signature index is out of range\")".to_string()
+    );
 }
 
 /// Used for generating test cases for the MultiEd25519 Move module.