faucet_core/

util.rs

//! Shared utilities used across faucet source and sink crates.

use std::collections::HashMap;

use crate::FaucetError;
use jsonpath_rust::JsonPath;
use serde_json::Value;

// ── SQL Utilities ───────────────────────────────────────────────────────────

/// Quote a SQL identifier to prevent SQL injection.
///
/// Wraps the name in double quotes and doubles any embedded double-quotes
/// per the SQL standard (ANSI SQL).
///
/// ```
/// use faucet_core::util::quote_ident;
/// assert_eq!(quote_ident("my_table"), "\"my_table\"");
/// assert_eq!(quote_ident("has\"quote"), "\"has\"\"quote\"");
/// ```
pub fn quote_ident(name: &str) -> String {
    format!("\"{}\"", name.replace('"', "\"\""))
}

// ── JSONPath Extraction ─────────────────────────────────────────────────────

/// Extract records from a JSON value using an optional JSONPath expression.
///
/// - If `path` is `Some`, queries the body with the JSONPath and returns
///   all matched values.
/// - If `path` is `None`, returns the body as-is: arrays are unpacked into
///   individual records, objects/scalars are returned as a single-element vec.
pub fn extract_records(body: &Value, path: Option<&str>) -> Result<Vec<Value>, FaucetError> {
    match path {
        Some(p) => {
            let results = body
                .query(p)
                .map_err(|e| FaucetError::JsonPath(format!("invalid JSONPath '{p}': {e}")))?;
            Ok(results.into_iter().cloned().collect())
        }
        None => match body {
            Value::Array(arr) => Ok(arr.clone()),
            other => Ok(vec![other.clone()]),
        },
    }
}

// ── HTTP Response Handling ──────────────────────────────────────────────────

/// Check an HTTP response status and return a [`FaucetError::HttpStatus`] on
/// non-success responses.
///
/// Reads the response body for error context, truncating to `max_body_len`
/// bytes (default: 2048) to avoid large error messages.
pub async fn check_http_response(
    resp: reqwest::Response,
    max_body_len: usize,
) -> Result<reqwest::Response, FaucetError> {
    if resp.status().is_success() {
        return Ok(resp);
    }

    let status = resp.status().as_u16();
    let url = resp.url().to_string();
    let body_text = resp.text().await.unwrap_or_default();

    let body = if body_text.len() > max_body_len {
        let end = body_text.floor_char_boundary(max_body_len);
        format!("{}...(truncated)", &body_text[..end])
    } else {
        body_text
    };

    Err(FaucetError::HttpStatus { status, url, body })
}

/// Default maximum body length for error responses.
pub const DEFAULT_ERROR_BODY_MAX_LEN: usize = 2048;

// ── Context Utilities ──────────────────────────────────────────────────────

/// Substitute `{key}` placeholders in a template string with values from context.
///
/// Value conversion rules:
/// - `String` -> raw string (no quotes)
/// - `Number` -> number as string
/// - `Bool` -> `"true"` / `"false"`
/// - `Null` -> `"null"`
/// - `Array` / `Object` -> JSON-serialized string
///
/// Unmatched placeholders are left as-is.
///
/// **Warning:** Do NOT use this for SQL queries (SQL injection risk) or for
/// substitution into serialized JSON (corruption risk with special characters).
/// Use [`substitute_context_bind_params`] for SQL and [`substitute_context_json`]
/// for serialized JSON.
pub fn substitute_context(template: &str, context: &HashMap<String, Value>) -> String {
    substitute_single_pass(template, context, |value| match value {
        Value::String(s) => s.clone(),
        Value::Number(n) => n.to_string(),
        Value::Bool(b) => b.to_string(),
        Value::Null => "null".to_string(),
        other => other.to_string(),
    })
}

/// Single left-to-right scan that replaces each recognised `{key}` placeholder
/// with `render(value)`. Unmatched placeholders are left verbatim; replacement
/// text is never re-scanned. Shared by [`substitute_context`] and
/// [`substitute_context_json`] so neither is O(template × context) (#78/#36).
fn substitute_single_pass(
    template: &str,
    context: &HashMap<String, Value>,
    render: impl Fn(&Value) -> String,
) -> String {
    if context.is_empty() {
        return template.to_string();
    }
    let mut result = String::with_capacity(template.len());
    let mut last_copied = 0;
    let mut search_from = 0;

    while search_from < template.len() {
        let Some(open_offset) = template[search_from..].find('{') else {
            break;
        };
        let open = search_from + open_offset;
        let Some(close_offset) = template[open + 1..].find('}') else {
            break;
        };
        let close = open + 1 + close_offset;
        let key = &template[open + 1..close];

        if let Some(value) = context.get(key) {
            result.push_str(&template[last_copied..open]);
            result.push_str(&render(value));
            last_copied = close + 1;
            search_from = close + 1;
        } else {
            search_from = open + 1;
        }
    }

    result.push_str(&template[last_copied..]);
    result
}

/// Replace `{key}` placeholders with SQL bind-parameter markers, returning
/// the rewritten query and an ordered list of values to bind.
///
/// Scans the template left-to-right; each recognised placeholder is replaced
/// with the marker produced by `marker_fn(index)`, and the corresponding
/// value is appended to the returned vector.  The same key appearing multiple
/// times produces one bind value per occurrence.
///
/// `start_index` is the 1-based index for the first parameter.
///
/// # Marker functions
///
/// - PostgreSQL: `|i| format!("${i}")`
/// - MySQL / SQLite: `|_| "?".to_string()`
///
/// Placeholders whose key is not present in `context` are left unchanged.
pub fn substitute_context_bind_params(
    template: &str,
    context: &HashMap<String, Value>,
    start_index: usize,
    marker_fn: impl Fn(usize) -> String,
) -> (String, Vec<Value>) {
    if context.is_empty() {
        return (template.to_string(), Vec::new());
    }

    let mut result = String::with_capacity(template.len());
    let mut values = Vec::new();
    let mut param_idx = start_index;
    let mut last_copied = 0;
    let mut search_from = 0;

    while search_from < template.len() {
        let Some(open_offset) = template[search_from..].find('{') else {
            break;
        };
        let open = search_from + open_offset;

        let Some(close_offset) = template[open + 1..].find('}') else {
            break;
        };
        let close = open + 1 + close_offset;
        let key = &template[open + 1..close];

        if let Some(value) = context.get(key) {
            result.push_str(&template[last_copied..open]);
            result.push_str(&marker_fn(param_idx));
            values.push(value.clone());
            param_idx += 1;
            last_copied = close + 1;
            search_from = close + 1;
        } else {
            search_from = open + 1;
        }
    }

    result.push_str(&template[last_copied..]);
    (result, values)
}

/// Substitute `{key}` placeholders within a serialized JSON string, escaping
/// string values so that the result remains valid JSON.
///
/// Use this instead of [`substitute_context`] when the template is a
/// `serde_json`-serialized value that will be deserialized back after
/// substitution.  String values are JSON-escaped (double-quotes, backslashes,
/// and control characters).  Numbers, bools, and null are substituted as-is.
pub fn substitute_context_json(template: &str, context: &HashMap<String, Value>) -> String {
    substitute_single_pass(template, context, |value| match value {
        Value::String(s) => json_escape_string(s),
        Value::Number(n) => n.to_string(),
        Value::Bool(b) => b.to_string(),
        Value::Null => "null".to_string(),
        other => other.to_string(),
    })
}

/// Escape a string for safe embedding inside a JSON string value.
///
/// Handles double-quotes, backslashes, and control characters per RFC 8259.
fn json_escape_string(s: &str) -> String {
    let mut escaped = String::with_capacity(s.len());
    for c in s.chars() {
        match c {
            '"' => escaped.push_str("\\\""),
            '\\' => escaped.push_str("\\\\"),
            '\n' => escaped.push_str("\\n"),
            '\r' => escaped.push_str("\\r"),
            '\t' => escaped.push_str("\\t"),
            c if c.is_control() => {
                escaped.push_str(&format!("\\u{:04x}", c as u32));
            }
            c => escaped.push(c),
        }
    }
    escaped
}

/// Strip credentials from a connection string so it can be used as a lineage
/// dataset URI without leaking secrets. Handles two shapes, best-effort:
///
/// - **URL userinfo** — `scheme://user:pass@host/...` → `scheme://host/...`
///   (the `user[:pass]@` between `://` and the authority terminator is removed).
/// - **Key/value (ADO.NET) connection strings** — any `Password=...` / `Pwd=...`
///   segment (case-insensitive key) has its value replaced with `***`.
///
/// Input with neither shape is returned unchanged.
pub fn redact_uri_credentials(uri: &str) -> String {
    let mut out = uri.to_string();
    // 1) URL userinfo: between "://" and the first '/' or '?' (the authority).
    if let Some(scheme_end) = out.find("://") {
        let after = scheme_end + 3;
        let auth_end = out[after..]
            .find(['/', '?'])
            .map(|i| after + i)
            .unwrap_or(out.len());
        if let Some(at_rel) = out[after..auth_end].find('@') {
            // Remove "user:pass@" inclusive of the '@'.
            out.replace_range(after..after + at_rel + 1, "");
        }
    }
    // 2) ADO.NET-style Password=/Pwd= tokens.
    if out.contains('=') {
        out = out
            .split(';')
            .map(|seg| match seg.find('=') {
                Some(eq)
                    if {
                        let k = seg[..eq].trim();
                        k.eq_ignore_ascii_case("password") || k.eq_ignore_ascii_case("pwd")
                    } =>
                {
                    format!("{}=***", &seg[..eq])
                }
                _ => seg.to_string(),
            })
            .collect::<Vec<_>>()
            .join(";");
    }
    out
}

/// Extract context values from a record using JSONPath expressions.
///
/// Each entry in `mapping` is `context_key -> json_path`. The function queries
/// the record with each JSONPath and stores the first matched value under the
/// corresponding context key.
///
/// Returns an error if any JSONPath matches nothing.
pub fn extract_context(
    record: &Value,
    mapping: &HashMap<String, String>,
) -> Result<HashMap<String, Value>, FaucetError> {
    let mut context = HashMap::with_capacity(mapping.len());
    for (context_key, json_path) in mapping {
        let results = record
            .query(json_path.as_str())
            .map_err(|e| FaucetError::JsonPath(format!("invalid JSONPath '{json_path}': {e}")))?;
        let value = results.first().ok_or_else(|| {
            FaucetError::JsonPath(format!(
                "JSONPath '{json_path}' matched nothing in record for context key '{context_key}'"
            ))
        })?;
        context.insert(context_key.clone(), (*value).clone());
    }
    Ok(context)
}

#[cfg(test)]
mod tests {
    use super::*;
    use serde_json::json;

    // ── quote_ident ─────────────────────────────────────────────────────

    #[test]
    fn quote_ident_simple() {
        assert_eq!(quote_ident("my_table"), "\"my_table\"");
    }

    #[test]
    fn quote_ident_with_embedded_quotes() {
        assert_eq!(quote_ident("has\"quote"), "\"has\"\"quote\"");
    }

    #[test]
    fn quote_ident_empty() {
        assert_eq!(quote_ident(""), "\"\"");
    }

    #[test]
    fn quote_ident_special_chars() {
        assert_eq!(quote_ident("table; DROP"), "\"table; DROP\"");
    }

    // ── extract_records ─────────────────────────────────────────────────

    #[test]
    fn extract_with_path() {
        let body = json!({"data": [{"id": 1}, {"id": 2}]});
        let records = extract_records(&body, Some("$.data[*]")).unwrap();
        assert_eq!(records.len(), 2);
        assert_eq!(records[0]["id"], 1);
    }

    #[test]
    fn extract_without_path_array() {
        let body = json!([{"id": 1}, {"id": 2}]);
        let records = extract_records(&body, None).unwrap();
        assert_eq!(records.len(), 2);
    }

    #[test]
    fn extract_without_path_object() {
        let body = json!({"id": 1});
        let records = extract_records(&body, None).unwrap();
        assert_eq!(records.len(), 1);
    }

    #[test]
    fn extract_empty_result() {
        let body = json!({"data": []});
        let records = extract_records(&body, Some("$.data[*]")).unwrap();
        assert!(records.is_empty());
    }

    #[test]
    fn extract_invalid_path_returns_error() {
        let body = json!({"data": 1});
        // jsonpath-rust handles most paths gracefully; test error propagation.
        let result = extract_records(&body, Some("$.data[*]"));
        // This should succeed (empty match) or fail; either is fine as long as
        // it doesn't panic.
        let _ = result;
    }

    // ── substitute_context ──────────────────────────────────────────────

    #[test]
    fn substitute_context_string_values() {
        let mut ctx = HashMap::new();
        ctx.insert("org".to_string(), json!("acme"));
        ctx.insert("repo".to_string(), json!("widgets"));
        let result = substitute_context("/orgs/{org}/repos/{repo}", &ctx);
        assert_eq!(result, "/orgs/acme/repos/widgets");
    }

    #[test]
    fn substitute_context_number_value() {
        let mut ctx = HashMap::new();
        ctx.insert("id".to_string(), json!(42));
        let result = substitute_context("/items/{id}", &ctx);
        assert_eq!(result, "/items/42");
    }

    #[test]
    fn substitute_context_bool_value() {
        let mut ctx = HashMap::new();
        ctx.insert("active".to_string(), json!(true));
        let result = substitute_context("/filter?active={active}", &ctx);
        assert_eq!(result, "/filter?active=true");
    }

    #[test]
    fn substitute_context_null_value() {
        let mut ctx = HashMap::new();
        ctx.insert("val".to_string(), json!(null));
        let result = substitute_context("/x/{val}", &ctx);
        assert_eq!(result, "/x/null");
    }

    #[test]
    fn substitute_context_array_value() {
        let mut ctx = HashMap::new();
        ctx.insert("ids".to_string(), json!([1, 2, 3]));
        let result = substitute_context("/x/{ids}", &ctx);
        assert_eq!(result, "/x/[1,2,3]");
    }

    #[test]
    fn substitute_context_unmatched_placeholder_left_as_is() {
        let ctx = HashMap::new();
        let result = substitute_context("/orgs/{org}/repos", &ctx);
        assert_eq!(result, "/orgs/{org}/repos");
    }

    #[test]
    fn substitute_context_empty_template() {
        let ctx = HashMap::new();
        let result = substitute_context("", &ctx);
        assert_eq!(result, "");
    }

    #[test]
    fn substitute_context_replaces_all_occurrences() {
        let mut ctx = HashMap::new();
        ctx.insert("id".to_string(), Value::String("42".to_string()));
        let result = substitute_context("/a/{id}/b/{id}", &ctx);
        assert_eq!(result, "/a/42/b/42");
    }

    #[test]
    fn substitute_context_does_not_rescan_replacement() {
        // Single-pass: a replacement value that itself looks like a placeholder
        // is emitted verbatim, never re-substituted (#78/#36).
        let mut ctx = HashMap::new();
        ctx.insert("a".to_string(), Value::String("{b}".to_string()));
        ctx.insert("b".to_string(), Value::String("SECRET".to_string()));
        let result = substitute_context("{a}", &ctx);
        assert_eq!(result, "{b}");
    }

    // ── extract_context ─────────────────────────────────────────────────

    #[test]
    fn extract_context_simple_paths() {
        let record = json!({"id": 1, "name": "alice"});
        let mut mapping = HashMap::new();
        mapping.insert("user_id".to_string(), "$.id".to_string());
        mapping.insert("user_name".to_string(), "$.name".to_string());
        let ctx = extract_context(&record, &mapping).unwrap();
        assert_eq!(ctx["user_id"], json!(1));
        assert_eq!(ctx["user_name"], json!("alice"));
    }

    #[test]
    fn extract_context_nested_path() {
        let record = json!({"data": {"info": {"id": 99}}});
        let mut mapping = HashMap::new();
        mapping.insert("deep_id".to_string(), "$.data.info.id".to_string());
        let ctx = extract_context(&record, &mapping).unwrap();
        assert_eq!(ctx["deep_id"], json!(99));
    }

    #[test]
    fn extract_context_missing_path_returns_error() {
        let record = json!({"id": 1});
        let mut mapping = HashMap::new();
        mapping.insert("missing".to_string(), "$.nonexistent".to_string());
        let result = extract_context(&record, &mapping);
        assert!(result.is_err());
    }

    #[test]
    fn extract_context_empty_mapping() {
        let record = json!({"id": 1});
        let mapping = HashMap::new();
        let ctx = extract_context(&record, &mapping).unwrap();
        assert!(ctx.is_empty());
    }

    // ── substitute_context_bind_params ──────────────────────────────────

    #[test]
    fn bind_params_postgres_style() {
        let mut ctx = HashMap::new();
        ctx.insert("org".to_string(), json!("acme"));
        ctx.insert("id".to_string(), json!(42));
        let (query, values) = substitute_context_bind_params(
            "SELECT * FROM t WHERE org = {org} AND id = {id}",
            &ctx,
            1,
            |i| format!("${i}"),
        );
        assert_eq!(query, "SELECT * FROM t WHERE org = $1 AND id = $2");
        assert_eq!(values.len(), 2);
        assert_eq!(values[0], json!("acme"));
        assert_eq!(values[1], json!(42));
    }

    #[test]
    fn bind_params_question_mark_style() {
        let mut ctx = HashMap::new();
        ctx.insert("name".to_string(), json!("test"));
        let (query, values) =
            substitute_context_bind_params("SELECT * FROM t WHERE name = {name}", &ctx, 1, |_| {
                "?".to_string()
            });
        assert_eq!(query, "SELECT * FROM t WHERE name = ?");
        assert_eq!(values, vec![json!("test")]);
    }

    #[test]
    fn bind_params_duplicate_key_produces_multiple_binds() {
        let mut ctx = HashMap::new();
        ctx.insert("id".to_string(), json!(5));
        let (query, values) = substitute_context_bind_params(
            "SELECT * FROM t WHERE a = {id} OR b = {id}",
            &ctx,
            3,
            |i| format!("${i}"),
        );
        assert_eq!(query, "SELECT * FROM t WHERE a = $3 OR b = $4");
        assert_eq!(values, vec![json!(5), json!(5)]);
    }

    #[test]
    fn bind_params_unknown_key_left_as_is() {
        let ctx = HashMap::new();
        let (query, values) =
            substitute_context_bind_params("SELECT * FROM t WHERE x = {unknown}", &ctx, 1, |i| {
                format!("${i}")
            });
        assert_eq!(query, "SELECT * FROM t WHERE x = {unknown}");
        assert!(values.is_empty());
    }

    #[test]
    fn bind_params_mixed_known_and_unknown() {
        let mut ctx = HashMap::new();
        ctx.insert("id".to_string(), json!(1));
        let (query, values) = substitute_context_bind_params(
            "SELECT * FROM t WHERE id = {id} AND x = {unknown}",
            &ctx,
            1,
            |i| format!("${i}"),
        );
        assert_eq!(query, "SELECT * FROM t WHERE id = $1 AND x = {unknown}");
        assert_eq!(values, vec![json!(1)]);
    }

    #[test]
    fn bind_params_empty_context() {
        let ctx = HashMap::new();
        let (query, values) =
            substitute_context_bind_params("SELECT 1", &ctx, 1, |i| format!("${i}"));
        assert_eq!(query, "SELECT 1");
        assert!(values.is_empty());
    }

    #[test]
    fn bind_params_start_index_offset() {
        let mut ctx = HashMap::new();
        ctx.insert("name".to_string(), json!("x"));
        let (query, values) =
            substitute_context_bind_params("SELECT * FROM t WHERE name = {name}", &ctx, 5, |i| {
                format!("${i}")
            });
        assert_eq!(query, "SELECT * FROM t WHERE name = $5");
        assert_eq!(values, vec![json!("x")]);
    }

    // ── substitute_context_json ─────────────────────────────────────────

    #[test]
    fn json_sub_escapes_double_quotes() {
        let mut ctx = HashMap::new();
        ctx.insert("name".to_string(), json!(r#"O'Brien "Bob""#));
        let template = r#"{"name":"{name}"}"#;
        let result = substitute_context_json(template, &ctx);
        let parsed: Value = serde_json::from_str(&result).unwrap();
        assert_eq!(parsed["name"], r#"O'Brien "Bob""#);
    }

    #[test]
    fn json_sub_escapes_backslashes() {
        let mut ctx = HashMap::new();
        ctx.insert("path".to_string(), json!("C:\\Users\\test"));
        let template = r#"{"path":"{path}"}"#;
        let result = substitute_context_json(template, &ctx);
        let parsed: Value = serde_json::from_str(&result).unwrap();
        assert_eq!(parsed["path"], "C:\\Users\\test");
    }

    #[test]
    fn json_sub_escapes_control_chars() {
        let mut ctx = HashMap::new();
        ctx.insert("text".to_string(), json!("line1\nline2\ttab"));
        let template = r#"{"text":"{text}"}"#;
        let result = substitute_context_json(template, &ctx);
        let parsed: Value = serde_json::from_str(&result).unwrap();
        assert_eq!(parsed["text"], "line1\nline2\ttab");
    }

    #[test]
    fn json_sub_number_value() {
        let mut ctx = HashMap::new();
        ctx.insert("id".to_string(), json!(42));
        let template = r#"{"user_id":"{id}"}"#;
        let result = substitute_context_json(template, &ctx);
        let parsed: Value = serde_json::from_str(&result).unwrap();
        assert_eq!(parsed["user_id"], "42");
    }

    #[test]
    fn json_sub_preserves_valid_json_without_special_chars() {
        let mut ctx = HashMap::new();
        ctx.insert("name".to_string(), json!("alice"));
        let template = r#"{"filter":{"name":"{name}"}}"#;
        let result = substitute_context_json(template, &ctx);
        let parsed: Value = serde_json::from_str(&result).unwrap();
        assert_eq!(parsed["filter"]["name"], "alice");
    }

    // ── json_escape_string ──────────────────────────────────────────────

    #[test]
    fn json_escape_plain_string() {
        assert_eq!(json_escape_string("hello"), "hello");
    }

    #[test]
    fn json_escape_quotes_and_backslashes() {
        assert_eq!(json_escape_string(r#"a"b\c"#), r#"a\"b\\c"#);
    }

    #[test]
    fn json_escape_newlines_and_tabs() {
        assert_eq!(json_escape_string("a\nb\tc"), "a\\nb\\tc");
    }

    // ── redact_uri_credentials ──────────────────────────────────────────

    #[test]
    fn redact_strips_url_userinfo() {
        assert_eq!(
            redact_uri_credentials("postgres://user:pass@host:5432/db"),
            "postgres://host:5432/db"
        );
        assert_eq!(
            redact_uri_credentials("mongodb://u:p@h/db?x=1"),
            "mongodb://h/db?x=1"
        );
    }

    #[test]
    fn redact_strips_user_only_userinfo() {
        assert_eq!(
            redact_uri_credentials("redis://user@127.0.0.1:6379"),
            "redis://127.0.0.1:6379"
        );
    }

    #[test]
    fn redact_handles_adonet_password_tokens() {
        assert_eq!(
            redact_uri_credentials("Server=tcp:h,1433;Database=db;User Id=sa;Password=secret;"),
            "Server=tcp:h,1433;Database=db;User Id=sa;Password=***;"
        );
        assert_eq!(
            redact_uri_credentials("server=h;pwd=secret"),
            "server=h;pwd=***"
        );
    }

    #[test]
    fn redact_passthrough_when_no_credentials() {
        assert_eq!(
            redact_uri_credentials("s3://bucket/prefix"),
            "s3://bucket/prefix"
        );
        assert_eq!(
            redact_uri_credentials("file:///tmp/x.csv"),
            "file:///tmp/x.csv"
        );
    }
}