aiway-plugin 0.3.5

//! WASM 插件导出宏
//!
//! 为插件开发者生成 C ABI 导出函数，供网关 wasmtime 运行时调用。

/// 格式化日志宏（ERROR 级别）
///
/// # 用法
/// ```ignore
/// log_error!(ctx, "request failed: {}", err);
/// ```
#[macro_export]
macro_rules! log_error {
    ($ctx:expr, $($arg:tt)*) => {
        $ctx.log_error(&format!($($arg)*))
    };
}

/// WARN日志
#[macro_export]
macro_rules! log_warn {
    ($ctx:expr, $($arg:tt)*) => {
        $ctx.log_warn(&format!($($arg)*))
    };
}

/// INFO日志
#[macro_export]
macro_rules! log_info {
    ($ctx:expr, $($arg:tt)*) => {
        $ctx.log_info(&format!($($arg)*))
    };
}

/// DEBUG日志
#[macro_export]
macro_rules! log_debug {
    ($ctx:expr, $($arg:tt)*) => {
        $ctx.log_debug(&format!($($arg)*))
    };
}

/// TRACE日志
#[macro_export]
macro_rules! log_trace {
    ($ctx:expr, $($arg:tt)*) => {
        $ctx.log_trace(&format!($($arg)*))
    };
}

/// 导出 WASM 插件
///
/// 生成以下导出函数：
/// - `plugin_info()` -> i64：返回插件元信息（bincode 编码）
/// - `aiway_call(hook_id, input_ptr, input_len)` -> i64：插件钩子调用入口
/// - `aiway_alloc(size)` -> i32：内存分配
/// - `aiway_dealloc(ptr, size)`：内存释放
///
/// # 用法
/// ```ignore
/// struct MyPlugin;
/// impl aiway_plugin::Plugin for MyPlugin { /* ... */ }
///
/// aiway_plugin::export_wasm!(MyPlugin);
/// ```
#[macro_export]
macro_rules! export_wasm {
    ($plugin_type:ty) => {
        /// 插件实例（全局单例，因为 WASM 插件应无状态）
        static PLUGIN: std::sync::LazyLock<$plugin_type> =
            std::sync::LazyLock::new(|| <$plugin_type>::new());

        /// 分配内存（供 Host 写入输入数据）
        #[unsafe(no_mangle)]
        pub extern "C" fn aiway_alloc(size: i32) -> i32 {
            let layout = std::alloc::Layout::from_size_align(size as usize, 1).unwrap();
            unsafe {
                let ptr = std::alloc::alloc(layout);
                ptr as i32
            }
        }

        /// 释放内存
        #[unsafe(no_mangle)]
        pub extern "C" fn aiway_dealloc(ptr: i32, size: i32) {
            let layout = std::alloc::Layout::from_size_align(size as usize, 1).unwrap();
            unsafe {
                std::alloc::dealloc(ptr as *mut u8, layout);
            }
        }

        /// 返回插件元信息
        ///
        /// 返回 i64，高 32 位 = 数据指针，低 32 位 = 数据长度
        #[unsafe(no_mangle)]
        pub extern "C" fn plugin_info() -> i64 {
            let info = aiway_plugin::wasm_types::WasmPluginInfo {
                name: PLUGIN.name().to_string(),
                version: PLUGIN.info().version.to_string(),
                description: PLUGIN.info().description.clone(),
                default_config: aiway_plugin::serde_json::to_string(&PLUGIN.info().default_config)
                    .unwrap_or_default(),
                readme:  PLUGIN.info().readme.clone(),
            };

            let bytes = $crate::bincode::serialize(&info).unwrap();
            let len = bytes.len();
            let ptr = bytes.as_ptr() as i32;
            std::mem::forget(bytes);

            ((ptr as i64) << 32) | (len as i64)
        }

        /// 插件钩子调用入口
        ///
        /// # 参数
        /// - `hook_id`: 钩子 ID（见 `wasm_types` 常量）
        /// - `input_ptr`: 输入数据在 WASM 内存中的指针
        /// - `input_len`: 输入数据长度
        ///
        /// # 返回
        /// i64，高 32 位 = 结果指针，低 32 位 = 结果长度。
        /// 若高 32 位为 0，表示错误，低 32 位是错误信息长度（数据在 ptr=1 处）。
        #[unsafe(no_mangle)]
        pub extern "C" fn aiway_call(hook_id: i32, input_ptr: i32, input_len: i32) -> i64 {
            // 读取输入数据
            let input_slice =
                unsafe { std::slice::from_raw_parts(input_ptr as *const u8, input_len as usize) };

            let input: $crate::wasm_types::WasmInput = match $crate::bincode::deserialize(input_slice)
            {
                Ok(v) => v,
                Err(e) => return encode_error(&format!("deserialize input failed: {}", e)),
            };

            // 根据 hook_id 分发
            let result: Result<aiway_plugin::wasm_types::WasmOutput, String> = match hook_id {
                aiway_plugin::wasm_types::HOOK_ON_REQUEST => handle_on_request(&PLUGIN, &input),
                aiway_plugin::wasm_types::HOOK_ON_REQUEST_BODY => {
                    handle_on_request_body(&PLUGIN, &input)
                }
                aiway_plugin::wasm_types::HOOK_ON_RESPONSE => handle_on_response(&PLUGIN, &input),
                aiway_plugin::wasm_types::HOOK_ON_RESPONSE_BODY => {
                    handle_on_response_body(&PLUGIN, &input)
                }
                aiway_plugin::wasm_types::HOOK_ON_LOGGING => handle_on_logging(&PLUGIN, &input),
                _ => Err(format!("unknown hook_id: {}", hook_id)),
            };

            match result {
                Ok(output) => encode_output(&output),
                Err(err_bytes) => encode_error(&err_bytes),
            }
        }

        /// 编码成功输出
        fn encode_output(output: &aiway_plugin::wasm_types::WasmOutput) -> i64 {
            match $crate::bincode::serialize(output) {
                Ok(bytes) => {
                    let len = bytes.len();
                    let ptr = bytes.as_ptr() as i32;
                    std::mem::forget(bytes);
                    ((ptr as i64) << 32) | (len as i64)
                }
                Err(e) => encode_error(&format!("serialize output failed: {}", e)),
            }
        }

        /// 编码错误信息：写入 ptr=1，返回 ptr=0 标记错误
        fn encode_error(msg: &str) -> i64 {
            let bytes = msg.as_bytes();
            // 将错误信息写入 ptr=1 处
            let dst = unsafe { std::slice::from_raw_parts_mut(1 as *mut u8, bytes.len()) };
            dst.copy_from_slice(bytes);
            (0i64 << 32) | (bytes.len() as i64)
        }

        /// 将 PluginError 编码为错误消息
        fn encode_plugin_error(e: aiway_plugin::PluginError) -> String {
            format!("{}", e)
        }

        /// 解析 JSON 配置字符串
        fn parse_config(config_str: &str) -> Result<aiway_plugin::serde_json::Value, String> {
            aiway_plugin::serde_json::from_str(config_str)
                .map_err(|e| format!("parse config failed: {}", e))
        }

        /// 处理 on_request
        fn handle_on_request(
            plugin: &$plugin_type,
            input: &aiway_plugin::wasm_types::WasmInput,
        ) -> Result<aiway_plugin::wasm_types::WasmOutput, String> {
            let config = parse_config(&input.config)?;
            let mut ctx = aiway_plugin::WasmHttpContext;

            let outcome = aiway_plugin::block_on(async {
                plugin.on_request(&config, &mut ctx).await
            })
            .map_err(encode_plugin_error)?;

            match outcome {
                aiway_plugin::Outcome::Continue => {
                    Ok(aiway_plugin::wasm_types::WasmOutput {
                        body: None,
                        respond: None,
                    })
                }
                aiway_plugin::Outcome::Respond(resp) => Ok(aiway_plugin::wasm_types::WasmOutput {
                    body: None,
                    respond: Some(aiway_plugin::wasm_types::WasmRespond {
                        status: resp.status,
                        headers: resp.headers,
                        body: resp.body,
                    }),
                }),
            }
        }

        /// 处理 on_request_body
        fn handle_on_request_body(
            plugin: &$plugin_type,
            input: &aiway_plugin::wasm_types::WasmInput,
        ) -> Result<aiway_plugin::wasm_types::WasmOutput, String> {
            let config = parse_config(&input.config)?;
            let mut body = input
                .body
                .as_ref()
                .map(|b| aiway_plugin::Bytes::from(b.clone()));
            let mut ctx = aiway_plugin::WasmHttpContext;

            let outcome = aiway_plugin::block_on(async {
                plugin.on_request_body(&config, &mut body, &mut ctx).await
            })
            .map_err(encode_plugin_error)?;

            match outcome {
                aiway_plugin::Outcome::Continue => Ok(aiway_plugin::wasm_types::WasmOutput {
                    body: body.map(|b| b.to_vec()),
                    respond: None,
                }),
                aiway_plugin::Outcome::Respond(resp) => Ok(
                    aiway_plugin::wasm_types::WasmOutput {
                        body: None,
                        respond: Some(aiway_plugin::wasm_types::WasmRespond {
                            status: resp.status,
                            headers: resp.headers,
                            body: resp.body,
                        }),
                    },
                ),
            }
        }

        /// 处理 on_response
        fn handle_on_response(
            plugin: &$plugin_type,
            input: &aiway_plugin::wasm_types::WasmInput,
        ) -> Result<aiway_plugin::wasm_types::WasmOutput, String> {
            let config = parse_config(&input.config)?;
            let mut ctx = aiway_plugin::WasmHttpContext;

            let outcome = aiway_plugin::block_on(async {
                plugin.on_response(&config, &mut ctx).await
            })
            .map_err(encode_plugin_error)?;

            match outcome {
                aiway_plugin::Outcome::Continue => {
                    Ok(aiway_plugin::wasm_types::WasmOutput {
                        body: None,
                        respond: None,
                    })
                }
                aiway_plugin::Outcome::Respond(resp) => Ok(aiway_plugin::wasm_types::WasmOutput {
                    body: None,
                    respond: Some(aiway_plugin::wasm_types::WasmRespond {
                        status: resp.status,
                        headers: resp.headers,
                        body: resp.body,
                    }),
                }),
            }
        }

        /// 处理 on_response_body
        fn handle_on_response_body(
            plugin: &$plugin_type,
            input: &aiway_plugin::wasm_types::WasmInput,
        ) -> Result<aiway_plugin::wasm_types::WasmOutput, String> {
            let config = parse_config(&input.config)?;
            let mut body = input
                .body
                .as_ref()
                .map(|b| aiway_plugin::Bytes::from(b.clone()));
            let mut ctx = aiway_plugin::WasmHttpContext;

            let outcome = aiway_plugin::block_on(async {
                plugin.on_response_body(&config, &mut body, &mut ctx).await
            })
            .map_err(encode_plugin_error)?;

            match outcome {
                aiway_plugin::Outcome::Continue => Ok(aiway_plugin::wasm_types::WasmOutput {
                    body: body.map(|b| b.to_vec()),
                    respond: None,
                }),
                aiway_plugin::Outcome::Respond(resp) => Ok(
                    aiway_plugin::wasm_types::WasmOutput {
                        body: None,
                        respond: Some(aiway_plugin::wasm_types::WasmRespond {
                            status: resp.status,
                            headers: resp.headers,
                            body: resp.body,
                        }),
                    },
                ),
            }
        }

        /// 处理 on_logging
        fn handle_on_logging(
            plugin: &$plugin_type,
            input: &aiway_plugin::wasm_types::WasmInput,
        ) -> Result<aiway_plugin::wasm_types::WasmOutput, String> {
            let config = parse_config(&input.config)?;
            let mut ctx = aiway_plugin::WasmHttpContext;

            aiway_plugin::block_on(async {
                plugin.on_logging(&config, &mut ctx).await;
            });

            Ok(aiway_plugin::wasm_types::WasmOutput {
                body: None,
                respond: None,
            })
        }
    };
}