use itertools::Itertools;
use pdk_core::log::debug;
use thiserror::Error;
mod parser;
mod validator;
use parser::{DataStream, JsonParser, ParserContext};
use validator::JsonConstraintValidator;
#[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>,
}
#[derive(Debug, PartialEq, Eq)]
pub enum ValidationResult {
Incomplete,
Complete,
}
#[non_exhaustive]
#[derive(Debug, Error, PartialEq, Eq)]
pub enum ValidationError {
#[error("Max depth exceeded.")]
MaxDepthExceeded,
#[error("Array exceeded max length")]
ArrayMaxLengthExceeded,
#[error("Object exceeded max length")]
ObjectMaxLengthExceeded,
#[error("String exceeded max length")]
StringMaxLengthExceeded,
#[error("Unable to process payload.")]
UnableToProcessPayload,
#[error("Unable to process payload. Could not complete token {0}")]
InvalidToken(String),
#[error("Error at idx {0}. Did not expect any more characters, found {1}")]
UnexpectedCharacter(usize, char),
#[error("Invalid json")]
InvalidJson,
#[error("Non-UTF8 payload")]
NonUtf8,
}
pub struct JsonValidatorBuilder {
constraints: JsonConstraints,
}
impl JsonValidatorBuilder {
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,
},
}
}
pub fn with_max_depth(mut self, depth: usize) -> Self {
self.constraints.max_container_depth = Some(depth);
self
}
pub fn with_max_array_length(mut self, length: usize) -> Self {
self.constraints.max_array_element_count = Some(length);
self
}
pub fn with_max_string_length(mut self, length: usize) -> Self {
self.constraints.max_string_value_length = Some(length);
self
}
pub fn with_max_object_entries(mut self, count: usize) -> Self {
self.constraints.max_object_entry_count = Some(count);
self
}
pub fn with_max_key_length(mut self, length: usize) -> Self {
self.constraints.max_object_entry_name_length = Some(length);
self
}
pub fn build(self) -> JsonValidator {
JsonValidator::new(self.constraints)
}
}
impl Default for JsonValidatorBuilder {
fn default() -> Self {
Self::new()
}
}
pub struct JsonValidator {
constraint_validator: JsonConstraintValidator,
parser: ParserContext,
data_stream_idx: (usize, usize), 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 {
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
}
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
}
pub fn validate_chunk(
&mut self,
chunk: &[u8],
end_of_stream: bool,
) -> Result<ValidationResult, ValidationError> {
self.buffer.extend_from_slice(chunk);
let (valid_utf8_ref, remainder_ref) = Self::split_utf8(&self.buffer);
if valid_utf8_ref.is_empty() && !end_of_stream {
return Ok(ValidationResult::Incomplete);
}
if valid_utf8_ref.is_empty() && end_of_stream {
return Err(ValidationError::NonUtf8);
}
let valid_utf8_vec = valid_utf8_ref.to_vec();
let remainder_vec = if !end_of_stream {
remainder_ref.to_vec()
} else {
Vec::new()
};
let json_str =
std::str::from_utf8(&valid_utf8_vec).map_err(|_| ValidationError::NonUtf8)?;
match self.validate_json(json_str, end_of_stream) {
Ok(()) if end_of_stream => {
self.buffer.clear();
Ok(ValidationResult::Complete)
}
Ok(()) => {
self.buffer = remainder_vec;
Ok(ValidationResult::Incomplete)
}
Err(e) => {
self.buffer.clear();
Err(e)
}
}
}
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();
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();
let result1 = validator.validate_chunk(b"{\"key\": \"val", false);
assert!(result1.is_ok());
assert_eq!(result1.unwrap(), ValidationResult::Incomplete);
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();
let _ = validator.validate_chunk(b"{\"key\": \"value\"}", true);
validator.reset();
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();
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() {
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));
}
}