rustywallet_taproot/
taproot.rs

1//! Taproot output and address generation
2//!
3//! Implements P2TR address generation from internal keys and script trees.
4
5use crate::error::TaprootError;
6use crate::tagged_hash::TapNodeHash;
7use crate::taptree::TapTree;
8use crate::tweak::tweak_public_key;
9use crate::xonly::{Parity, XOnlyPublicKey};
10use bech32::Hrp;
11
12/// Network for address generation
13#[derive(Clone, Copy, PartialEq, Eq, Debug)]
14pub enum Network {
15    /// Bitcoin mainnet
16    Mainnet,
17    /// Bitcoin testnet
18    Testnet,
19    /// Bitcoin signet
20    Signet,
21    /// Bitcoin regtest
22    Regtest,
23}
24
25impl Network {
26    /// Get the bech32 human-readable part
27    pub fn hrp(&self) -> &'static str {
28        match self {
29            Network::Mainnet => "bc",
30            Network::Testnet | Network::Signet => "tb",
31            Network::Regtest => "bcrt",
32        }
33    }
34}
35
36/// Taproot output key and spending info
37#[derive(Clone, Debug)]
38pub struct TaprootOutput {
39    /// The output public key (tweaked)
40    pub output_key: XOnlyPublicKey,
41    /// Parity of the output key
42    pub parity: Parity,
43    /// Internal key (before tweaking)
44    pub internal_key: XOnlyPublicKey,
45    /// Merkle root (None for key-path only)
46    pub merkle_root: Option<TapNodeHash>,
47}
48
49impl TaprootOutput {
50    /// Create a key-path only output (no script tree)
51    pub fn key_path_only(internal_key: XOnlyPublicKey) -> Result<Self, TaprootError> {
52        let (output_key, parity) = tweak_public_key(&internal_key, None)?;
53
54        Ok(Self {
55            output_key,
56            parity,
57            internal_key,
58            merkle_root: None,
59        })
60    }
61
62    /// Create an output with a script tree
63    pub fn with_script_tree(
64        internal_key: XOnlyPublicKey,
65        tree: &TapTree,
66    ) -> Result<Self, TaprootError> {
67        let merkle_root = tree.root_hash();
68        let (output_key, parity) = tweak_public_key(&internal_key, Some(&merkle_root))?;
69
70        Ok(Self {
71            output_key,
72            parity,
73            internal_key,
74            merkle_root: Some(merkle_root),
75        })
76    }
77
78    /// Create from output key directly (for parsing existing outputs)
79    pub fn from_output_key(output_key: XOnlyPublicKey) -> Self {
80        Self {
81            output_key,
82            parity: Parity::Even, // Unknown, default to even
83            internal_key: output_key, // Unknown, use output key
84            merkle_root: None,
85        }
86    }
87
88    /// Get the P2TR script pubkey (OP_1 <32-byte-key>)
89    pub fn script_pubkey(&self) -> Vec<u8> {
90        let mut script = Vec::with_capacity(34);
91        script.push(0x51); // OP_1 (witness version 1)
92        script.push(0x20); // Push 32 bytes
93        script.extend_from_slice(&self.output_key.serialize());
94        script
95    }
96
97    /// Get the bech32m address
98    pub fn address(&self, network: Network) -> Result<String, TaprootError> {
99        let hrp = Hrp::parse(network.hrp())
100            .map_err(|e| TaprootError::InvalidScript(e.to_string()))?;
101
102        // Witness version 1 + 32-byte key
103        let mut data = Vec::with_capacity(33);
104        data.push(1); // witness version
105        data.extend_from_slice(&self.output_key.serialize());
106
107        bech32::segwit::encode(hrp, bech32::segwit::VERSION_1, &self.output_key.serialize())
108            .map_err(|e| TaprootError::InvalidScript(e.to_string()))
109    }
110
111    /// Check if this is a key-path only output
112    pub fn is_key_path_only(&self) -> bool {
113        self.merkle_root.is_none()
114    }
115}
116
117/// Parse a P2TR address
118pub fn parse_address(address: &str) -> Result<XOnlyPublicKey, TaprootError> {
119    let (_hrp, version, program) = bech32::segwit::decode(address)
120        .map_err(|e| TaprootError::InvalidScript(e.to_string()))?;
121
122    // Check witness version
123    if version != bech32::segwit::VERSION_1 {
124        return Err(TaprootError::InvalidScript(format!(
125            "Expected witness version 1, got {:?}",
126            version
127        )));
128    }
129
130    // Check program length
131    if program.len() != 32 {
132        return Err(TaprootError::InvalidLength {
133            expected: 32,
134            got: program.len(),
135        });
136    }
137
138    XOnlyPublicKey::from_slice(&program)
139}
140
141/// Check if an address is a valid P2TR address
142pub fn is_p2tr_address(address: &str) -> bool {
143    parse_address(address).is_ok()
144}
145
146/// Create a P2TR address from an x-only public key
147pub fn create_address(pubkey: &XOnlyPublicKey, network: Network) -> Result<String, TaprootError> {
148    let output = TaprootOutput::from_output_key(*pubkey);
149    output.address(network)
150}
151
152#[cfg(test)]
153mod tests {
154    use super::*;
155    use secp256k1::{Secp256k1, SecretKey};
156
157    fn get_test_internal_key() -> XOnlyPublicKey {
158        let secp = Secp256k1::new();
159        let secret = [
160            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
161            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
162            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
163            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01,
164        ];
165        let sk = SecretKey::from_slice(&secret).unwrap();
166        let pk = sk.public_key(&secp);
167        let (xonly, _) = pk.x_only_public_key();
168        XOnlyPublicKey::from_inner(xonly)
169    }
170
171    #[test]
172    fn test_key_path_only() {
173        let internal_key = get_test_internal_key();
174        let output = TaprootOutput::key_path_only(internal_key).unwrap();
175
176        assert!(output.is_key_path_only());
177        assert!(output.merkle_root.is_none());
178        assert_ne!(output.output_key, internal_key); // Should be tweaked
179    }
180
181    #[test]
182    fn test_script_pubkey() {
183        let internal_key = get_test_internal_key();
184        let output = TaprootOutput::key_path_only(internal_key).unwrap();
185
186        let script = output.script_pubkey();
187        assert_eq!(script.len(), 34);
188        assert_eq!(script[0], 0x51); // OP_1
189        assert_eq!(script[1], 0x20); // Push 32 bytes
190    }
191
192    #[test]
193    fn test_address_mainnet() {
194        let internal_key = get_test_internal_key();
195        let output = TaprootOutput::key_path_only(internal_key).unwrap();
196
197        let address = output.address(Network::Mainnet).unwrap();
198        assert!(address.starts_with("bc1p"));
199    }
200
201    #[test]
202    fn test_address_testnet() {
203        let internal_key = get_test_internal_key();
204        let output = TaprootOutput::key_path_only(internal_key).unwrap();
205
206        let address = output.address(Network::Testnet).unwrap();
207        assert!(address.starts_with("tb1p"));
208    }
209
210    #[test]
211    fn test_address_roundtrip() {
212        let internal_key = get_test_internal_key();
213        let output = TaprootOutput::key_path_only(internal_key).unwrap();
214
215        let address = output.address(Network::Mainnet).unwrap();
216        let parsed = parse_address(&address).unwrap();
217
218        assert_eq!(output.output_key, parsed);
219    }
220
221    #[test]
222    fn test_with_script_tree() {
223        use crate::taptree::two_leaf_tree;
224
225        let internal_key = get_test_internal_key();
226        let tree = two_leaf_tree(vec![0x51], vec![0x52]);
227
228        let output = TaprootOutput::with_script_tree(internal_key, &tree).unwrap();
229
230        assert!(!output.is_key_path_only());
231        assert!(output.merkle_root.is_some());
232    }
233
234    #[test]
235    fn test_is_p2tr_address() {
236        let internal_key = get_test_internal_key();
237        let output = TaprootOutput::key_path_only(internal_key).unwrap();
238        let address = output.address(Network::Mainnet).unwrap();
239
240        assert!(is_p2tr_address(&address));
241        assert!(!is_p2tr_address("bc1qw508d6qejxtdg4y5r3zarvary0c5xw7kv8f3t4")); // P2WPKH
242        assert!(!is_p2tr_address("invalid"));
243    }
244}