use std::panic::AssertUnwindSafe;
use std::panic::catch_unwind;
use qubit_budget::ResourceBudget;
use qubit_budget::ResourceLimit;
use qubit_budget::StructureLimits;
use qubit_budget::json::JsonDecodeLimits;
use qubit_budget::json::JsonDecodeSession;
use qubit_budget::json::JsonEncodeLimits;
use qubit_budget::json::JsonEncodeSession;
use qubit_budget::json::JsonMeasurement;
use qubit_budget::json::JsonResource;
use qubit_budget::json::JsonValueBudget;
use qubit_budget::json::JsonValueLimits;
use qubit_json::decode::JsonDecodeError;
use qubit_json::decode::JsonDecoder;
use qubit_json::decode::NormalizingJsonDecodePolicy;
use qubit_json::decode::NormalizingJsonDecoder;
use serde::de::DeserializeOwned;
use serde::de::IgnoredAny;
fn run_with_session<'a, T>(
decoder: &NormalizingJsonDecoder<'_>,
input: &str,
session: &mut JsonDecodeSession<'a, JsonResource>,
) -> Result<T, JsonDecodeError>
where
T: DeserializeOwned,
{
let owned_session = std::mem::replace(session, JsonDecodeSession::from_limits(JsonDecodeLimits::new()));
let mut stateful = NormalizingJsonDecoder::new(decoder.policy().clone(), owned_session);
let result = stateful.decode_str(input);
*session = stateful.into_session();
result
}
#[test]
fn test_decode_and_encode_sessions_have_independent_directional_resources() {
let decode = JsonDecodeSession::from_limits(
JsonDecodeLimits::<JsonResource, usize>::builder()
.input_bytes_limit(ResourceLimit::new(JsonResource::InputBytes, 8))
.build(),
);
let encode = JsonEncodeSession::from_limits(
JsonEncodeLimits::<JsonResource, usize>::builder()
.output_bytes_limit(ResourceLimit::new(JsonResource::OutputBytes, 8))
.build(),
);
assert_eq!(decode.max_input_bytes(), Some(8));
assert_eq!(encode.max_output_bytes(), Some(8));
}
#[test]
fn test_decode_attempt_consumes_input_bytes_atomically() {
let mut session = JsonDecodeSession::from_limits(
JsonDecodeLimits::<JsonResource, usize>::builder()
.input_bytes_limit(ResourceLimit::new(JsonResource::InputBytes, 3))
.build(),
);
let mut attempt = session.begin_value();
attempt.try_consume_input_bytes(3).expect("exact input fits");
let error = attempt
.try_consume_input_bytes(1)
.expect_err("input budget is exhausted");
assert_eq!(*error.resource(), JsonResource::InputBytes);
}
#[test]
fn test_decode_session_borrowing_reuses_caller_owned_budgets() {
let mut input = ResourceBudget::new(JsonResource::InputBytes, 16_usize);
let mut value = JsonValueBudget::new(
JsonValueLimits::<JsonResource, usize>::builder()
.payload_bytes_limit(ResourceLimit::new(JsonResource::PayloadBytes, 3_usize))
.build(),
);
{
let session = JsonDecodeSession::borrowing_input(&mut input, &mut value);
let decoder = JsonDecoder::new(session);
let mut decoder = decoder;
decoder
.decode_utf8::<IgnoredAny>(br#"{"a":1}"#)
.expect("borrowed session should admit the document");
assert_eq!(decoder.session().max_input_bytes(), Some(16_usize));
assert_eq!(
decoder.session().input_budget().map(|budget| budget.limit()),
Some(16_usize)
);
}
assert_eq!(input.remaining(), 9_usize);
assert_eq!(value.used_payload_bytes(), Some(2_usize));
}
#[test]
fn test_decode_session_preserves_embedded_value_limits() {
let value_limits = JsonValueLimits::<JsonResource, usize>::builder()
.string_bytes_limit(ResourceLimit::new(JsonResource::StringBytes, 2))
.payload_bytes_limit(ResourceLimit::new(JsonResource::PayloadBytes, 3))
.structure_limits(StructureLimits::builder().nodes_limit(ResourceLimit::new(JsonResource::Nodes, 2)))
.build();
let mut session = JsonDecodeSession::from_limits(
JsonDecodeLimits::<JsonResource, usize>::builder()
.value_limits(value_limits)
.build(),
);
let mut attempt = session.begin_value();
attempt
.try_admit(JsonMeasurement::String { depth: 1, bytes: 2 })
.expect("exact string limit fits");
let first_error = attempt
.try_admit(JsonMeasurement::String { depth: 1, bytes: 3 })
.expect_err("overlong string poisons the attempt");
assert_eq!(*first_error.resource(), JsonResource::StringBytes);
let repeated_error = attempt
.try_admit(JsonMeasurement::Number { depth: 1, bytes: 1 })
.expect_err("poisoned attempt rejects later values");
assert_eq!(repeated_error.resource(), first_error.resource());
let commit_error = attempt.commit().expect_err("poisoned attempt cannot commit");
assert_eq!(commit_error.resource(), first_error.resource());
assert_eq!(session.value_budget().used_nodes(), Some(0));
assert_eq!(session.value_budget().used_payload_bytes(), Some(0));
}
#[test]
fn test_failed_second_value_preserves_first_commit_and_accumulates_input() {
let first = b"null";
let second = br#"[null,null]"#;
let third = b"null";
let session = JsonDecodeSession::from_limits(
JsonDecodeLimits::<JsonResource, usize>::builder()
.max_input_bytes(64)
.max_nodes(3)
.build(),
);
let mut decoder = JsonDecoder::new(session);
decoder.decode_utf8::<IgnoredAny>(first).expect("first value must fit");
assert!(decoder.decode_utf8::<IgnoredAny>(second).is_err());
assert_eq!(decoder.session().value_budget().used_nodes(), Some(1));
assert_eq!(
decoder.session().input_budget().expect("input budget").used(),
first.len() + second.len(),
);
decoder
.decode_utf8::<IgnoredAny>(third)
.expect("rolled-back second value must leave room for the third");
assert_eq!(decoder.session().value_budget().used_nodes(), Some(2));
}
#[test]
fn test_decode_attempt_panic_retains_input_and_reuses_value_capacity() {
let mut session = JsonDecodeSession::from_limits(
JsonDecodeLimits::<JsonResource, usize>::builder()
.max_input_bytes(8)
.max_nodes(1)
.build(),
);
let result = catch_unwind(AssertUnwindSafe(|| {
let mut attempt = session.begin_value();
attempt
.try_consume_input_bytes(4)
.expect("input must fit before the panic");
attempt
.try_admit(JsonMeasurement::Null { depth: 1 })
.expect("staged value must fit before the panic");
panic!("intentional decode-attempt panic");
}));
assert!(result.is_err());
assert_eq!(session.input_budget().expect("input budget").used(), 4,);
assert_eq!(session.value_budget().used_nodes(), Some(0));
let mut decoder = JsonDecoder::new(session);
decoder
.decode_utf8::<IgnoredAny>(b"null")
.expect("rolled-back value capacity must remain reusable");
assert_eq!(decoder.session().value_budget().used_nodes(), Some(1));
}
#[test]
fn test_lenient_typed_failure_retains_normalized_input_and_reuses_value_capacity() {
let rejected = "```json\nnull\n```";
let accepted = "null";
let mut session = JsonDecodeSession::from_limits(
JsonDecodeLimits::<JsonResource, usize>::builder()
.max_input_bytes(rejected.len() + accepted.len())
.max_normalized_input_bytes(8)
.max_nodes(1)
.build(),
);
let decoder = NormalizingJsonDecoder::with_limits(
NormalizingJsonDecodePolicy::default(),
JsonDecodeLimits::<JsonResource, usize>::default(),
);
assert!(run_with_session::<u8>(&decoder, rejected, &mut session).is_err());
assert_eq!(session.input_budget().expect("input budget").used(), rejected.len(),);
assert_eq!(
session
.normalized_input_budget()
.expect("normalized input budget")
.used(),
accepted.len(),
);
assert_eq!(session.value_budget().used_nodes(), Some(0));
run_with_session::<IgnoredAny>(&decoder, accepted, &mut session)
.expect("typed failure must leave value capacity for the next document");
assert_eq!(session.value_budget().used_nodes(), Some(1));
}