qail_core/

sanitize.rs

//! AST structural sanitization for untrusted input.
//!
//! The text parser enforces identifier constraints (alphanumeric + `_` + `.`),
//! but externally provided ASTs (binary endpoints, APIs, generated payloads, etc.)
//! may bypass parser-level identifier checks and still reach execution.
//!
//! Call [`validate_ast`](crate::sanitize::validate_ast) on any `Qail` obtained from an untrusted source
//! (binary endpoint, external API, etc.) before execution.

use crate::ast::{
    Action, ConflictAction, Expr, MergeAction, MergeSource, Qail, TableConstraint, Value,
};
use std::fmt;

/// Error returned when AST structural validation fails.
#[derive(Debug, Clone)]
pub struct SanitizeError {
    pub field: String,
    pub value: String,
    pub reason: String,
}

impl fmt::Display for SanitizeError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "AST validation failed: {} '{}' — {}",
            self.field, self.value, self.reason
        )
    }
}

impl std::error::Error for SanitizeError {}

/// Maximum identifier length (PostgreSQL NAMEDATALEN - 1).
const MAX_IDENT_LEN: usize = 63;

/// Validate that an identifier matches the parser grammar: `[a-zA-Z0-9_.]`.
///
/// Also rejects empty identifiers and those exceeding PostgreSQL's 63-char limit.
fn is_safe_identifier(s: &str) -> bool {
    !s.is_empty()
        && s.len() <= MAX_IDENT_LEN
        && s.bytes()
            .all(|b| b.is_ascii_alphanumeric() || b == b'_' || b == b'.')
}

/// Validate an identifier, returning a `SanitizeError` if invalid.
fn check_ident(field: &str, value: &str) -> Result<(), SanitizeError> {
    if is_safe_identifier(value) {
        Ok(())
    } else {
        Err(SanitizeError {
            field: field.to_string(),
            value: value.chars().take(40).collect(),
            reason: "identifiers must match [a-zA-Z0-9_.] and be ≤63 chars".to_string(),
        })
    }
}

/// Validate an `Expr` node for unsafe patterns.
///
/// - `Expr::Named` must be a safe identifier.
/// - Recursive variants (Cast, Binary, etc.) are validated recursively.
fn check_expr(field: &str, expr: &Expr) -> Result<(), SanitizeError> {
    match expr {
        Expr::Star => Ok(()),
        Expr::Named(name) => check_ident(field, name),
        Expr::Aliased { name, alias } => {
            check_ident(field, name)?;
            check_ident(&format!("{field}.alias"), alias)
        }
        Expr::Aggregate {
            col, alias, filter, ..
        } => {
            if col != "*" {
                check_ident(field, col)?;
            }
            if let Some(a) = alias {
                check_ident(&format!("{field}.alias"), a)?;
            }
            if let Some(conditions) = filter {
                for cond in conditions {
                    check_expr(&format!("{field}.filter"), &cond.left)?;
                    check_value(&format!("{field}.filter"), &cond.value)?;
                }
            }
            Ok(())
        }
        Expr::FunctionCall { name, args, alias } => {
            check_ident(field, name)?;
            if let Some(a) = alias {
                check_ident(&format!("{field}.alias"), a)?;
            }
            for arg in args {
                check_expr(&format!("{field}.arg"), arg)?;
            }
            Ok(())
        }
        Expr::Cast {
            expr,
            target_type,
            alias,
        } => {
            check_expr(field, expr)?;
            check_ident(&format!("{field}.cast_type"), target_type)?;
            if let Some(a) = alias {
                check_ident(&format!("{field}.alias"), a)?;
            }
            Ok(())
        }
        Expr::Binary {
            left, right, alias, ..
        } => {
            check_expr(field, left)?;
            check_expr(field, right)?;
            if let Some(a) = alias {
                check_ident(&format!("{field}.alias"), a)?;
            }
            Ok(())
        }
        Expr::Literal(_) => Ok(()),
        Expr::JsonAccess {
            column,
            alias,
            path_segments,
            ..
        } => {
            check_ident(field, column)?;
            for (key, _) in path_segments {
                // Integer indices are fine; string keys must be safe identifiers
                if key.parse::<i64>().is_err() && !is_safe_identifier(key) {
                    return Err(SanitizeError {
                        field: format!("{field}.json_path"),
                        value: key.chars().take(40).collect(),
                        reason: "JSON path key must be a safe identifier or integer".to_string(),
                    });
                }
            }
            if let Some(a) = alias {
                check_ident(&format!("{field}.alias"), a)?;
            }
            Ok(())
        }
        Expr::Subquery { query, alias } => {
            validate_ast(query)?;
            if let Some(a) = alias {
                check_ident(&format!("{field}.alias"), a)?;
            }
            Ok(())
        }
        Expr::Exists { query, alias, .. } => {
            validate_ast(query)?;
            if let Some(a) = alias {
                check_ident(&format!("{field}.alias"), a)?;
            }
            Ok(())
        }
        // For all other complex Expr variants, validate aliases where present
        Expr::Window {
            name,
            func,
            partition,
            params,
            order,
            ..
        } => {
            if !name.is_empty() {
                check_ident(&format!("{field}.window_alias"), name)?;
            }
            check_ident(&format!("{field}.window_func"), func)?;
            for p in partition {
                check_ident(&format!("{field}.partition"), p)?;
            }
            for p in params {
                check_expr(&format!("{field}.window_param"), p)?;
            }
            for cage in order {
                for cond in &cage.conditions {
                    check_expr(&format!("{field}.window_order"), &cond.left)?;
                    check_value(&format!("{field}.window_order"), &cond.value)?;
                }
            }
            Ok(())
        }
        Expr::Case {
            when_clauses,
            else_value,
            alias,
        } => {
            for (cond, val) in when_clauses {
                check_expr(&format!("{field}.case_when"), &cond.left)?;
                check_value(&format!("{field}.case_when"), &cond.value)?;
                check_expr(&format!("{field}.case_then"), val)?;
            }
            if let Some(e) = else_value {
                check_expr(&format!("{field}.case_else"), e)?;
            }
            if let Some(a) = alias {
                check_ident(&format!("{field}.alias"), a)?;
            }
            Ok(())
        }
        Expr::SpecialFunction { args, alias, name } => {
            check_ident(&format!("{field}.special_func"), name)?;
            for (_, arg) in args {
                check_expr(&format!("{field}.special_func_arg"), arg)?;
            }
            if let Some(a) = alias {
                check_ident(&format!("{field}.alias"), a)?;
            }
            Ok(())
        }
        Expr::ArrayConstructor { elements, alias } | Expr::RowConstructor { elements, alias } => {
            for elem in elements {
                check_expr(&format!("{field}.element"), elem)?;
            }
            if let Some(a) = alias {
                check_ident(&format!("{field}.alias"), a)?;
            }
            Ok(())
        }
        Expr::Subscript { expr, index, alias } => {
            check_expr(&format!("{field}.subscript_expr"), expr)?;
            check_expr(&format!("{field}.subscript_index"), index)?;
            if let Some(a) = alias {
                check_ident(&format!("{field}.alias"), a)?;
            }
            Ok(())
        }
        Expr::Collate {
            expr,
            collation,
            alias,
        } => {
            check_expr(&format!("{field}.collate_expr"), expr)?;
            check_ident(&format!("{field}.collation"), collation)?;
            if let Some(a) = alias {
                check_ident(&format!("{field}.alias"), a)?;
            }
            Ok(())
        }
        Expr::FieldAccess {
            expr,
            field: f,
            alias,
        } => {
            check_expr(&format!("{field}.field_access_expr"), expr)?;
            check_ident(&format!("{field}.field"), f)?;
            if let Some(a) = alias {
                check_ident(&format!("{field}.alias"), a)?;
            }
            Ok(())
        }
        Expr::Def { name, .. } => check_ident(field, name),
        Expr::Mod { col, .. } => check_expr(field, col),
    }
}

/// Check a `Value` for embedded subqueries.
fn check_value(field: &str, value: &Value) -> Result<(), SanitizeError> {
    match value {
        Value::Subquery(q) => validate_ast(q),
        Value::Array(vals) => {
            for v in vals {
                check_value(field, v)?;
            }
            Ok(())
        }
        Value::Expr(expr) => check_expr(field, expr),
        _ => Ok(()),
    }
}

/// Validate a `Qail` AST from an untrusted source.
///
/// Checks all identifier fields against the parser grammar (`[a-zA-Z0-9_.]`)
/// and rejects dangerous procedural actions.
///
/// # Errors
///
/// Returns `SanitizeError` on the first violation found.
pub fn validate_ast(cmd: &Qail) -> Result<(), SanitizeError> {
    // ── Block dangerous actions from binary path ─────────────────────
    match cmd.action {
        Action::Call | Action::Do | Action::SessionSet | Action::SessionReset => {
            return Err(SanitizeError {
                field: "action".to_string(),
                value: format!("{:?}", cmd.action),
                reason: "procedural/session actions are not allowed via binary AST".to_string(),
            });
        }
        _ => {}
    }

    // ── Table name ───────────────────────────────────────────────────
    if !cmd.table.is_empty() {
        check_ident("table", &cmd.table)?;
    }

    // ── Columns ──────────────────────────────────────────────────────
    for (i, col) in cmd.columns.iter().enumerate() {
        check_expr(&format!("columns[{i}]"), col)?;
    }

    // ── Table Constraints ────────────────────────────────────────────
    for (i, constraint) in cmd.table_constraints.iter().enumerate() {
        match constraint {
            TableConstraint::Unique(cols) | TableConstraint::PrimaryKey(cols) => {
                for col in cols {
                    check_ident(&format!("table_constraints[{i}].column"), col)?;
                }
            }
            TableConstraint::ForeignKey {
                name,
                columns,
                ref_table,
                ref_columns,
            } => {
                if let Some(name) = name {
                    check_ident(&format!("table_constraints[{i}].name"), name)?;
                }
                for col in columns {
                    check_ident(&format!("table_constraints[{i}].column"), col)?;
                }
                check_ident(&format!("table_constraints[{i}].ref_table"), ref_table)?;
                for col in ref_columns {
                    check_ident(&format!("table_constraints[{i}].ref_column"), col)?;
                }
            }
        }
    }

    // ── Joins ────────────────────────────────────────────────────────
    for (i, join) in cmd.joins.iter().enumerate() {
        // Join table may include alias: "users u"
        // Validate each space-separated token
        for token in join.table.split_whitespace() {
            check_ident(&format!("joins[{i}].table"), token)?;
        }
        if let Some(ref conditions) = join.on {
            for cond in conditions {
                check_expr(&format!("joins[{i}].on"), &cond.left)?;
                check_value(&format!("joins[{i}].on"), &cond.value)?;
            }
        }
    }

    // ── Cages (filters, sorts, etc.) ─────────────────────────────────
    for cage in &cmd.cages {
        for cond in &cage.conditions {
            check_expr("cage.condition.left", &cond.left)?;
            check_value("cage.condition.value", &cond.value)?;
        }
    }

    // ── CTEs ─────────────────────────────────────────────────────────
    for cte in &cmd.ctes {
        check_ident("cte.name", &cte.name)?;
        for col in &cte.columns {
            check_ident("cte.column", col)?;
        }
        validate_ast(&cte.base_query)?;
        if let Some(ref rq) = cte.recursive_query {
            validate_ast(rq)?;
        }
    }

    // ── DISTINCT ON ──────────────────────────────────────────────────
    for expr in &cmd.distinct_on {
        check_expr("distinct_on", expr)?;
    }

    // ── RETURNING ────────────────────────────────────────────────────
    if let Some(ref cols) = cmd.returning {
        for col in cols {
            check_expr("returning", col)?;
        }
    }

    // ── ON CONFLICT ──────────────────────────────────────────────────
    if let Some(ref oc) = cmd.on_conflict {
        for col in &oc.columns {
            check_ident("on_conflict.column", col)?;
        }
        if let ConflictAction::DoUpdate { assignments } = &oc.action {
            for (col, expr) in assignments {
                check_ident("on_conflict.assignment.column", col)?;
                check_expr("on_conflict.assignment.expr", expr)?;
            }
        }
    }

    // ── MERGE ────────────────────────────────────────────────────────
    if let Some(ref merge) = cmd.merge {
        if let Some(alias) = &merge.target_alias {
            check_ident("merge.target_alias", alias)?;
        }
        match &merge.source {
            MergeSource::Table { name, alias } => {
                check_ident("merge.source.table", name)?;
                if let Some(alias) = alias {
                    check_ident("merge.source.alias", alias)?;
                }
            }
            MergeSource::Query { query, alias } => {
                validate_ast(query)?;
                if let Some(alias) = alias {
                    check_ident("merge.source.alias", alias)?;
                }
            }
        }
        for cond in &merge.on {
            check_expr("merge.on.left", &cond.left)?;
            check_value("merge.on.value", &cond.value)?;
        }
        for clause in &merge.clauses {
            for cond in &clause.condition {
                check_expr("merge.clause.condition.left", &cond.left)?;
                check_value("merge.clause.condition.value", &cond.value)?;
            }
            match &clause.action {
                MergeAction::Update { assignments } => {
                    for (col, expr) in assignments {
                        check_ident("merge.update.column", col)?;
                        check_expr("merge.update.expr", expr)?;
                    }
                }
                MergeAction::Insert { columns, values } => {
                    for col in columns {
                        check_ident("merge.insert.column", col)?;
                    }
                    for expr in values {
                        check_expr("merge.insert.expr", expr)?;
                    }
                }
                MergeAction::Delete | MergeAction::DoNothing => {}
            }
        }
    }

    // ── FROM / USING tables ──────────────────────────────────────────
    for t in &cmd.from_tables {
        check_ident("from_tables", t)?;
    }
    for t in &cmd.using_tables {
        check_ident("using_tables", t)?;
    }

    // ── SET ops ──────────────────────────────────────────────────────
    for (_, sub) in &cmd.set_ops {
        validate_ast(sub)?;
    }

    // ── Source query (INSERT … SELECT) ───────────────────────────────
    if let Some(ref sq) = cmd.source_query {
        validate_ast(sq)?;
    }

    // ── HAVING ───────────────────────────────────────────────────────
    for cond in &cmd.having {
        check_expr("having", &cond.left)?;
        check_value("having", &cond.value)?;
    }

    // ── Channel (LISTEN/NOTIFY) ──────────────────────────────────────
    if let Some(ref ch) = cmd.channel {
        check_ident("channel", ch)?;
    }

    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::ast::Qail;

    #[test]
    fn valid_simple_query_passes() {
        let cmd = Qail::get("users").columns(["id", "name"]);
        assert!(validate_ast(&cmd).is_ok());
    }

    #[test]
    fn sql_injection_in_table_rejected() {
        let cmd = Qail::get("users; DROP TABLE users; --");
        let err = validate_ast(&cmd).unwrap_err();
        assert_eq!(err.field, "table");
    }

    #[test]
    fn call_action_rejected() {
        let cmd = Qail {
            action: Action::Call,
            table: "my_proc()".to_string(),
            ..Default::default()
        };
        let err = validate_ast(&cmd).unwrap_err();
        assert_eq!(err.field, "action");
    }

    #[test]
    fn do_action_rejected() {
        let cmd = Qail {
            action: Action::Do,
            table: "plpgsql".to_string(),
            ..Default::default()
        };
        let err = validate_ast(&cmd).unwrap_err();
        assert_eq!(err.field, "action");
    }

    #[test]
    fn valid_qualified_name_passes() {
        let cmd = Qail::get("public.users").columns(["users.id", "users.name"]);
        assert!(validate_ast(&cmd).is_ok());
    }

    #[test]
    fn injection_in_join_table_rejected() {
        use crate::ast::JoinKind;
        let cmd = Qail::get("users").join(
            JoinKind::Left,
            "orders; DROP TABLE x",
            "users.id",
            "orders.user_id",
        );
        let err = validate_ast(&cmd).unwrap_err();
        assert!(err.field.contains("joins"));
    }

    #[test]
    fn injection_in_column_rejected() {
        let cmd = Qail::get("users").columns(["id", "name; DROP TABLE x"]);
        let err = validate_ast(&cmd).unwrap_err();
        assert!(err.field.contains("columns"));
    }

    #[test]
    fn on_conflict_update_assignment_expression_injection_rejected() {
        let cmd = Qail::add("users")
            .set_value("id", 1)
            .set_value("name", "Alice")
            .on_conflict_update(
                &["id"],
                &[(
                    "name",
                    Expr::Named("EXCLUDED.name, is_admin = true".to_string()),
                )],
            );

        let err = validate_ast(&cmd).unwrap_err();

        assert_eq!(err.field, "on_conflict.assignment.expr");
        assert!(
            err.value.contains("EXCLUDED.name"),
            "unexpected rejected value: {}",
            err.value
        );
    }

    #[test]
    fn on_conflict_update_assignment_column_injection_rejected() {
        let cmd = Qail::add("users")
            .set_value("id", 1)
            .set_value("name", "Alice")
            .on_conflict_update(
                &["id"],
                &[("name, is_admin", Expr::Named("EXCLUDED.name".to_string()))],
            );

        let err = validate_ast(&cmd).unwrap_err();

        assert_eq!(err.field, "on_conflict.assignment.column");
    }

    #[test]
    fn aggregate_filter_value_expression_injection_rejected() {
        use crate::ast::{AggregateFunc, Condition, Operator, Value};

        let mut cmd = Qail::get("events");
        cmd.columns.push(Expr::Aggregate {
            col: "id".to_string(),
            func: AggregateFunc::Count,
            distinct: false,
            filter: Some(vec![Condition {
                left: Expr::Named("direction".to_string()),
                op: Operator::Eq,
                value: Value::Expr(Box::new(Expr::Named("bad;DROP".to_string()))),
                is_array_unnest: false,
            }]),
            alias: None,
        });

        let err = validate_ast(&cmd).unwrap_err();
        assert_eq!(err.field, "columns[0].filter");
    }

    #[test]
    fn count_star_aggregate_passes_sanitizer() {
        use crate::ast::AggregateFunc;

        let mut cmd = Qail::get("events");
        cmd.columns.push(Expr::Aggregate {
            col: "*".to_string(),
            func: AggregateFunc::Count,
            distinct: false,
            filter: None,
            alias: Some("total".to_string()),
        });

        assert!(validate_ast(&cmd).is_ok());
    }

    #[test]
    fn case_when_complex_condition_expression_passes_sanitizer() {
        use crate::ast::{Condition, Operator, Value};

        let mut cmd = Qail::get("users");
        cmd.columns.push(Expr::Case {
            when_clauses: vec![(
                Condition {
                    left: Expr::Cast {
                        expr: Box::new(Expr::JsonAccess {
                            column: "profile".to_string(),
                            path_segments: vec![("active".to_string(), true)],
                            alias: None,
                        }),
                        target_type: "integer".to_string(),
                        alias: None,
                    },
                    op: Operator::Gt,
                    value: Value::Int(0),
                    is_array_unnest: false,
                },
                Box::new(Expr::Literal(Value::String("active".to_string()))),
            )],
            else_value: Some(Box::new(Expr::Literal(Value::String(
                "inactive".to_string(),
            )))),
            alias: Some("status_label".to_string()),
        });

        assert!(validate_ast(&cmd).is_ok());
    }

    #[test]
    fn empty_table_name_passes() {
        // Some actions like TxnStart have empty table
        let cmd = Qail {
            action: Action::TxnStart,
            table: String::new(),
            ..Default::default()
        };
        assert!(validate_ast(&cmd).is_ok());
    }

    #[test]
    fn oversized_identifier_rejected() {
        let long_name = "a".repeat(64);
        let cmd = Qail::get(&long_name);
        let err = validate_ast(&cmd).unwrap_err();
        assert!(err.reason.contains("63"));
    }
}