Call framework host functions from RustScript examples

Author: fffonion

README.md

@@ -6,7 +6,10 @@ Standalone Pingora integration demo for `pd-vm` / RustScript.
 
 A Pingora gateway can keep the framework and proxy code compiled while moving request policy into RustScript:
 
-- input from `pingora::http::RequestHeader`: method, path, `x-user-tier`
+- input from real `pingora::http::RequestHeader`: method, path, `x-user-tier`
+- RustScript calls host functions backed by Pingora request data:
+  - `pingora_header(name) -> string`
+  - `pingora_method_is(expected) -> bool`
 - scripted output: `allow`, `route:<upstream>`, or `deny:<status>:<reason>`
 - optional conversion of deny decisions back into `pingora::http::ResponseHeader`
 

scripts/gateway_policy.rss

@@ -1,7 +1,10 @@
+fn pingora_header(name) -> string;
+fn pingora_method_is(expected) -> bool;
 fn gateway_decision(decision) -> string;
 
+let tier = pingora_header("x-user-tier");
 let blocked_admin = path == "/admin" && tier != "pro";
-let canary = path == "/canary" && method == "GET";
+let canary = path == "/canary" && pingora_method_is("GET");
 let decision = if blocked_admin => {
     "deny:403:upgrade required"
 } else if canary => {

src/lib.rs

@@ -16,19 +16,14 @@ pub struct ScriptedGatewayPolicy {
 impl ScriptedGatewayPolicy {
     pub fn from_source(source: impl Into<String>) -> Result<Self, String> {
         let source = source.into();
-        run_string(&wrap_request_source(&source, "GET", "/", ""))?;
+        let request = RequestHeader::build("GET", b"/", None)
+            .map_err(|err| format!("failed to build Pingora request: {err}"))?;
+        run_string(&wrap_request_source(&source, &request), &request)?;
         Ok(Self { source })
     }
 
     pub fn evaluate_request(&self, request: &RequestHeader) -> Result<GatewayDecision, String> {
-        let method = request.method.as_str();
-        let path = String::from_utf8_lossy(request.raw_path());
-        let tier = request
-            .headers
-            .get("x-user-tier")
-            .and_then(|value| value.to_str().ok())
-            .unwrap_or("");
-        let output = run_string(&wrap_request_source(&self.source, method, &path, tier))?;
+        let output = run_string(&wrap_request_source(&self.source, request), request)?;
         parse_decision(&output)
     }
 
@@ -69,23 +64,57 @@ fn parse_decision(output: &str) -> Result<GatewayDecision, String> {
     Err(format!("unknown gateway decision '{output}'"))
 }
 
-fn wrap_request_source(policy: &str, method: &str, path: &str, tier: &str) -> String {
+fn wrap_request_source(policy: &str, request: &RequestHeader) -> String {
+    let method = request.method.as_str();
+    let path = String::from_utf8_lossy(request.raw_path());
     format!(
-        "let method = {};\nlet path = {};\nlet tier = {};\n{}",
+        "let method = {};\nlet path = {};\n{}",
         rss_string(method),
-        rss_string(path),
-        rss_string(tier),
+        rss_string(&path),
         policy
     )
 }
 
-fn run_string(source: &str) -> Result<String, String> {
-    match run_value(source)? {
+fn run_string(source: &str, request: &RequestHeader) -> Result<String, String> {
+    match run_value(source, request)? {
         Value::String(value) => Ok(value.as_str().to_string()),
         other => Err(format!("script returned {other:?}; expected string")),
     }
 }
 
+#[derive(Debug, Clone)]
+struct PingoraRequestSnapshot {
+    method: String,
+    headers: Vec<(String, String)>,
+}
+
+impl PingoraRequestSnapshot {
+    fn from_request(request: &RequestHeader) -> Self {
+        let headers = request
+            .headers
+            .iter()
+            .filter_map(|(name, value)| {
+                Some((
+                    name.as_str().to_ascii_lowercase(),
+                    value.to_str().ok()?.to_string(),
+                ))
+            })
+            .collect();
+        Self {
+            method: request.method.as_str().to_string(),
+            headers,
+        }
+    }
+
+    fn header(&self, name: &str) -> String {
+        let needle = name.to_ascii_lowercase();
+        self.headers
+            .iter()
+            .find_map(|(header_name, value)| (header_name == &needle).then(|| value.clone()))
+            .unwrap_or_default()
+    }
+}
+
 struct GatewayDecisionHost;
 
 impl HostFunction for GatewayDecisionHost {
@@ -99,10 +128,51 @@ impl HostFunction for GatewayDecisionHost {
     }
 }
 
-fn run_value(source: &str) -> Result<Value, String> {
+struct PingoraHeaderHost {
+    request: PingoraRequestSnapshot,
+}
+
+impl HostFunction for PingoraHeaderHost {
+    fn call(&mut self, _vm: &mut Vm, args: &[Value]) -> Result<CallOutcome, VmError> {
+        match args {
+            [Value::String(name)] => Ok(CallOutcome::Return(CallReturn::one(Value::string(
+                self.request.header(name.as_str()),
+            )))),
+            _ => Err(VmError::TypeMismatch("header name string")),
+        }
+    }
+}
+
+struct PingoraMethodIsHost {
+    request: PingoraRequestSnapshot,
+}
+
+impl HostFunction for PingoraMethodIsHost {
+    fn call(&mut self, _vm: &mut Vm, args: &[Value]) -> Result<CallOutcome, VmError> {
+        match args {
+            [Value::String(expected)] => Ok(CallOutcome::Return(CallReturn::one(Value::Bool(
+                self.request.method == expected.as_str(),
+            )))),
+            _ => Err(VmError::TypeMismatch("method string")),
+        }
+    }
+}
+
+fn run_value(source: &str, request: &RequestHeader) -> Result<Value, String> {
     let compiled = compile_source(source).map_err(|err| err.to_string())?;
+    let snapshot = PingoraRequestSnapshot::from_request(request);
     let mut vm = Vm::new(compiled.program);
     vm.bind_function("gateway_decision", Box::new(GatewayDecisionHost));
+    vm.bind_function(
+        "pingora_header",
+        Box::new(PingoraHeaderHost {
+            request: snapshot.clone(),
+        }),
+    );
+    vm.bind_function(
+        "pingora_method_is",
+        Box::new(PingoraMethodIsHost { request: snapshot }),
+    );
     let status = vm.run().map_err(|err| err.to_string())?;
     if status != VmStatus::Halted {
         return Err(format!("script did not halt: {status:?}"));

tests/gateway_policy.rs

@@ -3,7 +3,7 @@ use pretty_assertions::assert_eq;
 use rustscript_pingora_gateway_policy::{GatewayDecision, ScriptedGatewayPolicy};
 
 #[test]
-fn rustscript_denies_pingora_admin_request_for_free_tier() {
+fn rustscript_calls_pingora_request_header_functions_for_free_tier() {
     let policy = ScriptedGatewayPolicy::from_source(include_str!("../scripts/gateway_policy.rss"))
         .expect("policy should compile");
     let mut request = RequestHeader::build("GET", b"/admin", None).expect("request should build");
@@ -25,7 +25,7 @@ fn rustscript_denies_pingora_admin_request_for_free_tier() {
 }
 
 #[test]
-fn rustscript_routes_pingora_canary_request_without_framework_fork() {
+fn rustscript_calls_pingora_method_function_for_canary_route() {
     let policy = ScriptedGatewayPolicy::from_source(include_str!("../scripts/gateway_policy.rss"))
         .expect("policy should compile");
     let request = RequestHeader::build("GET", b"/canary", None).expect("request should build");
@@ -39,3 +39,43 @@ fn rustscript_routes_pingora_canary_request_without_framework_fork() {
         GatewayDecision::Route("canary-upstream".to_string())
     );
 }
+
+#[test]
+fn rustscript_can_call_pingora_host_functions_from_inline_policy() {
+    let policy = ScriptedGatewayPolicy::from_source(
+        r#"
+fn pingora_header(name) -> string;
+fn pingora_method_is(expected) -> bool;
+fn gateway_decision(decision) -> string;
+
+let tier = pingora_header("x-user-tier");
+let blocked_admin = path == "/admin" && tier != "pro";
+let canary = path == "/canary" && pingora_method_is("GET");
+let decision = if blocked_admin => {
+    "deny:403:upgrade required"
+} else if canary => {
+    "route:canary-upstream"
+} else => {
+    "allow"
+};
+gateway_decision(decision);
+"#,
+    )
+    .expect("policy should compile");
+    let mut request = RequestHeader::build("GET", b"/admin", None).expect("request should build");
+    request
+        .insert_header("x-user-tier", "free")
+        .expect("header should insert");
+
+    let decision = policy
+        .evaluate_request(&request)
+        .expect("policy should evaluate");
+
+    assert_eq!(
+        decision,
+        GatewayDecision::Deny {
+            status: 403,
+            reason: "upgrade required".to_string(),
+        }
+    );
+}