use serde::{Deserialize, de::{DeserializeSeed, MapAccess, SeqAccess, Visitor}};
use serde_json::value::RawValue;
use std::{fmt, ops::Range};
const LIMIT: usize = crate::grouped_context::LIMIT;
#[derive(Clone, Copy, Default)]
struct Key { start: usize, end: usize, object: usize }
struct Keys { bytes: [u8; LIMIT], keys: [Key; LIMIT / 4], used: usize, count: usize }
#[derive(Default)]
pub(super) struct Fields { pub rpc: Option<Range<usize>>, pub trace: Option<Range<usize>> }
struct KeySeed<'a>(&'a mut Keys);
impl<'de> DeserializeSeed<'de> for KeySeed<'_> {
type Value = Range<usize>;
fn deserialize<D: serde::Deserializer<'de>>(self, d: D) -> Result<Self::Value, D::Error> {
d.deserialize_str(self)
}
}
impl<'de> Visitor<'de> for KeySeed<'_> {
type Value = Range<usize>;
fn expecting(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.write_str("JSON object key") }
fn visit_str<E: serde::de::Error>(self, value: &str) -> Result<Self::Value, E> {
let start = self.0.used;
let end = start.checked_add(value.len()).filter(|end| *end <= LIMIT).ok_or_else(|| E::custom("context too large"))?;
self.0.bytes[start..end].copy_from_slice(value.as_bytes()); self.0.used = end;
Ok(start..end)
}
}
struct Walk<'a, 'b> { keys: &'a mut Keys, raw: &'b str, fields: &'a mut Fields, depth: usize }
impl<'de> Visitor<'de> for Walk<'_, '_> {
type Value = ();
fn expecting(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.write_str("JSON container") }
fn visit_map<M: MapAccess<'de>>(self, mut map: M) -> Result<(), M::Error> {
while let Some(range) = map.next_key_seed(KeySeed(self.keys))? {
let object = self.raw.as_ptr() as usize;
if self.keys.keys[..self.keys.count].iter().any(|key| key.object == object && self.keys.bytes[key.start..key.end] == self.keys.bytes[range.clone()]) {
return Err(serde::de::Error::custom("duplicate context key"));
}
if self.keys.count == self.keys.keys.len() { return Err(serde::de::Error::custom("context too large")); }
self.keys.keys[self.keys.count] = Key { start: range.start, end: range.end, object }; self.keys.count += 1;
let raw = map.next_value::<&'de RawValue>()?.get();
if self.depth == 0 {
let offset = raw.as_ptr() as usize - object;
match &self.keys.bytes[range] {
b"rpcId" => self.fields.rpc = Some(offset..offset + raw.len()),
b"traceId" => self.fields.trace = Some(offset..offset + raw.len()),
_ => (),
}
}
walk(raw, self.keys, self.fields, self.depth + 1).map_err(serde::de::Error::custom)?;
}
Ok(())
}
fn visit_seq<M: SeqAccess<'de>>(self, mut seq: M) -> Result<(), M::Error> {
while let Some(raw) = seq.next_element::<&'de RawValue>()? {
walk(raw.get(), self.keys, self.fields, self.depth + 1).map_err(serde::de::Error::custom)?;
}
Ok(())
}
}
fn walk(raw: &str, keys: &mut Keys, fields: &mut Fields, depth: usize) -> Result<(), serde_json::Error> {
use serde::Deserializer;
if depth > 127 { return Err(<serde_json::Error as serde::de::Error>::custom("context nesting limit")); }
let mut d = serde_json::Deserializer::from_str(raw);
match raw.as_bytes().first() {
Some(b'{') => d.deserialize_map(Walk { keys, raw, fields, depth })?,
Some(b'[') => d.deserialize_seq(Walk { keys, raw, fields, depth })?,
_ => { let _ = <&RawValue>::deserialize(&mut d)?; }
}
d.end()
}
pub(super) fn object(raw: &str) -> Result<Fields, serde_json::Error> {
if !raw.starts_with('{') { return Err(<serde_json::Error as serde::de::Error>::custom("context must be an object")); }
let mut keys = Keys { bytes: [0; LIMIT], keys: [Key::default(); LIMIT / 4], used: 0, count: 0 };
let mut fields = Fields::default();
walk(raw, &mut keys, &mut fields, 0)?;
Ok(fields)
}
pub(super) fn present_raw<'de, D: serde::Deserializer<'de>>(d: D) -> Result<Option<&'de RawValue>, D::Error> {
<&RawValue>::deserialize(d).map(Some)
}