Skip to main content

kaspa_txscript/
script_builder.rs

1use std::iter::once;
2
3use crate::{
4    EngineFlags, MAX_SCRIPT_ELEMENT_SIZE, MAX_SCRIPTS_SIZE,
5    data_stack::OpcodeData,
6    opcodes::{OP_1_NEGATE_VAL, OP_DATA_MAX_VAL, OP_DATA_MIN_VAL, OP_SMALL_INT_MAX_VAL, OP_SMALL_INT_MIN_VAL, codes::*},
7};
8use hexplay::{HexView, HexViewBuilder};
9use kaspa_txscript_errors::SerializationError;
10use thiserror::Error;
11
12/// DEFAULT_SCRIPT_ALLOC is the default size used for the backing array
13/// for a script being built by the ScriptBuilder. The array will
14/// dynamically grow as needed, but this figure is intended to provide
15/// enough space for vast majority of scripts without needing to grow the
16/// backing array multiple times.
17const DEFAULT_SCRIPT_ALLOC: usize = 512;
18
19#[derive(Error, PartialEq, Eq, Debug, Clone, Copy)]
20pub enum ScriptBuilderError {
21    #[error("adding opcode {0} would exceed the maximum allowed canonical script length of {1}")]
22    OpCodeRejected(u8, usize),
23
24    #[error("adding {0} opcodes would exceed the maximum allowed canonical script length of {1}")]
25    OpCodesRejected(usize, usize),
26
27    #[error("adding {0} bytes of data would exceed the maximum allowed canonical script length of {1}")]
28    DataRejected(usize, usize),
29
30    #[error("adding a data element of {0} bytes exceed the maximum allowed script element size of {1}")]
31    ElementExceedsMaxSize(usize, usize),
32
33    #[error("adding integer {0} would exceed the maximum allowed canonical script length of {1}")]
34    IntegerRejected(i64, usize),
35
36    #[error(transparent)]
37    Serialization(#[from] SerializationError),
38}
39pub type ScriptBuilderResult<T> = std::result::Result<T, ScriptBuilderError>;
40
41/// ScriptBuilder provides a facility for building custom scripts. It allows
42/// you to push opcodes, ints, and data while respecting canonical encoding. In
43/// general it does not ensure the script will execute correctly, however any
44/// data pushes which would exceed the maximum allowed script engine limits and
45/// are therefore guaranteed not to execute will not be pushed and will result in
46/// the Script function returning an error.
47///
48/// For example, the following would build a 2-of-3 multisig script for usage in
49/// a pay-to-script-hash (although in this situation MultiSigScript() would be a
50/// better choice to generate the script):
51///
52/// ```
53/// use kaspa_txscript::opcodes::codes::*;
54/// use kaspa_txscript::script_builder::{ScriptBuilderResult, ScriptBuilder};
55/// fn build_multisig_script(pub_key1: &[u8], pub_key2: &[u8], pub_key3: &[u8]) -> ScriptBuilderResult<Vec<u8>> {
56///     Ok(ScriptBuilder::new()
57///         .add_op(Op2)?
58///         .add_data(pub_key1)?.add_data(pub_key2)?.add_data(pub_key3)?
59///         .add_op(Op3)?
60///         .add_op(OpCheckMultiSig)?
61///         .drain())
62/// }
63/// ```
64pub struct ScriptBuilder {
65    script: Vec<u8>,
66    flags: EngineFlags,
67}
68
69impl ScriptBuilder {
70    pub fn new() -> Self {
71        Self::with_flags(Default::default())
72    }
73
74    pub fn with_flags(flags: EngineFlags) -> Self {
75        Self { script: Vec::with_capacity(DEFAULT_SCRIPT_ALLOC), flags }
76    }
77
78    pub fn flags(&self) -> EngineFlags {
79        self.flags
80    }
81
82    pub fn script(&self) -> &[u8] {
83        &self.script
84    }
85
86    pub fn script_mut(&mut self) -> &mut Vec<u8> {
87        &mut self.script
88    }
89
90    pub fn drain(&mut self) -> Vec<u8> {
91        // Note that the internal script, when taken, is replaced by
92        // vector with no predefined capacity because the script
93        // builder is not supposed to be reused after a call
94        // to drain.
95        std::mem::take(&mut self.script)
96    }
97
98    /// Pushes the passed opcode to the end of the script. The script will not
99    /// be modified if pushing the opcode would cause the script to exceed the
100    /// maximum allowed script engine size.
101    pub fn add_op(&mut self, opcode: u8) -> ScriptBuilderResult<&mut Self> {
102        // Pushes that would cause the script to exceed the largest allowed
103        // script size would result in a non-canonical script.
104        let max_scripts_size = MAX_SCRIPTS_SIZE;
105        if self.script.len() >= max_scripts_size {
106            return Err(ScriptBuilderError::OpCodeRejected(opcode, max_scripts_size));
107        }
108
109        self.script.push(opcode);
110        Ok(self)
111    }
112
113    pub fn add_ops(&mut self, opcodes: &[u8]) -> ScriptBuilderResult<&mut Self> {
114        // Pushes that would cause the script to exceed the largest allowed
115        // script size would result in a non-canonical script.
116        let max_scripts_size = MAX_SCRIPTS_SIZE;
117        if self.script.len() + opcodes.len() > max_scripts_size {
118            return Err(ScriptBuilderError::OpCodesRejected(opcodes.len(), max_scripts_size));
119        }
120
121        self.script.extend_from_slice(opcodes);
122        Ok(self)
123    }
124
125    /// Returns the number of bytes the canonical encoding of the data will take.
126    pub fn canonical_data_size(data: &[u8]) -> usize {
127        let data_len = data.len();
128
129        // Empty data and single-byte small integers have one-byte canonical encodings.
130        if matches!(data, [] | [OP_SMALL_INT_MIN_VAL..=OP_SMALL_INT_MAX_VAL] | [OP_1_NEGATE_VAL]) {
131            return 1;
132        }
133
134        data_len
135            + if data_len <= OP_DATA_MAX_VAL as usize {
136                1 // length encoded as OpData#
137            } else if data_len <= u8::MAX as usize {
138                2 // length encoded as OpPushData1 + 1 byte for value
139            } else if data_len <= u16::MAX as usize {
140                3 // length encoded as OpPushData2 + 2 bytes for value
141            } else {
142                5 // length encoded as OpPushData4 + 4 bytes for value
143            }
144    }
145
146    /// Internal function that actually pushes the passed data to the
147    /// end of the script. It automatically chooses canonical opcodes depending on
148    /// the length of the data. A zero length buffer will lead to a push of empty
149    /// data onto the stack (OP_0). No data limits are enforced with this function.
150    fn add_raw_data(&mut self, data: &[u8]) -> &mut Self {
151        match data {
152            [OP_1_NEGATE_VAL] => {
153                self.script.push(Op1Negate);
154                self
155            }
156            // When the data consists of a single number that can be represented
157            // by one of the "small integer" opcodes, use that opcode instead of
158            // a data push opcode followed by the number.
159            [OP_SMALL_INT_MIN_VAL..=OP_SMALL_INT_MAX_VAL] => {
160                self.script.push((Op1 - 1) + data[0]);
161                self
162            }
163            // For all other data, choose the appropriate push opcode based on data length.
164            _ => self.add_raw_data_with_data_opcode(data),
165        }
166    }
167
168    /// This function should not typically be used by ordinary users as it does not
169    /// include the checks which prevent data pushes larger than the maximum allowed
170    /// sizes which leads to scripts that can't be executed. This is provided for
171    /// testing purposes such as tests where sizes are intentionally made larger
172    /// than allowed.
173    ///
174    /// Use add_data instead.
175    #[cfg(test)]
176    pub fn add_data_unchecked(&mut self, data: &[u8]) -> &mut Self {
177        self.add_raw_data(data)
178    }
179
180    fn validate_data_push(&self, data: &[u8], data_size: usize) -> ScriptBuilderResult<()> {
181        // Pushes that would cause the script to exceed the largest allowed
182        // script size would result in a non-canonical script.
183        let max_scripts_size = MAX_SCRIPTS_SIZE;
184        if self.script.len() + data_size > max_scripts_size {
185            return Err(ScriptBuilderError::DataRejected(data_size, max_scripts_size));
186        }
187
188        // Pushes larger than the max script element size would result in a
189        // script that is not canonical.
190        let data_len = data.len();
191        let max_script_element_size = MAX_SCRIPT_ELEMENT_SIZE;
192        if data_len > max_script_element_size {
193            return Err(ScriptBuilderError::ElementExceedsMaxSize(data_len, max_script_element_size));
194        }
195
196        Ok(())
197    }
198
199    /// AddData pushes the passed data to the end of the script. It automatically
200    /// chooses canonical opcodes depending on the length of the data.
201    ///
202    /// A zero length buffer will lead to a push of empty data onto the stack (Op0 = OpFalse)
203    /// and any push of data greater than the maximum script element size will not modify
204    /// the script since that is not allowed by the script engine.
205    ///
206    /// Also, the script will not be modified if pushing the data would cause the script to
207    /// exceed the maximum allowed script engine size.
208    pub fn add_data(&mut self, data: &[u8]) -> ScriptBuilderResult<&mut Self> {
209        self.validate_data_push(data, Self::canonical_data_size(data))?;
210
211        Ok(self.add_raw_data(data))
212    }
213
214    /// Adds `data` using an explicit push-data opcode chosen only by payload size.
215    ///
216    /// Unlike `add_data()`, this method does not canonicalize single-byte small
217    /// integers into `Op1Negate` or `Op1..Op16`.
218    ///
219    /// Empty data is encoded as `Op0`, and all other values are emitted using
220    /// one of the push-data forms (`OpDataN`, `OpPushData1`, `OpPushData2`, or
221    /// `OpPushData4`) according to the length of `data`.
222    pub fn add_data_with_push_opcode(&mut self, data: &[u8]) -> ScriptBuilderResult<&mut Self> {
223        self.validate_data_push(data, Self::explicit_push_encoded_size(data.len()))?;
224        Ok(self.add_raw_data_with_data_opcode(data))
225    }
226
227    /// Returns the number of bytes emitted by the explicit push-data encoding.
228    fn explicit_push_encoded_size(data_len: usize) -> usize {
229        data_len
230            + if data_len <= OP_DATA_MAX_VAL as usize {
231                1 // Op0 for empty data, otherwise length encoded as OpData#
232            } else if data_len <= u8::MAX as usize {
233                2 // length encoded as OpPushData1 + 1 byte for value
234            } else if data_len <= u16::MAX as usize {
235                3 // length encoded as OpPushData2 + 2 bytes for value
236            } else {
237                5 // length encoded as OpPushData4 + 4 bytes for value
238            }
239    }
240
241    /// Adds `data` using an explicit push-data opcode chosen only by payload size.
242    ///
243    /// It's an internal function that is used by `add_raw_data` if we can't apply
244    /// small integer optimization, or by `add_data_with_push_opcode`, when the
245    /// caller is not interested in small integer optimization and wants the
246    /// push data prefix to be determined only by the payload length.
247    fn add_raw_data_with_data_opcode(&mut self, data: &[u8]) -> &mut Self {
248        // Empty data can be pushed using Op0.
249        let data_len = data.len();
250        if data_len == 0 {
251            self.script.push(Op0);
252            return self;
253        }
254
255        // Use one of the OpData# opcodes if the length of the data is small
256        // enough so the data push instruction is only a single byte.
257        // Otherwise, choose the smallest possible OpPushData# opcode that
258        // can represent the length of the data.
259        if data_len <= OP_DATA_MAX_VAL as usize {
260            self.script.push((OP_DATA_MIN_VAL - 1) + data_len as u8);
261        } else if data_len <= u8::MAX as usize {
262            self.script.extend(once(OpPushData1).chain(once(data_len as u8)));
263        } else if data_len <= u16::MAX as usize {
264            self.script.extend(once(OpPushData2).chain((data_len as u16).to_le_bytes()));
265        } else {
266            self.script.extend(once(OpPushData4).chain((data_len as u32).to_le_bytes()));
267        }
268
269        // Append the actual data.
270        self.script.extend(data);
271        self
272    }
273
274    pub fn add_i64(&mut self, val: i64) -> ScriptBuilderResult<&mut Self> {
275        // Pushes that would cause the script to exceed the largest allowed
276        // script size would result in a non-canonical script.
277        let max_scripts_size = MAX_SCRIPTS_SIZE;
278        if self.script.len() + 1 > max_scripts_size {
279            return Err(ScriptBuilderError::IntegerRejected(val, max_scripts_size));
280        }
281
282        // Fast path for small integers and Op1Negate.
283        if val == 0 {
284            self.script.push(Op0);
285            return Ok(self);
286        }
287        if val == -1 || (1..=16).contains(&val) {
288            self.script.push(((Op1 as i64 - 1) + val) as u8);
289            return Ok(self);
290        }
291
292        let bytes: crate::data_stack::StackEntry = OpcodeData::<i64>::serialize(&val)?;
293        self.add_data(&bytes)
294    }
295
296    // This value is outside the range of numbers allowed in the script. Bitcoind tests utilizes this function
297    pub(crate) fn add_i64_min(&mut self) -> ScriptBuilderResult<&mut Self> {
298        let bytes: crate::data_stack::StackEntry =
299            OpcodeData::serialize(&crate::data_stack::SizedEncodeInt::<9>(i64::MIN)).expect("infallible");
300        self.add_data(&bytes)
301    }
302
303    /// Gets a u64 lock time, converts it to byte array in little-endian, and then used the add_data function.
304    pub fn add_lock_time(&mut self, lock_time: u64) -> ScriptBuilderResult<&mut Self> {
305        self.add_u64(lock_time)
306    }
307
308    /// Gets a u64 sequence, converts it to byte array in little-endian, and then used the add_data function.
309    pub fn add_sequence(&mut self, sequence: u64) -> ScriptBuilderResult<&mut Self> {
310        self.add_u64(sequence)
311    }
312
313    /// Gets a u64 lock time or sequence, converts it to byte array in little-endian, and then used the add_data function.
314    fn add_u64(&mut self, val: u64) -> ScriptBuilderResult<&mut Self> {
315        let buffer: [u8; 8] = val.to_le_bytes();
316        let trimmed_size = 8 - buffer.iter().rev().position(|x| *x != 0u8).unwrap_or(8);
317        let trimmed = &buffer[0..trimmed_size];
318        self.add_data(trimmed)
319    }
320
321    /// Return [`HexViewBuilder`] for the script
322    pub fn hex_view_builder(&self) -> HexViewBuilder<'_> {
323        HexViewBuilder::new(&self.script)
324    }
325
326    /// Return ready to use [`HexView`] for the script
327    pub fn hex_view(&self, offset: usize, width: usize) -> HexView<'_> {
328        HexViewBuilder::new(&self.script).address_offset(offset).row_width(width).finish()
329    }
330}
331
332impl Default for ScriptBuilder {
333    fn default() -> Self {
334        Self::new()
335    }
336}
337
338#[cfg(test)]
339mod tests {
340    use super::*;
341    use std::iter::{once, repeat_n};
342
343    // Tests that pushing opcodes to a script via the ScriptBuilder API works as expected.
344    #[test]
345    fn test_add_op() {
346        struct Test {
347            name: &'static str,
348            opcodes: Vec<u8>,
349            expected: Vec<u8>,
350        }
351
352        let tests = [
353            Test { name: "push OP_FALSE", opcodes: vec![OpFalse], expected: vec![OpFalse] },
354            Test { name: "push OP_TRUE", opcodes: vec![OpTrue], expected: vec![OpTrue] },
355            Test { name: "push OP_0", opcodes: vec![Op0], expected: vec![Op0] },
356            Test { name: "push OP_1 OP_2", opcodes: vec![Op1, Op2], expected: vec![Op1, Op2] },
357            Test { name: "push OP_BLAKE2B OP_EQUAL", opcodes: vec![OpBlake2b, OpEqual], expected: vec![OpBlake2b, OpEqual] },
358        ];
359
360        // Run tests and individually add each op via AddOp.
361        for test in tests.iter() {
362            let mut builder = ScriptBuilder::new();
363            test.opcodes.iter().for_each(|opcode| {
364                builder.add_op(*opcode).expect("the script is canonical");
365            });
366            let result = builder.script();
367            assert_eq!(result, &test.expected, "{} wrong result using add_op", test.name);
368        }
369
370        // Run tests and bulk add ops via AddOps.
371        for test in tests.iter() {
372            let mut builder = ScriptBuilder::new();
373            let result = builder.add_ops(&test.opcodes).expect("the script is canonical").script();
374            assert_eq!(result, &test.expected, "{} wrong result using add_ops", test.name);
375        }
376    }
377
378    /// Tests that pushing signed integers to a script via the ScriptBuilder API works as expected.
379    #[test]
380    fn test_add_i64() {
381        struct Test {
382            name: &'static str,
383            val: i64,
384            expected: Vec<u8>,
385        }
386
387        let tests = vec![
388            Test { name: "push -1", val: -1, expected: vec![Op1Negate] },
389            Test { name: "push small int 0", val: 0, expected: vec![Op0] },
390            Test { name: "push small int 1", val: 1, expected: vec![Op1] },
391            Test { name: "push small int 2", val: 2, expected: vec![Op2] },
392            Test { name: "push small int 3", val: 3, expected: vec![Op3] },
393            Test { name: "push small int 4", val: 4, expected: vec![Op4] },
394            Test { name: "push small int 5", val: 5, expected: vec![Op5] },
395            Test { name: "push small int 6", val: 6, expected: vec![Op6] },
396            Test { name: "push small int 7", val: 7, expected: vec![Op7] },
397            Test { name: "push small int 8", val: 8, expected: vec![Op8] },
398            Test { name: "push small int 9", val: 9, expected: vec![Op9] },
399            Test { name: "push small int 10", val: 10, expected: vec![Op10] },
400            Test { name: "push small int 11", val: 11, expected: vec![Op11] },
401            Test { name: "push small int 12", val: 12, expected: vec![Op12] },
402            Test { name: "push small int 13", val: 13, expected: vec![Op13] },
403            Test { name: "push small int 14", val: 14, expected: vec![Op14] },
404            Test { name: "push small int 15", val: 15, expected: vec![Op15] },
405            Test { name: "push small int 16", val: 16, expected: vec![Op16] },
406            Test { name: "push 17", val: 17, expected: vec![OpData1, 0x11] },
407            Test { name: "push 65", val: 65, expected: vec![OpData1, 0x41] },
408            Test { name: "push 127", val: 127, expected: vec![OpData1, 0x7f] },
409            Test { name: "push 128", val: 128, expected: vec![OpData2, 0x80, 0] },
410            Test { name: "push 255", val: 255, expected: vec![OpData2, 0xff, 0] },
411            Test { name: "push 256", val: 256, expected: vec![OpData2, 0, 0x01] },
412            Test { name: "push 32767", val: 32767, expected: vec![OpData2, 0xff, 0x7f] },
413            Test { name: "push 32768", val: 32768, expected: vec![OpData3, 0, 0x80, 0] },
414            Test { name: "push -2", val: -2, expected: vec![OpData1, 0x82] },
415            Test { name: "push -3", val: -3, expected: vec![OpData1, 0x83] },
416            Test { name: "push -4", val: -4, expected: vec![OpData1, 0x84] },
417            Test { name: "push -5", val: -5, expected: vec![OpData1, 0x85] },
418            Test { name: "push -17", val: -17, expected: vec![OpData1, 0x91] },
419            Test { name: "push -65", val: -65, expected: vec![OpData1, 0xc1] },
420            Test { name: "push -127", val: -127, expected: vec![OpData1, 0xff] },
421            Test { name: "push -128", val: -128, expected: vec![OpData2, 0x80, 0x80] },
422            Test { name: "push -255", val: -255, expected: vec![OpData2, 0xff, 0x80] },
423            Test { name: "push -256", val: -256, expected: vec![OpData2, 0x00, 0x81] },
424            Test { name: "push -32767", val: -32767, expected: vec![OpData2, 0xff, 0xff] },
425            Test { name: "push -32768", val: -32768, expected: vec![OpData3, 0x00, 0x80, 0x80] },
426            Test {
427                name: "push 9223372036854775807",
428                val: 9223372036854775807,
429                expected: vec![OpData8, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0x7F],
430            },
431        ];
432
433        for test in tests {
434            let mut builder = ScriptBuilder::new();
435            let result = builder.add_i64(test.val).expect("the script is canonical").script();
436            assert_eq!(result, test.expected, "{} wrong result", test.name);
437        }
438
439        // special case that used in bitcoind test
440        let mut builder = ScriptBuilder::new();
441        let result = builder.add_i64_min().expect("the script is canonical").script();
442        assert_eq!(
443            result,
444            vec![OpData9, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x80, 0x80],
445            "push -9223372036854775808 wrong result"
446        );
447    }
448
449    /// Tests that pushing data to a script via the ScriptBuilder API works as expected and conforms to BIP0062.
450    #[test]
451    fn test_add_data() {
452        struct Test {
453            name: &'static str,
454            data: Vec<u8>,
455            expected: ScriptBuilderResult<Vec<u8>>,
456            /// use add_data_unchecked instead of add_data
457            unchecked: bool,
458        }
459
460        let tests = vec![
461            // BIP0062: Pushing an empty byte sequence must use OP_0.
462            Test { name: "push empty byte sequence", data: vec![], expected: Ok(vec![Op0]), unchecked: false },
463            // BIP0062: Pushing a 1-byte sequence of byte 0x01 through 0x10 must use OP_n.
464            Test { name: "push 1 byte 0x01", data: vec![0x01], expected: Ok(vec![Op1]), unchecked: false },
465            Test { name: "push 1 byte 0x02", data: vec![0x02], expected: Ok(vec![Op2]), unchecked: false },
466            Test { name: "push 1 byte 0x03", data: vec![0x03], expected: Ok(vec![Op3]), unchecked: false },
467            Test { name: "push 1 byte 0x04", data: vec![0x04], expected: Ok(vec![Op4]), unchecked: false },
468            Test { name: "push 1 byte 0x05", data: vec![0x05], expected: Ok(vec![Op5]), unchecked: false },
469            Test { name: "push 1 byte 0x06", data: vec![0x06], expected: Ok(vec![Op6]), unchecked: false },
470            Test { name: "push 1 byte 0x07", data: vec![0x07], expected: Ok(vec![Op7]), unchecked: false },
471            Test { name: "push 1 byte 0x08", data: vec![0x08], expected: Ok(vec![Op8]), unchecked: false },
472            Test { name: "push 1 byte 0x09", data: vec![0x09], expected: Ok(vec![Op9]), unchecked: false },
473            Test { name: "push 1 byte 0x0a", data: vec![0x0a], expected: Ok(vec![Op10]), unchecked: false },
474            Test { name: "push 1 byte 0x0b", data: vec![0x0b], expected: Ok(vec![Op11]), unchecked: false },
475            Test { name: "push 1 byte 0x0c", data: vec![0x0c], expected: Ok(vec![Op12]), unchecked: false },
476            Test { name: "push 1 byte 0x0d", data: vec![0x0d], expected: Ok(vec![Op13]), unchecked: false },
477            Test { name: "push 1 byte 0x0e", data: vec![0x0e], expected: Ok(vec![Op14]), unchecked: false },
478            Test { name: "push 1 byte 0x0f", data: vec![0x0f], expected: Ok(vec![Op15]), unchecked: false },
479            Test { name: "push 1 byte 0x10", data: vec![0x10], expected: Ok(vec![Op16]), unchecked: false },
480            // BIP0062: Pushing the byte 0x81 must use OP_1NEGATE.
481            Test { name: "push 1 byte 0x81", data: vec![0x81], expected: Ok(vec![Op1Negate]), unchecked: false },
482            // BIP0062: Pushing any other byte sequence up to 75 bytes must
483            // use the normal data push (opcode byte n, with n the number of
484            // bytes, followed n bytes of data being pushed).
485            Test { name: "push 1 byte 0x00", data: vec![0x00], expected: Ok(vec![OpData1, 0x00]), unchecked: false },
486            Test { name: "push 1 byte 0x11", data: vec![0x11], expected: Ok(vec![OpData1, 0x11]), unchecked: false },
487            Test { name: "push 1 byte 0x80", data: vec![0x80], expected: Ok(vec![OpData1, 0x80]), unchecked: false },
488            Test { name: "push 1 byte 0x82", data: vec![0x82], expected: Ok(vec![OpData1, 0x82]), unchecked: false },
489            Test { name: "push 1 byte 0xff", data: vec![0xff], expected: Ok(vec![OpData1, 0xff]), unchecked: false },
490            Test {
491                name: "push data len 17",
492                data: vec![0x49; 17],
493                expected: Ok(once(OpData17).chain(repeat_n(0x49, 17)).collect()),
494                unchecked: false,
495            },
496            Test {
497                name: "push data len 75",
498                data: vec![0x49; 75],
499                expected: Ok(once(OpData75).chain(repeat_n(0x49, 75)).collect()),
500                unchecked: false,
501            },
502            // BIP0062: Pushing 76 to 255 bytes must use OP_PUSHDATA1.
503            Test {
504                name: "push data len 76",
505                data: vec![0x49; 76],
506                expected: Ok(once(OpPushData1).chain(once(76)).chain(repeat_n(0x49, 76)).collect()),
507                unchecked: false,
508            },
509            Test {
510                name: "push data len 255",
511                data: vec![0x49; 255],
512                expected: Ok(once(OpPushData1).chain(once(255)).chain(repeat_n(0x49, 255)).collect()),
513                unchecked: false,
514            },
515            // // BIP0062: Pushing 256 to 520 bytes must use OP_PUSHDATA2.
516            Test {
517                name: "push data len 256",
518                data: vec![0x49; 256],
519                expected: Ok(once(OpPushData2).chain([0, 1]).chain(repeat_n(0x49, 256)).collect()),
520                unchecked: false,
521            },
522            Test {
523                name: "push data len 520",
524                data: vec![0x49; 520],
525                expected: Ok(once(OpPushData2).chain([8, 2]).chain(repeat_n(0x49, 520)).collect()),
526                unchecked: false,
527            },
528            Test {
529                name: "push data len 300001 (canonical)",
530                data: vec![0x1; 3000001],
531                expected: Err(ScriptBuilderError::DataRejected(3000006, MAX_SCRIPTS_SIZE)),
532                unchecked: false,
533            },
534            // // Additional tests for the add_data_unchecked function that
535            // // intentionally allows data pushes to exceed the limit for
536            // // testing purposes.
537
538            // 3-byte data push via OP_PUSHDATA_2.
539            Test {
540                name: "push data len 32767 (non-canonical)",
541                data: vec![0x49; 32767],
542                expected: Ok(once(OpPushData2).chain([255, 127]).chain(repeat_n(0x49, 32767)).collect()),
543                unchecked: true,
544            },
545            // 5-byte data push via OP_PUSHDATA_4.
546            Test {
547                name: "push data len 65536 (non-canonical)",
548                data: vec![0x49; 65536],
549                expected: Ok(once(OpPushData4).chain([0, 0, 1, 0]).chain(repeat_n(0x49, 65536)).collect()),
550                unchecked: true,
551            },
552        ];
553
554        for test in tests {
555            let mut builder = ScriptBuilder::new();
556            let result = match test.unchecked {
557                false => builder.add_data(&test.data).map(|x| x.drain()),
558                true => {
559                    builder.add_data_unchecked(&test.data);
560                    Ok(builder.drain())
561                }
562            };
563            assert_eq!(result, test.expected, "{} wrong result", test.name);
564        }
565    }
566
567    #[test]
568    fn test_add_data_with_push_opcode() {
569        let tests = [
570            ("empty", vec![], vec![Op0]),
571            ("small int", vec![0x01], vec![OpData1, 0x01]),
572            ("one negate", vec![0x81], vec![OpData1, 0x81]),
573            ("normal byte", vec![0x11], vec![OpData1, 0x11]),
574            ("pushdata1", vec![0x49; 76], once(OpPushData1).chain(once(76)).chain(repeat_n(0x49, 76)).collect()),
575        ];
576
577        for (name, data, expected) in tests {
578            let result = ScriptBuilder::new().add_data_with_push_opcode(&data).expect("explicit push is allowed").drain();
579            assert_eq!(result, expected, "{name} wrong result");
580        }
581    }
582
583    #[test]
584    fn test_explicit_push_encoded_size_matches_canonical_on_boundaries() {
585        // Len 1 is excluded because canonical size is data-dependent there:
586        // small ints and Op1Negate encode as a single opcode. The other
587        // boundary lengths below are length-only.
588        let boundary_lengths = [
589            0,
590            2,
591            OP_DATA_MAX_VAL as usize - 1,
592            OP_DATA_MAX_VAL as usize,
593            OP_DATA_MAX_VAL as usize + 1,
594            u8::MAX as usize - 1,
595            u8::MAX as usize,
596            u8::MAX as usize + 1,
597            u16::MAX as usize - 1,
598            u16::MAX as usize,
599            u16::MAX as usize + 1,
600        ];
601
602        for data_len in boundary_lengths {
603            let data = vec![0x49; data_len];
604            assert_eq!(
605                ScriptBuilder::explicit_push_encoded_size(data_len),
606                ScriptBuilder::canonical_data_size(&data),
607                "wrong encoded size for len {data_len}"
608            );
609        }
610    }
611
612    #[test]
613    fn test_add_data_with_push_opcode_validates_explicit_size() {
614        let flags = Default::default();
615        let max_scripts_size = MAX_SCRIPTS_SIZE;
616
617        let mut canonical_builder = ScriptBuilder::with_flags(flags);
618        canonical_builder.script_mut().resize(max_scripts_size - 1, OpTrue);
619        assert_eq!(canonical_builder.add_data(&[0x01]).map(|_| ()), Ok(()));
620        assert_eq!(canonical_builder.script().len(), max_scripts_size);
621
622        let mut explicit_builder = ScriptBuilder::with_flags(flags);
623        explicit_builder.script_mut().resize(max_scripts_size - 1, OpTrue);
624        let original_script = Vec::from(explicit_builder.script());
625
626        let result = explicit_builder.add_data_with_push_opcode(&[0x01]).map(|_| ());
627        assert_eq!(result, Err(ScriptBuilderError::DataRejected(2, max_scripts_size)));
628        assert_eq!(explicit_builder.script(), &original_script);
629    }
630
631    #[test]
632    fn test_add_data_validates_zero_byte_canonical_size() {
633        let flags = Default::default();
634        let max_scripts_size = MAX_SCRIPTS_SIZE;
635
636        let mut valid_builder = ScriptBuilder::with_flags(flags);
637        valid_builder.script_mut().resize(max_scripts_size - 2, OpTrue);
638        valid_builder.add_data(&[0x00]).expect("two-byte zero push should fit");
639        assert_eq!(valid_builder.script().len(), max_scripts_size);
640        assert_eq!(&valid_builder.script()[max_scripts_size - 2..], &[OpData1, 0x00]);
641
642        let mut builder = ScriptBuilder::with_flags(flags);
643        builder.script_mut().resize(max_scripts_size - 1, OpTrue);
644        let original_script = Vec::from(builder.script());
645
646        let result = builder.add_data(&[0x00]).map(|_| ());
647        assert_eq!(result, Err(ScriptBuilderError::DataRejected(2, max_scripts_size)));
648        assert_eq!(builder.script(), &original_script);
649    }
650
651    #[test]
652    fn test_u64() {
653        struct Test {
654            name: &'static str,
655            value: u64,
656            expected: Vec<u8>,
657        }
658
659        let tests = vec![
660            Test { name: "0x00", value: 0x00, expected: vec![Op0] },
661            Test { name: "0x01", value: 0x01, expected: vec![Op1] },
662            Test { name: "0xff", value: 0xff, expected: vec![OpData1, 0xff] },
663            Test { name: "0xffee", value: 0xffee, expected: vec![OpData2, 0xee, 0xff] },
664            Test { name: "0xffeedd", value: 0xffeedd, expected: vec![OpData3, 0xdd, 0xee, 0xff] },
665            Test { name: "0xffeeddcc", value: 0xffeeddcc, expected: vec![OpData4, 0xcc, 0xdd, 0xee, 0xff] },
666            Test { name: "0xffeeddccbb", value: 0xffeeddccbb, expected: vec![OpData5, 0xbb, 0xcc, 0xdd, 0xee, 0xff] },
667            Test { name: "0xffeeddccbbaa", value: 0xffeeddccbbaa, expected: vec![OpData6, 0xaa, 0xbb, 0xcc, 0xdd, 0xee, 0xff] },
668            Test {
669                name: "0xffeeddccbbaa99",
670                value: 0xffeeddccbbaa99,
671                expected: vec![OpData7, 0x99, 0xaa, 0xbb, 0xcc, 0xdd, 0xee, 0xff],
672            },
673            Test {
674                name: "0xffeeddccbbaa9988",
675                value: 0xffeeddccbbaa9988,
676                expected: vec![OpData8, 0x88, 0x99, 0xaa, 0xbb, 0xcc, 0xdd, 0xee, 0xff],
677            },
678            Test { name: "0xffffffffffffffff", value: u64::MAX, expected: once(OpData8).chain(repeat_n(0xff, 8)).collect() },
679        ];
680
681        for test in tests {
682            let result = ScriptBuilder::new().add_u64(test.value).expect("the script is canonical").drain();
683            assert_eq!(result, test.expected, "{} wrong result", test.name);
684            let result = ScriptBuilder::new().add_lock_time(test.value).expect("the script is canonical").drain();
685            assert_eq!(result, test.expected, "{} wrong lock time result", test.name);
686            let result = ScriptBuilder::new().add_sequence(test.value).expect("the script is canonical").drain();
687            assert_eq!(result, test.expected, "{} wrong sequence result", test.name);
688        }
689    }
690
691    /// Ensures that all of the functions that can be used to add data to a script don't allow
692    /// the script to exceed the max allowed size.
693    #[test]
694    fn test_exceed_max_script_size() {
695        fn full_builder() -> ScriptBuilder {
696            let mut builder = ScriptBuilder::new();
697            builder.add_data_unchecked(&vec![0u8; MAX_SCRIPTS_SIZE - 3]);
698            builder
699        }
700        // Start off by constructing a max size script.
701        let mut builder = full_builder();
702        let original_result: Vec<u8> = Vec::from(builder.script());
703
704        // Ensure adding data that would exceed the maximum size of the script
705        // does not add the data.
706        let result = builder.add_data(&[0u8]).map(|_| ());
707        assert_eq!(
708            result,
709            Err(ScriptBuilderError::DataRejected(2, MAX_SCRIPTS_SIZE)),
710            "adding data that would exceed the maximum size of the script must fail"
711        );
712        assert_eq!(builder.script(), &original_result, "unexpected modified script");
713
714        // Ensure adding an opcode that would exceed the maximum size of the
715        // script does not add the data.
716        let result = builder.add_op(Op0).map(|_| ());
717        assert_eq!(
718            result,
719            Err(ScriptBuilderError::OpCodeRejected(Op0, MAX_SCRIPTS_SIZE)),
720            "adding an opcode that would exceed the maximum size of the script must fail"
721        );
722        assert_eq!(builder.script(), &original_result, "unexpected modified script");
723
724        // Ensure adding an opcode array that would exceed the maximum size of the
725        // script does not add the data.
726        let result = builder.add_ops(&[OpCheckSig]).map(|_| ());
727        assert_eq!(
728            result,
729            Err(ScriptBuilderError::OpCodesRejected(1, MAX_SCRIPTS_SIZE)),
730            "adding an opcode array that would exceed the maximum size of the script must fail"
731        );
732        assert_eq!(builder.script(), &original_result, "unexpected modified script");
733
734        // Ensure adding an integer that would exceed the maximum size of the
735        // script does not add the data.
736        let result = builder.add_i64(0).map(|_| ());
737        assert_eq!(
738            result,
739            Err(ScriptBuilderError::IntegerRejected(0, MAX_SCRIPTS_SIZE)),
740            "adding an integer that would exceed the maximum size of the script must fail"
741        );
742        assert_eq!(builder.script(), &original_result, "unexpected modified script");
743
744        // Ensure adding a lock time that would exceed the maximum size of the
745        // script does not add the data.
746        let result = builder.add_lock_time(0).map(|_| ());
747        assert_eq!(
748            result,
749            Err(ScriptBuilderError::DataRejected(1, MAX_SCRIPTS_SIZE)),
750            "adding a lock time that would exceed the maximum size of the script must fail"
751        );
752        assert_eq!(builder.script(), &original_result, "unexpected modified script");
753
754        // Ensure adding a sequence that would exceed the maximum size of the
755        // script does not add the data.
756        let result = builder.add_sequence(0).map(|_| ());
757        assert_eq!(
758            result,
759            Err(ScriptBuilderError::DataRejected(1, MAX_SCRIPTS_SIZE)),
760            "adding a sequence that would exceed the maximum size of the script must fail"
761        );
762        assert_eq!(builder.script(), &original_result, "unexpected modified script");
763    }
764}