1use std::collections::BTreeMap;
4use crate::error::PsbtError;
5use crate::types::{KeySource, PsbtSighashType, ProprietaryFields, UnknownFields};
6
7#[derive(Debug, Clone, PartialEq, Eq)]
9pub struct TxOut {
10 pub value: u64,
12 pub script_pubkey: Vec<u8>,
14}
15
16impl TxOut {
17 pub fn new(value: u64, script_pubkey: Vec<u8>) -> Self {
19 Self { value, script_pubkey }
20 }
21
22 pub fn to_bytes(&self) -> Vec<u8> {
24 let mut bytes = Vec::new();
25 bytes.extend_from_slice(&self.value.to_le_bytes());
26 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 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#[derive(Debug, Clone, PartialEq, Eq, Default)]
68pub struct Witness {
69 pub items: Vec<Vec<u8>>,
71}
72
73impl Witness {
74 pub fn new() -> Self {
76 Self { items: Vec::new() }
77 }
78
79 pub fn from_items(items: Vec<Vec<u8>>) -> Self {
81 Self { items }
82 }
83
84 pub fn push(&mut self, item: Vec<u8>) {
86 self.items.push(item);
87 }
88
89 pub fn is_empty(&self) -> bool {
91 self.items.is_empty()
92 }
93
94 pub fn to_bytes(&self) -> Vec<u8> {
96 let mut bytes = Vec::new();
97 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 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#[derive(Debug, Clone, Default)]
129pub struct InputMap {
130 pub non_witness_utxo: Option<Vec<u8>>,
132 pub witness_utxo: Option<TxOut>,
134 pub partial_sigs: BTreeMap<Vec<u8>, Vec<u8>>,
136 pub sighash_type: Option<PsbtSighashType>,
138 pub redeem_script: Option<Vec<u8>>,
140 pub witness_script: Option<Vec<u8>>,
142 pub bip32_derivation: BTreeMap<Vec<u8>, KeySource>,
144 pub final_script_sig: Option<Vec<u8>>,
146 pub final_script_witness: Option<Witness>,
148 pub tap_key_sig: Option<Vec<u8>>,
150 pub tap_script_sigs: BTreeMap<Vec<u8>, Vec<u8>>,
152 pub tap_leaf_scripts: BTreeMap<Vec<u8>, Vec<u8>>,
154 pub tap_bip32_derivation: BTreeMap<Vec<u8>, Vec<u8>>,
156 pub tap_internal_key: Option<Vec<u8>>,
158 pub tap_merkle_root: Option<Vec<u8>>,
160 pub previous_txid: Option<[u8; 32]>,
162 pub output_index: Option<u32>,
164 pub sequence: Option<u32>,
166 pub required_time_locktime: Option<u32>,
168 pub required_height_locktime: Option<u32>,
170 pub proprietary: ProprietaryFields,
172 pub unknown: UnknownFields,
174}
175
176impl InputMap {
177 pub fn new() -> Self {
179 Self::default()
180 }
181
182 pub fn is_finalized(&self) -> bool {
184 self.final_script_sig.is_some() || self.final_script_witness.is_some()
185 }
186
187 pub fn has_utxo(&self) -> bool {
189 self.witness_utxo.is_some() || self.non_witness_utxo.is_some()
190 }
191
192 pub fn utxo_value(&self) -> Option<u64> {
194 self.witness_utxo.as_ref().map(|utxo| utxo.value)
195 }
196
197 pub fn utxo_script(&self) -> Option<&[u8]> {
199 self.witness_utxo.as_ref().map(|utxo| utxo.script_pubkey.as_slice())
200 }
201
202 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
218pub 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
253pub 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], vec![0x02, 0x33], ]);
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}