ckb_ssri_std_proc_macro/
lib.rs

1#![no_std]
2
3extern crate alloc;
4extern crate proc_macro;
5
6use core::panic;
7
8use ckb_hash::blake2b_256;
9use proc_macro::TokenStream;
10use quote::quote;
11use syn::{parse::Parse, parse_macro_input, Expr, ExprLit, Ident, Lit, Token};
12
13use alloc::vec;
14use alloc::vec::Vec;
15
16fn encode_u64_vector(val: impl AsRef<[u64]>) -> Vec<u8> {
17    let val = val.as_ref();
18    u32::to_le_bytes(val.len() as u32)
19        .into_iter()
20        .chain(val.iter().flat_map(|v| u64::to_le_bytes(*v)))
21        .collect()
22}
23
24fn method_path(name: impl AsRef<[u8]>) -> u64 {
25    u64::from_le_bytes(blake2b_256(name)[0..8].try_into().unwrap())
26}
27
28struct Methods {
29    argv: Expr,
30    invalid_method: Expr,
31    invalid_args: Expr,
32    method_keys: Vec<u64>,
33    method_bodies: Vec<Expr>,
34}
35
36impl Parse for Methods {
37    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
38        input.parse::<Ident>()?;
39        input.parse::<Token![:]>()?;
40        let argv = input.parse::<Expr>()?;
41        input.parse::<Token![,]>()?;
42        input.parse::<Ident>()?;
43        input.parse::<Token![:]>()?;
44        let invalid_method = input.parse::<Expr>()?;
45        input.parse::<Token![,]>()?;
46        input.parse::<Ident>()?;
47        input.parse::<Token![:]>()?;
48        let invalid_args = input.parse::<Expr>()?;
49        input.parse::<Token![,]>()?;
50
51        let mut method_keys = vec![];
52        let mut method_bodies = vec![];
53        while !input.is_empty() {
54            let name = match input.parse::<Expr>()? {
55                Expr::Lit(ExprLit {
56                    lit: Lit::Str(v), ..
57                }) => v.value(),
58                _ => panic!("method name should be a string"),
59            };
60            input.parse::<Token![=>]>()?;
61            let body = input.parse::<Expr>()?;
62            input.parse::<Token![,]>()?;
63
64            method_keys.push(method_path(name));
65            method_bodies.push(body);
66        }
67
68        Ok(Methods {
69            argv,
70            invalid_method,
71            invalid_args,
72            method_keys,
73            method_bodies,
74        })
75    }
76}
77
78#[proc_macro]
79pub fn ssri_methods(input: TokenStream) -> TokenStream {
80    let Methods {
81        argv,
82        invalid_method,
83        invalid_args,
84        method_keys,
85        method_bodies,
86    } = parse_macro_input!(input as Methods);
87
88    let version_path = method_path("SSRI.version");
89    let get_methods_path = method_path("SSRI.get_methods");
90    let has_methods_path = method_path("SSRI.has_methods");
91
92    let raw_methods = encode_u64_vector(
93        [version_path, get_methods_path, has_methods_path]
94            .iter()
95            .chain(method_keys.iter())
96            .copied()
97            .collect::<Vec<_>>(),
98    );
99    let raw_methods_len = raw_methods.len();
100
101    TokenStream::from(quote! {
102        {
103            use alloc::{borrow::Cow, vec::Vec};
104            use ckb_std::high_level::decode_hex;
105            const RAW_METHODS: [u8; #raw_methods_len] = [#(#raw_methods,)*];
106            let res: Result<Cow<'static, [u8]>, Error> = match u64::from_le_bytes(
107                decode_hex(&(#argv)[0])?.try_into().map_err(|_| #invalid_method)?,
108            ) {
109                #version_path => Ok(Cow::from(&[0][..])),
110                #get_methods_path => {
111                    let offset = usize::min((u64::from_le_bytes(
112                        decode_hex(&(#argv)[1])?
113                            .try_into()
114                            .map_err(|_| #invalid_args)?
115                    ) as usize * 8), #raw_methods_len - 4);
116
117                    // If second argument is 0, take all remaining methods
118                    let second_arg = u64::from_le_bytes(
119                        decode_hex(&(#argv)[2])?
120                            .try_into()
121                            .map_err(|_| #invalid_args)?
122                    );
123
124                    let mut raw_result: Vec<u8>;
125                    if second_arg == 0 {
126                        // Take all remaining methods from offset
127                        raw_result = RAW_METHODS[(offset + 4)..].to_vec();
128                        let method_count = (RAW_METHODS.len() - (offset + 4)) / 8;
129                        let mut result = (method_count as u32).to_le_bytes().to_vec();
130                        result.extend_from_slice(&raw_result);
131                        Ok(Cow::from(result))
132                    } else {
133                        let limit = usize::min((offset + (second_arg as usize * 8)), #raw_methods_len - 4);
134                        raw_result = RAW_METHODS[(offset + 4)..(limit + 4)].to_vec();
135                        let method_count = (limit - offset) / 8;
136                        let mut result = (method_count as u32).to_le_bytes().to_vec();
137                        result.extend_from_slice(&raw_result);
138                        Ok(Cow::from(result))
139                    }
140                },
141                #has_methods_path => {
142                    let mut result = Vec::new();
143                    let matches = decode_hex(&(#argv)[1])?[4..].chunks(8).map(|path| {
144                        match RAW_METHODS[4..]
145                            .chunks(8)
146                            .find(|v| v == &path) {
147                                Some(_) => 1,
148                                None => 0,
149                            }
150                    }).collect::<Vec<_>>();
151                    result.extend_from_slice(&(matches.len() as u32).to_le_bytes());
152                    result.extend_from_slice(&matches);
153                    Ok(Cow::from(result))
154                },
155                #(
156                    #method_keys => #method_bodies,
157                )*
158                _ => Err(#invalid_method),
159            };
160            res
161        }
162    })
163}
ckb_ssri_std_proc_macro/lib.rs

ckb_ssri_std_proc_macro/
lib.rs
