pdk_json_validator_lib/

lib.rs

// Copyright (c) 2026, Salesforce, Inc.,
// All rights reserved.
// For full license text, see the LICENSE.txt file

//! PDK JSON Validator Library
//!
//! Provides JSON validation functionality with configurable constraints,
//! supporting incremental/streaming validation of JSON payloads.
//!
//! - Incremental JSON validation (chunk-based streaming)
//! - Configurable constraints (depth, array length, object entries, string/key lengths)
//! - UTF-8 boundary handling for incomplete chunks
//!
//! ## Primary types
//!
//! - [`JsonValidator`]: validator for JSON chunks incrementally
//! - [`ValidationResult`]: result indicating validation status
//! - [`ValidationError`]: error type for validation failures
//! - [`JsonValidatorBuilder`]: builder for creating validator instances
//!
//! ## Example
//!
//! ```rust,no_run
//! use pdk::json_validator::{JsonValidatorBuilder, ValidationResult, ValidationError};
//!
//! let mut validator = JsonValidatorBuilder::new()
//!     .with_max_depth(10)
//!     .with_max_array_length(1000)
//!     .with_max_string_length(10000)
//!     .build();
//!
//! // Validate first chunk
//! match validator.validate_chunk(b"{\"key\": \"val", false) {
//!     Ok(ValidationResult::Incomplete) => {
//!         // Wait for more chunks
//!     }
//!     Err(e) => {
//!         eprintln!("Validation error: {}", e);
//!     }
//!     _ => {}
//! }
//!
//! // Validate final chunk
//! match validator.validate_chunk(b"ue\"}", true) {
//!     Ok(ValidationResult::Complete) => {
//!         // JSON valid
//!     }
//!     Err(ValidationError::MaxDepthExceeded) => {
//!         // Handle specific constraint violation
//!     }
//!     Err(e) => {
//!         eprintln!("Validation error: {}", e);
//!     }
//! }
//! ```

use itertools::Itertools;
use pdk_core::log::debug;
use thiserror::Error;

mod parser;
mod validator;

use parser::{DataStream, JsonParser, ParserContext};
use validator::JsonConstraintValidator;

// JsonConstraints is private - only used internally
#[derive(Debug, Clone)]
struct JsonConstraints {
    max_container_depth: Option<usize>,
    max_object_entry_count: Option<usize>,
    max_object_entry_name_length: Option<usize>,
    max_array_element_count: Option<usize>,
    max_string_value_length: Option<usize>,
}

/// Result of validating a JSON chunk
#[derive(Debug, PartialEq, Eq)]
pub enum ValidationResult {
    /// JSON is valid but incomplete (waiting for more chunks)
    Incomplete,

    /// JSON is complete and valid
    Complete,
}

/// Errors that can occur during JSON validation
#[non_exhaustive]
#[derive(Debug, Error, PartialEq, Eq)]
pub enum ValidationError {
    /// Maximum depth constraint violated
    #[error("Max depth exceeded.")]
    MaxDepthExceeded,

    /// Array length constraint violated
    #[error("Array exceeded max length")]
    ArrayMaxLengthExceeded,

    /// Object entry count constraint violated
    #[error("Object exceeded max length")]
    ObjectMaxLengthExceeded,

    /// String length constraint violated
    #[error("String exceeded max length")]
    StringMaxLengthExceeded,

    /// Unable to process payload due to invalid syntax
    #[error("Unable to process payload.")]
    UnableToProcessPayload,

    /// Invalid token encountered
    #[error("Unable to process payload. Could not complete token {0}")]
    InvalidToken(String),

    /// Unexpected character at specific index
    #[error("Error at idx {0}. Did not expect any more characters, found {1}")]
    UnexpectedCharacter(usize, char),

    /// Invalid JSON structure
    #[error("Invalid json")]
    InvalidJson,

    /// Payload is not valid UTF-8
    #[error("Non-UTF8 payload")]
    NonUtf8,
}

/// Builder for creating JSON validation instances
///
/// The builder receives configuration values directly and constructs
/// constraints internally.
pub struct JsonValidatorBuilder {
    constraints: JsonConstraints,
}

impl JsonValidatorBuilder {
    /// Creates a new builder without restrictions (only validates JSON syntax)
    pub fn new() -> Self {
        Self {
            constraints: JsonConstraints {
                max_container_depth: None,
                max_object_entry_count: None,
                max_object_entry_name_length: None,
                max_array_element_count: None,
                max_string_value_length: None,
            },
        }
    }

    /// Sets the maximum depth of nested containers
    pub fn with_max_depth(mut self, depth: usize) -> Self {
        self.constraints.max_container_depth = Some(depth);
        self
    }

    /// Sets the maximum number of elements in arrays
    pub fn with_max_array_length(mut self, length: usize) -> Self {
        self.constraints.max_array_element_count = Some(length);
        self
    }

    /// Sets the maximum length of string values
    pub fn with_max_string_length(mut self, length: usize) -> Self {
        self.constraints.max_string_value_length = Some(length);
        self
    }

    /// Sets the maximum number of entries in objects
    pub fn with_max_object_entries(mut self, count: usize) -> Self {
        self.constraints.max_object_entry_count = Some(count);
        self
    }

    /// Sets the maximum length of key names in objects
    pub fn with_max_key_length(mut self, length: usize) -> Self {
        self.constraints.max_object_entry_name_length = Some(length);
        self
    }

    /// Builds the validation instance with the specified configuration
    pub fn build(self) -> JsonValidator {
        JsonValidator::new(self.constraints)
    }
}

impl Default for JsonValidatorBuilder {
    fn default() -> Self {
        Self::new()
    }
}

/// JSON validator with configurable constraints.
///
/// Handles incremental/streaming validation of JSON payloads and maintains internal
/// state between chunk validations.
///
/// Instances should be created using [`JsonValidatorBuilder`].
pub struct JsonValidator {
    constraint_validator: JsonConstraintValidator,
    parser: ParserContext,
    data_stream_idx: (usize, usize), // (idx, offset)
    buffer: Vec<u8>,
}

impl JsonValidator {
    fn new(constraints: JsonConstraints) -> Self {
        let constraint_validator = JsonConstraintValidator::new(
            constraints.max_array_element_count,
            constraints.max_string_value_length,
            constraints.max_object_entry_count,
            constraints.max_object_entry_name_length,
            constraints.max_container_depth,
        );

        Self {
            constraint_validator,
            parser: ParserContext::new(),
            data_stream_idx: (0, 0),
            buffer: Vec::new(),
        }
    }

    fn validate_json(&mut self, json: &str, end_of_stream: bool) -> Result<(), ValidationError> {
        let chars = json.as_bytes();
        let (idx, offset) = self.data_stream_idx;
        let mut stream = DataStream::new(chars, idx + offset);

        while !stream.is_empty() {
            if self.parser.is_empty() {
                if stream.consume().map(parser::is_whitespace).unwrap_or(false) {
                    continue;
                }
                return Err(ValidationError::UnexpectedCharacter(
                    stream.idx(),
                    stream.peek().unwrap_or_default(),
                ));
            }
            let current_parser = self.parser.current().unwrap();
            let result =
                current_parser.parse(&mut stream, &mut self.parser, &self.constraint_validator);
            if result.is_err() {
                debug!(
                    "Error at : {}{}",
                    stream.peek_range(-20).unwrap_or(""),
                    stream.peek_range(20).unwrap_or("")
                )
            }
            result?;
        }

        let stream_idx = stream.idx();
        self.data_stream_idx = (stream_idx, 0);

        let processing_valid_number = !self.parser.is_empty() && self.is_processing_valid_number();
        let is_valid = self.parser.is_empty() || processing_valid_number;
        let is_empty = self.data_stream_idx.0 + self.data_stream_idx.1 == 0;

        if !is_empty && is_valid {
            Ok(())
        } else if !end_of_stream {
            // JSON incomplete but valid so far
            Ok(())
        } else {
            return Err(ValidationError::InvalidJson);
        }
    }

    fn is_processing_valid_number(&self) -> bool {
        if let Some(JsonParser::Number { expected }) = self.parser.current() {
            return matches!(expected, parser::NumberStateMachine::DigitPeriodOrExponent);
        }
        false
    }

    /// Returns the longest UTF-8 valid prefix and the remainder
    ///
    /// This method handles UTF-8 boundaries correctly, ensuring that
    /// incomplete UTF-8 characters are not split across chunks.
    fn split_utf8(bytes: &[u8]) -> (&[u8], &[u8]) {
        if bytes.is_empty() {
            return (&[], &[]);
        }
        if Self::is_ascii(*bytes.last().unwrap()) {
            return (bytes, &[]);
        }

        if let Some((reversed_idx, b)) = bytes
            .iter()
            .rev()
            .find_position(|&&b| Self::is_utf8_lead(b))
        {
            let utf8_char_expected_len = Self::utf8_char_size(*b);
            let utf8_char_actual_len = reversed_idx + 1;
            let idx = if utf8_char_expected_len as usize == utf8_char_actual_len {
                bytes.len()
            } else {
                bytes
                    .len()
                    .checked_sub(reversed_idx + 1)
                    .unwrap_or_default()
            };
            (&bytes[..idx], &bytes[idx..])
        } else {
            (&[], &[])
        }
    }

    #[inline]
    fn is_ascii(b: u8) -> bool {
        !Self::is_not_ascii(b)
    }

    #[inline]
    fn is_utf8_lead(b: u8) -> bool {
        (b & 0b11000000_u8) == 0b11000000_u8
    }

    #[inline]
    fn is_not_ascii(byte: u8) -> bool {
        byte & 0b10000000_u8 == 0b10000000_u8
    }

    fn utf8_char_size(mut lead_byte: u8) -> u8 {
        let mut counter = 0;
        while Self::is_not_ascii(lead_byte) {
            counter += 1;
            lead_byte <<= 1;
        }
        counter
    }

    /// Validates a JSON chunk incrementally
    ///
    /// This method can be called multiple times with different chunks
    /// of the same JSON. The validator maintains internal state to handle
    /// incremental validation and UTF-8 boundaries.
    ///
    /// # Arguments
    ///
    /// * `chunk` - Bytes chunk to validate. Can be a fragment of the complete JSON.
    /// * `end_of_stream` - `true` if this is the last chunk of the stream, `false` if there are more chunks pending.
    ///
    /// # Returns
    ///
    /// * `Ok(ValidationResult::Incomplete)` - The chunk is valid but the JSON is incomplete.
    ///   `validate_chunk()` should be called again with the next chunk.
    /// * `Ok(ValidationResult::Complete)` - The JSON is complete and valid (only when `end_of_stream = true`).
    /// * `Err(ValidationError)` - Validation error (invalid JSON or constraint violation).
    pub fn validate_chunk(
        &mut self,
        chunk: &[u8],
        end_of_stream: bool,
    ) -> Result<ValidationResult, ValidationError> {
        // Accumulate chunk in buffer
        self.buffer.extend_from_slice(chunk);

        // Handle UTF-8 boundaries (split into valid prefix + remainder)
        let (valid_utf8_ref, remainder_ref) = Self::split_utf8(&self.buffer);

        // If no valid UTF-8 and not end_of_stream, wait for more data
        if valid_utf8_ref.is_empty() && !end_of_stream {
            return Ok(ValidationResult::Incomplete);
        }

        // If no valid UTF-8 and end_of_stream, error
        if valid_utf8_ref.is_empty() && end_of_stream {
            return Err(ValidationError::NonUtf8);
        }

        // Copy slices before mutating self (needed for borrow checker)
        // Only copy remainder if it can be needed (not end_of_stream)
        let valid_utf8_vec = valid_utf8_ref.to_vec();
        let remainder_vec = if !end_of_stream {
            remainder_ref.to_vec()
        } else {
            Vec::new()
        };

        // Convert to UTF-8 string
        let json_str =
            std::str::from_utf8(&valid_utf8_vec).map_err(|_| ValidationError::NonUtf8)?;

        // Validate JSON with incremental parser
        match self.validate_json(json_str, end_of_stream) {
            Ok(()) if end_of_stream => {
                // Clear buffer
                self.buffer.clear();
                Ok(ValidationResult::Complete)
            }
            Ok(()) => {
                // Save remainder for next chunk
                self.buffer = remainder_vec;
                Ok(ValidationResult::Incomplete)
            }
            Err(e) => {
                // Clear buffer on error
                self.buffer.clear();
                Err(e)
            }
        }
    }

    /// Resets the validator state to start a new validation
    pub fn reset(&mut self) {
        self.parser = ParserContext::new();
        self.buffer.clear();
        self.data_stream_idx = (0, 0);
    }
}

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

    #[test]
    fn builder_creates_validator_without_constraints() {
        let mut validator = JsonValidatorBuilder::new().build();
        let result = validator.validate_chunk(b"{\"key\": \"value\"}", true);
        assert!(result.is_ok());
        assert_eq!(result.unwrap(), ValidationResult::Complete);
    }

    #[test]
    fn builder_configures_max_depth() {
        let mut validator = JsonValidatorBuilder::new().with_max_depth(2).build();

        // JSON with depth 3 should fail
        let deep_json = r#"{"a": {"b": {"c": "value"}}}"#;
        let result = validator.validate_chunk(deep_json.as_bytes(), true);
        assert!(result.is_err());
        assert!(matches!(
            result.unwrap_err(),
            ValidationError::MaxDepthExceeded
        ));
    }

    #[test]
    fn validator_handles_incremental_chunks() {
        let mut validator = JsonValidatorBuilder::new().build();

        // First incomplete chunk
        let result1 = validator.validate_chunk(b"{\"key\": \"val", false);
        assert!(result1.is_ok());
        assert_eq!(result1.unwrap(), ValidationResult::Incomplete);

        // Second complete chunk
        let result2 = validator.validate_chunk(b"ue\"}", true);
        assert!(result2.is_ok());
        assert_eq!(result2.unwrap(), ValidationResult::Complete);
    }

    #[test]
    fn validator_rejects_invalid_json() {
        let mut validator = JsonValidatorBuilder::new().build();
        let result = validator.validate_chunk(b"{invalid", true);
        assert!(result.is_err());
        assert!(matches!(
            result.unwrap_err(),
            ValidationError::InvalidJson
                | ValidationError::UnableToProcessPayload
                | ValidationError::InvalidToken(_)
        ));
    }

    #[test]
    fn validator_reset_clears_state() {
        let mut validator = JsonValidatorBuilder::new().build();

        // Validate a JSON
        let _ = validator.validate_chunk(b"{\"key\": \"value\"}", true);

        // Reset
        validator.reset();

        // Should be able to validate another JSON from scratch
        let result = validator.validate_chunk(b"{\"another\": \"json\"}", true);
        assert!(result.is_ok());
        assert_eq!(result.unwrap(), ValidationResult::Complete);
    }

    #[test]
    fn validator_rejects_non_utf8_payload() {
        let mut validator = JsonValidatorBuilder::new().build();
        // Invalid UTF-8 bytes
        let invalid_utf8: &[u8] = &[0xFF, 0xFE, 0xFD];
        let result = validator.validate_chunk(invalid_utf8, true);
        assert!(result.is_err());
        assert!(matches!(result.unwrap_err(), ValidationError::NonUtf8));
    }

    #[test]
    fn split_utf8_bytes_are_not_changed() {
        let original = "my very ∆˚∆ utf8 string".as_bytes();
        let (prefix, remainder) = JsonValidator::split_utf8(original);
        assert_eq!(original, prefix);
        assert_eq!(&[] as &[u8], remainder);
    }

    #[test]
    fn split_utf8_string_with_incomplete_char_splits_correctly() {
        let original: &[u8] = &[
            0b01101101, 0b01111001, 0b00100000, 0b01110110, 0b01100101, 0b01110010, 0b01111001,
            0b00100000, 0b11110010, 0b10100010, 0b10100010,
        ];
        let expected_prefix: &[u8] = &[
            0b01101101, 0b01111001, 0b00100000, 0b01110110, 0b01100101, 0b01110010, 0b01111001,
            0b00100000,
        ];
        let expected_remainder: &[u8] = &[0b11110010, 0b10100010, 0b10100010];

        let (prefix, remainder) = JsonValidator::split_utf8(original);
        assert_eq!(expected_prefix, prefix);
        assert_eq!(expected_remainder, remainder);
    }

    #[test]
    fn split_utf8_empty_bytes_returns_empty() {
        let (prefix, remainder) = JsonValidator::split_utf8(&[]);
        assert_eq!(&[] as &[u8], prefix);
        assert_eq!(&[] as &[u8], remainder);
    }

    #[test]
    fn split_utf8_ascii_string_returns_full_string() {
        let ascii = b"hello world";
        let (prefix, remainder) = JsonValidator::split_utf8(ascii);
        assert_eq!(ascii, prefix);
        assert_eq!(&[] as &[u8], remainder);
    }

    #[test]
    fn plain_positive_integer_is_not_blocked() {
        // pre-migration InvalidJsonTest: numericPayloadIsNotBlocked
        let mut validator = JsonValidatorBuilder::new().build();
        let result = validator.validate_chunk(b"1234", true);
        assert_eq!(result, Ok(ValidationResult::Complete));
    }

    #[test]
    fn plain_negative_integer_is_not_blocked() {
        let mut validator = JsonValidatorBuilder::new().build();
        let result = validator.validate_chunk(b"-42", true);
        assert_eq!(result, Ok(ValidationResult::Complete));
    }

    #[test]
    fn exponential_integer_notation_is_not_blocked() {
        let mut validator = JsonValidatorBuilder::new().build();
        let result = validator.validate_chunk(b"1e10", true);
        assert_eq!(result, Ok(ValidationResult::Complete));
    }

    #[test]
    fn exponential_float_notation_uppercase_is_not_blocked() {
        let mut validator = JsonValidatorBuilder::new().build();
        let result = validator.validate_chunk(b"8.8E9", true);
        assert_eq!(result, Ok(ValidationResult::Complete));
    }

    #[test]
    fn exponential_with_negative_exponent_is_not_blocked() {
        let mut validator = JsonValidatorBuilder::new().build();
        let result = validator.validate_chunk(b"2.5e-3", true);
        assert_eq!(result, Ok(ValidationResult::Complete));
    }

    #[test]
    fn exponential_with_positive_exponent_is_not_blocked() {
        let mut validator = JsonValidatorBuilder::new().build();
        let result = validator.validate_chunk(b"-14.15e+13", true);
        assert_eq!(result, Ok(ValidationResult::Complete));
    }

    #[test]
    fn exponential_notation_in_object_is_not_blocked() {
        let mut validator = JsonValidatorBuilder::new().build();
        let result = validator.validate_chunk(br#"{"value": 1e10}"#, true);
        assert_eq!(result, Ok(ValidationResult::Complete));
    }

    #[test]
    fn exponential_notation_in_array_is_not_blocked() {
        let mut validator = JsonValidatorBuilder::new().build();
        let result = validator.validate_chunk(b"[1e10, 8.8E9, 2.5e-3]", true);
        assert_eq!(result, Ok(ValidationResult::Complete));
    }

    #[test]
    fn incomplete_number_missing_exponent_digits_in_array_is_blocked() {
        let mut validator = JsonValidatorBuilder::new().build();
        let result = validator.validate_chunk(b"[1e, 2]", true);
        assert_eq!(result, Err(ValidationError::UnableToProcessPayload));
    }

    #[test]
    fn incomplete_number_missing_fraction_digits_in_array_is_blocked() {
        let mut validator = JsonValidatorBuilder::new().build();
        let result = validator.validate_chunk(b"[1., 2]", true);
        assert_eq!(result, Err(ValidationError::UnableToProcessPayload));
    }

    #[test]
    fn incomplete_number_bare_minus_in_array_is_blocked() {
        let mut validator = JsonValidatorBuilder::new().build();
        let result = validator.validate_chunk(b"[-]", true);
        assert_eq!(result, Err(ValidationError::UnableToProcessPayload));
    }

    #[test]
    fn incomplete_number_exponent_with_sign_only_in_array_is_blocked() {
        let mut validator = JsonValidatorBuilder::new().build();
        let result = validator.validate_chunk(b"[1e+]", true);
        assert_eq!(result, Err(ValidationError::UnableToProcessPayload));
    }

    #[test]
    fn incomplete_standalone_number_missing_exponent_digits_is_blocked() {
        let mut validator = JsonValidatorBuilder::new().build();
        let result = validator.validate_chunk(b"1e", true);
        assert_eq!(result, Err(ValidationError::InvalidJson));
    }

    #[test]
    fn incomplete_standalone_number_missing_fraction_digits_is_blocked() {
        let mut validator = JsonValidatorBuilder::new().build();
        let result = validator.validate_chunk(b"1.", true);
        assert_eq!(result, Err(ValidationError::InvalidJson));
    }

    #[test]
    fn incomplete_standalone_number_bare_minus_is_blocked() {
        let mut validator = JsonValidatorBuilder::new().build();
        let result = validator.validate_chunk(b"-", true);
        assert_eq!(result, Err(ValidationError::InvalidJson));
    }

    #[test]
    fn incomplete_standalone_number_exponent_with_sign_only_is_blocked() {
        let mut validator = JsonValidatorBuilder::new().build();
        let result = validator.validate_chunk(b"1e+", true);
        assert_eq!(result, Err(ValidationError::InvalidJson));
    }

    #[test]
    fn valid_string_json_is_not_blocked() {
        let mut validator = JsonValidatorBuilder::new().build();
        let result = validator.validate_chunk(b"\"string is valid json\"", true);
        assert_eq!(result, Ok(ValidationResult::Complete));
    }

    #[test]
    fn two_objects_payload_is_blocked() {
        let mut validator = JsonValidatorBuilder::new().build();
        let result = validator.validate_chunk(br#"{"menu": "file"}{"foo":"bar"}"#, true);
        assert!(matches!(
            result,
            Err(ValidationError::UnexpectedCharacter(..))
        ));
    }

    #[test]
    fn two_numbers_payload_is_blocked() {
        let mut validator = JsonValidatorBuilder::new().build();
        let result = validator.validate_chunk(b"123 456", true);
        assert!(matches!(
            result,
            Err(ValidationError::UnexpectedCharacter(..))
        ));
    }

    #[test]
    fn string_and_number_payload_is_blocked() {
        let mut validator = JsonValidatorBuilder::new().build();
        let result = validator.validate_chunk(b"\"string1\" 123", true);
        assert!(matches!(
            result,
            Err(ValidationError::UnexpectedCharacter(..))
        ));
    }

    #[test]
    fn array_payload_is_not_blocked() {
        let mut validator = JsonValidatorBuilder::new().build();
        let result = validator.validate_chunk(b"[1, 2, \"a\", \"b\"]", true);
        assert_eq!(result, Ok(ValidationResult::Complete));
    }

    #[test]
    fn nested_object_in_array_is_not_blocked() {
        let mut validator = JsonValidatorBuilder::new().build();
        let result = validator.validate_chunk(br#"[1, 2, "a", "b", {"menu": "file"}]"#, true);
        assert_eq!(result, Ok(ValidationResult::Complete));
    }

    #[test]
    fn payload_with_whitespace_is_not_blocked() {
        let mut validator = JsonValidatorBuilder::new().build();
        let result = validator.validate_chunk("{\"menu\": \r\n  \"file\" \t}".as_bytes(), true);
        assert_eq!(result, Ok(ValidationResult::Complete));
    }
}