drain_macros/
lib.rs

1use proc_macro::TokenStream;
2use regex::Regex;
3
4#[proc_macro_attribute]
5pub fn drain_endpoint(attr: TokenStream, item: TokenStream) -> TokenStream {
6    let item_str = item.to_string();
7    let attr_str = attr.to_string();
8    let attr_str = attr_str.trim_matches('"');
9
10    let page_fn_split = item_str.split_once('(').unwrap();
11    let page_fn_body = page_fn_split.1.split_once('{').unwrap().1;
12
13    let attr_regex = Regex::new(
14        r#"^((([A-Za-z0-9\-_]*\.[[:alnum:]]+/?)+)+|([A-Za-z0-9\-_]+/?)+)+$"#
15    ).unwrap();
16
17    if !attr_regex.is_match(&attr_str) {
18        panic!("Endpoint's name must consist only of alphanumeric characters, hyphens, underscores, slashes and dots.");
19    }
20
21    let attr_str_prepared = attr_str
22        .trim_end_matches(|x| x == '/' || x == '\\')
23        .replace(|x| x == '/' || x == '\\', "::");
24
25    format!("#[unsafe(export_name = \"{attr_str_prepared}\")]{}(REQUEST_DATA: drain_common::RequestData, \
26                REQUEST_HEADERS: &std::collections::HashMap<String, String>, \
27                RESPONSE_HEADERS: &mut std::collections::HashMap<String, String>, \
28                SET_COOKIE: &mut std::collections::HashMap<String, drain_common::cookies::SetCookie>,\
29                HTTP_STATUS_CODE: &mut u16) -> Result<Option<Vec<u8>>, Box<dyn std::any::Any + Send>> {{\
30                    std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {{
31                        tokio::runtime::Builder::new_multi_thread()\
32                            .enable_all()\
33                            .build()\
34                            .unwrap()\
35                            .block_on(async {{\
36                                {page_fn_body}\
37                            )\
38                    }}))
39                }}",
40            page_fn_split.0)
41        .parse()
42        .unwrap()
43}
44
45#[proc_macro]
46pub fn set_header(input: TokenStream) -> TokenStream {
47    let in_str = input.to_string();
48    let args_split = in_str.split_once(',').unwrap();
49    format!("RESPONSE_HEADERS.insert({}.to_lowercase(), String::from({}))",
50            args_split.0,
51            args_split.1)
52        .parse()
53        .unwrap()
54}
55
56#[proc_macro]
57pub fn header(input: TokenStream) -> TokenStream {
58    let in_str = input.to_string();
59    format!("REQUEST_HEADERS.get({})", in_str)
60        .parse()
61        .unwrap()
62}
63
64#[proc_macro]
65pub fn cookies(_input: TokenStream) -> TokenStream {
66    "drain_common::cookies::cookies(REQUEST_HEADERS)"
67        .parse()
68        .unwrap()
69}
70
71#[proc_macro]
72pub fn start_session(_input: TokenStream) -> TokenStream {
73    "drain_common::sessions::start_session(REQUEST_HEADERS, SET_COOKIE)"
74        .parse()
75        .unwrap()
76}
77
78#[proc_macro_derive(SessionValue)]
79pub fn derive_session_value(input: TokenStream) -> TokenStream {
80    let input_str = input.to_string();
81    let impl_decl = input_str
82        .trim_start_matches("pub")
83        .trim_start_matches("(crate)")
84        .trim_start_matches("(super)")
85        .split_once(|p| p == '{' || p == ';' || p == '(')
86        .unwrap().0
87        .trim()
88        .split_once(' ')
89        .unwrap();
90
91    if !impl_decl.0.eq("struct") && !impl_decl.0.eq("enum") {
92        panic!("Implementor of SessionValue must be either a struct or an enum.");
93    }
94
95    let impl_name = impl_decl.1;
96
97    format!("\nimpl SessionValue for {impl_name} {{ fn as_any(&self) -> &dyn std::any::Any {{ self }} }}")
98        .parse()
99        .unwrap()
100}
drain_macros/lib.rs

drain_macros/
lib.rs
