rustywallet_psbt/
input.rs

1//! PSBT input map
2
3use std::collections::BTreeMap;
4use crate::error::PsbtError;
5use crate::types::{KeySource, PsbtSighashType, ProprietaryFields, UnknownFields};
6
7/// Transaction output (for witness_utxo)
8#[derive(Debug, Clone, PartialEq, Eq)]
9pub struct TxOut {
10    /// Value in satoshis
11    pub value: u64,
12    /// Script pubkey
13    pub script_pubkey: Vec<u8>,
14}
15
16impl TxOut {
17    /// Create a new TxOut
18    pub fn new(value: u64, script_pubkey: Vec<u8>) -> Self {
19        Self { value, script_pubkey }
20    }
21
22    /// Serialize to bytes
23    pub fn to_bytes(&self) -> Vec<u8> {
24        let mut bytes = Vec::new();
25        bytes.extend_from_slice(&self.value.to_le_bytes());
26        // Compact size for script length
27        let script_len = self.script_pubkey.len();
28        if script_len < 0xfd {
29            bytes.push(script_len as u8);
30        } else if script_len <= 0xffff {
31            bytes.push(0xfd);
32            bytes.extend_from_slice(&(script_len as u16).to_le_bytes());
33        } else {
34            bytes.push(0xfe);
35            bytes.extend_from_slice(&(script_len as u32).to_le_bytes());
36        }
37        bytes.extend_from_slice(&self.script_pubkey);
38        bytes
39    }
40
41    /// Parse from bytes
42    pub fn from_bytes(bytes: &[u8]) -> Result<(Self, usize), PsbtError> {
43        if bytes.len() < 9 {
44            return Err(PsbtError::InvalidFormat("TxOut too short".into()));
45        }
46
47        let value = u64::from_le_bytes([
48            bytes[0], bytes[1], bytes[2], bytes[3],
49            bytes[4], bytes[5], bytes[6], bytes[7],
50        ]);
51
52        let (script_len, offset) = read_compact_size(&bytes[8..])?;
53        let start = 8 + offset;
54        let end = start + script_len;
55
56        if bytes.len() < end {
57            return Err(PsbtError::InvalidFormat("TxOut script truncated".into()));
58        }
59
60        let script_pubkey = bytes[start..end].to_vec();
61
62        Ok((Self { value, script_pubkey }, end))
63    }
64}
65
66/// Witness stack
67#[derive(Debug, Clone, PartialEq, Eq, Default)]
68pub struct Witness {
69    /// Witness stack items
70    pub items: Vec<Vec<u8>>,
71}
72
73impl Witness {
74    /// Create a new empty witness
75    pub fn new() -> Self {
76        Self { items: Vec::new() }
77    }
78
79    /// Create witness from items
80    pub fn from_items(items: Vec<Vec<u8>>) -> Self {
81        Self { items }
82    }
83
84    /// Add an item to the witness
85    pub fn push(&mut self, item: Vec<u8>) {
86        self.items.push(item);
87    }
88
89    /// Check if witness is empty
90    pub fn is_empty(&self) -> bool {
91        self.items.is_empty()
92    }
93
94    /// Serialize to bytes
95    pub fn to_bytes(&self) -> Vec<u8> {
96        let mut bytes = Vec::new();
97        // Number of items
98        write_compact_size(&mut bytes, self.items.len());
99        for item in &self.items {
100            write_compact_size(&mut bytes, item.len());
101            bytes.extend_from_slice(item);
102        }
103        bytes
104    }
105
106    /// Parse from bytes
107    pub fn from_bytes(bytes: &[u8]) -> Result<Self, PsbtError> {
108        let mut offset = 0;
109        let (count, size) = read_compact_size(&bytes[offset..])?;
110        offset += size;
111
112        let mut items = Vec::with_capacity(count);
113        for _ in 0..count {
114            let (item_len, size) = read_compact_size(&bytes[offset..])?;
115            offset += size;
116            if bytes.len() < offset + item_len {
117                return Err(PsbtError::InvalidFormat("Witness item truncated".into()));
118            }
119            items.push(bytes[offset..offset + item_len].to_vec());
120            offset += item_len;
121        }
122
123        Ok(Self { items })
124    }
125}
126
127/// PSBT input map
128#[derive(Debug, Clone, Default)]
129pub struct InputMap {
130    /// Non-witness UTXO (full previous transaction)
131    pub non_witness_utxo: Option<Vec<u8>>,
132    /// Witness UTXO (just the output being spent)
133    pub witness_utxo: Option<TxOut>,
134    /// Partial signatures (pubkey -> signature)
135    pub partial_sigs: BTreeMap<Vec<u8>, Vec<u8>>,
136    /// Sighash type
137    pub sighash_type: Option<PsbtSighashType>,
138    /// Redeem script (for P2SH)
139    pub redeem_script: Option<Vec<u8>>,
140    /// Witness script (for P2WSH)
141    pub witness_script: Option<Vec<u8>>,
142    /// BIP32 derivation paths (pubkey -> key source)
143    pub bip32_derivation: BTreeMap<Vec<u8>, KeySource>,
144    /// Final scriptSig
145    pub final_script_sig: Option<Vec<u8>>,
146    /// Final witness
147    pub final_script_witness: Option<Witness>,
148    /// Taproot key spend signature
149    pub tap_key_sig: Option<Vec<u8>>,
150    /// Taproot script spend signatures
151    pub tap_script_sigs: BTreeMap<Vec<u8>, Vec<u8>>,
152    /// Taproot leaf scripts
153    pub tap_leaf_scripts: BTreeMap<Vec<u8>, Vec<u8>>,
154    /// Taproot BIP32 derivation
155    pub tap_bip32_derivation: BTreeMap<Vec<u8>, Vec<u8>>,
156    /// Taproot internal key
157    pub tap_internal_key: Option<Vec<u8>>,
158    /// Taproot merkle root
159    pub tap_merkle_root: Option<Vec<u8>>,
160    /// PSBT v2: Previous txid
161    pub previous_txid: Option<[u8; 32]>,
162    /// PSBT v2: Output index
163    pub output_index: Option<u32>,
164    /// PSBT v2: Sequence
165    pub sequence: Option<u32>,
166    /// PSBT v2: Required time locktime
167    pub required_time_locktime: Option<u32>,
168    /// PSBT v2: Required height locktime
169    pub required_height_locktime: Option<u32>,
170    /// Proprietary fields
171    pub proprietary: ProprietaryFields,
172    /// Unknown fields
173    pub unknown: UnknownFields,
174}
175
176impl InputMap {
177    /// Create a new empty input map
178    pub fn new() -> Self {
179        Self::default()
180    }
181
182    /// Check if this input is finalized
183    pub fn is_finalized(&self) -> bool {
184        self.final_script_sig.is_some() || self.final_script_witness.is_some()
185    }
186
187    /// Check if this input has UTXO information
188    pub fn has_utxo(&self) -> bool {
189        self.witness_utxo.is_some() || self.non_witness_utxo.is_some()
190    }
191
192    /// Get the value of the UTXO being spent
193    pub fn utxo_value(&self) -> Option<u64> {
194        self.witness_utxo.as_ref().map(|utxo| utxo.value)
195    }
196
197    /// Get the script pubkey of the UTXO being spent
198    pub fn utxo_script(&self) -> Option<&[u8]> {
199        self.witness_utxo.as_ref().map(|utxo| utxo.script_pubkey.as_slice())
200    }
201
202    /// Clear non-final fields (called after finalization)
203    pub fn clear_for_finalization(&mut self) {
204        self.partial_sigs.clear();
205        self.sighash_type = None;
206        self.redeem_script = None;
207        self.witness_script = None;
208        self.bip32_derivation.clear();
209        self.tap_key_sig = None;
210        self.tap_script_sigs.clear();
211        self.tap_leaf_scripts.clear();
212        self.tap_bip32_derivation.clear();
213        self.tap_internal_key = None;
214        self.tap_merkle_root = None;
215    }
216}
217
218/// Read compact size from bytes
219pub fn read_compact_size(bytes: &[u8]) -> Result<(usize, usize), PsbtError> {
220    if bytes.is_empty() {
221        return Err(PsbtError::InvalidFormat("Empty compact size".into()));
222    }
223
224    match bytes[0] {
225        0..=0xfc => Ok((bytes[0] as usize, 1)),
226        0xfd => {
227            if bytes.len() < 3 {
228                return Err(PsbtError::InvalidFormat("Truncated compact size".into()));
229            }
230            let value = u16::from_le_bytes([bytes[1], bytes[2]]) as usize;
231            Ok((value, 3))
232        }
233        0xfe => {
234            if bytes.len() < 5 {
235                return Err(PsbtError::InvalidFormat("Truncated compact size".into()));
236            }
237            let value = u32::from_le_bytes([bytes[1], bytes[2], bytes[3], bytes[4]]) as usize;
238            Ok((value, 5))
239        }
240        0xff => {
241            if bytes.len() < 9 {
242                return Err(PsbtError::InvalidFormat("Truncated compact size".into()));
243            }
244            let value = u64::from_le_bytes([
245                bytes[1], bytes[2], bytes[3], bytes[4],
246                bytes[5], bytes[6], bytes[7], bytes[8],
247            ]) as usize;
248            Ok((value, 9))
249        }
250    }
251}
252
253/// Write compact size to buffer
254pub fn write_compact_size(buf: &mut Vec<u8>, value: usize) {
255    if value < 0xfd {
256        buf.push(value as u8);
257    } else if value <= 0xffff {
258        buf.push(0xfd);
259        buf.extend_from_slice(&(value as u16).to_le_bytes());
260    } else if value <= 0xffffffff {
261        buf.push(0xfe);
262        buf.extend_from_slice(&(value as u32).to_le_bytes());
263    } else {
264        buf.push(0xff);
265        buf.extend_from_slice(&(value as u64).to_le_bytes());
266    }
267}
268
269#[cfg(test)]
270mod tests {
271    use super::*;
272
273    #[test]
274    fn test_txout_roundtrip() {
275        let txout = TxOut::new(
276            100_000_000,
277            vec![0x00, 0x14, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08,
278                 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f, 0x10, 0x11, 0x12,
279                 0x13, 0x14],
280        );
281
282        let bytes = txout.to_bytes();
283        let (parsed, _) = TxOut::from_bytes(&bytes).unwrap();
284
285        assert_eq!(txout, parsed);
286    }
287
288    #[test]
289    fn test_witness_roundtrip() {
290        let witness = Witness::from_items(vec![
291            vec![0x30, 0x44], // signature
292            vec![0x02, 0x33], // pubkey
293        ]);
294
295        let bytes = witness.to_bytes();
296        let parsed = Witness::from_bytes(&bytes).unwrap();
297
298        assert_eq!(witness, parsed);
299    }
300
301    #[test]
302    fn test_compact_size() {
303        let test_cases = [
304            (0usize, vec![0x00]),
305            (252, vec![0xfc]),
306            (253, vec![0xfd, 0xfd, 0x00]),
307            (0xffff, vec![0xfd, 0xff, 0xff]),
308            (0x10000, vec![0xfe, 0x00, 0x00, 0x01, 0x00]),
309        ];
310
311        for (value, expected) in test_cases {
312            let mut buf = Vec::new();
313            write_compact_size(&mut buf, value);
314            assert_eq!(buf, expected);
315
316            let (parsed, _) = read_compact_size(&buf).unwrap();
317            assert_eq!(parsed, value);
318        }
319    }
320}