1use crate::{
2 builder::{BuilderError, CallFlags, InteropService},
3 codec::Encoder,
4 crypto::Secp256r1PublicKey,
5 neo_crypto::utils::{FromBase64String, FromHexString, ToHexString},
6 Bytes, ContractParameter, ContractParameterType, OpCode, ParameterValue, ScriptHashExtension,
7};
8use futures_util::future::ok;
9use getset::{Getters, Setters};
10use num_bigint::BigInt;
11use num_traits::{Signed, ToPrimitive};
12use primitive_types::H160;
13use serde::{Deserialize, Serialize};
14use std::{
15 cmp::PartialEq,
16 collections::HashMap,
17 convert::TryInto,
18 fmt::{Debug, Formatter},
19 str::FromStr,
20};
21use tokio::io::AsyncWriteExt;
22
23#[derive(Debug, PartialEq, Eq, Hash, Getters, Setters)]
44pub struct ScriptBuilder {
45 #[getset(get = "pub")]
46 pub script: Encoder,
47}
48
49impl ScriptBuilder {
50 pub fn new() -> Self {
60 Self { script: Encoder::new() }
61 }
62
63 pub fn op_code(&mut self, op_codes: &[OpCode]) -> &mut Self {
83 for opcode in op_codes {
84 self.script.write_u8(opcode.opcode());
85 }
86 self
87 }
88
89 pub fn op_code_with_arg(&mut self, opcode: OpCode, argument: Bytes) -> &mut Self {
110 self.script.write_u8(opcode.opcode());
111 let _ = self.script.write_bytes(&argument);
112 self
113 }
114
115 pub fn contract_call(
147 &mut self,
148 hash160: &H160,
149 method: &str,
150 params: &[ContractParameter],
151 call_flags: Option<CallFlags>,
152 ) -> Result<&mut Self, BuilderError> {
153 if params.is_empty() {
154 self.op_code(&[OpCode::NewArray0]);
155 } else {
156 self.push_params(params);
157 }
158
159 Ok(self
160 .push_integer(BigInt::from(match call_flags {
161 Some(flags) => flags.value(),
162 None => CallFlags::All.value(),
163 }))
164 .push_data(method.as_bytes().to_vec())
165 .push_data(hash160.to_vec())
166 .sys_call(InteropService::SystemContractCall))
167 }
168
169 pub fn sys_call(&mut self, operation: InteropService) -> &mut Self {
189 self.push_opcode_bytes(
190 OpCode::Syscall,
191 operation
192 .hash()
193 .from_hex_string()
194 .map_err(|e| {
195 BuilderError::IllegalArgument(format!("Invalid operation hash: {}", e))
196 })
197 .expect("InteropService hash should always be valid hex"),
198 )
199 }
200
201 pub fn push_params(&mut self, params: &[ContractParameter]) -> Result<&mut Self, BuilderError> {
225 for param in params {
226 self.push_param(param).map_err(|e| {
227 BuilderError::IllegalArgument(format!("Failed to push parameter: {}", e))
228 })?;
229 }
230
231 Ok(self.push_integer(BigInt::from(params.len())).op_code(&[OpCode::Pack]))
232 }
233
234 pub fn push_param(&mut self, param: &ContractParameter) -> Result<&mut Self, BuilderError> {
255 if param.get_type() == ContractParameterType::Any {
256 self.op_code(&[OpCode::PushNull]);
257 return Ok(self);
258 }
259 match ¶m
260 .value
261 .clone()
262 .ok_or_else(|| BuilderError::IllegalArgument("Parameter value is None".to_string()))?
263 {
264 ParameterValue::Boolean(b) => self.push_bool(*b),
265 ParameterValue::Integer(i) => self.push_integer(BigInt::from(i.clone())),
266 ParameterValue::ByteArray(b) => {
267 let bytes = b.from_base64_string().map_err(|e| {
269 BuilderError::IllegalArgument(format!(
270 "Failed to decode base64 ByteArray: {}",
271 e
272 ))
273 })?;
274 self.push_data(bytes)
275 },
276 ParameterValue::Signature(b) | ParameterValue::PublicKey(b) =>
277 self.push_data(b.as_bytes().to_vec()),
278 ParameterValue::H160(h) => self.push_data(h.as_bytes().to_vec()),
279 ParameterValue::H256(h) => self.push_data(h.as_bytes().to_vec()),
280 ParameterValue::String(s) => self.push_data(s.as_bytes().to_vec()),
281 ParameterValue::Array(arr) => self.push_array(arr).map_err(|e| {
282 BuilderError::IllegalArgument(format!("Failed to push array: {}", e))
283 })?,
284 ParameterValue::Map(map) => self
285 .push_map(&map.0)
286 .map_err(|e| BuilderError::IllegalArgument(format!("Failed to push map: {}", e)))?,
287 _ =>
288 return Err(BuilderError::IllegalArgument("Unsupported parameter type".to_string())),
289 };
290
291 Ok(self)
292 }
293
294 pub fn push_integer(&mut self, i: BigInt) -> &mut Self {
318 if i >= BigInt::from(-1) && i <= BigInt::from(16) {
319 self.op_code(
320 vec![OpCode::try_from(i.to_i32().unwrap() as u8 + OpCode::Push0 as u8).unwrap()]
321 .as_slice(),
322 );
323 } else {
324 let mut bytes = i.to_signed_bytes_le();
325
326 while bytes.len() > 1 && bytes[bytes.len() - 1] == 0 && !i.is_negative() {
332 bytes.pop();
333 }
334
335 let len = bytes.len();
336
337 match len {
340 1 => self.push_opcode_bytes(OpCode::PushInt8, bytes),
341 2 => self.push_opcode_bytes(OpCode::PushInt16, bytes),
342 len if len <= 4 => self.push_opcode_bytes(
343 OpCode::PushInt32,
344 Self::pad_right(&bytes, 4, i.is_negative()),
345 ),
346 len if len <= 8 => self.push_opcode_bytes(
347 OpCode::PushInt64,
348 Self::pad_right(&bytes, 8, i.is_negative()),
349 ),
350 len if len <= 16 => self.push_opcode_bytes(
351 OpCode::PushInt128,
352 Self::pad_right(&bytes, 16, i.is_negative()),
353 ),
354 len if len <= 32 => self.push_opcode_bytes(
355 OpCode::PushInt256,
356 Self::pad_right(&bytes, 32, i.is_negative()),
357 ),
358 _ => {
359 eprintln!("Warning: Integer too large, truncating to 32 bytes");
362 self.push_opcode_bytes(
363 OpCode::PushInt256,
364 Self::pad_right(&bytes[..32.min(bytes.len())], 32, i.is_negative()),
365 )
366 },
367 };
368 }
369
370 self
371 }
372
373 pub fn push_opcode_bytes(&mut self, opcode: OpCode, argument: Vec<u8>) -> &mut ScriptBuilder {
394 self.script.write_u8(opcode as u8);
395 self.script.write_bytes(&argument);
396
397 self
398 }
399
400 fn pad_right(bytes: &[u8], size: usize, negative: bool) -> Vec<u8> {
401 let pad_value = if negative { 0xFF } else { 0 };
402
403 let mut padded = vec![0; size];
404 padded[0..bytes.len()].copy_from_slice(bytes);
405 padded[bytes.len()..].fill(pad_value);
406 padded
407 }
408
409 pub fn push_data(&mut self, data: Vec<u8>) -> &mut Self {
430 match data.len() {
431 0..=0xff => {
432 self.op_code(&[OpCode::PushData1]);
433 self.script.write_u8(data.len() as u8);
434 let _ = self.script.write_bytes(&data);
435 },
436 0x100..=0xffff => {
437 self.op_code(&[OpCode::PushData2]);
438 self.script.write_u16(data.len() as u16);
439 let _ = self.script.write_bytes(&data);
440 },
441 _ => {
442 self.op_code(&[OpCode::PushData4]);
443 self.script.write_u32(data.len() as u32);
444 let _ = self.script.write_bytes(&data);
445 }, }
447 self
448 }
449
450 pub fn push_bool(&mut self, b: bool) -> &mut Self {
469 if b {
470 self.op_code(&[OpCode::PushTrue])
471 } else {
472 self.op_code(&[OpCode::PushFalse])
473 };
474 self
475 }
476
477 pub fn push_array(&mut self, arr: &[ContractParameter]) -> Result<&mut Self, BuilderError> {
489 if arr.is_empty() {
490 self.op_code(&[OpCode::NewArray0]);
491 } else {
492 self.push_params(arr);
493 };
494 Ok(self)
495 }
496
497 pub fn push_map(
508 &mut self,
509 map: &HashMap<ContractParameter, ContractParameter>,
510 ) -> Result<&mut Self, BuilderError> {
511 for (k, v) in map {
512 let kk: ContractParameter = k.clone().into();
513 let vv: ContractParameter = v.clone().into();
514 self.push_param(&vv).unwrap();
515 self.push_param(&kk).unwrap();
516 }
517
518 Ok(self.push_integer(BigInt::from(map.len())).op_code(&[OpCode::PackMap]))
519 }
520
521 pub fn pack(&mut self) -> &mut Self {
527 self.op_code(&[OpCode::Pack])
528 }
529
530 pub fn to_bytes(&self) -> Bytes {
532 self.script.to_bytes()
533 }
534
535 pub fn build_verification_script(pub_key: &Secp256r1PublicKey) -> Bytes {
545 let mut sb = ScriptBuilder::new();
546 sb.push_data(pub_key.get_encoded(true))
547 .sys_call(InteropService::SystemCryptoCheckSig);
548 sb.to_bytes()
549 }
550
551 pub fn build_multi_sig_script(
563 pubkeys: &mut [Secp256r1PublicKey],
564 threshold: u8,
565 ) -> Result<Bytes, BuilderError> {
566 let mut sb = ScriptBuilder::new();
567 sb.push_integer(BigInt::from(threshold));
568 pubkeys.sort_by(|a, b| a.get_encoded(true).cmp(&b.get_encoded(true)));
569 for pk in pubkeys.iter() {
570 sb.push_data(pk.get_encoded(true));
571 }
572 sb.push_integer(BigInt::from(pubkeys.len()));
573 sb.sys_call(InteropService::SystemCryptoCheckMultiSig);
574 Ok(sb.to_bytes())
575 }
576
577 pub fn build_contract_script(
613 sender: &H160,
614 nef_checksum: u32,
615 name: &str,
616 ) -> Result<Bytes, BuilderError> {
617 let mut sb = ScriptBuilder::new();
618 sb.op_code(&[OpCode::Abort])
619 .push_data(sender.to_vec())
620 .push_integer(BigInt::from(nef_checksum))
621 .push_data(name.as_bytes().to_vec());
622 Ok(sb.to_bytes())
623 }
624
625 pub fn build_contract_call_and_unwrap_iterator(
674 contract_hash: &H160,
675 method: &str,
676 params: &[ContractParameter],
677 max_items: u32,
678 call_flags: Option<CallFlags>,
679 ) -> Result<Bytes, BuilderError> {
680 let mut sb = Self::new();
681 sb.push_integer(BigInt::from(max_items));
682
683 sb.contract_call(contract_hash, method, params, call_flags).unwrap();
684
685 sb.op_code(&[OpCode::NewArray]);
686
687 let cycle_start = sb.len();
688 sb.op_code(&[OpCode::Over]);
689 sb.sys_call(InteropService::SystemIteratorNext);
690
691 let jmp_if_not = sb.len();
692 sb.op_code_with_arg(OpCode::JmpIf, vec![0]);
693
694 sb.op_code(&[OpCode::Dup, OpCode::Push2, OpCode::Pick])
695 .sys_call(InteropService::SystemIteratorValue)
696 .op_code(&[
697 OpCode::Append,
698 OpCode::Dup,
699 OpCode::Size,
700 OpCode::Push3,
701 OpCode::Pick,
702 OpCode::Ge,
703 ]);
704
705 let jmp_if_max = sb.len();
706 sb.op_code_with_arg(OpCode::JmpIf, vec![0]);
707
708 let jmp_offset = sb.len();
709 let jmp_bytes = (cycle_start as i32 - jmp_offset as i32) as i8;
711 sb.op_code_with_arg(OpCode::Jmp, vec![jmp_bytes as u8]);
712
713 let load_result = sb.len();
714 sb.op_code(&[OpCode::Nip, OpCode::Nip]);
715
716 let mut script = sb.to_bytes();
717 let jmp_not_bytes = (load_result - jmp_if_not) as i8;
718 script[jmp_if_not + 1] = jmp_not_bytes as u8;
719
720 let jmp_max_bytes = (load_result - jmp_if_max) as i8;
721 script[jmp_if_max + 1] = jmp_max_bytes as u8;
722
723 Ok(script)
724 }
725
726 pub fn len(&self) -> usize {
746 self.script().size()
747 }
748}
749
750#[cfg(test)]
751mod tests {
752 use super::*;
753 use crate::{
754 neo_types::{contract::ContractParameterMap, ContractParameter, ContractParameterType},
755 prelude::Bytes,
756 };
757 use num_bigint::BigInt;
758 use std::collections::HashMap;
759
760 #[test]
761 fn test_push_empty_array() {
762 let mut builder = ScriptBuilder::new();
763 builder.push_array(&[]).unwrap();
764 assert_builder(&builder, &[OpCode::NewArray0 as u8]);
765 }
766
767 #[test]
768 fn test_push_byte_array() {
769 let mut builder = ScriptBuilder::new();
770 let data = vec![0x01, 0x02, 0x03];
771 builder.push_data(data.clone());
772
773 let mut expected = vec![OpCode::PushData1 as u8, data.len() as u8];
774 expected.extend(data);
775 assert_builder(&builder, &expected);
776 }
777
778 #[test]
779 fn test_push_string() {
780 let mut builder = ScriptBuilder::new();
781 let string_data = "Hello, Neo!";
782 builder.push_data(string_data.as_bytes().to_vec());
783
784 let mut expected = vec![OpCode::PushData1 as u8, string_data.len() as u8];
785 expected.extend(string_data.as_bytes());
786 assert_builder(&builder, &expected);
787 }
788
789 #[test]
790 fn test_push_integer() {
791 let mut builder = ScriptBuilder::new();
792
793 builder.push_integer(BigInt::from(0));
795 assert_builder(&builder, &[OpCode::Push0 as u8]);
796
797 let mut builder = ScriptBuilder::new();
798 builder.push_integer(BigInt::from(16));
799 assert_builder(&builder, &[OpCode::Push16 as u8]);
800
801 let mut builder = ScriptBuilder::new();
803 builder.push_integer(BigInt::from(255));
804 assert_builder(&builder, &[OpCode::PushInt8 as u8, 0xff]);
805
806 let mut builder = ScriptBuilder::new();
807 builder.push_integer(BigInt::from(65535));
808 assert_builder(&builder, &[OpCode::PushInt16 as u8, 0xff, 0xff]);
809
810 let mut builder = ScriptBuilder::new();
812 builder.push_integer(BigInt::from(-1000000000000000i64));
813 let expected_bytes = vec![OpCode::PushInt64 as u8, 0, 128, 57, 91, 129, 114, 252, 255];
816 assert_builder(&builder, &expected_bytes);
817
818 let mut builder = ScriptBuilder::new();
819 builder.push_integer(BigInt::from(1000000000000000i64));
820 let pos_big_int = BigInt::from(1000000000000000i64);
823 let pos_bytes = pos_big_int.to_signed_bytes_le();
824 let mut expected_pos = vec![OpCode::PushInt64 as u8];
825 let padded_pos = ScriptBuilder::pad_right(&pos_bytes, 8, false);
826 expected_pos.extend(padded_pos);
827 assert_builder(&builder, &expected_pos);
828 }
829
830 #[test]
831 fn test_verification_script() {
832 let public_key = Secp256r1PublicKey::from_bytes(
833 &hex::decode("035fdb1d1f06759547020891ae97c729327853aeb1256b6fe0473bc2e9fa42ff50")
834 .unwrap(),
835 )
836 .unwrap();
837
838 let script = ScriptBuilder::build_verification_script(&public_key);
839
840 let mut expected = vec![0x0c, 0x21]; expected.extend_from_slice(
844 &hex::decode("035fdb1d1f06759547020891ae97c729327853aeb1256b6fe0473bc2e9fa42ff50")
845 .unwrap(),
846 );
847 expected.push(0x41); expected
849 .extend_from_slice(&hex::decode(&InteropService::SystemCryptoCheckSig.hash()).unwrap());
850
851 assert_eq!(script.to_vec(), expected);
852 }
853
854 #[test]
855 fn test_map() {
856 let mut map = HashMap::new();
857 map.insert(
858 ContractParameter::string("first".to_string()),
859 ContractParameter::byte_array(hex::decode("7365636f6e64").unwrap()),
860 );
861
862 let mut builder = ScriptBuilder::new();
863 builder.push_map(&map).unwrap();
864
865 let expected = builder.to_bytes().to_hex_string();
866
867 let mut builder2 = ScriptBuilder::new();
868 builder2
869 .push_data(hex::decode("7365636f6e64").unwrap())
870 .push_data("first".as_bytes().to_vec())
871 .push_integer(BigInt::from(1))
872 .op_code(&[OpCode::PackMap]);
873
874 let expected2 = builder2.to_bytes().to_hex_string();
875
876 let mut builder3 = ScriptBuilder::new().push_map(&map).unwrap().to_bytes().to_hex_string();
877
878 assert_eq!(expected, expected2);
879 assert_eq!(expected, builder3);
880 }
881
882 #[test]
883 fn test_map_nested() {
884 let mut inner = ContractParameterMap::new();
885 inner.0.insert(
886 ContractParameter::string("inner_key".to_string()),
887 ContractParameter::integer(42),
888 );
889
890 let mut outer = ContractParameterMap::new();
891 outer.0.insert(
892 ContractParameter::string("outer_key".to_string()),
893 ContractParameter::map(inner),
894 );
895
896 let expected = ScriptBuilder::new().push_map(&outer.0).unwrap().to_bytes().to_hex_string();
897
898 let mut manual_builder = ScriptBuilder::new();
900 manual_builder
901 .push_integer(BigInt::from(42))
902 .push_data("inner_key".as_bytes().to_vec())
903 .push_integer(BigInt::from(1))
904 .op_code(&[OpCode::PackMap])
905 .push_data("outer_key".as_bytes().to_vec())
906 .push_integer(BigInt::from(1))
907 .op_code(&[OpCode::PackMap]);
908
909 let manual_expected = manual_builder.to_bytes().to_hex_string();
910
911 assert_eq!(expected, manual_expected);
912 }
913
914 fn assert_builder(builder: &ScriptBuilder, expected: &[u8]) {
915 assert_eq!(builder.to_bytes().to_vec(), expected);
916 }
917
918 fn byte_array(size: usize) -> Vec<u8> {
919 (0..size).map(|i| (i % 256) as u8).collect()
920 }
921}