featherbit/plugins/native/

exit_transformer.rs

//! The `exit-transformer` node — reshapes gateway-generated responses
//! ("exits": auth rejections, rate-limit denials, upstream failures, ...)
//! with a status-code remap and a body template. Reinterpreted subset of
//! APISIX's `exit-transformer` plugin (response-phase: place after
//! `upstream`, before `client`).
//!
//! **This is a deliberate reinterpretation, not a faithful port.** APISIX's
//! plugin registers user-supplied *Lua functions* that receive
//! `(code, body, headers)` and return replacements. Arbitrary scripted
//! transformation belongs to featherbit's `script` node; this node keeps the
//! declarative core of the idea:
//! - `status_map` remaps exit status codes (e.g. hide a 502 as a 503), and
//! - `body` renders a `$var` template (including `$status`) as the new body,
//!
//! applied only to responses the gateway itself generated (heuristic:
//! `Context.errors` non-empty), or unconditionally with `always: true`.

use async_trait::async_trait;
use bytes::Bytes;
use std::collections::HashMap;

use crate::context::Context;
use crate::plugins::{Plugin, PluginOutput, PluginResult};
use crate::vars;

/// Remaps `Context.response.status_code` via `status_map` and/or replaces
/// the body with an interpolated `$var` template — but only when the
/// response was generated by the gateway (`Context.errors` non-empty),
/// unless `always` is set. This plugin never fails at execution time.
pub struct ExitTransformerPlugin {
    status_map: HashMap<u16, u16>,
    body: Option<String>,
    always: bool,
}

/// Validates one status code (map key or value): an integer within 100-599.
fn parse_status(what: &str, raw: &str, value: Option<u64>) -> Result<u16, String> {
    let n = value
        .or_else(|| raw.parse().ok())
        .ok_or_else(|| format!("status_map {} '{}' must be an integer", what, raw))?;
    if !(100..=599).contains(&n) {
        return Err(format!(
            "status_map {} '{}' must be within 100-599",
            what, raw
        ));
    }
    Ok(n as u16)
}

impl ExitTransformerPlugin {
    /// Builds the plugin from node config.
    ///
    /// Accepted keys (all optional):
    /// - `status_map` (map of status → status, e.g. `"502": 503`): when the
    ///   response status matches a key, it is replaced by the value. Keys
    ///   are strings (YAML/JSON object keys), values integers; both must be
    ///   within 100-599.
    /// - `body` (string): replacement body template with `$var` / `${var}`
    ///   interpolation — `$status` (the status *after* remapping), `$uri`,
    ///   `$host`, `$http_<name>`, `$msg_<key>`, and every other variable
    ///   from [`crate::vars`].
    /// - `always` (bool, default `false`): apply to every response. When
    ///   false, only responses generated by the gateway itself — heuristic:
    ///   `Context.errors` is non-empty — are transformed; clean upstream
    ///   responses pass through untouched.
    ///
    /// ```yaml
    /// type: exit-transformer
    /// config:
    ///   status_map:
    ///     "401": 403
    ///     "502": 503
    ///   body: '{"status": $status, "path": "$uri"}'
    /// ```
    pub fn from_config(config: &HashMap<String, serde_json::Value>) -> Result<Self, String> {
        let mut status_map = HashMap::new();
        if let Some(raw) = config.get("status_map") {
            let obj = raw
                .as_object()
                .ok_or("status_map must be a map of status -> status".to_string())?;
            for (key, value) in obj {
                let from = parse_status("key", key, None)?;
                let to = parse_status("value", &value.to_string(), value.as_u64())?;
                status_map.insert(from, to);
            }
        }

        let body = match config.get("body") {
            None => None,
            Some(v) => Some(
                v.as_str()
                    .ok_or("body must be a string".to_string())?
                    .to_string(),
            ),
        };

        let always = config
            .get("always")
            .and_then(|v| v.as_bool())
            .unwrap_or(false);

        Ok(Self {
            status_map,
            body,
            always,
        })
    }
}

#[async_trait]
impl Plugin for ExitTransformerPlugin {
    fn plugin_type(&self) -> &str {
        "exit-transformer"
    }

    async fn execute(
        &self,
        mut ctx: Context,
        _named_inputs: &HashMap<String, serde_json::Value>,
    ) -> PluginResult {
        // Gate: gateway-generated exits only, unless `always`.
        let applies = self.always || !ctx.errors.is_empty();

        if applies {
            // Remap first so a $status in the body template reflects the
            // final, client-visible status code.
            if let Some(&new_status) = self.status_map.get(&ctx.response.status_code) {
                ctx.response.status_code = new_status;
            }

            if let Some(template) = &self.body {
                let body = vars::interpolate(&ctx, template);
                ctx.response.body = Bytes::from(body);
                // Body-mutation convention: recompute length; the rendered
                // template is not encoded.
                ctx.response.headers.remove("content-length");
                ctx.response.headers.remove("content-encoding");
            }
        }

        Ok(PluginOutput {
            context: ctx,
            named_outputs: HashMap::new(),
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::context::{GatewayError, GatewayRequest, GatewayResponse, Protocol};

    fn test_context(status: u16, gateway_generated: bool) -> Context {
        let mut response_headers = HashMap::new();
        response_headers.insert("content-length".to_string(), vec!["8".to_string()]);
        let errors = if gateway_generated {
            vec![GatewayError {
                node_id: "auth-1".to_string(),
                code: "UNAUTHORIZED".to_string(),
                message: "missing credentials".to_string(),
                metadata: HashMap::new(),
            }]
        } else {
            Vec::new()
        };
        Context {
            request: GatewayRequest {
                method: "GET".to_string(),
                path: "/api/users".to_string(),
                host: "localhost".to_string(),
                scheme: "http".to_string(),
                headers: HashMap::new(),
                query_params: HashMap::new(),
                body: Bytes::new(),
                remote_addr: "127.0.0.1:12345".to_string(),
                protocol: Protocol::Http1,
            },
            response: GatewayResponse {
                status_code: status,
                headers: response_headers,
                body: Bytes::from_static(b"original"),
            },
            message: HashMap::new(),
            errors,
        }
    }

    fn plugin(config: serde_json::Value) -> ExitTransformerPlugin {
        let map: HashMap<String, serde_json::Value> =
            serde_json::from_value(config).expect("test config must be an object");
        ExitTransformerPlugin::from_config(&map).expect("config should be valid")
    }

    #[tokio::test]
    async fn test_exit_transformer_remaps_status_and_body() {
        let p = plugin(serde_json::json!({
            "status_map": { "502": 503 },
            "body": "{\"status\": $status, \"path\": \"$uri\"}"
        }));
        let out = p
            .execute(test_context(502, true), &HashMap::new())
            .await
            .unwrap();
        assert_eq!(out.context.response.status_code, 503);
        // $status reflects the remapped status.
        assert_eq!(
            out.context.response.body.as_ref(),
            b"{\"status\": 503, \"path\": \"/api/users\"}"
        );
        assert!(!out.context.response.headers.contains_key("content-length"));
    }

    #[tokio::test]
    async fn test_exit_transformer_skips_upstream_responses() {
        let p = plugin(serde_json::json!({
            "status_map": { "502": 503 },
            "body": "transformed"
        }));
        // 502 from the upstream itself: no gateway errors → passthrough.
        let out = p
            .execute(test_context(502, false), &HashMap::new())
            .await
            .unwrap();
        assert_eq!(out.context.response.status_code, 502);
        assert_eq!(out.context.response.body.as_ref(), b"original");
        assert!(out.context.response.headers.contains_key("content-length"));
    }

    #[tokio::test]
    async fn test_exit_transformer_always_applies_unconditionally() {
        let p = plugin(serde_json::json!({
            "status_map": { "502": 503 },
            "always": true
        }));
        let out = p
            .execute(test_context(502, false), &HashMap::new())
            .await
            .unwrap();
        assert_eq!(out.context.response.status_code, 503);
        // No body template configured → body untouched.
        assert_eq!(out.context.response.body.as_ref(), b"original");
        assert!(out.context.response.headers.contains_key("content-length"));
    }

    #[tokio::test]
    async fn test_exit_transformer_unmapped_status_kept() {
        let p = plugin(serde_json::json!({
            "status_map": { "502": 503 }
        }));
        let out = p
            .execute(test_context(401, true), &HashMap::new())
            .await
            .unwrap();
        assert_eq!(out.context.response.status_code, 401);
    }

    #[test]
    fn test_exit_transformer_config_validation() {
        for bad in [
            serde_json::json!({ "status_map": ["not", "a", "map"] }),
            serde_json::json!({ "status_map": { "abc": 503 } }),
            serde_json::json!({ "status_map": { "99": 503 } }),
            serde_json::json!({ "status_map": { "502": 600 } }),
            serde_json::json!({ "status_map": { "502": "not a number" } }),
            serde_json::json!({ "body": 42 }),
        ] {
            let map: HashMap<String, serde_json::Value> =
                serde_json::from_value(bad.clone()).unwrap();
            assert!(
                ExitTransformerPlugin::from_config(&map).is_err(),
                "should reject: {bad}"
            );
        }
    }
}