windmill/backend/windmill-common/src/schema.rs

use anyhow::anyhow;
use itertools::Itertools;
use quick_cache::sync::Cache;
use regex::Regex;
use serde::{Deserialize, Serialize};
use std::{collections::HashMap, str::FromStr, sync::Arc};

use serde_json::{value::RawValue, Value};

use crate::{error::Error, scripts::ScriptLang};

#[derive(PartialEq, Serialize, Deserialize, Debug, Clone)]
pub enum JsonPrimitiveType {
    String,
    Number,
    Integer,
    Object,
    Array,
    Boolean,
    Null,
}

#[derive(PartialEq, Serialize, Deserialize, Debug, Clone)]
pub enum SchemaValidationRule {
    StrictEnum(Vec<Value>),
    IsNull,
    IsInteger,
    IsString,
    IsBool,
    IsDatetime,
    IsNumber,
    IsEmail,
    IsObject(Vec<(String, Vec<SchemaValidationRule>)>),
    IsArray(Vec<SchemaValidationRule>),
    IsUnionType(Vec<Vec<SchemaValidationRule>>),
    IsOneOf(HashMap<String, Vec<SchemaValidationRule>>),
    IsBytes,
}

impl SchemaValidationRule {
    fn from_primitive(p: &JsonPrimitiveType, val: &Value) -> Result<Vec<Self>, anyhow::Error> {
        let mut schema_rules = vec![];

        match p {
            JsonPrimitiveType::String => {
                schema_rules.push(SchemaValidationRule::IsString);

                if let Some(format) = val.get("format").and_then(|f| f.as_str()) {
                    if format == "date" || format == "date-time" {
                        schema_rules.push(SchemaValidationRule::IsDatetime);
                    }

                    if format == "email" {
                        schema_rules.push(SchemaValidationRule::IsEmail);
                    }
                }

                if let Some(encoding) = val.get("contentEncoding").and_then(|e| e.as_str()) {
                    if encoding == "base64" {
                        schema_rules.push(SchemaValidationRule::IsBytes);
                    }
                }
            }

            JsonPrimitiveType::Number => {
                schema_rules.push(SchemaValidationRule::IsNumber);
            }
            JsonPrimitiveType::Integer => {
                schema_rules.push(SchemaValidationRule::IsInteger);
            }
            JsonPrimitiveType::Object => {
                let mut obj_rules = vec![];

                if let Some(properties) = val.get("properties") {
                    let properties = properties
                        .as_object()
                        .ok_or(anyhow!("Field properties should be an object"))?;

                    for (key, v) in properties {
                        obj_rules.push((key.clone(), SchemaValidationRule::from_value(v)?))
                    }

                    schema_rules.push(SchemaValidationRule::IsObject(obj_rules));
                } else if let Some(one_of) = val.get("oneOf") {
                    let one_of = one_of
                        .as_array()
                        .ok_or(anyhow!("`oneOf` needs to be an array"))?;
                    let mut rules_map: HashMap<String, Vec<SchemaValidationRule>> = HashMap::new();

                    for variant in one_of {
                        let variant_label = variant
                            .get("title")
                            .ok_or(anyhow!(
                                "oneOf variant definition should have a `title` field"
                            ))?
                            .as_str()
                            .ok_or(anyhow!(
                                "oneOf variant definition `title` field should be a string"
                            ))?;
                        if !rules_map.contains_key(variant_label) {
                            rules_map.insert(
                                variant_label.to_string(),
                                SchemaValidationRule::from_value(variant)?,
                            );
                        } else {
                            return Err(anyhow!(
                                "oneOf definition has a duplicate variant `{variant_label}`"
                            ));
                        }
                    }

                    schema_rules.push(SchemaValidationRule::IsOneOf(rules_map))
                } else {
                    let is_resource = val
                        .get("format")
                        .and_then(|f| f.as_str())
                        .map(|f| f.starts_with("resource"))
                        .unwrap_or(false);
                    if !is_resource {
                        return Err(anyhow!(
                        "Object type should have a `properties` or `anyOf` field, or be a resource"
                        ));
                    }
                }
            }
            JsonPrimitiveType::Array => {
                let items = val
                    .get("items")
                    .ok_or(anyhow!("Array type should have field `items`"))?;

                let arr_rules = SchemaValidationRule::from_value(items)?;

                schema_rules.push(SchemaValidationRule::IsArray(arr_rules));
            }
            JsonPrimitiveType::Boolean => {
                schema_rules.push(SchemaValidationRule::IsBool);
            }
            JsonPrimitiveType::Null => {
                schema_rules.push(SchemaValidationRule::IsNull);
            }
        }

        Ok(schema_rules)
    }

    fn from_value(val: &Value) -> Result<Vec<Self>, Error> {
        if let Some(any_of) = val.get("anyOf").and_then(|any_of| any_of.as_array()) {
            let mut r = vec![];

            for variant in any_of {
                r.push(SchemaValidationRule::from_value(variant)?);
            }
            return Ok(vec![SchemaValidationRule::IsUnionType(r)]);
        }

        let mut schema_rules = vec![];

        let typ = val.get("type").ok_or(anyhow!("Missing `type` field"))?;

        if let Some(typ) = typ.as_str() {
            schema_rules.append(&mut SchemaValidationRule::from_primitive(
                &JsonPrimitiveType::from_str(typ)?,
                val,
            )?);
        } else if let Some(typ_arr) = typ.as_array() {
            let typ_arr = typ_arr
                .into_iter()
                .map(|v| {
                    SchemaValidationRule::from_primitive(
                        &JsonPrimitiveType::from_str(
                            v.as_str()
                                .ok_or(anyhow!("Expected array of strings for `type` field"))?,
                        )?,
                        v,
                    )
                })
                .collect::<Result<Vec<Vec<SchemaValidationRule>>, anyhow::Error>>()?;

            schema_rules.push(SchemaValidationRule::IsUnionType(typ_arr));
        } else {
            return Err(anyhow!(
                "Unsupported value for type field, expected string or string array"
            )
            .into());
        }

        if let Some(enum_variants) = val.get("enum") {
            let variants = enum_variants
                .as_array()
                .ok_or(anyhow!("enum variants are not in an array"))?
                .clone();
            schema_rules.push(SchemaValidationRule::StrictEnum(variants));
        }

        Ok(schema_rules)
    }

    fn apply_rule(&self, key: &str, val: &Value, required: bool) -> Result<(), Error> {
        if val.is_null() {
            if !required {
                return Ok(());
            }
            return Err(Error::ArgumentErr(format!("Argument {key} cannot be null")));
        }
        match self {
            SchemaValidationRule::IsNull => {
                if !val.is_null() {
                    return Err(Error::ArgumentErr(format!(
                        "Argument `{key}` should be null"
                    )));
                }
            }
            SchemaValidationRule::StrictEnum(vec) => {
                if !vec.contains(val) {
                    let options = vec.iter().map(|s| s.to_string()).join(", ");
                    return Err(Error::ArgumentErr(format!(
                        "Enum type argument `{key}` expected one of `[{options}]` but received {}",
                        val.to_string()
                    )));
                }
            }
            SchemaValidationRule::IsNumber => {
                if !val.is_number() {
                    return Err(Error::ArgumentErr(format!(
                        "Argument `{key}` should be a numeric value"
                    )));
                }
            }
            SchemaValidationRule::IsInteger => {
                if !val.is_i64() && !val.is_u64() {
                    return Err(Error::ArgumentErr(format!(
                        "Argument `{key}` should be an integer"
                    )));
                }
            }
            SchemaValidationRule::IsString => {
                if !val.is_string() {
                    return Err(Error::ArgumentErr(format!(
                        "Argument `{key}` should be a string"
                    )));
                }
            }
            SchemaValidationRule::IsBool => {
                if !val.is_boolean() {
                    return Err(Error::ArgumentErr(format!(
                        "Argument `{key}` should be a boolean"
                    )));
                }
            }
            SchemaValidationRule::IsObject(o) => {
                if !val.is_object() {
                    return Err(Error::ArgumentErr(format!(
                        "Argument `{key}` should be an object"
                    )));
                }

                for (s, rules) in o {
                    let v = val
                        .get(&s)
                        .ok_or(Error::ArgumentErr(format!("Missing field {s} in {key}")))?;
                    for r in rules {
                        r.apply_rule(&format!("{key}.{s}"), v, true)?;
                    }
                }
            }
            SchemaValidationRule::IsArray(vec) => {
                if let Some(arr) = val.as_array() {
                    for (i, el) in arr.iter().enumerate() {
                        for r in vec {
                            r.apply_rule(&format!("{key}[{i}]"), el, true)?;
                        }
                    }
                } else {
                    return Err(Error::ArgumentErr(format!(
                        "Argument `{key}` should be an array"
                    )));
                }
            }
            // TODO: For better error messages on OneOf, make a dedicated OneOf type that matches the label instead of trying the whole type.
            SchemaValidationRule::IsUnionType(vec) => {
                let mut match_count = 0;

                let mut errors = String::new();
                for typ in vec {
                    if let Some(e) = typ
                        .iter()
                        .map(|r| r.apply_rule(key, val, true))
                        .find_map(Result::err)
                    {
                        errors.push_str(&format!("- {e}\n"));
                    } else {
                        match_count += 1;
                    }
                }

                if match_count == 0 {
                    return Err(Error::ArgumentErr(format!(
                        "Argument `{key}` is not valid, failed matching to one of the expected types. Here is a list of possible errors:\n{errors}"
                    )));
                }
            }
            SchemaValidationRule::IsOneOf(vec) => {
                let variant_label = val
                    .get("label")
                    .ok_or(Error::ArgumentErr(format!(
                        "oneOf Variant  for argument `{key}` should have a label field"
                    )))?
                    .as_str()
                    .ok_or(Error::ArgumentErr(format!(
                        "Argument `{key}` of type oneOf expected the label to be a string"
                    )))?;

                let variant_rules = vec
                    .get(variant_label)
                    .ok_or_else(|| Error::ArgumentErr(format!(
                        "Argument `{key}` of type oneOf expected one of the following variants {}, but received `{variant_label}`", vec.keys().join(", ")
                    )))?;

                for r in variant_rules {
                    r.apply_rule(key, val, true).map_err(|e| Error::ArgumentErr(format!("Argument `{key}`: The schema for the selected oneOf variant `{variant_label}` was not respected: {e}")))?;
                }
            }
            // TODO: Implement validation on these
            SchemaValidationRule::IsDatetime => (),
            SchemaValidationRule::IsEmail => (),
            SchemaValidationRule::IsBytes => (),
        }

        Ok(())
    }
}

fn find_annotation(comm_lit: &str, annotation: &str, code: &str) -> bool {
    let a = format!("{comm_lit} {annotation}");
    for l in code.lines() {
        if !l.starts_with(comm_lit) {
            break;
        }

        if l.trim_end() == a {
            return true;
        }
    }

    false
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CiTestedItem {
    pub path: String,
    pub kind: String,
}

/// Parse a CI test annotation from the top of a script.
///
/// Multi-line format:
/// ```
/// // test:
/// // script/u/admin/my_script
/// // flow/u/admin/my_flow
/// // resource/u/admin/my_resource
/// ```
///
/// One-line format:
/// ```
/// // test: script/u/admin/my_script
/// ```
pub fn parse_ci_test_annotation(code: &str, comment_prefix: &str) -> Option<Vec<CiTestedItem>> {
    let test_marker = format!("{comment_prefix} test:");
    let mut lines = code.lines();

    let first = lines.next()?;
    let first_trimmed = first.trim_end();
    if !first_trimmed.starts_with(&test_marker) {
        return None;
    }

    let mut items = Vec::new();

    // Check for one-line format: "// test: script/path"
    let inline = first_trimmed[test_marker.len()..].trim();
    if !inline.is_empty() {
        if let Some(item) = parse_ci_test_item(inline) {
            items.push(item);
        }
    }

    for line in lines {
        let Some(after_prefix) = line.strip_prefix(comment_prefix) else {
            break;
        };
        let stripped = after_prefix.trim();
        if stripped.is_empty() {
            continue;
        }

        if let Some(item) = parse_ci_test_item(stripped) {
            items.push(item);
        } else {
            break;
        }
    }

    if items.is_empty() {
        None
    } else {
        Some(items)
    }
}

fn parse_ci_test_item(s: &str) -> Option<CiTestedItem> {
    if let Some(path) = s.strip_prefix("script/") {
        Some(CiTestedItem { path: path.trim().to_string(), kind: "script".to_string() })
    } else if let Some(path) = s.strip_prefix("flow/") {
        Some(CiTestedItem { path: path.trim().to_string(), kind: "flow".to_string() })
    } else if let Some(path) = s.strip_prefix("resource/") {
        Some(CiTestedItem { path: path.trim().to_string(), kind: "resource".to_string() })
    } else {
        None
    }
}

lazy_static::lazy_static! {
    static ref CI_TEST_PATTERN_CACHE: Cache<String, Arc<Regex>> = Cache::new(256);
}

/// Does `path` match an annotation pattern?
/// - `**` matches any characters including `/`
/// - `*` matches any characters within a single path segment (no `/`)
/// - All other regex metacharacters are treated literally
pub fn ci_test_path_matches(path: &str, pattern: &str) -> bool {
    // Fast path: literal pattern with no wildcard.
    if !pattern.contains('*') {
        return path == pattern;
    }

    let regex = if let Some(cached) = CI_TEST_PATTERN_CACHE.get(pattern) {
        cached
    } else {
        let mut re = String::with_capacity(pattern.len() + 4);
        re.push('^');
        let mut chars = pattern.chars().peekable();
        while let Some(c) = chars.next() {
            if c == '*' {
                if chars.peek() == Some(&'*') {
                    chars.next();
                    re.push_str(".*");
                } else {
                    re.push_str("[^/]*");
                }
            } else {
                // regex::escape would reallocate per char; inline the ASCII metachar set.
                if matches!(
                    c,
                    '.' | '+' | '?' | '(' | ')' | '[' | ']' | '{' | '}' | '|' | '^' | '$' | '\\'
                ) {
                    re.push('\\');
                }
                re.push(c);
            }
        }
        re.push('$');
        match Regex::new(&re) {
            Ok(r) => {
                let arc = Arc::new(r);
                CI_TEST_PATTERN_CACHE.insert(pattern.to_string(), arc.clone());
                arc
            }
            Err(e) => {
                tracing::warn!(%e, pattern, "invalid ci test pattern, skipping");
                return false;
            }
        }
    };

    regex.is_match(path)
}

pub fn should_validate_schema(code: &str, lang: &ScriptLang) -> bool {
    let annotation = "schema_validation";
    find_annotation(&lang.as_comment_lit(), annotation, code)
}

#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct SchemaValidator {
    pub required: Vec<String>,
    pub rules: Vec<(String, Vec<SchemaValidationRule>)>,
}

impl SchemaValidator {
    pub fn validate(&self, args: &HashMap<String, Box<RawValue>>) -> Result<(), Error> {
        for key in &self.required {
            if !args.contains_key(key) {
                return Err(Error::ArgumentErr(format!("Argument {key} is required")));
            }
        }

        for (key, rules) in &self.rules {
            if let Some(raw_val) = args.get(key) {
                let parsed_val = Value::from_str(raw_val.get()).map_err(|e| {
                    Error::ArgumentErr(format!("Failed to parse `{key}` argument: {e}"))
                })?;
                for rule in rules {
                    rule.apply_rule(key, &parsed_val, self.required.contains(key))?;
                }
            }
        }

        Ok(())
    }

    pub fn from_schema(schema: &str) -> Result<Self, Error> {
        let schema: Value = serde_json::from_str(schema)?;

        if let Some(draft_version) = schema.get("$schema") {
            match draft_version.as_str() {
                Some("https://json-schema.org/draft/2020-12/schema") => (),
                _ => return Err(anyhow!("Supplied schema draft version is unsuported").into()),
            }
        } else {
            return Err(anyhow!("No draft version supplied").into());
        }

        let required: Vec<String> = schema
            .get("required")
            .ok_or(anyhow!("Missing `required` field on schema"))?
            .as_array()
            .ok_or(anyhow!("`required` field should be an array of strings"))?
            .into_iter()
            .map(|v| {
                v.as_str()
                    .map(|s| s.to_string())
                    .ok_or(anyhow!("required field key is not a string"))
            })
            .collect::<Result<Vec<String>, anyhow::Error>>()?;

        let properties = schema
            .get("properties")
            .ok_or(anyhow!("Missing `properties` field on schema"))?
            .as_object()
            .ok_or(anyhow!("`properties` field should be an object"))?;

        let mut rules = vec![];

        for (key, val) in properties {
            rules.push((
                key.clone(),
                SchemaValidationRule::from_value(val)
                    .map_err(|e| anyhow!("Problem making rule for {key}: {e}"))?,
            ));
        }

        Ok(Self { required, rules })
    }
}

impl JsonPrimitiveType {
    fn from_str(typ: &str) -> Result<Self, anyhow::Error> {
        match typ {
            "string" => {
                return Ok(JsonPrimitiveType::String);
            }
            "number" => {
                return Ok(JsonPrimitiveType::Number);
            }
            "integer" => {
                return Ok(JsonPrimitiveType::Integer);
            }
            "object" => {
                return Ok(JsonPrimitiveType::Object);
            }
            "array" => {
                return Ok(JsonPrimitiveType::Array);
            }
            "boolean" => {
                return Ok(JsonPrimitiveType::Boolean);
            }
            "null" => {
                return Ok(JsonPrimitiveType::Null);
            }
            other => return Err(anyhow!("Received unsupported type `{other}`").into()),
        }
    }
}

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

    use super::*;

    fn value_to_rawvalue_map(
        value: Value,
    ) -> Result<HashMap<String, Box<RawValue>>, anyhow::Error> {
        match value {
            Value::Object(map) => {
                let mut result = HashMap::new();
                for (key, val) in map {
                    let raw = serde_json::to_string(&val)?; // Serialize the Value to a string
                    let raw_value: Box<RawValue> = serde_json::from_str(&raw)?; // Convert string to Box<RawValue>
                    result.insert(key, raw_value);
                }
                Ok(result)
            }
            _ => Err(anyhow!("Expected a JSON object")),
        }
    }
    #[test]
    fn test_parse_and_validate_schema() {
        let schema = r#"{
    "$schema": "https://json-schema.org/draft/2020-12/schema",
    "properties": {
        "a": {
            "contentEncoding": "base64",
            "default": null,
            "description": "",
            "originalType": "bytes",
            "type": "string"
        },
        "b": {
            "default": null,
            "description": "",
            "enum": [
                "my",
                "enum"
            ],
            "originalType": "enum",
            "type": "string"
        },
        "e": {
            "default": "inferred type string from default arg",
            "description": "",
            "originalType": "string",
            "type": "string"
        },
        "f": {
            "default": {
                "nested": "object"
            },
            "description": "",
            "properties": {
                "nested": {
                    "description": "",
                    "type": "string",
                    "originalType": "string"
                }
            },
            "type": "object"
        },
        "g": {
            "default": null,
            "description": "",
            "oneOf": [
                {
                    "type": "object",
                    "title": "Variant 1",
                    "properties": {
                        "label": {
                            "description": "",
                            "type": "string",
                            "originalType": "enum",
                            "enum": [
                                "Variant 1"
                            ]
                        },
                        "foo": {
                            "description": "",
                            "type": "string",
                            "originalType": "string"
                        }
                    }
                },
                {
                    "type": "object",
                    "title": "Variant 2",
                    "properties": {
                        "label": {
                            "description": "",
                            "type": "string",
                            "originalType": "enum",
                            "enum": [
                                "Variant 2"
                            ]
                        },
                        "bar": {
                            "description": "",
                            "type": "number"
                        }
                    }
                }
            ],
            "type": "object"
        }
    },
    "required": [
        "a",
        "b",
        "g"
    ],
    "type": "object"
}
"#;

        let validator = SchemaValidator::from_schema(schema)
            .expect("Schema couldn't be built from a valid schema");

        let args = json!(
                    {
                        "g": {
                                "label": "Variant 1",
                                "foo": ""
                        },
                        "f": {
                                "nested": "object"
                        },
                        "e": "inferred type string from default arg",
                        "b": "my",
                        "a": null
                    }
        );

        validator
            .validate(&value_to_rawvalue_map(args).unwrap())
            .err()
            .expect("Validation should not work for this");

        let args = json!(
                    {
                        "g": {
                                "label": "Variant 1",
                                "foo": ""
                        },
                        "f": {
                                "nested": "object"
                        },
                        "e": "inferred type string from default arg",
                        "b": "not_enum",
                        "a": "123"
                    }
        );

        validator
            .validate(&value_to_rawvalue_map(args).unwrap())
            .err()
            .expect("Validation should not work for this");

        let args = json!(
                    {
                        "g": {
                                "label": "Variant 1",
                                "foo": ""
                        },
                        "f": {
                                "nested": "object"
                        },
                        "e": "inferred type string from default arg",
                        "b": "my",
                        "a": "123"
                    }
        );

        validator
            .validate(&value_to_rawvalue_map(args).unwrap())
            .expect("Validation should work for this");
    }

    #[test]
    fn ci_test_path_matches_literal() {
        assert!(ci_test_path_matches("u/user/foo", "u/user/foo"));
        assert!(!ci_test_path_matches("u/user/foo", "u/user/bar"));
        assert!(!ci_test_path_matches("u/user/foo/deep", "u/user/foo"));
    }

    #[test]
    fn ci_test_path_matches_single_star_stays_in_segment() {
        assert!(ci_test_path_matches("u/user/foo", "u/user/*"));
        assert!(ci_test_path_matches("u/user/bar", "u/user/*"));
        assert!(!ci_test_path_matches("u/user/foo/deep", "u/user/*"));
        assert!(!ci_test_path_matches("u/user", "u/user/*"));
    }

    #[test]
    fn ci_test_path_matches_double_star_crosses_slashes() {
        assert!(ci_test_path_matches("u/user/foo", "u/user/**"));
        assert!(ci_test_path_matches("u/user/foo/deep", "u/user/**"));
        assert!(ci_test_path_matches("u/user/a/b/c", "u/user/**"));
        assert!(!ci_test_path_matches("u/other/foo", "u/user/**"));
    }

    #[test]
    fn ci_test_path_matches_escapes_regex_metachars() {
        // `.` is a regex metachar; must be matched literally.
        assert!(ci_test_path_matches("u/user/file.txt", "u/user/file.txt"));
        assert!(!ci_test_path_matches("u/user/fileXtxt", "u/user/file.txt"));
        // `+` / `?` / parens must also match literally.
        assert!(ci_test_path_matches("u/user/a+b", "u/user/a+b"));
        assert!(ci_test_path_matches("u/user/a(b)", "u/user/a(b)"));
    }

    #[test]
    fn ci_test_path_matches_star_at_boundaries() {
        assert!(ci_test_path_matches("u/user/foo", "*/user/foo"));
        assert!(ci_test_path_matches("u/user/foo", "u/*/foo"));
        assert!(ci_test_path_matches("prefix_suffix", "prefix_*"));
        assert!(ci_test_path_matches("", ""));
        assert!(!ci_test_path_matches("u/user/a/b", "u/user/a"));
    }

    #[test]
    fn ci_test_path_matches_underscore_is_literal() {
        // Underscore is a SQL LIKE metachar but plain text for regex; must match literally.
        assert!(ci_test_path_matches("u/user/my_script", "u/user/my_script"));
        assert!(!ci_test_path_matches(
            "u/user/myXscript",
            "u/user/my_script"
        ));
    }
}