featherbit/plugins/script/
lua_runtime.rs

1//! Lua scripting runtime for the `script` plugin, built on mlua's Luau VM.
2//!
3//! Marshals the gateway `Context` into a Lua table, calls the script's
4//! global `execute(ctx)` function, and marshals the returned table back into
5//! a `Context`. Also installs a sandboxed `require` restricted to a
6//! configured modules directory.
7
8use mlua::prelude::*;
9use std::collections::HashMap;
10use std::path::PathBuf;
11
12use crate::context::{Context, GatewayError, GatewayRequest, GatewayResponse, Protocol};
13use crate::plugins::PluginExecutionError;
14
15/// Holds a validated Lua script and executes it against a `Context`.
16///
17/// A fresh Lua VM is created for every execution, so scripts cannot leak
18/// state between requests; only the source text is retained between calls.
19pub struct LuaRuntime {
20    /// The full script source, re-loaded into a fresh VM per execution.
21    source: String,
22    /// Directory the sandboxed `require` resolves modules from; `None`
23    /// disables `require`.
24    modules_path: Option<PathBuf>,
25    /// Configured execution timeout in milliseconds (currently stored but
26    /// not enforced by the VM).
27    #[allow(dead_code)] // see roadmap: script execution timeouts
28    timeout_ms: u64,
29}
30
31impl LuaRuntime {
32    /// Compiles and validates the script in a throwaway VM, failing early if
33    /// the source has syntax errors, its top level errors on load, or it does
34    /// not define a global `execute` function. This runs once at
35    /// policy-compile time, not per request.
36    pub fn new(
37        source: &str,
38        timeout_ms: u64,
39        modules_path: Option<PathBuf>,
40    ) -> Result<Self, String> {
41        // Validate the script compiles
42        let lua = Lua::new();
43        setup_module_loader(&lua, &modules_path);
44        lua.load(source)
45            .exec()
46            .map_err(|e| format!("Lua compilation error: {}", e))?;
47
48        // Verify execute function exists
49        lua.globals()
50            .get::<LuaFunction>("execute")
51            .map_err(|_| "Lua script must define an 'execute(ctx)' function".to_string())?;
52
53        Ok(Self {
54            source: source.to_string(),
55            modules_path,
56            timeout_ms,
57        })
58    }
59
60    /// Runs the script's `execute(ctx)` against the given context in a fresh
61    /// VM and returns the context rebuilt from the table the script returned.
62    ///
63    /// Every failure mode (load, marshalling either way, missing `execute`,
64    /// or a runtime error raised by the script) returns a
65    /// `PluginExecutionError` carrying the original context, with a
66    /// distinguishing error code (`LUA_LOAD_ERROR`, `LUA_MARSHAL_ERROR`,
67    /// `LUA_MISSING_EXECUTE`, `LUA_EXECUTION_ERROR`, `LUA_UNMARSHAL_ERROR`),
68    /// so the graph engine routes through the error port exactly like a
69    /// native plugin failure.
70    pub fn execute(&self, ctx: Context) -> Result<Context, PluginExecutionError> {
71        let lua = Lua::new();
72        setup_module_loader(&lua, &self.modules_path);
73
74        if let Err(e) = lua.load(&self.source).exec() {
75            return Err(PluginExecutionError {
76                context: ctx,
77                error: GatewayError {
78                    node_id: String::new(),
79                    code: "LUA_LOAD_ERROR".to_string(),
80                    message: format!("Failed to load Lua script: {}", e),
81                    metadata: HashMap::new(),
82                },
83            });
84        }
85
86        let ctx_table = match context_to_lua(&lua, &ctx) {
87            Ok(t) => t,
88            Err(e) => {
89                return Err(PluginExecutionError {
90                    context: ctx,
91                    error: GatewayError {
92                        node_id: String::new(),
93                        code: "LUA_MARSHAL_ERROR".to_string(),
94                        message: format!("Failed to marshal context to Lua: {}", e),
95                        metadata: HashMap::new(),
96                    },
97                });
98            }
99        };
100
101        let execute_fn: LuaFunction = match lua.globals().get("execute") {
102            Ok(f) => f,
103            Err(e) => {
104                return Err(PluginExecutionError {
105                    context: ctx,
106                    error: GatewayError {
107                        node_id: String::new(),
108                        code: "LUA_MISSING_EXECUTE".to_string(),
109                        message: format!("Missing execute function: {}", e),
110                        metadata: HashMap::new(),
111                    },
112                });
113            }
114        };
115
116        let result_table: LuaTable = match execute_fn.call(ctx_table) {
117            Ok(t) => t,
118            Err(e) => {
119                return Err(PluginExecutionError {
120                    context: ctx,
121                    error: GatewayError {
122                        node_id: String::new(),
123                        code: "LUA_EXECUTION_ERROR".to_string(),
124                        message: format!("Lua execution error: {}", e),
125                        metadata: HashMap::new(),
126                    },
127                });
128            }
129        };
130
131        // Carry over the fields scripts never see: the wire protocol and the
132        // errors accumulated by earlier nodes must survive a script node.
133        let protocol = ctx.request.protocol.clone();
134        let errors = ctx.errors.clone();
135        match lua_to_context(&result_table, protocol, errors) {
136            Ok(new_ctx) => Ok(new_ctx),
137            Err(e) => Err(PluginExecutionError {
138                context: ctx,
139                error: GatewayError {
140                    node_id: String::new(),
141                    code: "LUA_UNMARSHAL_ERROR".to_string(),
142                    message: format!("Failed to unmarshal context from Lua: {}", e),
143                    metadata: HashMap::new(),
144                },
145            }),
146        }
147    }
148}
149
150/// Registers a custom `require()` loader that reads `.lua` files from the modules directory.
151///
152/// The loader is sandboxed: module names containing `..`, `/`, or `\` are
153/// rejected, so only files directly inside `modules_path` can be loaded
154/// (resolved as `<modules_path>/<name>.lua`). When `modules_path` is `None`,
155/// no `require` is installed. Note: modules are re-evaluated on every
156/// `require`; results are not cached.
157fn setup_module_loader(lua: &Lua, modules_path: &Option<PathBuf>) {
158    let Some(base_path) = modules_path.clone() else {
159        return;
160    };
161
162    // In Luau, we use the require override approach
163    let loader = lua
164        .create_function(move |lua, module_name: String| {
165            // Sanitize: no path traversal
166            if module_name.contains("..") || module_name.contains('/') || module_name.contains('\\')
167            {
168                return Err(LuaError::runtime(format!(
169                    "Invalid module name '{}': path traversal not allowed",
170                    module_name
171                )));
172            }
173
174            // Try module_name.lua
175            let file_path = base_path.join(format!("{}.lua", module_name));
176            let source = std::fs::read_to_string(&file_path).map_err(|e| {
177                LuaError::runtime(format!(
178                    "Cannot load module '{}' from {:?}: {}",
179                    module_name, file_path, e
180                ))
181            })?;
182
183            // Execute the module and return its result
184            lua.load(&source).eval::<LuaValue>().map_err(|e| {
185                LuaError::runtime(format!("Error loading module '{}': {}", module_name, e))
186            })
187        })
188        .expect("Failed to create module loader");
189
190    lua.globals()
191        .set("require", loader)
192        .expect("Failed to set require");
193}
194
195/// Marshals a `Context` into a Lua table with `request`, `response`, and
196/// `message` sub-tables. Header and query-param values become 1-indexed
197/// arrays of strings; bodies become Lua strings; `message` values are
198/// converted from JSON. `errors` is not exposed to scripts.
199fn context_to_lua(lua: &Lua, ctx: &Context) -> LuaResult<LuaTable> {
200    let table = lua.create_table()?;
201
202    // request
203    let req = lua.create_table()?;
204    req.set("method", ctx.request.method.as_str())?;
205    req.set("path", ctx.request.path.as_str())?;
206    req.set("host", ctx.request.host.as_str())?;
207    req.set("scheme", ctx.request.scheme.as_str())?;
208    req.set("remote_addr", ctx.request.remote_addr.as_str())?;
209
210    let headers = lua.create_table()?;
211    for (k, v) in &ctx.request.headers {
212        let vals = lua.create_table()?;
213        for (i, val) in v.iter().enumerate() {
214            vals.set(i + 1, val.as_str())?;
215        }
216        headers.set(k.as_str(), vals)?;
217    }
218    req.set("headers", headers)?;
219
220    let query = lua.create_table()?;
221    for (k, v) in &ctx.request.query_params {
222        let vals = lua.create_table()?;
223        for (i, val) in v.iter().enumerate() {
224            vals.set(i + 1, val.as_str())?;
225        }
226        query.set(k.as_str(), vals)?;
227    }
228    req.set("query_params", query)?;
229
230    req.set("body", lua.create_string(&ctx.request.body)?)?;
231    table.set("request", req)?;
232
233    // response
234    let resp = lua.create_table()?;
235    resp.set("status_code", ctx.response.status_code)?;
236    let resp_headers = lua.create_table()?;
237    for (k, v) in &ctx.response.headers {
238        let vals = lua.create_table()?;
239        for (i, val) in v.iter().enumerate() {
240            vals.set(i + 1, val.as_str())?;
241        }
242        resp_headers.set(k.as_str(), vals)?;
243    }
244    resp.set("headers", resp_headers)?;
245    resp.set("body", lua.create_string(&ctx.response.body)?)?;
246    table.set("response", resp)?;
247
248    // message
249    let msg = lua.create_table()?;
250    for (k, v) in &ctx.message {
251        let lua_val = json_to_lua(lua, v)?;
252        msg.set(k.as_str(), lua_val)?;
253    }
254    table.set("message", msg)?;
255
256    Ok(table)
257}
258
259/// Rebuilds a `Context` from the table returned by the script's `execute`.
260///
261/// `request` and `response` (including their `headers` and bodies) are
262/// required and fail unmarshalling if malformed; `query_params` and
263/// `message` are optional. `protocol` and `errors` are not exposed to Lua,
264/// so the caller passes the original context's values through unchanged.
265fn lua_to_context(
266    table: &LuaTable,
267    protocol: Protocol,
268    errors: Vec<GatewayError>,
269) -> LuaResult<Context> {
270    let req_table: LuaTable = table.get("request")?;
271    let resp_table: LuaTable = table.get("response")?;
272
273    let mut request_headers = HashMap::new();
274    let headers_table: LuaTable = req_table.get("headers")?;
275    for pair in headers_table.pairs::<String, LuaTable>() {
276        let (k, v) = pair?;
277        let mut vals = Vec::new();
278        for val in v.sequence_values::<String>() {
279            vals.push(val?);
280        }
281        request_headers.insert(k, vals);
282    }
283
284    let mut query_params = HashMap::new();
285    if let Ok(qp_table) = req_table.get::<LuaTable>("query_params") {
286        for pair in qp_table.pairs::<String, LuaTable>() {
287            let (k, v) = pair?;
288            let mut vals = Vec::new();
289            for val in v.sequence_values::<String>() {
290                vals.push(val?);
291            }
292            query_params.insert(k, vals);
293        }
294    }
295
296    let body_str: mlua::String = req_table.get("body")?;
297    let request = GatewayRequest {
298        method: req_table.get("method")?,
299        path: req_table.get("path")?,
300        host: req_table.get("host")?,
301        scheme: req_table.get("scheme")?,
302        headers: request_headers,
303        query_params,
304        body: bytes::Bytes::from(body_str.as_bytes().to_vec()),
305        remote_addr: req_table.get("remote_addr")?,
306        protocol,
307    };
308
309    let mut response_headers = HashMap::new();
310    let resp_headers_table: LuaTable = resp_table.get("headers")?;
311    for pair in resp_headers_table.pairs::<String, LuaTable>() {
312        let (k, v) = pair?;
313        let mut vals = Vec::new();
314        for val in v.sequence_values::<String>() {
315            vals.push(val?);
316        }
317        response_headers.insert(k, vals);
318    }
319
320    let resp_body_str: mlua::String = resp_table.get("body")?;
321    let response = GatewayResponse {
322        status_code: resp_table.get("status_code")?,
323        headers: response_headers,
324        body: bytes::Bytes::from(resp_body_str.as_bytes().to_vec()),
325    };
326
327    let mut message = HashMap::new();
328    if let Ok(msg_table) = table.get::<LuaTable>("message") {
329        for pair in msg_table.pairs::<String, LuaValue>() {
330            let (k, v) = pair?;
331            message.insert(k, lua_to_json(&v));
332        }
333    }
334
335    Ok(Context {
336        request,
337        response,
338        message,
339        errors,
340    })
341}
342
343/// Converts a JSON value to the corresponding Lua value (arrays become
344/// 1-indexed tables, objects become string-keyed tables).
345fn json_to_lua(lua: &Lua, value: &serde_json::Value) -> LuaResult<LuaValue> {
346    match value {
347        serde_json::Value::Null => Ok(LuaValue::Nil),
348        serde_json::Value::Bool(b) => Ok(LuaValue::Boolean(*b)),
349        serde_json::Value::Number(n) => {
350            if let Some(i) = n.as_i64() {
351                Ok(LuaValue::Integer(i as _))
352            } else {
353                Ok(LuaValue::Number(n.as_f64().unwrap_or(0.0)))
354            }
355        }
356        serde_json::Value::String(s) => Ok(LuaValue::String(lua.create_string(s)?)),
357        serde_json::Value::Array(arr) => {
358            let table = lua.create_table()?;
359            for (i, v) in arr.iter().enumerate() {
360                table.set(i + 1, json_to_lua(lua, v)?)?;
361            }
362            Ok(LuaValue::Table(table))
363        }
364        serde_json::Value::Object(obj) => {
365            let table = lua.create_table()?;
366            for (k, v) in obj {
367                table.set(k.as_str(), json_to_lua(lua, v)?)?;
368            }
369            Ok(LuaValue::Table(table))
370        }
371    }
372}
373
374/// Converts a Lua value back to JSON. Tables with a non-zero sequence length
375/// become JSON arrays, other tables become objects with string keys;
376/// unconvertible values (functions, userdata, non-UTF-8 strings) degrade to
377/// null or empty strings rather than erroring.
378fn lua_to_json(value: &LuaValue) -> serde_json::Value {
379    match value {
380        LuaValue::Nil => serde_json::Value::Null,
381        LuaValue::Boolean(b) => serde_json::Value::Bool(*b),
382        LuaValue::Integer(i) => serde_json::json!(*i),
383        LuaValue::Number(n) => serde_json::json!(*n),
384        LuaValue::String(s) => {
385            serde_json::Value::String(std::str::from_utf8(&s.as_bytes()).unwrap_or("").to_string())
386        }
387        LuaValue::Table(t) => {
388            let len = t.raw_len();
389            if len > 0 {
390                let arr: Vec<serde_json::Value> = (1..=len)
391                    .filter_map(|i| t.get::<LuaValue>(i).ok().map(|v| lua_to_json(&v)))
392                    .collect();
393                serde_json::Value::Array(arr)
394            } else {
395                let mut map = serde_json::Map::new();
396                if let Ok(pairs) = t
397                    .clone()
398                    .pairs::<String, LuaValue>()
399                    .collect::<Result<Vec<_>, _>>()
400                {
401                    for (k, v) in pairs {
402                        map.insert(k, lua_to_json(&v));
403                    }
404                }
405                serde_json::Value::Object(map)
406            }
407        }
408        _ => serde_json::Value::Null,
409    }
410}
411
412#[cfg(test)]
413mod tests {
414    use super::*;
415    use crate::context::Protocol;
416
417    fn test_context() -> Context {
418        Context {
419            request: GatewayRequest {
420                method: "GET".to_string(),
421                path: "/test".to_string(),
422                host: "localhost".to_string(),
423                scheme: "http".to_string(),
424                headers: HashMap::new(),
425                query_params: HashMap::new(),
426                body: bytes::Bytes::new(),
427                remote_addr: "127.0.0.1:1234".to_string(),
428                protocol: Protocol::Http1,
429            },
430            response: GatewayResponse {
431                status_code: 0,
432                headers: HashMap::new(),
433                body: bytes::Bytes::new(),
434            },
435            message: HashMap::new(),
436            errors: Vec::new(),
437        }
438    }
439
440    #[test]
441    fn test_lua_modify_path() {
442        let rt = LuaRuntime::new(
443            r#"
444            function execute(ctx)
445                ctx.request.path = "/modified"
446                return ctx
447            end
448            "#,
449            5000,
450            None,
451        )
452        .unwrap();
453
454        let ctx = test_context();
455        let result = rt.execute(ctx).unwrap();
456        assert_eq!(result.request.path, "/modified");
457    }
458
459    #[test]
460    fn test_lua_set_message() {
461        let rt = LuaRuntime::new(
462            r#"
463            function execute(ctx)
464                ctx.message.enriched = true
465                ctx.message.user_id = "abc123"
466                return ctx
467            end
468            "#,
469            5000,
470            None,
471        )
472        .unwrap();
473
474        let ctx = test_context();
475        let result = rt.execute(ctx).unwrap();
476        assert_eq!(
477            result.message.get("enriched"),
478            Some(&serde_json::json!(true))
479        );
480        assert_eq!(
481            result.message.get("user_id"),
482            Some(&serde_json::json!("abc123"))
483        );
484    }
485
486    #[test]
487    fn test_lua_add_header() {
488        let rt = LuaRuntime::new(
489            r#"
490            function execute(ctx)
491                ctx.request.headers["x-custom"] = {"hello"}
492                return ctx
493            end
494            "#,
495            5000,
496            None,
497        )
498        .unwrap();
499
500        let ctx = test_context();
501        let result = rt.execute(ctx).unwrap();
502        assert_eq!(
503            result.request.headers.get("x-custom"),
504            Some(&vec!["hello".to_string()])
505        );
506    }
507
508    #[test]
509    fn test_lua_preserves_errors_and_protocol() {
510        let rt = LuaRuntime::new(
511            r#"
512            function execute(ctx)
513                return ctx
514            end
515            "#,
516            5000,
517            None,
518        )
519        .unwrap();
520
521        let mut ctx = test_context();
522        ctx.request.protocol = Protocol::Http2;
523        ctx.errors.push(GatewayError {
524            node_id: "upstream".to_string(),
525            code: "UPSTREAM_CONNECTION_ERROR".to_string(),
526            message: "boom".to_string(),
527            metadata: HashMap::new(),
528        });
529
530        let result = rt.execute(ctx).unwrap();
531        assert_eq!(result.request.protocol, Protocol::Http2);
532        assert_eq!(result.errors.len(), 1);
533        assert_eq!(result.errors[0].code, "UPSTREAM_CONNECTION_ERROR");
534    }
535
536    #[test]
537    fn test_lua_error_handling() {
538        let rt = LuaRuntime::new(
539            r#"
540            function execute(ctx)
541                error("something went wrong")
542                return ctx
543            end
544            "#,
545            5000,
546            None,
547        )
548        .unwrap();
549
550        let ctx = test_context();
551        let result = rt.execute(ctx);
552        assert!(result.is_err());
553        let err = result.unwrap_err();
554        assert_eq!(err.error.code, "LUA_EXECUTION_ERROR");
555    }
556
557    #[test]
558    fn test_lua_require_module() {
559        // Create a temp directory with a module
560        let tmp = std::env::temp_dir().join("gw_lua_test");
561        let _ = std::fs::create_dir_all(&tmp);
562        std::fs::write(
563            tmp.join("helpers.lua"),
564            r#"
565            local M = {}
566            function M.greet(name)
567                return "hello " .. name
568            end
569            return M
570            "#,
571        )
572        .unwrap();
573
574        let rt = LuaRuntime::new(
575            r#"
576            local helpers = require("helpers")
577
578            function execute(ctx)
579                ctx.message.greeting = helpers.greet("world")
580                return ctx
581            end
582            "#,
583            5000,
584            Some(tmp.clone()),
585        )
586        .unwrap();
587
588        let ctx = test_context();
589        let result = rt.execute(ctx).unwrap();
590        assert_eq!(
591            result.message.get("greeting"),
592            Some(&serde_json::json!("hello world"))
593        );
594
595        let _ = std::fs::remove_dir_all(&tmp);
596    }
597}
featherbit/plugins/script/lua_runtime.rs

featherbit/plugins/script/
lua_runtime.rs
