Add isolated Luau rule engine service

crates/rule-engine/src/lib.rs

@@ -0,0 +1,250 @@
+use mlua::{Function, Lua, LuaSerdeExt, Value as LuaValue};
+use serde::{Deserialize, Serialize};
+use serde_json::Value;
+use uuid::Uuid;
+
+pub const CONTRACT_VERSION: u16 = 1;
+pub const EVALUATION_SUBJECT: &str = "ai-spel.rules.evaluate.v1";
+pub const MAX_RULE_BYTES: usize = 16 * 1024;
+pub const MAX_SNAPSHOT_BYTES: usize = 32 * 1024;
+pub const MAX_RESULT_BYTES: usize = 16 * 1024;
+pub const MAX_EFFECTS: usize = 16;
+pub const MAX_MEMORY_BYTES: usize = 2 * 1024 * 1024;
+pub const MAX_INTERRUPTS: usize = 10_000;
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+#[serde(deny_unknown_fields)]
+pub struct WorldSnapshot {
+    pub world_id: Uuid,
+    pub version: i64,
+    pub world_time: i64,
+    pub facts: Value,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
+#[serde(deny_unknown_fields)]
+pub struct ProposedEffect {
+    pub kind: String,
+    pub payload: Value,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+#[serde(deny_unknown_fields)]
+pub struct RuleInvocation {
+    pub rule_id: Uuid,
+    pub rule_version: i32,
+    pub source: String,
+    pub snapshot: WorldSnapshot,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+#[serde(deny_unknown_fields)]
+pub struct EvaluationRequest {
+    pub contract_version: u16,
+    pub correlation_id: Uuid,
+    pub invocation: RuleInvocation,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+#[serde(tag = "status", rename_all = "snake_case", deny_unknown_fields)]
+pub enum EvaluationOutcome {
+    Evaluated { effects: Vec<ProposedEffect> },
+    Rejected { code: String, detail: String },
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+#[serde(deny_unknown_fields)]
+pub struct EvaluationResponse {
+    pub contract_version: u16,
+    pub correlation_id: Uuid,
+    pub outcome: EvaluationOutcome,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub enum RuleError {
+    RuleTooLarge,
+    SnapshotTooLarge,
+    ResultTooLarge,
+    TooManyEffects,
+    InvalidEffect,
+    Sandbox(String),
+}
+
+impl RuleError {
+    pub fn code(&self) -> &'static str {
+        match self {
+            Self::RuleTooLarge => "rule_too_large",
+            Self::SnapshotTooLarge => "snapshot_too_large",
+            Self::ResultTooLarge => "result_too_large",
+            Self::TooManyEffects => "too_many_effects",
+            Self::InvalidEffect => "invalid_effect",
+            Self::Sandbox(_) => "sandbox_rejected",
+        }
+    }
+}
+
+#[derive(Debug, Default, Clone, Copy)]
+pub struct LuauRuntime;
+
+impl LuauRuntime {
+    pub fn evaluate(&self, invocation: &RuleInvocation) -> Result<Vec<ProposedEffect>, RuleError> {
+        if invocation.source.len() > MAX_RULE_BYTES {
+            return Err(RuleError::RuleTooLarge);
+        }
+        let snapshot_json = serde_json::to_vec(&invocation.snapshot)
+            .map_err(|error| RuleError::Sandbox(error.to_string()))?;
+        if snapshot_json.len() > MAX_SNAPSHOT_BYTES {
+            return Err(RuleError::SnapshotTooLarge);
+        }
+
+        let lua = Lua::new();
+        lua.set_memory_limit(MAX_MEMORY_BYTES)
+            .map_err(|error| RuleError::Sandbox(error.to_string()))?;
+        lua.sandbox(true)
+            .map_err(|error| RuleError::Sandbox(error.to_string()))?;
+        let interrupt_count = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0));
+        lua.set_interrupt({
+            let interrupt_count = interrupt_count.clone();
+            move |_| {
+                let count = interrupt_count.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
+                if count >= MAX_INTERRUPTS {
+                    return Err(mlua::Error::RuntimeError(
+                        "rule execution budget exceeded".into(),
+                    ));
+                }
+                Ok(mlua::VmState::Continue)
+            }
+        });
+
+        let result = (|| -> mlua::Result<LuaValue> {
+            let rule: Function = lua.load(&invocation.source).set_name("game-rule").eval()?;
+            let snapshot = lua.to_value(&invocation.snapshot)?;
+            rule.call(snapshot)
+        })()
+        .map_err(|error| RuleError::Sandbox(error.to_string()))?;
+
+        let effects: Vec<ProposedEffect> = lua
+            .from_value(result)
+            .map_err(|error| RuleError::Sandbox(error.to_string()))?;
+        if effects.len() > MAX_EFFECTS {
+            return Err(RuleError::TooManyEffects);
+        }
+        let encoded = serde_json::to_vec(&effects)
+            .map_err(|error| RuleError::Sandbox(error.to_string()))?;
+        if encoded.len() > MAX_RESULT_BYTES {
+            return Err(RuleError::ResultTooLarge);
+        }
+        if effects.iter().any(|effect| {
+            effect.kind.trim().is_empty()
+                || effect.kind.len() > 64
+                || !effect.payload.is_object()
+                || serde_json::to_vec(&effect.payload)
+                    .map(|payload| payload.len() > MAX_RESULT_BYTES)
+                    .unwrap_or(true)
+        }) {
+            return Err(RuleError::InvalidEffect);
+        }
+        Ok(effects)
+    }
+}
+
+pub fn evaluate_request(request: EvaluationRequest) -> EvaluationResponse {
+    let correlation_id = request.correlation_id;
+    let outcome = if request.contract_version != CONTRACT_VERSION {
+        EvaluationOutcome::Rejected {
+            code: "unsupported_contract_version".into(),
+            detail: "The requested rule evaluation contract version is not supported.".into(),
+        }
+    } else {
+        match LuauRuntime.evaluate(&request.invocation) {
+            Ok(effects) => EvaluationOutcome::Evaluated { effects },
+            Err(error) => EvaluationOutcome::Rejected {
+                code: error.code().into(),
+                detail: match error {
+                    RuleError::Sandbox(detail) => detail,
+                    _ => format!("Rule evaluation rejected: {}", error.code()),
+                },
+            },
+        }
+    };
+    EvaluationResponse {
+        contract_version: CONTRACT_VERSION,
+        correlation_id,
+        outcome,
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn request(source: &str) -> EvaluationRequest {
+        EvaluationRequest {
+            contract_version: CONTRACT_VERSION,
+            correlation_id: Uuid::nil(),
+            invocation: RuleInvocation {
+                rule_id: Uuid::nil(),
+                rule_version: 1,
+                source: source.into(),
+                snapshot: WorldSnapshot {
+                    world_id: Uuid::nil(),
+                    version: 7,
+                    world_time: 12,
+                    facts: serde_json::json!({"minutes":5}),
+                },
+            },
+        }
+    }
+
+    #[test]
+    fn evaluates_luau_to_typed_proposals() {
+        let result = evaluate_request(request(
+            r#"return function(snapshot)
+                return {{ kind = "advance_time", payload = { minutes = snapshot.facts.minutes } }}
+            end"#,
+        ));
+        let EvaluationOutcome::Evaluated { effects } = result.outcome else {
+            panic!("expected a typed effect proposal");
+        };
+        assert_eq!(effects[0].kind, "advance_time");
+        assert_eq!(effects[0].payload["minutes"], 5);
+    }
+
+    #[test]
+    fn rejects_ambient_io_and_non_effect_results() {
+        let result = evaluate_request(request(
+            "return function() return io.open('/etc/passwd') end",
+        ));
+        assert!(matches!(result.outcome, EvaluationOutcome::Rejected { .. }));
+        let result = evaluate_request(request(
+            "return function() return {{ kind = 'advance_time', payload = 'not an object' }} end",
+        ));
+        assert!(matches!(result.outcome, EvaluationOutcome::Rejected { .. }));
+        let result = evaluate_request(request(
+            "return function() return 'not an effect list' end",
+        ));
+        assert!(matches!(result.outcome, EvaluationOutcome::Rejected { .. }));
+    }
+
+    #[test]
+    fn rejects_execution_budget_exhaustion() {
+        let result = evaluate_request(request("return function() while true do end end"));
+        assert!(matches!(result.outcome, EvaluationOutcome::Rejected { .. }));
+    }
+
+    #[test]
+    fn rejects_wrong_contract_and_oversized_source() {
+        let mut wrong_version = request("return function() return {} end");
+        wrong_version.contract_version = CONTRACT_VERSION + 1;
+        let result = evaluate_request(wrong_version);
+        assert!(matches!(
+            result.outcome,
+            EvaluationOutcome::Rejected { ref code, .. } if code == "unsupported_contract_version"
+        ));
+        let result = evaluate_request(request(&" ".repeat(MAX_RULE_BYTES + 1)));
+        assert!(matches!(
+            result.outcome,
+            EvaluationOutcome::Rejected { ref code, .. } if code == "rule_too_large"
+        ));
+    }
+}

crates/rule-engine/src/main.rs

@@ -0,0 +1,171 @@
+use ai_spel_rule_engine::{
+    CONTRACT_VERSION, EVALUATION_SUBJECT, EvaluationRequest, EvaluationResponse, evaluate_request,
+};
+use async_nats::{Client, ConnectOptions, Subscriber, Subject};
+use axum::{Router, extract::State, http::StatusCode, routing::get};
+use futures_util::StreamExt;
+use std::{env, sync::Arc, time::Duration};
+use tokio::{task::JoinSet, time};
+
+#[derive(Clone)]
+struct Health {
+    nats_connected: bool,
+}
+
+async fn live() -> StatusCode {
+    StatusCode::NO_CONTENT
+}
+
+async fn ready(State(health): State<Arc<tokio::sync::RwLock<Health>>>) -> StatusCode {
+    if health.read().await.nats_connected {
+        StatusCode::NO_CONTENT
+    } else {
+        StatusCode::SERVICE_UNAVAILABLE
+    }
+}
+
+async fn connect_nats() -> Result<Client, Box<dyn std::error::Error>> {
+    let url = env::var("NATS_URL").unwrap_or_else(|_| "nats://ai-spel-nats:4222".into());
+    let username = env::var("NATS_USERNAME")?;
+    let password = env::var("NATS_PASSWORD")?;
+    Ok(ConnectOptions::new()
+        .user_and_password(username, password)
+        .connect(url)
+        .await?)
+}
+
+async fn smoke_nats() -> Result<(), Box<dyn std::error::Error>> {
+    let client = connect_nats().await?;
+    let correlation_id = uuid::Uuid::new_v4();
+    let request = EvaluationRequest {
+        contract_version: CONTRACT_VERSION,
+        correlation_id,
+        invocation: ai_spel_rule_engine::RuleInvocation {
+            rule_id: uuid::Uuid::nil(),
+            rule_version: 1,
+            source: "return function() return {} end".into(),
+            snapshot: ai_spel_rule_engine::WorldSnapshot {
+                world_id: uuid::Uuid::nil(),
+                version: 0,
+                world_time: 0,
+                facts: serde_json::json!({}),
+            },
+        },
+    };
+    let bytes = serde_json::to_vec(&request)?;
+    let response = time::timeout(
+        Duration::from_secs(5),
+        client.request(EVALUATION_SUBJECT, bytes.into()),
+    )
+    .await??;
+    let response: EvaluationResponse = serde_json::from_slice(&response.payload)?;
+    if response.contract_version != CONTRACT_VERSION || response.correlation_id != correlation_id {
+        return Err("rule service smoke response did not match the request contract".into());
+    }
+    match response.outcome {
+        ai_spel_rule_engine::EvaluationOutcome::Evaluated { effects } if effects.is_empty() => {
+            println!("Rule engine NATS request/reply smoke passed")
+        }
+        _ => return Err("rule service smoke evaluation was rejected or returned effects".into()),
+    }
+    Ok(())
+}
+
+fn response_for_payload(payload: &[u8]) -> Option<Vec<u8>> {
+    let request: EvaluationRequest = match serde_json::from_slice(payload) {
+        Ok(request) => request,
+        Err(_) => return None,
+    };
+    serde_json::to_vec(&evaluate_request(request)).ok()
+}
+
+async fn respond(client: Client, reply: Subject, payload: Vec<u8>) {
+    if let Err(error) = client.publish(reply, payload.into()).await {
+        eprintln!(
+            "{{\"level\":\"error\",\"event\":\"rule.reply_publish_failed\",\"error_category\":\"nats_publish\",\"detail\":{:?}}}",
+            error.to_string()
+        );
+    }
+}
+
+async fn serve_rule_requests(
+    client: Client,
+    mut subscriber: Subscriber,
+    max_concurrency: usize,
+) -> Result<(), Box<dyn std::error::Error>> {
+    let mut tasks = JoinSet::new();
+    loop {
+        tokio::select! {
+            message = subscriber.next() => {
+                let Some(message) = message else { break; };
+                let Some(reply) = message.reply else { continue; };
+                while tasks.len() >= max_concurrency {
+                    let _ = tasks.join_next().await;
+                }
+                let client = client.clone();
+                tasks.spawn(async move {
+                    let payload = tokio::task::spawn_blocking(move || response_for_payload(&message.payload)).await;
+                    if let Ok(Some(payload)) = payload {
+                        respond(client, reply, payload).await;
+                    }
+                });
+            }
+            Some(_) = tasks.join_next(), if !tasks.is_empty() => {}
+        }
+    }
+    while tasks.join_next().await.is_some() {}
+    Ok(())
+}
+
+#[tokio::main]
+async fn main() -> Result<(), Box<dyn std::error::Error>> {
+    if env::args().any(|argument| argument == "--smoke") {
+        return smoke_nats().await;
+    }
+    let health = Arc::new(tokio::sync::RwLock::new(Health {
+        nats_connected: false,
+    }));
+    let health_router = Router::new()
+        .route("/health/live", get(live))
+        .route("/health/ready", get(ready))
+        .with_state(health.clone());
+    let health_address = env::var("HEALTH_ADDR").unwrap_or_else(|_| "0.0.0.0:8080".into());
+    let health_listener = tokio::net::TcpListener::bind(health_address).await?;
+    tokio::spawn(async move { let _ = axum::serve(health_listener, health_router).await; });
+
+    let client = loop {
+        match connect_nats().await {
+            Ok(client) => break client,
+            Err(error) => {
+                eprintln!("{{\"level\":\"warn\",\"event\":\"nats.connect_retry\",\"error_category\":\"nats_connect\",\"detail\":{:?}}}", error.to_string());
+                time::sleep(Duration::from_secs(2)).await;
+            }
+        }
+    };
+    let queue_group = env::var("NATS_QUEUE_GROUP").unwrap_or_else(|_| "ai-spel-rule-engine".into());
+    let subscriber = client
+        .queue_subscribe(EVALUATION_SUBJECT, queue_group)
+        .await?;
+    health.write().await.nats_connected = true;
+    {
+        let health = health.clone();
+        let health_client = client.clone();
+        tokio::spawn(async move {
+            loop {
+                let connected = matches!(
+                    time::timeout(Duration::from_secs(2), health_client.flush()).await,
+                    Ok(Ok(()))
+                );
+                health.write().await.nats_connected = connected;
+                time::sleep(Duration::from_secs(5)).await;
+            }
+        });
+    }
+    let concurrency = env::var("RULE_MAX_CONCURRENCY")
+        .ok()
+        .and_then(|value| value.parse::<usize>().ok())
+        .filter(|value| (1..=32).contains(value))
+        .unwrap_or(8);
+    eprintln!("{{\"level\":\"info\",\"event\":\"rule_engine.ready\",\"contract_version\":{},\"max_concurrency\":{}}}", CONTRACT_VERSION, concurrency);
+    serve_rule_requests(client, subscriber, concurrency).await
+}