Merge #932: wallet_policy: fix set_key_info silently dropping keys for repeated placeholders

1f655a9 wallet_policy: add regression test for repeated placeholder key drop (Sanket Kanjalkar)
716f0ce wallet_policy: fix set_key_info accepting extra keys for repeated placeholders (Sanket Kanjalkar)

Pull request description:

  `set_key_info()` counted key-expression occurrences instead of unique placeholder indices. A template reusing `@0` with disjoint paths (e.g. `@0/**,@0/<2;3>/*`) would accept extra keys that are silently dropped during descriptor translation. Fix by counting unique placeholder indices via `BTreeSet`.
  
  I'm not super familiar with wallet_policy/BIP-388, but this looks like a bug to me. Credit to jmeccom for reporting this.


ACKs for top commit:
  apoelstra:
    ACK 1f655a9; successfully ran local tests


Tree-SHA512: c1399f6c3225482d796bd32c457b096a2b69ed75129a4ea7c6250ceef2b760a8e583e52c1e0e15e0874dabeca8adf7f15c0c4ec0ea2491b1611d886b5229ffad

Author: apoelstra

src/descriptor/wallet_policy/mod.rs

@@ -5,7 +5,7 @@ use core::str::FromStr;
 
 use super::key::XKeyParseError;
 use super::{DerivPaths, DescriptorKeyParseError, Wildcard};
-use crate::{Descriptor, DescriptorPublicKey, String, Translator, Vec};
+use crate::{BTreeSet, Descriptor, DescriptorPublicKey, String, Translator, Vec};
 
 mod key_expression;
 
@@ -119,7 +119,9 @@ impl WalletPolicy {
     /// Sets the key information so that `WalletPolicy::into_descriptor` can be
     /// called successfully. Errors when there are not enough keys for the template.
     pub fn set_key_info(&mut self, keys: &[DescriptorPublicKey]) -> Result<(), WalletPolicyError> {
-        if keys.len() != self.template.iter_pk().count() {
+        let unique_placeholders: BTreeSet<u32> =
+            self.template.iter_pk().map(|k| k.index.0).collect();
+        if keys.len() != unique_placeholders.len() {
             return Err(WalletPolicyError::WalletPolicyInvalidKeyInfo);
         }
         self.key_info = keys.to_vec();
@@ -377,4 +379,27 @@ mod tests {
         template_only.set_key_info(&keys).unwrap();
         assert!(template_only.clone().into_descriptor().is_ok());
     }
+
+    // Regression test for a bug where set_key_info() counted key-expression
+    // occurrences instead of unique placeholder indices. A template reusing @0
+    // with disjoint paths (e.g. @0/**,@0/<2;3>/*) has 2 occurrences but only 1
+    // unique placeholder. The old code accepted 2 keys, silently dropping the
+    // second one during descriptor translation since both @0 resolve to
+    // key_info[0].
+    // Reported by jm@squareup.com
+    #[test]
+    fn set_key_info_rejects_extra_keys_for_repeated_placeholder() {
+        let mut policy = WalletPolicy::from_str("wsh(sortedmulti(2,@0/**,@0/<2;3>/*))").unwrap();
+        let attacker_key: DescriptorPublicKey = "[6738736c/48'/0'/0'/2']xpub6FC1fXFP1GXLX5TKtcjHGT4q89SDRehkQLtbKJ2PzWcvbBHtyDsJPLtpLtkGqYNYZdVVAjRQ5kug9CsapegmmeRutpP7PW4u4wVF9JfkDhw".parse().unwrap();
+        let victim_key: DescriptorPublicKey = "[b2b1f0cf/48'/0'/0'/2']xpub6EWhjpPa6FqrcaPBuGBZRJVjzGJ1ZsMygRF26RwN932Vfkn1gyCiTbECVitBjRCkexEvetLdiqzTcYimmzYxyR1BZ79KNevgt61PDcukmC7".parse().unwrap();
+
+        // Must reject: 2 keys provided but only 1 unique placeholder (@0)
+        assert_eq!(
+            policy.set_key_info(&[attacker_key.clone(), victim_key]),
+            Err(WalletPolicyError::WalletPolicyInvalidKeyInfo),
+        );
+
+        // Must accept: 1 key for 1 unique placeholder
+        assert!(policy.set_key_info(&[attacker_key]).is_ok());
+    }
 }