89 lines
2.5 KiB
Rust
89 lines
2.5 KiB
Rust
//! WASM Guest Logging Example for sgl-model-gateway
|
|
//!
|
|
//! This example demonstrates logging and tracing middleware
|
|
//! for sgl-model-gateway using the WebAssembly Component Model.
|
|
//!
|
|
//! Features:
|
|
//! - Request tracking and tracing headers
|
|
//! - Response status code conversion
|
|
|
|
wit_bindgen::generate!({
|
|
path: "../../../src/wasm/interface",
|
|
world: "sgl-model-gateway",
|
|
});
|
|
|
|
use exports::sgl::model_gateway::{
|
|
middleware_on_request::Guest as OnRequestGuest,
|
|
middleware_on_response::Guest as OnResponseGuest,
|
|
};
|
|
use sgl::model_gateway::middleware_types::{Action, Header, ModifyAction, Request, Response};
|
|
|
|
/// Main middleware implementation
|
|
struct Middleware;
|
|
|
|
// Helper function to create header
|
|
fn create_header(name: &str, value: &str) -> Header {
|
|
Header {
|
|
name: name.to_string(),
|
|
value: value.to_string(),
|
|
}
|
|
}
|
|
|
|
// Implement on-request interface
|
|
impl OnRequestGuest for Middleware {
|
|
fn on_request(req: Request) -> Action {
|
|
let mut modify_action = ModifyAction {
|
|
status: None,
|
|
headers_set: vec![],
|
|
headers_add: vec![],
|
|
headers_remove: vec![],
|
|
body_replace: None,
|
|
};
|
|
|
|
// Request Logging and Tracing
|
|
// Add tracing headers with request ID
|
|
modify_action
|
|
.headers_add
|
|
.push(create_header("x-request-id", &req.request_id));
|
|
modify_action
|
|
.headers_add
|
|
.push(create_header("x-wasm-processed", "true"));
|
|
modify_action.headers_add.push(create_header(
|
|
"x-processed-at",
|
|
&req.now_epoch_ms.to_string(),
|
|
));
|
|
|
|
// Add custom header for API requests
|
|
if req.path.starts_with("/api") || req.path.starts_with("/v1") {
|
|
modify_action
|
|
.headers_add
|
|
.push(create_header("x-api-route", "true"));
|
|
}
|
|
|
|
Action::Modify(modify_action)
|
|
}
|
|
}
|
|
|
|
// Implement on-response interface
|
|
impl OnResponseGuest for Middleware {
|
|
fn on_response(resp: Response) -> Action {
|
|
// Status code conversion: Convert 500 to 503 for better client handling
|
|
if resp.status == 500 {
|
|
let modify_action = ModifyAction {
|
|
status: Some(503),
|
|
headers_set: vec![],
|
|
headers_add: vec![],
|
|
headers_remove: vec![],
|
|
body_replace: None,
|
|
};
|
|
Action::Modify(modify_action)
|
|
} else {
|
|
// No modification needed
|
|
Action::Continue
|
|
}
|
|
}
|
|
}
|
|
|
|
// Export the component
|
|
export!(Middleware);
|