Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 22 additions & 1 deletion crates/fiber-lib/src/ckb/config.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@ use crate::Result;
use ckb_jsonrpc_types::{OutPoint as OutPointWrapper, Script as ScriptWrapper};
use ckb_sdk::{traits::DefaultCellCollector, CkbRpcAsyncClient};
use ckb_types::core::ScriptHashType;
use ckb_types::prelude::Builder;
use ckb_types::prelude::{Builder, Pack};
use ckb_types::H256;
use ckb_types::{
core::DepType,
Expand Down Expand Up @@ -270,6 +270,27 @@ pub struct UdtArgInfo {
#[derive(Serialize, Deserialize, Clone, Debug, Default, Eq, PartialEq, Hash)]
pub struct UdtCfgInfos(pub Vec<UdtArgInfo>);

impl UdtCfgInfos {
/// Find a matching UDT info by script (code_hash, hash_type, and args pattern)
pub fn find_matching_udt(&self, udt_script: &Script) -> Option<&UdtArgInfo> {
use regex::Regex;
for udt in &self.0 {
if let Ok(hash_type) = udt_script.hash_type().try_into() {
if udt.script.code_hash.pack() == udt_script.code_hash()
&& udt.script.hash_type == hash_type
{
let args = format!("0x{:x}", udt_script.args().raw_data());
let pattern = Regex::new(&udt.script.args).expect("invalid expression");
if pattern.is_match(&args) {
return Some(udt);
}
}
}
}
None
}
}

impl FromStr for UdtCfgInfos {
type Err = serde_json::Error;

Expand Down
18 changes: 2 additions & 16 deletions crates/fiber-lib/src/ckb/contracts.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,6 @@ use ckb_types::{
prelude::{Builder, Entity, Pack, PackVec},
};
use once_cell::sync::OnceCell;
use regex::Regex;
use serde::{Deserialize, Serialize};
use std::{collections::HashMap, vec};
use thiserror::Error;
Expand Down Expand Up @@ -377,20 +376,7 @@ impl ContractsContext {
}

pub(crate) fn get_udt_info(&self, udt_script: &Script) -> Option<&UdtArgInfo> {
for udt in &self.get_udt_whitelist().0 {
if let Ok(_type) = udt_script.hash_type().try_into() {
if udt.script.code_hash.pack() == udt_script.code_hash()
&& udt.script.hash_type == _type
{
let args = format!("0x{:x}", udt_script.args().raw_data());
let pattern = Regex::new(&udt.script.args).expect("invalid expression");
if pattern.is_match(&args) {
return Some(udt);
}
}
}
}
None
self.get_udt_whitelist().find_matching_udt(udt_script)
}
}

Expand Down Expand Up @@ -445,7 +431,7 @@ pub fn get_cell_deps_count_by_contracts(contracts: Vec<Contract>) -> usize {
get_contracts_context().get_cell_deps_count(contracts)
}

fn get_udt_info(script: &Script) -> Option<UdtArgInfo> {
pub fn get_udt_info(script: &Script) -> Option<UdtArgInfo> {
get_contracts_context().get_udt_info(script).cloned()
}

Expand Down
193 changes: 193 additions & 0 deletions crates/fiber-lib/src/ckb/tests/config.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,10 @@ use crate::ckb::config::{UdtArgInfo, UdtCellDep, UdtCfgInfos, UdtDep, UdtScript}
use crate::fiber::gen::fiber::UdtCfgInfos as MoleculeUdtCfgInfos;
use ckb_jsonrpc_types::OutPoint;
use ckb_types::core::{DepType, ScriptHashType};
use ckb_types::packed::Script;
use ckb_types::prelude::{Builder, Pack};
use ckb_types::H256;
use hex;
use molecule::prelude::Entity;

#[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
Expand Down Expand Up @@ -30,3 +33,193 @@ fn test_udt_whitelist() {
UdtCfgInfos::from(MoleculeUdtCfgInfos::from_slice(&serialized).expect("invalid mol"));
assert_eq!(udt_whitelist, deserialized);
}

#[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
#[cfg_attr(not(target_arch = "wasm32"), test)]
fn test_find_matching_udt_exact_match() {
let code_hash = H256::from([1u8; 32]);
let args = "0x1234".to_string();
let udt_whitelist = UdtCfgInfos(vec![UdtArgInfo {
name: "TestUDT".to_string(),
script: UdtScript {
code_hash: code_hash.clone(),
hash_type: ScriptHashType::Data,
args: args.clone(),
},
auto_accept_amount: Some(100),
cell_deps: vec![],
}]);

let args_bytes = hex::decode(&args[2..]).unwrap();
let script = Script::new_builder()
.code_hash(code_hash.pack())
.hash_type(ScriptHashType::Data.into())
.args(args_bytes.pack())
.build();

let found = udt_whitelist.find_matching_udt(&script);
assert!(found.is_some());
assert_eq!(found.unwrap().name, "TestUDT");
}

#[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
#[cfg_attr(not(target_arch = "wasm32"), test)]
fn test_find_matching_udt_regex_pattern() {
let code_hash = H256::from([2u8; 32]);
let udt_whitelist = UdtCfgInfos(vec![UdtArgInfo {
name: "RegexUDT".to_string(),
script: UdtScript {
code_hash: code_hash.clone(),
hash_type: ScriptHashType::Data,
args: "0x[0-9a-f]{4}".to_string(), // Regex pattern matching 4 hex digits
},
auto_accept_amount: Some(200),
cell_deps: vec![],
}]);

// Test with matching args
let args_bytes = hex::decode("abcd").unwrap();
let script = Script::new_builder()
.code_hash(code_hash.pack())
.hash_type(ScriptHashType::Data.into())
.args(args_bytes.pack())
.build();

let found = udt_whitelist.find_matching_udt(&script);
assert!(found.is_some());
assert_eq!(found.unwrap().name, "RegexUDT");

// Test with non-matching args
let args_bytes_no_match = hex::decode("ab").unwrap(); // Only 2 hex digits (1 byte), doesn't match pattern
let script_no_match = Script::new_builder()
.code_hash(code_hash.pack())
.hash_type(ScriptHashType::Data.into())
.args(args_bytes_no_match.pack())
.build();

let found = udt_whitelist.find_matching_udt(&script_no_match);
assert!(found.is_none());
}

#[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
#[cfg_attr(not(target_arch = "wasm32"), test)]
fn test_find_matching_udt_wrong_code_hash() {
let code_hash1 = H256::from([3u8; 32]);
let code_hash2 = H256::from([4u8; 32]);
let udt_whitelist = UdtCfgInfos(vec![UdtArgInfo {
name: "TestUDT".to_string(),
script: UdtScript {
code_hash: code_hash1.clone(),
hash_type: ScriptHashType::Data,
args: "0x00".to_string(),
},
auto_accept_amount: Some(100),
cell_deps: vec![],
}]);

let args_bytes = hex::decode("00").unwrap();
let script = Script::new_builder()
.code_hash(code_hash2.pack()) // Different code_hash
.hash_type(ScriptHashType::Data.into())
.args(args_bytes.pack())
.build();

let found = udt_whitelist.find_matching_udt(&script);
assert!(found.is_none());
}

#[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
#[cfg_attr(not(target_arch = "wasm32"), test)]
fn test_find_matching_udt_wrong_hash_type() {
let code_hash = H256::from([5u8; 32]);
let udt_whitelist = UdtCfgInfos(vec![UdtArgInfo {
name: "TestUDT".to_string(),
script: UdtScript {
code_hash: code_hash.clone(),
hash_type: ScriptHashType::Data,
args: "0x00".to_string(),
},
auto_accept_amount: Some(100),
cell_deps: vec![],
}]);

let args_bytes = hex::decode("00").unwrap();
let script = Script::new_builder()
.code_hash(code_hash.pack())
.hash_type(ScriptHashType::Type.into()) // Different hash_type
.args(args_bytes.pack())
.build();

let found = udt_whitelist.find_matching_udt(&script);
assert!(found.is_none());
}

#[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
#[cfg_attr(not(target_arch = "wasm32"), test)]
fn test_find_matching_udt_multiple_udts() {
let code_hash1 = H256::from([6u8; 32]);
let code_hash2 = H256::from([7u8; 32]);
let udt_whitelist = UdtCfgInfos(vec![
UdtArgInfo {
name: "FirstUDT".to_string(),
script: UdtScript {
code_hash: code_hash1.clone(),
hash_type: ScriptHashType::Data,
args: "0x00".to_string(),
},
auto_accept_amount: Some(100),
cell_deps: vec![],
},
UdtArgInfo {
name: "SecondUDT".to_string(),
script: UdtScript {
code_hash: code_hash2.clone(),
hash_type: ScriptHashType::Data,
args: "0x01".to_string(),
},
auto_accept_amount: Some(200),
cell_deps: vec![],
},
]);

// Test finding first UDT
let args_bytes1 = hex::decode("00").unwrap();
let script1 = Script::new_builder()
.code_hash(code_hash1.pack())
.hash_type(ScriptHashType::Data.into())
.args(args_bytes1.pack())
.build();

let found = udt_whitelist.find_matching_udt(&script1);
assert!(found.is_some());
assert_eq!(found.unwrap().name, "FirstUDT");

// Test finding second UDT
let args_bytes2 = hex::decode("01").unwrap();
let script2 = Script::new_builder()
.code_hash(code_hash2.pack())
.hash_type(ScriptHashType::Data.into())
.args(args_bytes2.pack())
.build();

let found = udt_whitelist.find_matching_udt(&script2);
assert!(found.is_some());
assert_eq!(found.unwrap().name, "SecondUDT");
}

#[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
#[cfg_attr(not(target_arch = "wasm32"), test)]
fn test_find_matching_udt_empty_whitelist() {
let udt_whitelist = UdtCfgInfos(vec![]);
let code_hash = H256::from([8u8; 32]);

let args_bytes = hex::decode("00").unwrap();
let script = Script::new_builder()
.code_hash(code_hash.pack())
.hash_type(ScriptHashType::Data.into())
.args(args_bytes.pack())
.build();

let found = udt_whitelist.find_matching_udt(&script);
assert!(found.is_none());
}
Loading
Loading