botnet_macros/
lib.rs

1extern crate proc_macro;
2
3use botnet_utils::type_id;
4use proc_macro::TokenStream;
5use proc_macro_error::proc_macro_error;
6use quote::{format_ident, quote};
7use syn::{parse_macro_input, Ident, ItemFn, ItemStruct};
8
9fn process_task(attrs: TokenStream, input: TokenStream) -> TokenStream {
10    let func = parse_macro_input!(input as ItemFn);
11    let ident = format_ident!("{}Task", parse_macro_input!(attrs as Ident));
12
13    let output = quote! {
14
15        struct #ident{}
16
17        #[async_trait::async_trait]
18        impl<K, D> Task<K, D> for #ident
19        where
20            K: DatabaseKey + 'static,
21            D: Database + Send + Sync + 'static
22        {
23            #[allow(unused)]
24            #func
25        }
26
27    };
28
29    TokenStream::from(output)
30}
31
32fn process_evaluator(attrs: TokenStream, input: TokenStream) -> TokenStream {
33    let func = parse_macro_input!(input as ItemFn);
34    let ident = format_ident!("{}Evaluator", parse_macro_input!(attrs as Ident));
35
36    let output = quote! {
37
38        struct #ident{}
39
40        #[async_trait::async_trait]
41        impl Evaluator for #ident {
42            #[allow(unused)]
43            #func
44        }
45
46    };
47
48    TokenStream::from(output)
49}
50
51fn process_key(input: TokenStream) -> TokenStream {
52    let item = parse_macro_input!(input as ItemStruct);
53    let ident = &item.ident;
54    let name = ident.to_string();
55    let ty_id = type_id(&*name);
56
57    let output = quote! {
58
59        #[derive(Debug, PartialEq, Eq, Hash, Clone)]
60        #item
61
62        impl DatabaseKey for #ident {}
63
64        impl Key for #ident {
65
66            const NAME: &'static str = #name;
67            const TYPE_ID: usize = #ty_id;
68
69            type Item = Field;
70            type Metadata = KeyMetadata;
71
72            fn builder() -> Self {
73                Self {
74                    fields: Vec::new(),
75                    metadata: KeyMetadata::default(),
76                }
77            }
78
79            fn metadata(&mut self, meta: KeyMetadata) -> &mut Self {
80                self.metadata = KeyMetadata::with_key_code(meta, Self::TYPE_ID);
81                self
82            }
83
84            fn field(&mut self, f: Self::Item) -> &mut Self {
85                self.fields.push(f);
86                self
87            }
88
89            fn build(&self) -> Self {
90                self.clone()
91            }
92
93            fn flatten(&self) -> Bytes {
94                let fields = Bytes::from(
95                    self.fields
96                        .iter()
97                        .map(|f| usize::to_le_bytes(f.type_id).to_vec())
98                        .collect::<Vec<Vec<u8>>>()
99                        .into_iter()
100                        .flatten()
101                        .collect::<Vec<u8>>(),
102                );
103
104                let id = Bytes::from(usize::to_le_bytes(Self::TYPE_ID).to_vec());
105                Bytes::from([id, fields].concat())
106            }
107
108            fn get_metadata(&self) -> Self::Metadata {
109                self.metadata.clone()
110            }
111
112            fn type_id(&self) -> usize {
113                Self::TYPE_ID
114            }
115
116            fn name(&self) -> &'static str {
117                Self::NAME
118            }
119        }
120
121        impl AsBytes for #ident {
122            fn as_bytes(&self) -> &[u8] {
123                <#ident as AsBytes>::as_bytes(self)
124            }
125        }
126
127        impl #ident {
128            pub fn from_input(value: Input, extractors: &Extractors, metadata: &Metadata) -> BotnetResult<Self> {
129                let meta = metadata.get(&Self::TYPE_ID);
130                let fields = extractors.items.iter().map(|e| e.1.call(&value).expect("Failed to call on input.")).collect::<Vec<Field>>();
131
132                // TODO: use builder pattern
133                Ok(#ident{ fields, metadata: meta.to_owned() })
134            }
135
136            pub fn from_bytes(b: Bytes, metadata: &Metadata) -> BotnetResult<#ident> {
137                let mut parts = b.chunks_exact(64);
138
139                let key_ty_id = parts.next().unwrap();
140                let meta = metadata.get(&Self::TYPE_ID);
141
142                // let ty_id = usize::from_le_bytes(parts[0]);
143                // TODO: finish
144                Ok(#ident::builder())
145            }
146        }
147
148    };
149
150    TokenStream::from(output)
151}
152
153#[proc_macro_error]
154#[proc_macro_attribute]
155pub fn task(attrs: TokenStream, input: TokenStream) -> TokenStream {
156    process_task(attrs, input)
157}
158
159#[proc_macro_error]
160#[proc_macro_attribute]
161pub fn evaluator(attrs: TokenStream, input: TokenStream) -> TokenStream {
162    process_evaluator(attrs, input)
163}
164
165#[proc_macro_error]
166#[proc_macro_attribute]
167pub fn key(_attrs: TokenStream, input: TokenStream) -> TokenStream {
168    process_key(input)
169}
botnet_macros/lib.rs

botnet_macros/
lib.rs
