Feature 11: WASM Plugin System

Feature 11: WASM Plugin System

Priority: High | Complexity: High | Phase: 4 (Differentiation)

---

Overview

Problem Statement

Every organization has unique requirements:

• Custom authentication schemes
• Domain-specific query transformations
• Proprietary caching strategies
• Integration with internal systems

Static proxy features can’t anticipate every use case. Plugin systems traditionally require:

• Native code (C/Rust) with security risks
• Scripting languages (Lua) with performance costs
• Custom builds per deployment

Solution

Implement a WebAssembly (WASM) plugin system that enables safe, sandboxed, high-performance extensibility:

                    ┌─────────────────────────────────────────────────┐
                    │              WASM PLUGIN RUNTIME                 │
                    │                                                  │
                    │  ┌──────────────────────────────────────────┐   │
                    │  │ Plugin Manager                           │   │
                    │  │ - Load/unload plugins                    │   │
                    │  │ - Version management                     │   │
                    │  │ - Health monitoring                      │   │
                    │  └──────────────────────────────────────────┘   │
                    │                    │                             │
                    │  ┌─────────────────┼─────────────────┐          │
                    │  │                 │                 │          │
                    │  ▼                 ▼                 ▼          │
                    │ ┌────────┐    ┌────────┐       ┌────────┐      │
                    │ │Plugin A│    │Plugin B│       │Plugin C│      │
                    │ │(Auth)  │    │(Cache) │       │(Rewrite│      │
                    │ │.wasm   │    │.wasm   │       │.wasm   │      │
                    │ └────────┘    └────────┘       └────────┘      │
                    │      │              │               │           │
                    │  ┌───┴──────────────┴───────────────┴───┐      │
                    │  │ Host Functions (Secure API)          │      │
                    │  │ - Query execution                    │      │
                    │  │ - Cache access                        │      │
                    │  │ - Metrics                            │      │
                    │  │ - Logging                            │      │
                    │  └──────────────────────────────────────┘      │
                    └─────────────────────────────────────────────────┘

---

Architecture

Plugin Runtime

use wasmtime::{Engine, Module, Store, Instance, Linker};

pub struct WasmPluginRuntime {
    /// WASM engine
    engine: Engine,

    /// Loaded plugins
    plugins: DashMap<String, LoadedPlugin>,

    /// Host function registry
    host_functions: HostFunctions,

    /// Configuration
    config: PluginRuntimeConfig,
}

pub struct LoadedPlugin {
    /// Plugin metadata
    pub metadata: PluginMetadata,

    /// Compiled WASM module
    module: Module,

    /// Plugin instance (per-thread)
    instances: ThreadLocal<Instance>,

    /// Plugin state
    state: PluginState,
}

#[derive(Debug, Clone)]
pub struct PluginMetadata {
    pub name: String,
    pub version: String,
    pub description: String,
    pub author: String,
    pub hooks: Vec<HookType>,
    pub permissions: Vec<Permission>,
}

#[derive(Debug, Clone, Copy)]
pub enum HookType {
    /// Before query execution
    PreQuery,

    /// After query execution
    PostQuery,

    /// Authentication
    Authenticate,

    /// Authorization
    Authorize,

    /// Cache lookup
    CacheGet,

    /// Cache store
    CacheSet,

    /// Routing decision
    Route,

    /// Query rewriting
    Rewrite,

    /// Metrics collection
    Metrics,
}

pub struct PluginRuntimeConfig {
    /// Memory limit per plugin (bytes)
    pub memory_limit: usize,

    /// Execution timeout per call
    pub timeout: Duration,

    /// Max plugins loaded
    pub max_plugins: usize,

    /// Plugin directory
    pub plugin_dir: PathBuf,

    /// Auto-reload on file change
    pub hot_reload: bool,
}

Host Functions

pub struct HostFunctions {
    /// Query executor
    query_executor: Arc<QueryExecutor>,

    /// Cache interface
    cache: Arc<QueryCache>,

    /// Metrics registry
    metrics: Arc<MetricsRegistry>,

    /// Logger
    logger: Arc<Logger>,
}

impl HostFunctions {
    /// Register host functions in WASM linker
    pub fn register(&self, linker: &mut Linker<PluginContext>) -> Result<()> {
        // Query execution
        linker.func_wrap("helios", "execute_query", |ctx: Caller<PluginContext>, query_ptr: i32, query_len: i32| -> i32 {
            let query = ctx.read_string(query_ptr, query_len);
            let result = ctx.data().host.query_executor.execute(&query);
            ctx.write_result(&result)
        })?;

        // Cache operations
        linker.func_wrap("helios", "cache_get", |ctx: Caller<PluginContext>, key_ptr: i32, key_len: i32| -> i32 {
            let key = ctx.read_string(key_ptr, key_len);
            let value = ctx.data().host.cache.get(&key);
            ctx.write_option(&value)
        })?;

        linker.func_wrap("helios", "cache_set", |ctx: Caller<PluginContext>, key_ptr: i32, key_len: i32, value_ptr: i32, value_len: i32, ttl: i64| {
            let key = ctx.read_string(key_ptr, key_len);
            let value = ctx.read_bytes(value_ptr, value_len);
            ctx.data().host.cache.set(&key, &value, Duration::from_secs(ttl as u64));
        })?;

        // Logging
        linker.func_wrap("helios", "log", |ctx: Caller<PluginContext>, level: i32, msg_ptr: i32, msg_len: i32| {
            let msg = ctx.read_string(msg_ptr, msg_len);
            let level = LogLevel::from_i32(level);
            ctx.data().host.logger.log(level, &msg);
        })?;

        // Metrics
        linker.func_wrap("helios", "metric_inc", |ctx: Caller<PluginContext>, name_ptr: i32, name_len: i32, value: f64| {
            let name = ctx.read_string(name_ptr, name_len);
            ctx.data().host.metrics.increment(&name, value);
        })?;

        Ok(())
    }
}

Plugin Development Kit

// Plugin SDK (compiled to WASM)
// This would be a separate crate: heliosproxy-plugin-sdk

/// Plugin trait that all plugins must implement
#[no_mangle]
pub trait HeliosPlugin {
    /// Plugin metadata
    fn metadata() -> PluginMetadata;

    /// Called when plugin is loaded
    fn on_load() -> Result<(), PluginError>;

    /// Called when plugin is unloaded
    fn on_unload();
}

/// Pre-query hook
#[no_mangle]
pub trait PreQueryHook {
    fn pre_query(ctx: &QueryContext) -> PreQueryResult;
}

/// Authentication hook
#[no_mangle]
pub trait AuthenticateHook {
    fn authenticate(request: &AuthRequest) -> AuthResult;
}

/// Cache hook
#[no_mangle]
pub trait CacheHook {
    fn cache_key(ctx: &QueryContext) -> Option<String>;
    fn should_cache(ctx: &QueryContext, result: &QueryResult) -> bool;
}

// Example plugin implementation
mod my_custom_auth {
    use heliosproxy_plugin_sdk::*;

    #[no_mangle]
    pub fn metadata() -> PluginMetadata {
        PluginMetadata {
            name: "custom-auth".to_string(),
            version: "1.0.0".to_string(),
            hooks: vec![HookType::Authenticate],
            permissions: vec![Permission::HttpFetch],
        }
    }

    #[no_mangle]
    pub fn authenticate(request: &AuthRequest) -> AuthResult {
        // Custom authentication logic
        let token = request.header("X-Custom-Token")?;

        // Validate token (plugin has HTTP access)
        let response = http_fetch("https://auth.internal/validate", &token)?;

        if response.status == 200 {
            AuthResult::Success(Identity {
                user_id: response.body["user_id"].clone(),
                roles: response.body["roles"].clone(),
            })
        } else {
            AuthResult::Denied("Invalid token")
        }
    }
}

---

API Specification

Configuration (heliosproxy.toml)

[plugins]
enabled = true
plugin_dir = "/etc/heliosproxy/plugins"
hot_reload = true

# Resource limits
memory_limit = "64MB"
timeout = "100ms"
max_plugins = 20

# Plugin-specific config
[plugins.custom-auth]
enabled = true
priority = 100
config = { auth_endpoint = "https://auth.internal/validate" }

[plugins.custom-cache]
enabled = true
priority = 50
config = { ttl_multiplier = 2.0 }

# Permissions per plugin
[plugins.permissions]
custom-auth = ["http_fetch", "cache_read"]
custom-cache = ["cache_read", "cache_write"]

Admin API

GET /plugins
{
  "plugins": [
    {
      "name": "custom-auth",
      "version": "1.0.0",
      "status": "running",
      "hooks": ["authenticate"],
      "memory_used": "2.3MB",
      "invocations": 15234,
      "avg_latency_us": 45
    }
  ]
}

POST /plugins/load
# Load new plugin
{ "path": "/etc/heliosproxy/plugins/new-plugin.wasm" }

POST /plugins/{name}/reload
# Reload plugin (hot reload)

POST /plugins/{name}/disable
# Disable plugin without unloading

DELETE /plugins/{name}
# Unload plugin

GET /plugins/{name}/logs
# Plugin-specific logs

GET /plugins/{name}/metrics
# Plugin-specific metrics

Plugin Manifest

# plugin.yaml - shipped with .wasm file
name: custom-auth
version: 1.0.0
description: Custom authentication using internal auth service
author: Acme Corp
license: MIT

# Required hooks
hooks:
  - authenticate
  - authorize

# Required permissions
permissions:
  - http_fetch
  - cache_read

# Configuration schema
config_schema:
  auth_endpoint:
    type: string
    required: true
    description: URL of internal auth service
  cache_ttl:
    type: integer
    default: 300
    description: Cache TTL for auth results

# Resource requirements
resources:
  min_memory: 16MB
  max_memory: 64MB

---

AI/Agent Innovations

1. Custom Embedding Plugins

Allow custom embedding models:

// Plugin for custom embeddings
#[no_mangle]
pub fn compute_embedding(text: &str) -> Vec<f32> {
    // Load custom model (cached in WASM memory)
    let model = get_or_load_model();

    // Compute embedding
    model.embed(text)
}

#[no_mangle]
pub fn similarity(a: &[f32], b: &[f32]) -> f32 {
    // Custom similarity metric
    cosine_similarity(a, b)
}

2. RAG Pipeline Plugins

Custom RAG orchestration:

// Plugin for custom RAG logic
#[no_mangle]
pub fn rag_retrieve(query: &str, config: &RagConfig) -> Vec<Chunk> {
    // Custom retrieval strategy
    let embeddings = compute_query_embedding(query);

    // Multi-index search
    let semantic = search_semantic_index(&embeddings);
    let keyword = search_keyword_index(query);

    // Custom fusion
    reciprocal_rank_fusion(semantic, keyword)
}

#[no_mangle]
pub fn rag_rerank(query: &str, chunks: &[Chunk]) -> Vec<Chunk> {
    // Custom reranking
    let scores = cross_encoder_score(query, chunks);
    sort_by_score(chunks, &scores)
}

3. Agent Tool Plugins

Define custom tools for AI agents:

// Plugin for custom agent tools
#[no_mangle]
pub fn list_tools() -> Vec<ToolDefinition> {
    vec![
        ToolDefinition {
            name: "internal_api_call",
            description: "Call internal API endpoint",
            parameters: json!({
                "endpoint": {"type": "string"},
                "method": {"type": "string", "enum": ["GET", "POST"]},
                "body": {"type": "object"}
            }),
        }
    ]
}

#[no_mangle]
pub fn execute_tool(name: &str, params: &Value) -> ToolResult {
    match name {
        "internal_api_call" => {
            let endpoint = params["endpoint"].as_str().unwrap();
            let method = params["method"].as_str().unwrap();
            let body = &params["body"];

            // Execute internal API call
            http_fetch(endpoint, method, body)
        }
        _ => ToolResult::Unknown,
    }
}

4. LLM Response Plugins

Custom response processing:

// Plugin for LLM response processing
#[no_mangle]
pub fn post_process_response(response: &str, context: &Context) -> String {
    // Custom post-processing
    let sanitized = remove_pii(response);
    let formatted = apply_formatting_rules(sanitized, &context.format);
    add_citations(formatted, &context.sources)
}

---

HeliosDB-Lite Integration

1. Branch Policy Plugins

Custom branch access rules:

// Plugin for custom branch policies
#[no_mangle]
pub fn can_access_branch(identity: &Identity, branch: &str) -> bool {
    // Custom access logic
    match branch {
        b if b.starts_with("feature_") => identity.has_role("developer"),
        "production" => identity.has_role("admin"),
        _ => true,
    }
}

#[no_mangle]
pub fn on_branch_merge(source: &str, target: &str, identity: &Identity) -> MergePolicy {
    // Custom merge approval
    if target == "production" {
        MergePolicy::RequireApproval {
            approvers: vec!["admin", "lead"],
            min_approvals: 2,
        }
    } else {
        MergePolicy::AutoMerge
    }
}

2. Vector Index Plugins

Custom vector indexing:

// Plugin for custom vector operations
#[no_mangle]
pub fn custom_distance(a: &[f32], b: &[f32]) -> f32 {
    // Custom distance metric (e.g., weighted cosine)
    weighted_cosine(a, b, &DOMAIN_WEIGHTS)
}

#[no_mangle]
pub fn custom_quantization(vectors: &[Vec<f32>]) -> QuantizedVectors {
    // Custom quantization strategy
    product_quantization(vectors, 8, 256)
}

3. Time-Travel Plugins

Custom history policies:

// Plugin for custom time-travel policies
#[no_mangle]
pub fn time_travel_policy(identity: &Identity, table: &str, as_of: Timestamp) -> TimePolicy {
    // Compliance-based access
    let retention = get_retention_policy(table);

    if as_of < now() - retention {
        TimePolicy::Denied("Exceeds retention period")
    } else if table == "audit_log" && !identity.has_role("auditor") {
        TimePolicy::Denied("Audit access required")
    } else {
        TimePolicy::Allowed
    }
}

4. Sync Mode Plugins

Custom replication decisions:

// Plugin for custom sync mode selection
#[no_mangle]
pub fn select_sync_mode(query: &str, context: &QueryContext) -> SyncMode {
    // Business-critical queries get sync
    if context.tables.iter().any(|t| CRITICAL_TABLES.contains(t)) {
        return SyncMode::Sync;
    }

    // Batch operations can be async
    if context.is_batch || context.rows_affected > 1000 {
        return SyncMode::Async;
    }

    SyncMode::SemiSync
}

---

Implementation Notes

File Locations

src/proxy/
├── plugins/
│   ├── mod.rs              # Public API
│   ├── runtime.rs          # WasmPluginRuntime
│   ├── loader.rs           # Plugin loading
│   ├── host_functions.rs   # Host function registry
│   ├── sandbox.rs          # Security sandbox
│   ├── hot_reload.rs       # Hot reload support
│   └── metrics.rs          # Plugin metrics

# Plugin SDK (separate crate)
crates/heliosproxy-plugin-sdk/
├── Cargo.toml
├── src/
│   ├── lib.rs
│   ├── hooks.rs
│   ├── types.rs
│   └── bindings.rs

Key Considerations

1. Security: WASM provides memory isolation. Limit permissions per plugin.

2. Performance: Pre-compile WASM to native code. Use instance pooling.

3. Debugging: Provide stack traces and logging from WASM.

4. Versioning: Support multiple versions of plugin API.

5. Distribution: Plugin registry for sharing community plugins.

WASM Runtime Choice

Cargo.toml

[dependencies]
wasmtime = "18"  # Production-ready, fast
# Alternative: wasmer = "4"

Security Sandbox

pub struct PluginSandbox {
    /// Allowed host functions
    allowed_functions: HashSet<String>,

    /// Resource limits
    limits: ResourceLimits,

    /// Network policy
    network_policy: NetworkPolicy,
}

impl PluginSandbox {
    pub fn validate_call(&self, function: &str) -> Result<(), SecurityError> {
        if !self.allowed_functions.contains(function) {
            return Err(SecurityError::FunctionNotAllowed(function.to_string()));
        }
        Ok(())
    }
}

Performance Targets

Metric	Target	Measurement
Plugin load time	<100ms	cold start
Hook invocation	<50μs	p99 (warm)
Memory per plugin	<64MB	max
WASM compilation	<500ms	per plugin

---

Related Features

• Query Rewriting - Rewriting via plugins
• Authentication Proxy - Auth via plugins
• Query Caching - Caching via plugins