1mod utils;
2
3use proc_macro::TokenStream;
4use proc_macro2::{Ident, Span};
5use quote::{quote, ToTokens};
6use rand::Rng;
7use serde_json::json;
8use std::{collections::HashSet, sync::LazyLock};
9use utils::StringExt as _;
10
11static ARG_TYPES: LazyLock<HashSet<&'static str>> = LazyLock::new(|| {
12 [
13 "String", "bool", "u8", "u16", "u32", "u64", "usize", "i8", "i16", "i32", "i64", "isize",
14 "f32", "f64",
15 ]
16 .into_iter()
17 .collect()
18});
19
20fn random_ident() -> Ident {
21 let mut rng = rand::thread_rng();
22 let value = format!("__potato_id_{}", rng.gen::<u64>());
23 Ident::new(&value, Span::call_site())
24}
25
26fn http_handler_macro(attr: TokenStream, input: TokenStream, req_name: &str) -> TokenStream {
27 let req_name = Ident::new(req_name, Span::call_site());
28 let (route_path, oauth_arg) = {
29 let mut oroute_path = syn::parse::<syn::LitStr>(attr.clone())
30 .ok()
31 .map(|path| path.value());
32 let mut oauth_arg = None;
33 if oroute_path.is_none() {
35 let http_parser = syn::meta::parser(|meta| {
36 if meta.path.is_ident("path") {
37 if let Ok(arg) = meta.value() {
38 if let Ok(route_path) = arg.parse::<syn::LitStr>() {
39 let route_path = route_path.value();
40 oroute_path = Some(route_path);
41 }
42 }
43 Ok(())
44 } else if meta.path.is_ident("auth_arg") {
45 if let Ok(arg) = meta.value() {
46 if let Ok(tmp_field) = arg.parse::<Ident>() {
47 oauth_arg = Some(tmp_field.to_string());
48 }
49 }
50 Ok(())
51 } else {
52 Err(meta.error("unsupported annotation property"))
53 }
54 });
55 syn::parse_macro_input!(attr with http_parser);
56 }
57 if oroute_path.is_none() {
58 panic!("`path` argument is required");
59 }
60 let route_path = oroute_path.unwrap();
61 if !route_path.starts_with('/') {
62 panic!("route path must start with '/'");
63 }
64 (route_path, oauth_arg)
65 };
66 let root_fn = syn::parse_macro_input!(input as syn::ItemFn);
67 let doc_show = {
68 let mut doc_show = true;
69 for attr in root_fn.attrs.iter() {
70 if attr.meta.path().get_ident().map(|p| p.to_string()) == Some("doc".to_string()) {
71 if let Ok(meta_list) = attr.meta.require_list() {
72 if meta_list.tokens.to_string() == "hidden" {
73 doc_show = false;
74 break;
75 }
76 }
77 }
78 }
79 doc_show
80 };
81 let doc_auth = oauth_arg.is_some();
82 let doc_summary = {
83 let mut docs = vec![];
84 for attr in root_fn.attrs.iter() {
85 if let Ok(attr) = attr.meta.require_name_value() {
86 if attr.path.get_ident().map(|p| p.to_string()) == Some("doc".to_string()) {
87 let mut doc = attr.value.to_token_stream().to_string();
88 if doc.starts_with('\"') {
89 doc.remove(0);
90 doc.pop();
91 }
92 docs.push(doc);
93 }
94 }
95 }
96 if docs.iter().all(|d| d.starts_with(' ')) {
97 for doc in docs.iter_mut() {
98 doc.remove(0);
99 }
100 }
101 docs.join("\n")
102 };
103 let doc_desp = "";
104 let fn_name = root_fn.sig.ident.clone();
105 let wrap_func_name = random_ident();
106 let mut args = vec![];
107 let mut doc_args = vec![];
108 let mut arg_auth_mark = false;
109 for arg in root_fn.sig.inputs.iter() {
110 if let syn::FnArg::Typed(arg) = arg {
111 let arg_type_str = arg
112 .ty
113 .as_ref()
114 .to_token_stream()
115 .to_string()
116 .type_simplify();
117 let arg_name_str = arg.pat.to_token_stream().to_string();
118 args.push(match &arg_type_str[..] {
119 "& mut HttpRequest" => quote! { req },
120 "PostFile" => {
121 doc_args.push(json!({ "name": arg_name_str, "type": arg_type_str }));
122 quote! {
123 match req.body_files.get(&potato::utils::refstr::RefOrString::from_str(#arg_name_str)).cloned() {
124 Some(file) => file,
125 None => return HttpResponse::error(format!("miss arg: {}", #arg_name_str)),
126 }
127 }
128 },
129 arg_type_str if ARG_TYPES.contains(arg_type_str) => {
130 let is_auth_arg = match oauth_arg.as_ref() {
131 Some(auth_arg) => auth_arg == &arg_name_str,
132 None => false,
133 };
134 if is_auth_arg {
135 if arg_type_str != "String" {
136 panic!("auth_arg argument is must String type");
137 }
138 arg_auth_mark = true;
139 quote! {
140 match req.headers
141 .get(&potato::utils::refstr::HeaderRefOrString::from_str("Authorization"))
142 .map(|v| v.to_str()) {
143 Some(mut auth) => {
144 if auth.starts_with("Bearer ") {
145 auth = &auth[7..];
146 }
147 match potato::ServerAuth::jwt_check(&auth).await {
148 Ok(payload) => payload,
149 Err(err) => return HttpResponse::error(format!("auth failed: {err:?}")),
150 }
151 }
152 None => return HttpResponse::error("miss header : Authorization"),
153 }
154 }
155 } else {
156 doc_args.push(json!({ "name": arg_name_str, "type": arg_type_str }));
157 let mut arg_value = quote! {
158 match req.body_pairs
159 .get(&potato::utils::refstr::RefOrString::from_str(#arg_name_str))
160 .map(|p| p.to_string()) {
161 Some(val) => val,
162 None => match req.url_query
163 .get(&potato::utils::refstr::RefOrString::from_str(#arg_name_str))
164 .map(|p| p.to_str().to_string()) {
165 Some(val) => val,
166 None => return HttpResponse::error(format!("miss arg: {}", #arg_name_str)),
167 },
168 }
169 };
170 if arg_type_str != "String" {
171 arg_value = quote! {
172 match #arg_value.parse() {
173 Ok(val) => val,
174 Err(err) => return HttpResponse::error(format!("arg[{}] is not {} type", #arg_name_str, #arg_type_str)),
175 }
176 }
177 }
178 arg_value
179 }
180 },
181 _ => panic!("unsupported arg type: [{arg_type_str}]"),
182 });
183 } else {
184 panic!("unsupported: {}", arg.to_token_stream().to_string());
185 }
186 }
187 if !arg_auth_mark && doc_auth {
188 panic!("`auth_arg` attribute is must point to an existing argument");
189 }
190 let wrap_func_name2 = random_ident();
191 let ret_type = root_fn
192 .sig
193 .output
194 .to_token_stream()
195 .to_string()
196 .type_simplify();
197 let wrap_func_body = match &ret_type[..] {
198 "Result < () >" => quote! {
199 match #fn_name(#(#args),*).await {
200 Ok(ret) => HttpResponse::text("ok"),
201 Err(err) => HttpResponse::error(format!("{err:?}")),
202 }
203 },
204 "Result < HttpResponse >" => quote! {
205 match #fn_name(#(#args),*).await {
206 Ok(ret) => ret,
207 Err(err) => HttpResponse::error(format!("{err:?}")),
208 }
209 },
210 "()" => quote! {
211 #fn_name(#(#args),*).await;
212 HttpResponse::text("ok")
213 },
214 "HttpResponse" => quote! {
215 #fn_name(#(#args),*).await
216 },
217 _ => panic!("unsupported ret type: {ret_type}"),
218 };
219 let doc_args = serde_json::to_string(&doc_args).unwrap();
220 quote! {
222 #root_fn
223
224 #[doc(hidden)]
225 async fn #wrap_func_name2(req: &mut potato::HttpRequest) -> potato::HttpResponse {
226 #wrap_func_body
227 }
228
229 #[doc(hidden)]
230 fn #wrap_func_name(req: &mut potato::HttpRequest) -> std::pin::Pin<Box<dyn std::future::Future<Output = potato::HttpResponse> + Send + '_>> {
231 Box::pin(#wrap_func_name2(req))
232 }
233
234 potato::inventory::submit!{potato::RequestHandlerFlag::new(
235 potato::HttpMethod::#req_name,
236 #route_path,
237 #wrap_func_name,
238 potato::RequestHandlerFlagDoc::new(#doc_show, #doc_auth, #doc_summary, #doc_desp, #doc_args)
239 )}
240 }.into()
241 }
245
246#[proc_macro_attribute]
247pub fn http_get(attr: TokenStream, input: TokenStream) -> TokenStream {
248 http_handler_macro(attr, input, "GET")
249}
250
251#[proc_macro_attribute]
252pub fn http_post(attr: TokenStream, input: TokenStream) -> TokenStream {
253 http_handler_macro(attr, input, "POST")
254}
255
256#[proc_macro_attribute]
257pub fn http_put(attr: TokenStream, input: TokenStream) -> TokenStream {
258 http_handler_macro(attr, input, "PUT")
259}
260
261#[proc_macro_attribute]
262pub fn http_delete(attr: TokenStream, input: TokenStream) -> TokenStream {
263 http_handler_macro(attr, input, "DELETE")
264}
265
266#[proc_macro_attribute]
267pub fn http_options(attr: TokenStream, input: TokenStream) -> TokenStream {
268 http_handler_macro(attr, input, "OPTIONS")
269}
270
271#[proc_macro_attribute]
272pub fn http_head(attr: TokenStream, input: TokenStream) -> TokenStream {
273 http_handler_macro(attr, input, "HEAD")
274}
275
276#[proc_macro]
277pub fn embed_dir(input: TokenStream) -> TokenStream {
278 let path = syn::parse_macro_input!(input as syn::LitStr).value();
279 quote! {{
280 #[derive(potato::rust_embed::Embed)]
281 #[folder = #path]
282 struct Asset;
283
284 potato::load_embed::<Asset>()
285 }}
286 .into()
287}
288
289#[proc_macro_derive(StandardHeader)]
290pub fn standard_header_derive(input: TokenStream) -> TokenStream {
291 let root_enum = syn::parse_macro_input!(input as syn::ItemEnum);
292 let enum_name = root_enum.ident;
293 let mut try_from_str_items = vec![];
294 let mut to_str_items = vec![];
295 let mut headers_items = vec![];
296 let mut headers_apply_items = vec![];
297 for root_field in root_enum.variants.iter() {
298 let name = root_field.ident.clone();
299 if root_field.fields.iter().next().is_some() {
300 panic!("unsupported enum type");
301 }
302 let str_name = name.to_string().replace("_", "-");
303 let len = str_name.len();
304 try_from_str_items
305 .push(quote! { #len if value.eq_ignore_ascii_case(#str_name) => Some(Self::#name), });
306 to_str_items.push(quote! { Self::#name => #str_name, });
307 headers_items.push(quote! { #name(String), });
308 headers_apply_items
309 .push(quote! { Headers::#name(s) => self.set_header(HeaderItem::#name.to_str(), s), });
310 }
311 quote! {
312 impl #enum_name {
313 pub fn try_from_str(value: &str) -> Option<Self> {
314 match value.len() {
315 #( #try_from_str_items )*
316 _ => None,
317 }
318 }
319
320 pub fn to_str(&self) -> &'static str {
321 match self {
322 #( #to_str_items )*
323 }
324 }
325 }
326
327 pub enum Headers {
328 #( #headers_items )*
329 Custom((String, String)),
330 }
331
332 impl HttpRequest {
333 pub fn apply_header(&mut self, header: Headers) {
334 match header {
335 #( #headers_apply_items )*
336 Headers::Custom((k, v)) => self.set_header(&k[..], v),
337 }
338 }
339 }
340 }
341 .into()
342}