Building a Custom Guardrail Pipeline for LLM API Gateways Using Rust and WASM — code snippets

snippet-1.txt

use log::{info, warn};
use proxy_wasm::traits::*;
use proxy_wasm::types::*;

proxy_wasm::main! {{
    proxy_wasm::set_log_level(LogLevel::Info);
    proxy_wasm::set_root_context(|_| -> Box<dyn RootContext> {
        Box::new(GuardrailRootContext {
            config: GuardrailConfig::default(),
        })
    });
}}

#[derive(Default, Clone)]
struct GuardrailConfig {
    blocked_patterns: Vec<String>,
    redaction_replacement: String,
    pii_redaction_enabled: bool,
}

struct GuardrailRootContext {
    config: GuardrailConfig,
}

impl RootContext for GuardrailRootContext {
    fn create_http_context(&self, _context_id: u32) -> Option<Box<dyn HttpContext>> {
        Some(Box::new(GuardrailHttpContext {
            config: self.config.clone(),
            request_buffer: Vec::new(),
            response_buffer: Vec::new(),
        }))
    }

    fn get_type(&self) -> Option<ContextType> {
        Some(ContextType::HttpContext)
    }
}

struct GuardrailHttpContext {
    config: GuardrailConfig,
    request_buffer: Vec<u8>,
    response_buffer: Vec<u8>,
}

impl HttpContext for GuardrailHttpContext {
    fn on_http_request_headers(&mut self, _num_headers: usize, _end_of_stream: bool) -> Action {
        // Strip downstream trace headers or inject telemetry metadata
        Action::Continue
    }
}

snippet-2.txt

use aho_corasick::{AhoCorasick, AhoCorasickBuilder, MatchKind};

pub struct PatternScanner {
    automaton: AhoCorasick,
}

impl PatternScanner {
    pub fn new(patterns: &[String]) -> Result<Self, String> {
        let automaton = AhoCorasickBuilder::new()
            .match_kind(MatchKind::LeftmostLongest)
            .ascii_case_insensitive(true)
            .build(patterns)
            .map_err(|e| format!("Failed to compile Aho-Corasick: {:?}", e))?;
        Ok(Self { automaton })
    }

    pub fn scan_and_sanitize(&self, input: &str, replacement: &str) -> (String, bool) {
        let mut output = String::with_capacity(input.len());
        let mut last_match_end = 0;
        let mut matched = false;

        for mat in self.automaton.find_iter(input) {
            matched = true;
            output.push_str(&input[last_match_end..mat.start()]);
            output.push_str(replacement);
            last_match_end = mat.end();
        }
        output.push_str(&input[last_match_end..]);
        (output, matched)
    }
}

snippet-3.txt

use serde::Deserialize;

#[derive(Deserialize, Default, Clone)]
struct RawGuardrailConfig {
    #[serde(default)]
    blocked_patterns: Vec<String>,
    #[serde(default = "default_replacement")]
    redaction_replacement: String,
    #[serde(default = "default_true")]
    pii_redaction_enabled: bool,
}

fn default_replacement() -> String {
    "[REDACTED]".to_string()
}

fn default_true() -> bool {
    true
}

impl RootContext for GuardrailRootContext {
    fn on_configure(&mut self, _plugin_configuration_size: usize) -> bool {
        if let Some(config_bytes) = self.get_plugin_configuration() {
            match serde_json::from_slice::<RawGuardrailConfig>(&config_bytes) {
                Ok(raw) => {
                    self.config = GuardrailConfig {
                        blocked_patterns: raw.blocked_patterns,
                        redaction_replacement: raw.redaction_replacement,
                        pii_redaction_enabled: raw.pii_redaction_enabled,
                    };
                    info!("WASM configuration successfully parsed and updated.");
                    true
                }
                Err(e) => {
                    warn!("Invalid WASM JSON configuration: {:?}", e);
                    false
                }
            }
        } else {
            warn!("No WASM configuration provided; using defaults.");
            true
        }
    }
}

snippet-4.txt

impl HttpContext for GuardrailHttpContext {
    fn on_http_response_body(&mut self, body_limit: usize, end_of_stream: bool) -> Action {
        if let Some(chunk) = self.get_http_response_body(0, body_limit) {
            self.response_buffer.extend_from_slice(&chunk);
            
            let raw_text = match std::str::from_utf8(&self.response_buffer) {
                Ok(t) => t,
                Err(_) => {
                    // Buffer is cut off in the middle of a multi-byte UTF-8 character.
                    // Yield execution and wait for next chunk.
                    return Action::Pause;
                }
            };

            // Build scanner with active configuration
            let scanner = match PatternScanner::new(&self.config.blocked_patterns) {
                Ok(s) => s,
                Err(_) => return Action::Continue,
            };

            let (sanitized_text, modified) = scanner.scan_and_sanitize(raw_text, &self.config.redaction_replacement);
            
            if modified {
                // Mutate the body inside Envoy's filter buffer
                self.set_http_response_body(0, self.response_buffer.len(), sanitized_text.as_bytes());
                self.response_buffer.clear();
            }
        }

        if end_of_stream {
            Action::Continue
        } else {
            // Keep proxy filter execution paused for this chunk until buffer is flushed
            Action::Continue
        }
    }
}

snippet-6.sh

#!/usr/bin/env bash
# Build script for WASM compilation and local deployment validation
set -euo pipefail

TARGET_ARCH="wasm32-wasi"
OUTPUT_DIR="./dist"

echo "[*] Adding WebAssembly build targets..."
rustup target add ${TARGET_ARCH}

echo "[*] Compiling guardrail binary to WebAssembly in release mode..."
cargo build --target ${TARGET_ARCH} --release

echo "[*] Copying binary output to destination directory..."
mkdir -p "${OUTPUT_DIR}"
cp target/${TARGET_ARCH}/release/guardrail_pipeline.wasm "${OUTPUT_DIR}/guardrail_pipeline.wasm"

echo "[*] Running local Envoy container with WASM volume mount for verification..."
docker run --rm -d \
  -p 8080:8080 \
  -v "$(pwd)/dist/guardrail_pipeline.wasm:/etc/envoy/wasm/guardrail_pipeline.wasm:ro" \
  -v "$(pwd)/envoy.yaml:/etc/envoy/envoy.yaml:ro" \
  --name envoy-guardrail-test \
  envoyproxy/envoy:v1.28.0

echo "[*] Verifying Envoy container startup health..."
sleep 2
if docker ps | grep -q "envoy-guardrail-test"; then
  echo "[+] Deployment verification succeeded. Envoy is running and listening on port 8080."
else
  echo "[!] Envoy container failed to start. Check docker logs."
  exit 1
fi