Skip to main content

jsonrpc_server_macro/
lib.rs

1use proc_macro::TokenStream;
2use quote::{quote, quote_spanned};
3use syn::punctuated::Punctuated;
4use syn::spanned::Spanned;
5use syn::{
6    Error, FnArg, ImplItem, ImplItemFn, Item, ItemFn, ItemImpl, Pat, ReturnType, Signature, Token,
7    Type, parse, parse_macro_input,
8};
9
10#[proc_macro_attribute]
11pub fn jsonrpc(_attr: TokenStream, input: TokenStream) -> TokenStream {
12    match parse_macro_input!(input as Item) {
13        Item::Fn(item) => expand_fn(item).unwrap_or_else(|e| e.to_compile_error().into()),
14        Item::Impl(item) => expand_impl(item).unwrap_or_else(|e| e.to_compile_error().into()),
15        item => Error::new_spanned(item, "#[jsonrpc]: expected fn or impl block")
16            .to_compile_error()
17            .into(),
18    }
19}
20
21fn expand_impl(item: ItemImpl) -> Result<TokenStream, Error> {
22    if let Some((_, ref path, ..)) = item.trait_ {
23        return Err(Error::new_spanned(
24            path,
25            "#[jsonrpc]: trait impl is not supported",
26        ));
27    }
28
29    if !item.generics.params.is_empty() {
30        return Err(Error::new_spanned(
31            item.generics,
32            "#[jsonrpc]: generic is not supported",
33        ));
34    }
35
36    let prefix = match *item.self_ty {
37        Type::Path(ref path) => snake_case(
38            path.path
39                .segments
40                .last()
41                .unwrap()
42                .ident
43                .to_string()
44                .as_bytes(),
45        ),
46        _ => {
47            return Err(Error::new_spanned(
48                item.self_ty,
49                "#[jsonrpc]: not supported",
50            ));
51        }
52    };
53
54    let mut names = Vec::new();
55    let mut methods = Vec::new();
56    for impl_item in &item.items {
57        if let ImplItem::Fn(item_fn) = impl_item {
58            names.push(format!("{}.{}", prefix, item_fn.sig.ident));
59            methods.push(generate_method(&item, item_fn)?);
60        }
61    }
62
63    let ty = &item.self_ty;
64    Ok(quote! {
65        #item
66        impl jsonrpc_server::Register for #ty {
67            fn register(&self, registry: &mut jsonrpc_server::Registry) {
68                #(registry.add(#names, #methods);)*
69            }
70        }
71    }
72    .into())
73}
74
75fn generate_method(
76    item: &ItemImpl,
77    item_fn: &ImplItemFn,
78) -> Result<proc_macro2::TokenStream, Error> {
79    if !item_fn.sig.generics.params.is_empty() {
80        return Err(Error::new_spanned(
81            &item_fn.sig.generics,
82            "#[jsonrpc]: generic is not supported",
83        ));
84    }
85    let ident = &item_fn.sig.ident;
86    let self_ty = &item.self_ty;
87    let ret_assert = ret_assert(&item_fn.sig)?;
88    let (arg_assert, args) = arg_assert(&item_fn.sig.inputs, true)?;
89    let wait = item_fn
90        .sig
91        .asyncness
92        .map(|_| quote!(let result = result.await;));
93    let set_method = set_method(ident);
94    let argc = 0..args.len();
95    let st = quote_spanned! {item_fn.span()=>
96        {
97            struct __Method(#self_ty);
98
99            impl jsonrpc_server::Method for __Method {
100                fn call(&self, args: jsonrpc_server::serde_json::Value) -> jsonrpc_server::BoxFuture<'_> {
101                    #arg_assert
102                    #ret_assert
103
104                    #[allow(unused)]
105                    macro_rules! arg {
106                        ($v:expr) => {
107                            jsonrpc_server::serde_json::from_value($v).map_err(|err| {
108                                jsonrpc_server::error!("deserialize parameter error: {}", err);
109                                jsonrpc_server::Error::invalid_params()
110                            })
111                        };
112                    }
113
114                    Box::pin(async move {
115                        #[allow(unused)]
116                        let result = match args {
117                            jsonrpc_server::serde_json::Value::Array(mut args) => {
118                                self.0.#ident(#(arg!(args.get_mut(#argc).map(jsonrpc_server::serde_json::Value::take).unwrap_or(jsonrpc_server::serde_json::Value::Null))?),*)
119                            }
120                            jsonrpc_server::serde_json::Value::Object(mut args) => {
121                                self.0.#ident(#(arg!(args.remove(#args).unwrap_or(jsonrpc_server::serde_json::Value::Null))?),*)
122                            }
123                            _ => return Err(jsonrpc_server::Error::invalid_params()),
124                        };
125                        #wait
126                        #set_method
127                        Ok(jsonrpc_server::serde_json::to_value(result?).expect("serialize error"))
128                    })
129                }
130            }
131
132            __Method(self.clone())
133        }
134    };
135    parse(st.into())
136}
137
138fn expand_fn(item: ItemFn) -> Result<TokenStream, Error> {
139    let ident = &item.sig.ident;
140    let vis = &item.vis;
141    if !item.sig.generics.params.is_empty() {
142        return Err(Error::new_spanned(
143            item.sig.generics,
144            "#[jsonrpc]: generic is not supported",
145        ));
146    }
147
148    let ret_assert = ret_assert(&item.sig)?;
149    let (arg_assert, args) = arg_assert(&item.sig.inputs, false)?;
150    let wait = item
151        .sig
152        .asyncness
153        .map(|_| quote!(let result = result.await;));
154    let set_method = set_method(ident);
155    let argc = 0..args.len();
156    let ts = quote! {
157        #vis fn #ident(args: jsonrpc_server::serde_json::Value) -> std::pin::Pin<Box<dyn std::future::Future<Output=std::result::Result<jsonrpc_server::serde_json::Value, jsonrpc_server::Error>> + Send>> {
158            #arg_assert
159            #ret_assert
160            #item
161
162            #[allow(unused)]
163            macro_rules! arg {
164                ($v:expr) => {
165                    jsonrpc_server::serde_json::from_value($v).map_err(|err| {
166                        jsonrpc_server::error!("deserialize parameter error: {}", err);
167                        jsonrpc_server::Error::invalid_params()
168                    })
169                };
170            }
171
172            Box::pin(async move {
173                #[allow(unused)]
174                let result = match args {
175                    jsonrpc_server::serde_json::Value::Array(mut args) => {
176                        #ident(#(arg!(args.get_mut(#argc).map(jsonrpc_server::serde_json::Value::take).unwrap_or(jsonrpc_server::serde_json::Value::Null))?),*)
177                    }
178                    jsonrpc_server::serde_json::Value::Object(mut args) => {
179                        #ident(#(arg!(args.remove(#args).unwrap_or(jsonrpc_server::serde_json::Value::Null))?),*)
180                    }
181                    _ => return Err(jsonrpc_server::Error::invalid_params()),
182                };
183                #wait
184                #set_method
185                Ok(jsonrpc_server::serde_json::to_value(result?).expect("serialize error"))
186            })
187        }
188    };
189    Ok(ts.into())
190}
191
192#[cfg(feature = "anyhow")]
193fn set_method(name: &proc_macro2::Ident) -> Option<proc_macro2::TokenStream> {
194    let name = name.to_string();
195    Some(quote! {
196        let result = match result {
197            Ok(v) => Ok(v),
198            Err(e) => {
199                let mut e = jsonrpc_server::Error::from(e);
200                e.method = Some(#name);
201                Err(e)
202            }
203        };
204    })
205}
206
207#[cfg(not(feature = "anyhow"))]
208fn set_method(_: &proc_macro2::Ident) -> Option<proc_macro2::TokenStream> {
209    None
210}
211
212fn ret_assert(sig: &Signature) -> Result<proc_macro2::TokenStream, Error> {
213    match sig.output {
214        ReturnType::Default => Err(Error::new_spanned(sig, "#[jsonrpc]: expected return value")),
215        ReturnType::Type(_, ref ty) => Ok(quote_spanned! {ty.span()=>
216            {
217                fn assert(_: Option<std::result::Result<impl jsonrpc_server::serde::Serialize, impl Into<jsonrpc_server::Error>>>) {}
218                assert(None::<#ty>);
219            }
220        }),
221    }
222}
223
224fn arg_assert(
225    inputs: &Punctuated<FnArg, Token![,]>,
226    is_method: bool,
227) -> Result<(proc_macro2::TokenStream, Vec<String>), Error> {
228    let mut assert = vec![];
229    let mut args = Vec::with_capacity(inputs.len());
230    for arg in inputs.iter().skip(if is_method { 1 } else { 0 }) {
231        match arg {
232            FnArg::Typed(arg) => match *arg.pat {
233                Pat::Ident(ref pat) => {
234                    args.push(pat.ident.to_string());
235                    let ty = &arg.ty;
236                    assert.push(quote_spanned! {ty.span()=>
237                        { struct _Assert where #ty: jsonrpc_server::serde::de::DeserializeOwned; }
238                    })
239                }
240                _ => return Err(Error::new_spanned(arg, "#[jsonrpc]: unsupported argument")),
241            },
242            FnArg::Receiver(_) => unreachable!(),
243        }
244    }
245    Ok((quote!(#(#assert)*), args))
246}
247
248fn snake_case(s: &[u8]) -> String {
249    let mut result = String::with_capacity(s.len());
250    for &b in s {
251        match b {
252            b'A'..=b'Z' => {
253                if !result.is_empty() {
254                    result.push('_');
255                }
256                result.push((b + 32) as char);
257            }
258            b => result.push(b as char),
259        }
260    }
261    result
262}