featherbit/plugins/native/

exit_transformer.rs

//! The `exit-transformer` node — reshapes gateway-generated responses
//! ("exits") with a status-code remap and a body template. Reinterpreted
//! subset of APISIX's `exit-transformer` plugin (response-phase: place after
//! `upstream`, before `client`).
//!
//! ## Which exits it applies to
//!
//! By default the node acts only on responses whose context carries an
//! **error record** (`Context.errors` non-empty) — a failed upstream, a failed
//! auth callout, an `error-handler`-rendered body. A deliberate rejection that
//! left its node on an *outcome* port (`denied`, `limited`, `broken`, `abort`,
//! `redirect`, ...) carries **no** error record and is therefore **not**
//! transformed unless `always: true` is set.
//!
//! So to reshape denials and throttles, set `always: true` and put the node on
//! the path the outcome port takes:
//!
//! ```yaml
//! nodes:
//!   - id: shape-exits
//!     type: exit-transformer
//!     config:
//!       always: true                 # required: a denial carries no error record
//!       status_map: { "401": 403 }
//!       body: '{"status": $status, "path": "$uri"}'
//! edges:
//!   - from: auth.denied
//!     to: shape-exits.in             # not straight to client
//!   - from: shape-exits.success
//!     to: client.in
//! ```
//!
//! With `always: true` the node transforms *every* response reaching it,
//! clean upstream replies included — so give it its own branch as above
//! rather than putting it on the main success path.
//!
//! **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 `{{namespace.path}}` template (plus legacy `$var`
//!   interpolation, 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::template::Template;

/// Remaps `Context.response.status_code` via `status_map` and/or replaces
/// the body with an interpolated `{{namespace.path}}` (plus legacy `$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 template: supports `{{namespace.path}}` references and legacy
    /// `$var` interpolation (see [`Template::render_with_legacy`]).
    body: Option<Template>,
    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. Supports
    ///   `{{namespace.path}}` references plus legacy `$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) => {
                let s = v.as_str().ok_or("body must be a string".to_string())?;
                // Discard warnings here — the compile-time walk (a later
                // task) reports well-formed-but-unknown references;
                // execution must not.
                Some(Template::parse(s).0)
            }
        };

        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) -> 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 = template.render_with_legacy(&ctx);
                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::success(ctx))
    }
}

#[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"),
                stream: None,
            },
            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)).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)).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)).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)).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}"
            );
        }
    }
}