sapio_psbt/
lib.rs

1// Copyright Judica, Inc 2022
2//
3// This Source Code Form is subject to the terms of the Mozilla Public
4//  License, v. 2.0. If a copy of the MPL was not distributed with this
5//  file, You can obtain one at https://mozilla.org/MPL/2.0/.
6
7use bitcoin::consensus::serialize;
8use bitcoin::schnorr::TapTweak;
9use bitcoin::secp256k1::rand::Rng;
10use bitcoin::secp256k1::{rand, Signing, Verification};
11use bitcoin::util::bip32::{ExtendedPubKey, Fingerprint, KeySource};
12use bitcoin::util::sighash::Prevouts;
13use bitcoin::util::taproot::TapLeafHash;
14use bitcoin::util::taproot::TapSighashHash;
15use bitcoin::XOnlyPublicKey;
16use bitcoin::{
17    psbt::PartiallySignedTransaction, secp256k1::Secp256k1, util::bip32::ExtendedPrivKey,
18};
19use bitcoin::{KeyPair, TxOut};
20use bitcoin::{Network, SchnorrSig};
21use std::collections::BTreeMap;
22use std::error::Error;
23use std::fmt::Display;
24pub mod external_api;
25
26pub struct SigningKey(pub Vec<ExtendedPrivKey>);
27
28impl SigningKey {
29    pub fn read_key_from_buf(buf: &[u8]) -> Result<Self, bitcoin::util::bip32::Error> {
30        ExtendedPrivKey::decode(buf).map(|k| SigningKey(vec![k]))
31    }
32    pub fn new_key(network: Network) -> Result<Self, bitcoin::util::bip32::Error> {
33        let seed: [u8; 32] = rand::thread_rng().gen();
34        let xpriv = ExtendedPrivKey::new_master(network, &seed)?;
35        Ok(SigningKey(vec![xpriv]))
36    }
37    pub fn merge(&mut self, other: SigningKey) -> &mut SigningKey {
38        self.0.extend(other.0);
39        self
40    }
41    pub fn pubkey<C: Signing>(&self, secp: &Secp256k1<C>) -> Vec<ExtendedPubKey> {
42        self.0
43            .iter()
44            .map(|s| ExtendedPubKey::from_priv(secp, s))
45            .collect()
46    }
47    pub fn sign(
48        &self,
49        mut psbt: PartiallySignedTransaction,
50        hash_ty: bitcoin::SchnorrSighashType,
51    ) -> Result<Vec<u8>, PSBTSigningError> {
52        self.sign_psbt_mut(&mut psbt, &Secp256k1::new(), hash_ty)?;
53        let bytes = serialize(&psbt);
54        Ok(bytes)
55    }
56    pub fn sign_psbt<C: Signing + Verification>(
57        &self,
58        mut psbt: PartiallySignedTransaction,
59        secp: &Secp256k1<C>,
60        hash_ty: bitcoin::SchnorrSighashType,
61    ) -> Result<PartiallySignedTransaction, (PartiallySignedTransaction, PSBTSigningError)> {
62        match self.sign_psbt_mut(&mut psbt, secp, hash_ty) {
63            Ok(()) => Ok(psbt),
64            Err(e) => Err((psbt, e)),
65        }
66    }
67    pub fn sign_psbt_mut<C: Signing + Verification>(
68        &self,
69        psbt: &mut PartiallySignedTransaction,
70        secp: &Secp256k1<C>,
71        hash_ty: bitcoin::SchnorrSighashType,
72    ) -> Result<(), PSBTSigningError> {
73        let l = psbt.inputs.len();
74        for idx in 0..l {
75            self.sign_psbt_input_mut(psbt, secp, idx, hash_ty)?;
76        }
77        Ok(())
78    }
79    pub fn sign_psbt_input<C: Signing + Verification>(
80        &self,
81        mut psbt: PartiallySignedTransaction,
82        secp: &Secp256k1<C>,
83        idx: usize,
84        hash_ty: bitcoin::SchnorrSighashType,
85    ) -> Result<PartiallySignedTransaction, (PartiallySignedTransaction, PSBTSigningError)> {
86        match self.sign_psbt_input_mut(&mut psbt, secp, idx, hash_ty) {
87            Ok(()) => Ok(psbt),
88            Err(e) => Err((psbt, e)),
89        }
90    }
91    pub fn sign_psbt_input_mut<C: Signing + Verification>(
92        &self,
93        psbt: &mut PartiallySignedTransaction,
94        secp: &Secp256k1<C>,
95        idx: usize,
96        hash_ty: bitcoin::SchnorrSighashType,
97    ) -> Result<(), PSBTSigningError> {
98        let tx = psbt.clone().extract_tx();
99        let utxos: Vec<TxOut> = psbt
100            .inputs
101            .iter()
102            .enumerate()
103            .map(|(i, o)| {
104                if let Some(ref utxo) = o.witness_utxo {
105                    Ok(utxo.clone())
106                } else {
107                    Err(i)
108                }
109            })
110            .collect::<Result<Vec<TxOut>, usize>>()
111            .map_err(PSBTSigningError::NoUTXOAtIndex)?;
112        let mut sighash = bitcoin::util::sighash::SighashCache::new(&tx);
113        let input = &mut psbt
114            .inputs
115            .get_mut(idx)
116            .ok_or(PSBTSigningError::NoInputAtIndex(idx))?;
117        let prevouts = &Prevouts::All(&utxos);
118        let fingerprints_map = self.compute_fingerprint_map(secp);
119        self.sign_taproot_top_key(
120            secp,
121            input,
122            &mut sighash,
123            prevouts,
124            hash_ty,
125            &fingerprints_map,
126        );
127        self.sign_all_tapleaf_branches(
128            secp,
129            input,
130            &mut sighash,
131            prevouts,
132            hash_ty,
133            &fingerprints_map,
134        );
135        Ok(())
136    }
137
138    fn sign_all_tapleaf_branches<C: Signing + Verification>(
139        &self,
140        secp: &Secp256k1<C>,
141        input: &mut bitcoin::psbt::Input,
142        sighash: &mut bitcoin::util::sighash::SighashCache<&bitcoin::Transaction>,
143        prevouts: &Prevouts<TxOut>,
144        hash_ty: bitcoin::SchnorrSighashType,
145        fingerprints_map: &Vec<(Fingerprint, &ExtendedPrivKey)>,
146    ) {
147        let signers = self.compute_matching_keys(secp, &input.tap_key_origins, fingerprints_map);
148        for (kp, vtlh) in signers {
149            for tlh in vtlh {
150                let sig = get_sig(
151                    sighash,
152                    prevouts,
153                    hash_ty,
154                    secp,
155                    &kp,
156                    &Some((*tlh, DEFAULT_CODESEP)),
157                );
158                input
159                    .tap_script_sigs
160                    .insert((kp.x_only_public_key().0, *tlh), sig);
161            }
162        }
163    }
164
165    fn sign_taproot_top_key<C: Signing + Verification>(
166        &self,
167        secp: &Secp256k1<C>,
168        input: &mut bitcoin::psbt::Input,
169        sighash: &mut bitcoin::util::sighash::SighashCache<&bitcoin::Transaction>,
170        prevouts: &Prevouts<TxOut>,
171        hash_ty: bitcoin::SchnorrSighashType,
172        fingerprints_map: &Vec<(Fingerprint, &ExtendedPrivKey)>,
173    ) -> Option<()> {
174        // first attempt to use derivations from the key source map
175        let key = input.tap_internal_key?;
176        let untweaked = self.find_internal_keypair(input, key, fingerprints_map, secp)?;
177        let tweaked = untweaked
178            .tap_tweak(secp, input.tap_merkle_root)
179            .into_inner();
180        input.tap_key_sig = Some(get_sig(sighash, prevouts, hash_ty, secp, &tweaked, &None));
181        Some(())
182    }
183
184    fn find_internal_keypair<C: Signing>(
185        &self,
186        input: &mut bitcoin::psbt::Input,
187        input_key: XOnlyPublicKey,
188        fingerprints_map: &Vec<(Fingerprint, &ExtendedPrivKey)>,
189        secp: &Secp256k1<C>,
190    ) -> Option<KeyPair> {
191        // Assume that the key is an exact, non derived, match for a key we know already
192        for kp in self.0.iter() {
193            let untweaked = kp.to_keypair(secp);
194            let pk = XOnlyPublicKey::from_keypair(&untweaked);
195            if input_key == pk.0 {
196                return Some(untweaked);
197            }
198        }
199        // Otherwise, try to derive a key
200        let (_, (f, path)) = input.tap_key_origins.get(&input_key)?;
201        let idx = fingerprints_map.partition_point(|(x, _)| *x < *f);
202        for (_, key) in fingerprints_map.iter().skip(idx).take_while(|k| k.0 == *f) {
203            if let Ok(sk) = key.derive_priv(secp, path) {
204                let untweaked = sk.to_keypair(secp);
205                let pk = untweaked.public_key().x_only_public_key().0;
206                if pk == input_key {
207                    return Some(untweaked);
208                }
209            }
210        }
211        None
212    }
213
214    /// Compute keypairs for all matching fingerprints
215    fn compute_matching_keys<'a, C: Signing>(
216        &'a self,
217        secp: &'a Secp256k1<C>,
218        input: &'a BTreeMap<XOnlyPublicKey, (Vec<TapLeafHash>, KeySource)>,
219        fingerprints_map: &'a Vec<(Fingerprint, &'a ExtendedPrivKey)>,
220    ) -> impl Iterator<Item = (KeyPair, &'a Vec<TapLeafHash>)> + 'a {
221        // TODO: Cache this on type creation?
222        input.iter().filter_map(move |(x, (vlth, (f, path)))| {
223            let idx = fingerprints_map.partition_point(|(x, _)| *x < *f);
224            for (_, key) in fingerprints_map
225                .iter()
226                .skip(idx)
227                .take_while(|(x, _)| *x == *f)
228            {
229                match key.derive_priv(secp, path).map(|k| k.to_keypair(secp)) {
230                    Ok(kp) => {
231                        if kp.public_key().x_only_public_key().0 == *x {
232                            return Some((kp, vlth));
233                        } else {
234                            return None;
235                        }
236                    }
237                    Err(_) => continue,
238                }
239            }
240            None
241        })
242    }
243
244    /// Computes a map of all fingerprints
245    // TODO: consider more memory efficient representations
246    fn compute_fingerprint_map<'a, C: Signing>(
247        &'a self,
248        secp: &Secp256k1<C>,
249    ) -> Vec<(Fingerprint, &'a ExtendedPrivKey)> {
250        let fingerprint = self.0.iter().map(|k| (k.fingerprint(secp), k));
251        let mut keyarr: Vec<(Fingerprint, &ExtendedPrivKey)> = fingerprint.collect();
252        keyarr.sort_by_key(|k| k.0);
253        keyarr
254    }
255}
256
257#[derive(Debug, Clone)]
258pub enum PSBTSigningError {
259    NoUTXOAtIndex(usize),
260    NoInputAtIndex(usize),
261}
262
263impl Display for PSBTSigningError {
264    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
265        write!(f, "{:?}", self)
266    }
267}
268impl Error for PSBTSigningError {}
269
270const DEFAULT_CODESEP: u32 = 0xffff_ffff;
271fn get_sig<C: Signing>(
272    sighash: &mut bitcoin::util::sighash::SighashCache<&bitcoin::Transaction>,
273    prevouts: &Prevouts<TxOut>,
274    hash_ty: bitcoin::SchnorrSighashType,
275    secp: &Secp256k1<C>,
276    kp: &bitcoin::KeyPair,
277    path: &Option<(TapLeafHash, u32)>,
278) -> SchnorrSig {
279    let annex = None;
280    let sighash: TapSighashHash = sighash
281        .taproot_signature_hash(0, prevouts, annex, *path, hash_ty)
282        .expect("Signature hash cannot fail...");
283    let msg = bitcoin::secp256k1::Message::from_slice(&sighash[..]).expect("Size must be correct.");
284    let sig = secp.sign_schnorr_no_aux_rand(&msg, kp);
285    SchnorrSig { sig, hash_ty }
286}
sapio_psbt/lib.rs

sapio_psbt/
lib.rs
