Skip to main content

aiway_plugin/
macros.rs

1//! WASM 插件导出宏
2//!
3//! 为插件开发者生成 C ABI 导出函数,供网关 wasmtime 运行时调用。
4
5/// 格式化日志宏(ERROR 级别)
6///
7/// # 用法
8/// ```ignore
9/// log_error!(ctx, "request failed: {}", err);
10/// ```
11#[macro_export]
12macro_rules! log_error {
13    ($ctx:expr, $($arg:tt)*) => {
14        $ctx.log_error(&format!($($arg)*))
15    };
16}
17
18/// WARN日志
19#[macro_export]
20macro_rules! log_warn {
21    ($ctx:expr, $($arg:tt)*) => {
22        $ctx.log_warn(&format!($($arg)*))
23    };
24}
25
26/// INFO日志
27#[macro_export]
28macro_rules! log_info {
29    ($ctx:expr, $($arg:tt)*) => {
30        $ctx.log_info(&format!($($arg)*))
31    };
32}
33
34/// DEBUG日志
35#[macro_export]
36macro_rules! log_debug {
37    ($ctx:expr, $($arg:tt)*) => {
38        $ctx.log_debug(&format!($($arg)*))
39    };
40}
41
42/// TRACE日志
43#[macro_export]
44macro_rules! log_trace {
45    ($ctx:expr, $($arg:tt)*) => {
46        $ctx.log_trace(&format!($($arg)*))
47    };
48}
49
50/// 导出 WASM 插件
51///
52/// 生成以下导出函数:
53/// - `plugin_info()` -> i64:返回插件元信息(bincode 编码)
54/// - `aiway_call(hook_id, input_ptr, input_len)` -> i64:插件钩子调用入口
55/// - `aiway_alloc(size)` -> i32:内存分配
56/// - `aiway_dealloc(ptr, size)`:内存释放
57///
58/// # 用法
59/// ```ignore
60/// struct MyPlugin;
61/// impl aiway_plugin::Plugin for MyPlugin { /* ... */ }
62///
63/// aiway_plugin::export_wasm!(MyPlugin);
64/// ```
65#[macro_export]
66macro_rules! export_wasm {
67    ($plugin_type:ty) => {
68        /// 插件实例(全局单例,因为 WASM 插件应无状态)
69        static PLUGIN: std::sync::LazyLock<$plugin_type> =
70            std::sync::LazyLock::new(|| <$plugin_type>::new());
71
72        /// 分配内存(供 Host 写入输入数据)
73        #[unsafe(no_mangle)]
74        pub extern "C" fn aiway_alloc(size: i32) -> i32 {
75            let layout = std::alloc::Layout::from_size_align(size as usize, 1).unwrap();
76            unsafe {
77                let ptr = std::alloc::alloc(layout);
78                ptr as i32
79            }
80        }
81
82        /// 释放内存
83        #[unsafe(no_mangle)]
84        pub extern "C" fn aiway_dealloc(ptr: i32, size: i32) {
85            let layout = std::alloc::Layout::from_size_align(size as usize, 1).unwrap();
86            unsafe {
87                std::alloc::dealloc(ptr as *mut u8, layout);
88            }
89        }
90
91        /// 返回插件元信息
92        ///
93        /// 返回 i64,高 32 位 = 数据指针,低 32 位 = 数据长度
94        #[unsafe(no_mangle)]
95        pub extern "C" fn plugin_info() -> i64 {
96            let info = aiway_plugin::wasm_types::WasmPluginInfo {
97                name: PLUGIN.name().to_string(),
98                version: PLUGIN.info().version.to_string(),
99                description: PLUGIN.info().description.clone(),
100                default_config: aiway_plugin::serde_json::to_string(&PLUGIN.info().default_config)
101                    .unwrap_or_default(),
102                readme:  PLUGIN.info().readme.clone(),
103            };
104
105            let bytes = $crate::bincode::serialize(&info).unwrap();
106            let len = bytes.len();
107            let ptr = bytes.as_ptr() as i32;
108            std::mem::forget(bytes);
109
110            ((ptr as i64) << 32) | (len as i64)
111        }
112
113        /// 插件钩子调用入口
114        ///
115        /// # 参数
116        /// - `hook_id`: 钩子 ID(见 `wasm_types` 常量)
117        /// - `input_ptr`: 输入数据在 WASM 内存中的指针
118        /// - `input_len`: 输入数据长度
119        ///
120        /// # 返回
121        /// i64,高 32 位 = 结果指针,低 32 位 = 结果长度。
122        /// 若高 32 位为 0,表示错误,低 32 位是错误信息长度(数据在 ptr=1 处)。
123        #[unsafe(no_mangle)]
124        pub extern "C" fn aiway_call(hook_id: i32, input_ptr: i32, input_len: i32) -> i64 {
125            // 读取输入数据
126            let input_slice =
127                unsafe { std::slice::from_raw_parts(input_ptr as *const u8, input_len as usize) };
128
129            let input: $crate::wasm_types::WasmInput = match $crate::bincode::deserialize(input_slice)
130            {
131                Ok(v) => v,
132                Err(e) => return encode_error(&format!("deserialize input failed: {}", e)),
133            };
134
135            // 根据 hook_id 分发
136            let result: Result<aiway_plugin::wasm_types::WasmOutput, String> = match hook_id {
137                aiway_plugin::wasm_types::HOOK_ON_REQUEST => handle_on_request(&PLUGIN, &input),
138                aiway_plugin::wasm_types::HOOK_ON_REQUEST_BODY => {
139                    handle_on_request_body(&PLUGIN, &input)
140                }
141                aiway_plugin::wasm_types::HOOK_ON_RESPONSE => handle_on_response(&PLUGIN, &input),
142                aiway_plugin::wasm_types::HOOK_ON_RESPONSE_BODY => {
143                    handle_on_response_body(&PLUGIN, &input)
144                }
145                aiway_plugin::wasm_types::HOOK_ON_LOGGING => handle_on_logging(&PLUGIN, &input),
146                _ => Err(format!("unknown hook_id: {}", hook_id)),
147            };
148
149            match result {
150                Ok(output) => encode_output(&output),
151                Err(err_bytes) => encode_error(&err_bytes),
152            }
153        }
154
155        /// 编码成功输出
156        fn encode_output(output: &aiway_plugin::wasm_types::WasmOutput) -> i64 {
157            match $crate::bincode::serialize(output) {
158                Ok(bytes) => {
159                    let len = bytes.len();
160                    let ptr = bytes.as_ptr() as i32;
161                    std::mem::forget(bytes);
162                    ((ptr as i64) << 32) | (len as i64)
163                }
164                Err(e) => encode_error(&format!("serialize output failed: {}", e)),
165            }
166        }
167
168        /// 编码错误信息:写入 ptr=1,返回 ptr=0 标记错误
169        fn encode_error(msg: &str) -> i64 {
170            let bytes = msg.as_bytes();
171            // 将错误信息写入 ptr=1 处
172            let dst = unsafe { std::slice::from_raw_parts_mut(1 as *mut u8, bytes.len()) };
173            dst.copy_from_slice(bytes);
174            (0i64 << 32) | (bytes.len() as i64)
175        }
176
177        /// 将 PluginError 编码为错误消息
178        fn encode_plugin_error(e: aiway_plugin::PluginError) -> String {
179            format!("{}", e)
180        }
181
182        /// 解析 JSON 配置字符串
183        fn parse_config(config_str: &str) -> Result<aiway_plugin::serde_json::Value, String> {
184            aiway_plugin::serde_json::from_str(config_str)
185                .map_err(|e| format!("parse config failed: {}", e))
186        }
187
188        /// 处理 on_request
189        fn handle_on_request(
190            plugin: &$plugin_type,
191            input: &aiway_plugin::wasm_types::WasmInput,
192        ) -> Result<aiway_plugin::wasm_types::WasmOutput, String> {
193            let config = parse_config(&input.config)?;
194            let mut ctx = aiway_plugin::WasmHttpContext;
195
196            let outcome = aiway_plugin::block_on(async {
197                plugin.on_request(&config, &mut ctx).await
198            })
199            .map_err(encode_plugin_error)?;
200
201            match outcome {
202                aiway_plugin::Outcome::Continue => {
203                    Ok(aiway_plugin::wasm_types::WasmOutput {
204                        body: None,
205                        respond: None,
206                    })
207                }
208                aiway_plugin::Outcome::Respond(resp) => Ok(aiway_plugin::wasm_types::WasmOutput {
209                    body: None,
210                    respond: Some(aiway_plugin::wasm_types::WasmRespond {
211                        status: resp.status,
212                        headers: resp.headers,
213                        body: resp.body,
214                    }),
215                }),
216            }
217        }
218
219        /// 处理 on_request_body
220        fn handle_on_request_body(
221            plugin: &$plugin_type,
222            input: &aiway_plugin::wasm_types::WasmInput,
223        ) -> Result<aiway_plugin::wasm_types::WasmOutput, String> {
224            let config = parse_config(&input.config)?;
225            let mut body = input
226                .body
227                .as_ref()
228                .map(|b| aiway_plugin::Bytes::from(b.clone()));
229            let mut ctx = aiway_plugin::WasmHttpContext;
230
231            let outcome = aiway_plugin::block_on(async {
232                plugin.on_request_body(&config, &mut body, &mut ctx).await
233            })
234            .map_err(encode_plugin_error)?;
235
236            match outcome {
237                aiway_plugin::Outcome::Continue => Ok(aiway_plugin::wasm_types::WasmOutput {
238                    body: body.map(|b| b.to_vec()),
239                    respond: None,
240                }),
241                aiway_plugin::Outcome::Respond(resp) => Ok(
242                    aiway_plugin::wasm_types::WasmOutput {
243                        body: None,
244                        respond: Some(aiway_plugin::wasm_types::WasmRespond {
245                            status: resp.status,
246                            headers: resp.headers,
247                            body: resp.body,
248                        }),
249                    },
250                ),
251            }
252        }
253
254        /// 处理 on_response
255        fn handle_on_response(
256            plugin: &$plugin_type,
257            input: &aiway_plugin::wasm_types::WasmInput,
258        ) -> Result<aiway_plugin::wasm_types::WasmOutput, String> {
259            let config = parse_config(&input.config)?;
260            let mut ctx = aiway_plugin::WasmHttpContext;
261
262            let outcome = aiway_plugin::block_on(async {
263                plugin.on_response(&config, &mut ctx).await
264            })
265            .map_err(encode_plugin_error)?;
266
267            match outcome {
268                aiway_plugin::Outcome::Continue => {
269                    Ok(aiway_plugin::wasm_types::WasmOutput {
270                        body: None,
271                        respond: None,
272                    })
273                }
274                aiway_plugin::Outcome::Respond(resp) => Ok(aiway_plugin::wasm_types::WasmOutput {
275                    body: None,
276                    respond: Some(aiway_plugin::wasm_types::WasmRespond {
277                        status: resp.status,
278                        headers: resp.headers,
279                        body: resp.body,
280                    }),
281                }),
282            }
283        }
284
285        /// 处理 on_response_body
286        fn handle_on_response_body(
287            plugin: &$plugin_type,
288            input: &aiway_plugin::wasm_types::WasmInput,
289        ) -> Result<aiway_plugin::wasm_types::WasmOutput, String> {
290            let config = parse_config(&input.config)?;
291            let mut body = input
292                .body
293                .as_ref()
294                .map(|b| aiway_plugin::Bytes::from(b.clone()));
295            let mut ctx = aiway_plugin::WasmHttpContext;
296
297            let outcome = aiway_plugin::block_on(async {
298                plugin.on_response_body(&config, &mut body, &mut ctx).await
299            })
300            .map_err(encode_plugin_error)?;
301
302            match outcome {
303                aiway_plugin::Outcome::Continue => Ok(aiway_plugin::wasm_types::WasmOutput {
304                    body: body.map(|b| b.to_vec()),
305                    respond: None,
306                }),
307                aiway_plugin::Outcome::Respond(resp) => Ok(
308                    aiway_plugin::wasm_types::WasmOutput {
309                        body: None,
310                        respond: Some(aiway_plugin::wasm_types::WasmRespond {
311                            status: resp.status,
312                            headers: resp.headers,
313                            body: resp.body,
314                        }),
315                    },
316                ),
317            }
318        }
319
320        /// 处理 on_logging
321        fn handle_on_logging(
322            plugin: &$plugin_type,
323            input: &aiway_plugin::wasm_types::WasmInput,
324        ) -> Result<aiway_plugin::wasm_types::WasmOutput, String> {
325            let config = parse_config(&input.config)?;
326            let mut ctx = aiway_plugin::WasmHttpContext;
327
328            aiway_plugin::block_on(async {
329                plugin.on_logging(&config, &mut ctx).await;
330            });
331
332            Ok(aiway_plugin::wasm_types::WasmOutput {
333                body: None,
334                respond: None,
335            })
336        }
337    };
338}