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 "HttpRequest" => quote! { req },
120 "SocketAddr" => quote! { client },
121 "& mut WebsocketContext" => quote! { wsctx },
122 "PostFile" => {
123 doc_args.push(json!({ "name": arg_name_str, "type": arg_type_str }));
124 quote! {
125 match req.body_files.get(&potato::utils::refstr::RefOrString::from_str(#arg_name_str)).cloned() {
126 Some(file) => file,
127 None => return HttpResponse::error(format!("miss arg: {}", #arg_name_str)),
128 }
129 }
130 },
131 arg_type_str if ARG_TYPES.contains(arg_type_str) => {
132 let is_auth_arg = match oauth_arg.as_ref() {
133 Some(auth_arg) => auth_arg == &arg_name_str,
134 None => false,
135 };
136 if is_auth_arg {
137 if arg_type_str != "String" {
138 panic!("auth_arg argument is must String type");
139 }
140 arg_auth_mark = true;
141 quote! {
142 match req.headers
143 .get(&potato::utils::refstr::HeaderRefOrString::from_str("Authorization"))
144 .map(|v| v.to_str()) {
145 Some(mut auth) => {
146 if auth.starts_with("Bearer ") {
147 auth = &auth[7..];
148 }
149 match potato::ServerAuth::jwt_check(&auth).await {
150 Ok(payload) => payload,
151 Err(err) => return HttpResponse::error(format!("auth failed: {:?}", err)),
152 }
153 }
154 None => return HttpResponse::error("miss header : Authorization"),
155 }
156 }
157 } else {
158 doc_args.push(json!({ "name": arg_name_str, "type": arg_type_str }));
159 let mut arg_value = quote! {
160 match req.body_pairs
161 .get(&potato::utils::refstr::RefOrString::from_str(#arg_name_str))
162 .map(|p| p.to_string()) {
163 Some(val) => val,
164 None => match req.url_query
165 .get(&potato::utils::refstr::RefOrString::from_str(#arg_name_str))
166 .map(|p| p.to_str().to_string()) {
167 Some(val) => val,
168 None => return HttpResponse::error(format!("miss arg: {}", #arg_name_str)),
169 },
170 }
171 };
172 if arg_type_str != "String" {
173 arg_value = quote! {
174 match #arg_value.parse() {
175 Ok(val) => val,
176 Err(err) => return HttpResponse::error(format!("arg[{}] is not {} type", #arg_name_str, #arg_type_str)),
177 }
178 }
179 }
180 arg_value
181 }
182 },
183 _ => panic!("unsupported arg type: [{arg_type_str}]"),
184 });
185 } else {
186 panic!("unsupported: {}", arg.to_token_stream().to_string());
187 }
188 }
189 if !arg_auth_mark && doc_auth {
190 panic!("`auth_arg` attribute is must point to an existing argument");
191 }
192 let wrap_func_name2 = random_ident();
193 let ret_type = root_fn
194 .sig
195 .output
196 .to_token_stream()
197 .to_string()
198 .type_simplify();
199 let wrap_func_body = match &ret_type[..] {
200 "Result < () >" => quote! {
201 match #fn_name(#(#args),*).await {
202 Ok(ret) => HttpResponse::text("ok"),
203 Err(err) => HttpResponse::error(format!("{err:?}")),
204 }
205 },
206 "Result < HttpResponse >" => quote! {
207 match #fn_name(#(#args),*).await {
208 Ok(ret) => ret,
209 Err(err) => HttpResponse::error(format!("{err:?}")),
210 }
211 },
212 "()" => quote! {
213 #fn_name(#(#args),*).await;
214 HttpResponse::text("ok")
215 },
216 "HttpResponse" => quote! {
217 #fn_name(#(#args),*).await
218 },
219 _ => panic!("unsupported ret type: {ret_type}"),
220 };
221 let doc_args = serde_json::to_string(&doc_args).unwrap();
222 quote! {
224 #root_fn
225
226 #[doc(hidden)]
227 async fn #wrap_func_name2(
228 req: potato::HttpRequest, client: std::net::SocketAddr, wsctx: &mut potato::WebsocketContext
229 ) -> potato::HttpResponse {
230 #wrap_func_body
231 }
232
233 #[doc(hidden)]
234 fn #wrap_func_name<'a>(
235 req: potato::HttpRequest, client: std::net::SocketAddr, wsctx: &'a mut potato::WebsocketContext
236 ) -> std::pin::Pin<Box<dyn std::future::Future<Output = potato::HttpResponse> + Send + 'a>> {
237 Box::pin(#wrap_func_name2(req, client, wsctx))
238 }
239
240 potato::inventory::submit!{potato::RequestHandlerFlag::new(
241 potato::HttpMethod::#req_name,
242 #route_path,
243 #wrap_func_name,
244 potato::RequestHandlerFlagDoc::new(#doc_show, #doc_auth, #doc_summary, #doc_desp, #doc_args)
245 )}
246 }.into()
247 }
251
252#[proc_macro_attribute]
253pub fn http_get(attr: TokenStream, input: TokenStream) -> TokenStream {
254 http_handler_macro(attr, input, "GET")
255}
256
257#[proc_macro_attribute]
258pub fn http_post(attr: TokenStream, input: TokenStream) -> TokenStream {
259 http_handler_macro(attr, input, "POST")
260}
261
262#[proc_macro_attribute]
263pub fn http_put(attr: TokenStream, input: TokenStream) -> TokenStream {
264 http_handler_macro(attr, input, "PUT")
265}
266
267#[proc_macro_attribute]
268pub fn http_delete(attr: TokenStream, input: TokenStream) -> TokenStream {
269 http_handler_macro(attr, input, "DELETE")
270}
271
272#[proc_macro_attribute]
273pub fn http_options(attr: TokenStream, input: TokenStream) -> TokenStream {
274 http_handler_macro(attr, input, "OPTIONS")
275}
276
277#[proc_macro_attribute]
278pub fn http_head(attr: TokenStream, input: TokenStream) -> TokenStream {
279 http_handler_macro(attr, input, "HEAD")
280}
281
282#[proc_macro]
283pub fn embed_dir(input: TokenStream) -> TokenStream {
284 let path = syn::parse_macro_input!(input as syn::LitStr).value();
285 quote! {{
286 #[derive(potato::rust_embed::Embed)]
287 #[folder = #path]
288 struct Asset;
289
290 potato::load_embed::<Asset>()
291 }}
292 .into()
293}
294
295#[proc_macro_derive(StandardHeader)]
296pub fn standard_header_derive(input: TokenStream) -> TokenStream {
297 let root_enum = syn::parse_macro_input!(input as syn::ItemEnum);
298 let enum_name = root_enum.ident;
299 let mut try_from_str_items = vec![];
300 let mut to_str_items = vec![];
301 let mut headers_items = vec![];
302 let mut headers_apply_items = vec![];
303 for root_field in root_enum.variants.iter() {
304 let name = root_field.ident.clone();
305 if root_field.fields.iter().next().is_some() {
306 panic!("unsupported enum type");
307 }
308 let str_name = name.to_string().replace("_", "-");
309 let len = str_name.len();
310 try_from_str_items
311 .push(quote! { #len if value.eq_ignore_ascii_case(#str_name) => Some(Self::#name), });
312 to_str_items.push(quote! { Self::#name => #str_name, });
313 headers_items.push(quote! { #name(String), });
314 headers_apply_items
315 .push(quote! { Headers::#name(s) => self.set_header(HeaderItem::#name.to_str(), s), });
316 }
317 quote! {
318 impl #enum_name {
319 pub fn try_from_str(value: &str) -> Option<Self> {
320 match value.len() {
321 #( #try_from_str_items )*
322 _ => None,
323 }
324 }
325
326 pub fn to_str(&self) -> &'static str {
327 match self {
328 #( #to_str_items )*
329 }
330 }
331 }
332
333 pub enum Headers {
334 #( #headers_items )*
335 Custom((String, String)),
336 }
337
338 impl HttpRequest {
339 pub fn apply_header(&mut self, header: Headers) {
340 match header {
341 #( #headers_apply_items )*
342 Headers::Custom((k, v)) => self.set_header(&k[..], v),
343 }
344 }
345 }
346 }
347 .into()
348}