Skip to main content

trezor_connect_rs/compose/
sorting.rs

1//! Transaction sorting strategies.
2//!
3//! Implements BIP-69, random, and no-op sorting for transaction inputs and outputs.
4
5use serde::{Deserialize, Serialize};
6
7/// Sorting strategy for transaction inputs and outputs.
8#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
9pub enum SortingStrategy {
10    /// BIP-69: deterministic lexicographic sorting
11    Bip69,
12    /// Random shuffle (better privacy, prevents fingerprinting)
13    #[default]
14    Random,
15    /// Keep original order
16    None,
17}
18
19/// A composed input (for sorting purposes).
20#[derive(Debug, Clone)]
21pub struct SortableInput {
22    /// Original index in the input array
23    pub index: usize,
24    /// Transaction ID (hex)
25    pub txid: String,
26    /// Output index
27    pub vout: u32,
28}
29
30/// A composed output (for sorting purposes).
31#[derive(Debug, Clone)]
32pub struct SortableOutput {
33    /// Original index in the output array
34    pub index: usize,
35    /// Amount in satoshis
36    pub amount: u64,
37    /// Raw scriptPubKey bytes for BIP69 comparison
38    pub script_pubkey: Vec<u8>,
39    /// Whether this is a change output (used by random sorting)
40    pub is_change: bool,
41}
42
43/// Derive a scriptPubKey from an address string.
44///
45/// Supports P2PKH (1...), P2SH (3...), P2WPKH (bc1q.../tb1q.../bcrt1q...),
46/// P2TR (bc1p.../tb1p.../bcrt1p...), and testnet/regtest variants.
47/// Returns empty vec for unrecognized addresses.
48pub fn address_to_script_pubkey(address: &str) -> Vec<u8> {
49    let addr_lower = address.to_lowercase();
50
51    // Bech32/Bech32m (P2WPKH or P2TR)
52    if addr_lower.starts_with("bc1")
53        || addr_lower.starts_with("tb1")
54        || addr_lower.starts_with("bcrt1")
55    {
56        // Find the separator '1' - for bech32, it's the last '1' in the string
57        if let Some(sep_pos) = addr_lower.rfind('1') {
58            let data_part = &addr_lower[sep_pos + 1..];
59            if let Some(decoded) = bech32_decode_data(data_part) {
60                if decoded.is_empty() {
61                    return Vec::new();
62                }
63                let witness_version = decoded[0];
64                let witness_program = convert_bits(&decoded[1..], 5, 8, false);
65                if let Some(program) = witness_program {
66                    let mut script = Vec::new();
67                    // Witness version: OP_0 (0x00) for v0, OP_1..OP_16 (0x51..0x60) for v1..v16
68                    if witness_version == 0 {
69                        script.push(0x00);
70                    } else {
71                        script.push(0x50 + witness_version);
72                    }
73                    script.push(program.len() as u8);
74                    script.extend_from_slice(&program);
75                    return script;
76                }
77            }
78        }
79        return Vec::new();
80    }
81
82    // Base58Check (P2PKH or P2SH)
83    if let Some(decoded) = base58check_decode(address) {
84        if decoded.len() == 21 {
85            let version = decoded[0];
86            let hash = &decoded[1..];
87            match version {
88                // P2PKH mainnet (0x00) or testnet (0x6f)
89                0x00 | 0x6f => {
90                    // OP_DUP OP_HASH160 <20> <hash> OP_EQUALVERIFY OP_CHECKSIG
91                    let mut script = vec![0x76, 0xa9, 0x14];
92                    script.extend_from_slice(hash);
93                    script.push(0x88);
94                    script.push(0xac);
95                    return script;
96                }
97                // P2SH mainnet (0x05) or testnet (0xc4)
98                0x05 | 0xc4 => {
99                    // OP_HASH160 <20> <hash> OP_EQUAL
100                    let mut script = vec![0xa9, 0x14];
101                    script.extend_from_slice(hash);
102                    script.push(0x87);
103                    return script;
104                }
105                _ => {}
106            }
107        }
108    }
109
110    Vec::new()
111}
112
113/// Derive a scriptPubKey for an OP_RETURN output from hex data.
114pub fn op_return_script_pubkey(data_hex: &str) -> Vec<u8> {
115    let data = hex::decode(data_hex).unwrap_or_default();
116    let mut script = vec![0x6a]; // OP_RETURN
117    if data.len() <= 75 {
118        script.push(data.len() as u8);
119    } else if data.len() <= 255 {
120        script.push(0x4c); // OP_PUSHDATA1
121        script.push(data.len() as u8);
122    } else {
123        script.push(0x4d); // OP_PUSHDATA2
124        script.push((data.len() & 0xff) as u8);
125        script.push((data.len() >> 8) as u8);
126    }
127    script.extend_from_slice(&data);
128    script
129}
130
131/// Decode bech32 data part (5-bit values).
132fn bech32_decode_data(data: &str) -> Option<Vec<u8>> {
133    const CHARSET: &str = "qpzry9x8gf2tvdw0s3jn54khce6mua7l";
134
135    let mut values = Vec::with_capacity(data.len());
136    for c in data.chars() {
137        if let Some(pos) = CHARSET.find(c) {
138            values.push(pos as u8);
139        } else {
140            return None;
141        }
142    }
143
144    // Strip the 6-character checksum
145    if values.len() < 7 {
146        return None;
147    }
148    values.truncate(values.len() - 6);
149
150    Some(values)
151}
152
153/// Convert between bit groups (e.g., 5-bit to 8-bit for bech32).
154fn convert_bits(data: &[u8], from_bits: u32, to_bits: u32, pad: bool) -> Option<Vec<u8>> {
155    let mut acc: u32 = 0;
156    let mut bits: u32 = 0;
157    let mut result = Vec::new();
158    let max_v = (1u32 << to_bits) - 1;
159
160    for &value in data {
161        if (value as u32) >> from_bits != 0 {
162            return None;
163        }
164        acc = (acc << from_bits) | value as u32;
165        bits += from_bits;
166        while bits >= to_bits {
167            bits -= to_bits;
168            result.push(((acc >> bits) & max_v) as u8);
169        }
170    }
171
172    if pad {
173        if bits > 0 {
174            result.push(((acc << (to_bits - bits)) & max_v) as u8);
175        }
176    } else if bits >= from_bits || ((acc << (to_bits - bits)) & max_v) != 0 {
177        return None;
178    }
179
180    Some(result)
181}
182
183/// Decode a Base58Check-encoded string.
184fn base58check_decode(input: &str) -> Option<Vec<u8>> {
185    const ALPHABET: &[u8] = b"123456789ABCDEFGHJKLMNPQRSTUVWXYZabcdefghijkmnopqrstuvwxyz";
186
187    let mut result = vec![0u8; 1]; // Start with capacity
188    // Actually, let's use a big number approach
189    let mut bytes = Vec::new();
190    let mut leading_zeros = 0;
191
192    // Count leading '1's (which represent leading zero bytes)
193    for c in input.chars() {
194        if c == '1' {
195            leading_zeros += 1;
196        } else {
197            break;
198        }
199    }
200
201    // Convert from base58 to base256 using big number arithmetic
202    let mut num = vec![0u8; 0];
203    for c in input.bytes() {
204        let digit = ALPHABET.iter().position(|&x| x == c)? as u32;
205
206        // Multiply num by 58 and add digit
207        let mut carry = digit;
208        for byte in num.iter_mut().rev() {
209            let val = (*byte as u32) * 58 + carry;
210            *byte = (val & 0xff) as u8;
211            carry = val >> 8;
212        }
213        while carry > 0 {
214            num.insert(0, (carry & 0xff) as u8);
215            carry >>= 8;
216        }
217    }
218
219    // Add leading zero bytes
220    for _ in 0..leading_zeros {
221        bytes.push(0);
222    }
223    bytes.extend_from_slice(&num);
224
225    // Verify checksum (last 4 bytes)
226    if bytes.len() < 4 {
227        return None;
228    }
229    let payload_len = bytes.len() - 4;
230    let payload = &bytes[..payload_len];
231    let checksum = &bytes[payload_len..];
232
233    use sha2::{Digest, Sha256};
234    let hash1 = Sha256::digest(payload);
235    let hash2 = Sha256::digest(hash1);
236
237    if &hash2[..4] != checksum {
238        return None;
239    }
240
241    result = payload.to_vec();
242    Some(result)
243}
244
245/// Sort inputs and outputs according to the given strategy.
246/// Returns the permutation indices for outputs.
247pub fn sort_transaction(
248    inputs: &mut Vec<SortableInput>,
249    outputs: &mut Vec<SortableOutput>,
250    strategy: SortingStrategy,
251) -> Vec<usize> {
252    match strategy {
253        SortingStrategy::Bip69 => {
254            // Sort inputs by txid (as raw bytes via hex), then by vout
255            inputs.sort_by(|a, b| a.txid.cmp(&b.txid).then(a.vout.cmp(&b.vout)));
256
257            // Sort outputs by amount, then by scriptPubKey bytes (BIP69 spec)
258            outputs.sort_by(|a, b| {
259                a.amount
260                    .cmp(&b.amount)
261                    .then(a.script_pubkey.cmp(&b.script_pubkey))
262            });
263        }
264        SortingStrategy::Random => {
265            use rand::seq::SliceRandom;
266            let mut rng = rand::thread_rng();
267            inputs.shuffle(&mut rng);
268
269            // Insert change outputs at random positions among payment outputs,
270            // matching JS behavior for better privacy.
271            // Payment outputs keep their user-specified order (JS randomSortingStrategy.ts).
272            let (mut payment_outputs, change_outputs): (Vec<_>, Vec<_>) =
273                outputs.drain(..).partition(|o| !o.is_change);
274
275            // Insert each change output at a random position
276            for change in change_outputs {
277                use rand::Rng;
278                let pos = rng.gen_range(0..=payment_outputs.len());
279                payment_outputs.insert(pos, change);
280            }
281
282            *outputs = payment_outputs;
283        }
284        SortingStrategy::None => {
285            // Keep original order
286        }
287    }
288
289    // Build output permutation: permutation[new_position] = original_index
290    outputs.iter().map(|o| o.index).collect()
291}
292
293#[cfg(test)]
294mod tests {
295    use super::*;
296
297    #[test]
298    fn test_bip69_sort_inputs() {
299        let mut inputs = vec![
300            SortableInput {
301                index: 0,
302                txid: "bbbb".into(),
303                vout: 1,
304            },
305            SortableInput {
306                index: 1,
307                txid: "aaaa".into(),
308                vout: 0,
309            },
310            SortableInput {
311                index: 2,
312                txid: "aaaa".into(),
313                vout: 2,
314            },
315        ];
316        let mut outputs = vec![
317            SortableOutput {
318                index: 0,
319                amount: 5000,
320                script_pubkey: vec![0xcc],
321                is_change: false,
322            },
323            SortableOutput {
324                index: 1,
325                amount: 1000,
326                script_pubkey: vec![0xaa],
327                is_change: false,
328            },
329        ];
330
331        let perm = sort_transaction(&mut inputs, &mut outputs, SortingStrategy::Bip69);
332
333        // Inputs should be sorted: aaaa:0, aaaa:2, bbbb:1
334        assert_eq!(inputs[0].txid, "aaaa");
335        assert_eq!(inputs[0].vout, 0);
336        assert_eq!(inputs[1].txid, "aaaa");
337        assert_eq!(inputs[1].vout, 2);
338        assert_eq!(inputs[2].txid, "bbbb");
339
340        // Outputs sorted by amount: 1000, 5000
341        assert_eq!(outputs[0].amount, 1000);
342        assert_eq!(outputs[1].amount, 5000);
343
344        // Permutation maps new positions to original indices
345        assert_eq!(perm, vec![1, 0]);
346    }
347
348    #[test]
349    fn test_none_sort_preserves_order() {
350        let mut inputs = vec![
351            SortableInput {
352                index: 0,
353                txid: "bb".into(),
354                vout: 1,
355            },
356            SortableInput {
357                index: 1,
358                txid: "aa".into(),
359                vout: 0,
360            },
361        ];
362        let mut outputs = vec![
363            SortableOutput {
364                index: 0,
365                amount: 5000,
366                script_pubkey: vec![0xcc],
367                is_change: false,
368            },
369            SortableOutput {
370                index: 1,
371                amount: 1000,
372                script_pubkey: vec![0xaa],
373                is_change: false,
374            },
375        ];
376
377        let perm = sort_transaction(&mut inputs, &mut outputs, SortingStrategy::None);
378
379        assert_eq!(inputs[0].txid, "bb");
380        assert_eq!(inputs[1].txid, "aa");
381        assert_eq!(perm, vec![0, 1]);
382    }
383
384    #[test]
385    fn test_bip69_sort_outputs_by_script_pubkey() {
386        // Two outputs with same amount but different scriptPubKeys
387        let mut inputs = vec![SortableInput {
388            index: 0,
389            txid: "aa".into(),
390            vout: 0,
391        }];
392        let mut outputs = vec![
393            SortableOutput {
394                index: 0,
395                amount: 1000,
396                script_pubkey: vec![0x76, 0xa9],
397                is_change: false,
398            },
399            SortableOutput {
400                index: 1,
401                amount: 1000,
402                script_pubkey: vec![0x00, 0x14],
403                is_change: false,
404            },
405        ];
406
407        let perm = sort_transaction(&mut inputs, &mut outputs, SortingStrategy::Bip69);
408
409        // scriptPubKey [0x00, 0x14] < [0x76, 0xa9]
410        assert_eq!(outputs[0].script_pubkey, vec![0x00, 0x14]);
411        assert_eq!(outputs[1].script_pubkey, vec![0x76, 0xa9]);
412        assert_eq!(perm, vec![1, 0]);
413    }
414
415    #[test]
416    fn test_address_to_script_pubkey_p2wpkh() {
417        // Known P2WPKH address: bc1qw508d6qejxtdg4y5r3zarvary0c5xw7kv8f3t4
418        // Expected scriptPubKey: 0014751e76e8199196d454941c45d1b3a323f1433bd6
419        let script = address_to_script_pubkey("bc1qw508d6qejxtdg4y5r3zarvary0c5xw7kv8f3t4");
420        assert!(!script.is_empty(), "Should decode P2WPKH address");
421        assert_eq!(script[0], 0x00, "Witness version 0");
422        assert_eq!(script[1], 0x14, "20-byte program push");
423        assert_eq!(script.len(), 22, "P2WPKH script is 22 bytes");
424        assert_eq!(
425            hex::encode(&script),
426            "0014751e76e8199196d454941c45d1b3a323f1433bd6"
427        );
428    }
429
430    #[test]
431    fn test_address_to_script_pubkey_p2tr() {
432        // Known P2TR address: bc1pqyqszqgpqyqszqgpqyqszqgpqyqszqgpqyqszqgpqyqszqgpqyqs3wf0qm
433        // Witness version 1, 32-byte program
434        let script = address_to_script_pubkey(
435            "bc1p0xlxvlhemja6c4dqv22uapctqupfhlxm9h8z3k2e72q4k9hcz7vqzk5jj0",
436        );
437        assert!(!script.is_empty(), "Should decode P2TR address");
438        assert_eq!(script[0], 0x51, "Witness version 1 = OP_1");
439        assert_eq!(script[1], 0x20, "32-byte program push");
440        assert_eq!(script.len(), 34, "P2TR script is 34 bytes");
441    }
442
443    #[test]
444    fn test_address_to_script_pubkey_p2pkh() {
445        // Known P2PKH address: 1BvBMSEYstWetqTFn5Au4m4GFg7xJaNVN2
446        let script = address_to_script_pubkey("1BvBMSEYstWetqTFn5Au4m4GFg7xJaNVN2");
447        assert!(!script.is_empty(), "Should decode P2PKH address");
448        assert_eq!(script[0], 0x76, "OP_DUP");
449        assert_eq!(script[1], 0xa9, "OP_HASH160");
450        assert_eq!(script[2], 0x14, "20-byte push");
451        assert_eq!(script.len(), 25, "P2PKH script is 25 bytes");
452    }
453
454    #[test]
455    fn test_address_to_script_pubkey_p2sh() {
456        // Known P2SH address: 3J98t1WpEZ73CNmQviecrnyiWrnqRhWNLy
457        let script = address_to_script_pubkey("3J98t1WpEZ73CNmQviecrnyiWrnqRhWNLy");
458        assert!(!script.is_empty(), "Should decode P2SH address");
459        assert_eq!(script[0], 0xa9, "OP_HASH160");
460        assert_eq!(script[1], 0x14, "20-byte push");
461        assert_eq!(script.len(), 23, "P2SH script is 23 bytes");
462    }
463
464    #[test]
465    fn test_address_to_script_pubkey_regtest() {
466        // Regtest P2WPKH address
467        let script = address_to_script_pubkey("bcrt1qj2gz3meule5mc4r4knv65vjds3g88rlxs0jlmq");
468        assert!(!script.is_empty(), "Should decode regtest bech32 address");
469        assert_eq!(script[0], 0x00, "Witness version 0");
470        assert_eq!(script[1], 0x14, "20-byte program push");
471        assert_eq!(script.len(), 22);
472    }
473
474    #[test]
475    fn test_op_return_script_pubkey() {
476        let script = op_return_script_pubkey("deadbeef");
477        assert_eq!(script[0], 0x6a, "OP_RETURN");
478        assert_eq!(script[1], 4, "4 bytes of data");
479        assert_eq!(&script[2..], &[0xde, 0xad, 0xbe, 0xef]);
480    }
481}