ser_hex/
lib.rs

1//! Serialization tracing library for visualizing where bytes came from.
2//!
3//! This crate provides tools to trace serialization operations and record which
4//! bytes were read/written during deserialization, helping answer "where did these
5//! bytes come from?" when examining binary data.
6//!
7//! # Usage
8//!
9//! There are two main approaches:
10//!
11//! ## Tracing with `tracing` instrumentation
12//!
13//! Use [`read`] or [`read_incremental`] with `tracing::instrument` annotations:
14//!
15//! ```no_run
16//! use std::io::{Cursor, Read};
17//!
18//! let mut input = Cursor::new([1, 2, 3]);
19//! ser_hex::read_incremental("trace.json", &mut input, read).unwrap();
20//!
21//! #[tracing::instrument(skip_all)]
22//! fn read<R: Read>(input: &mut R) -> std::io::Result<()> {
23//!     input.read_exact(&mut [0; 3])?;
24//!     Ok(())
25//! }
26//! ```
27//!
28//! ## Using [`TraceStream`] directly
29//!
30//! Wrap any `Read + Seek` stream with [`TraceStream`]:
31//!
32//! ```no_run
33//! use std::io::{Cursor, Read};
34//! use ser_hex::TraceStream;
35//!
36//! let input = Cursor::new([1, 2, 3]);
37//! let mut tracer = TraceStream::new("trace.json", input);
38//! // ... perform reads on tracer ...
39//! ```
40
41use serde::{Deserialize, Serialize};
42use tracing::{
43    Event, Id, Metadata,
44    span::{self, EnteredSpan},
45    subscriber::{self, DefaultGuard, Subscriber},
46};
47use tracing_core::span::Current;
48
49use std::{
50    collections::HashMap,
51    fs,
52    io::{self, Cursor, Read, Seek, SeekFrom, Write},
53    path::{Path, PathBuf},
54    sync::{Arc, Mutex},
55};
56
57/// Build a stream (Cursor<Vec<u8>>) mirroring all the data in the underlying stream and cursor position
58fn build_mirror<S: Read + Seek>(stream: &mut S) -> Result<Cursor<Vec<u8>>, io::Error> {
59    let pos = stream.stream_position()?;
60    stream.seek(SeekFrom::Start(0))?;
61    let mut data = vec![];
62    stream.read_to_end(&mut data)?;
63    let mut cursor = Cursor::new(data);
64    stream.seek(SeekFrom::Start(pos))?;
65    cursor.seek(SeekFrom::Start(pos))?;
66    Ok(cursor)
67}
68
69pub fn read<'t, 'r: 't, P: AsRef<Path>, R: Read + Seek + 'r, F, T>(
70    out_path: P,
71    reader: &'r mut R,
72    f: F,
73) -> T
74where
75    F: FnOnce(&mut TraceStream<&'r mut R>) -> T,
76{
77    let cursor = build_mirror(reader).unwrap();
78    CounterSubscriber::read(out_path.as_ref().to_owned(), Some(cursor), reader, f)
79}
80
81pub fn read_incremental<'t, 'r: 't, P: AsRef<Path>, R: Read + 'r, F, T>(
82    out_path: P,
83    reader: &'r mut R,
84    f: F,
85) -> T
86where
87    F: FnOnce(&mut TraceStream<&'r mut R>) -> T,
88{
89    CounterSubscriber::read(out_path.as_ref().to_owned(), None, reader, f)
90}
91
92pub struct TraceStream<S> {
93    stream: S,
94
95    // first drop span
96    #[allow(unused)]
97    scope_guard: EnteredSpan,
98
99    // then drop subscriber guard
100    #[allow(unused)]
101    guard: Option<DefaultGuard>,
102
103    // finally drop subscriber which writes trace
104    subscriber: CounterSubscriber,
105}
106
107impl<S: Read + Seek> TraceStream<S> {
108    pub fn new<P: Into<PathBuf>>(trace_path: P, mut inner_stream: S) -> Self {
109        let cursor = build_mirror(&mut inner_stream).unwrap();
110        let subscriber = CounterSubscriber::new(trace_path.into(), cursor);
111        let guard = Some(tracing::subscriber::set_default(subscriber.clone()));
112        Self::new_internal(inner_stream, subscriber, guard)
113    }
114}
115impl<S> TraceStream<S> {
116    pub fn new_incremental<P: Into<PathBuf>>(trace_path: P, inner_stream: S) -> Self {
117        let subscriber = CounterSubscriber::new(trace_path.into(), Cursor::new(vec![]));
118        let guard = Some(tracing::subscriber::set_default(subscriber.clone()));
119        Self::new_internal(inner_stream, subscriber, guard)
120    }
121}
122impl<S> TraceStream<S> {
123    fn new_internal(stream: S, subscriber: CounterSubscriber, guard: Option<DefaultGuard>) -> Self {
124        Self {
125            stream,
126            scope_guard: tracing::info_span!("root").entered(),
127            guard,
128            subscriber,
129        }
130    }
131}
132impl<R: Seek> Seek for TraceStream<R> {
133    fn seek(&mut self, pos: std::io::SeekFrom) -> std::io::Result<u64> {
134        self.stream
135            .seek(pos)
136            .inspect(|&to| self.subscriber.seek_action(to))
137    }
138}
139impl<R: Read> Read for TraceStream<R> {
140    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
141        self.stream
142            .read(buf)
143            .inspect(|&s| self.subscriber.read_action(buf, s))
144    }
145}
146impl<R: Write> Write for TraceStream<R> {
147    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
148        self.stream
149            .write(buf)
150            .inspect(|&s| self.subscriber.read_action(buf, s))
151    }
152
153    fn flush(&mut self) -> io::Result<()> {
154        self.stream.flush()
155    }
156}
157
158#[derive(Debug, Serialize, Deserialize)]
159pub enum Action<S> {
160    Read(usize),
161    Seek(usize),
162    Span(S),
163}
164
165#[derive(Debug, Serialize, Deserialize)]
166pub struct ReadSpan<S = TreeSpan> {
167    pub name: std::borrow::Cow<'static, str>,
168    pub actions: Vec<Action<S>>,
169}
170impl<S> ReadSpan<S> {
171    fn new(name: &'static str) -> Self {
172        Self {
173            name: name.into(),
174            actions: vec![],
175        }
176    }
177}
178
179struct CounterSubscriberInner {
180    out_path: PathBuf,
181    start_index: usize,
182    data: Cursor<Vec<u8>>,
183    last_id: u64,
184    root_span: Option<Id>,
185    spans: HashMap<Id, ReadSpan<Id>>,
186    metadata: HashMap<Id, &'static Metadata<'static>>,
187    stack: Vec<Id>,
188}
189impl CounterSubscriberInner {
190    fn new(out_path: PathBuf, mut data: Cursor<Vec<u8>>) -> Self {
191        Self {
192            out_path,
193            start_index: data.stream_position().unwrap() as usize,
194            data,
195            last_id: Default::default(),
196            root_span: Default::default(),
197            spans: Default::default(),
198            metadata: Default::default(),
199            stack: Default::default(),
200        }
201    }
202}
203
204#[derive(Debug, Serialize, Deserialize)]
205pub struct Trace<D: AsRef<[u8]> = Vec<u8>> {
206    #[serde(
207        serialize_with = "base64::serialize",
208        deserialize_with = "base64::deserialize",
209        bound(deserialize = "D: From<Vec<u8>>")
210    )]
211    pub data: D,
212    pub start_index: usize,
213    pub root: Action<TreeSpan>,
214}
215impl<D: AsRef<[u8]>> Trace<D> {
216    pub fn save(&self, path: impl AsRef<Path>) -> Result<(), std::io::Error> {
217        let json = serde_json::to_string(&self).unwrap();
218        fs::write(path, json)
219    }
220}
221
222mod base64 {
223    use base64::prelude::*;
224    use serde::{Deserialize, Serialize};
225    use serde::{Deserializer, Serializer};
226
227    pub fn serialize<V, S: Serializer>(v: V, s: S) -> Result<S::Ok, S::Error>
228    where
229        V: AsRef<[u8]>,
230    {
231        let base64 = BASE64_STANDARD.encode(v.as_ref());
232        String::serialize(&base64, s)
233    }
234
235    pub fn deserialize<'de, V: From<Vec<u8>>, D: Deserializer<'de>>(d: D) -> Result<V, D::Error> {
236        let base64 = String::deserialize(d)?;
237        BASE64_STANDARD
238            .decode(base64.as_bytes())
239            .map_err(serde::de::Error::custom)
240            .map(|v| v.into())
241    }
242}
243
244#[derive(Debug, Serialize, Deserialize)]
245#[repr(transparent)]
246pub struct TreeSpan(pub ReadSpan);
247impl TreeSpan {
248    fn into_tree(id: Id, spans: &mut HashMap<Id, ReadSpan<Id>>) -> Self {
249        let read_span = spans.remove(&id).unwrap();
250        Self(ReadSpan {
251            name: read_span.name,
252            actions: read_span
253                .actions
254                .into_iter()
255                .map(|a| match a {
256                    Action::Read(i) => Action::Read(i),
257                    Action::Seek(i) => Action::Seek(i),
258                    Action::Span(id) => Action::Span(Self::into_tree(id, spans)),
259                })
260                .collect(),
261        })
262    }
263}
264
265impl Drop for CounterSubscriberInner {
266    fn drop(&mut self) {
267        let tree = TreeSpan::into_tree(self.root_span.as_ref().cloned().unwrap(), &mut self.spans);
268        Trace {
269            data: std::mem::take(&mut self.data).into_inner(),
270            start_index: self.start_index,
271            root: Action::Span(tree),
272        }
273        .save(&self.out_path)
274        .unwrap()
275    }
276}
277
278#[derive(Clone)]
279struct CounterSubscriber {
280    inner: Arc<Mutex<CounterSubscriberInner>>,
281}
282impl CounterSubscriber {
283    fn new(out_path: PathBuf, data: Cursor<Vec<u8>>) -> Self {
284        Self {
285            inner: Arc::new(Mutex::new(CounterSubscriberInner::new(out_path, data))),
286        }
287    }
288    fn read<'d, 't, 'r: 't, R: Read + 'r, P, F, T>(
289        out_path: P,
290        data: Option<Cursor<Vec<u8>>>,
291        reader: &'r mut R,
292        f: F,
293    ) -> T
294    where
295        F: FnOnce(&mut TraceStream<&'r mut R>) -> T,
296        P: Into<PathBuf>,
297    {
298        let sub = Self::new(out_path.into(), data.unwrap_or_default());
299        tracing::subscriber::with_default(sub.clone(), || {
300            // must build TraceStream after defualt subscriber is set because it enters root span
301            f(&mut TraceStream::new_internal(reader, sub, None))
302        })
303    }
304    fn read_action(&self, buf: &[u8], size: usize) {
305        let mut lock = self.inner.lock().unwrap();
306        let current = lock.stack.last().cloned().unwrap();
307        lock.data.write_all(&buf[..size]).unwrap();
308        lock.spans
309            .get_mut(&current)
310            .unwrap()
311            .actions
312            .push(Action::Read(size));
313    }
314    fn seek_action(&self, to: u64) {
315        let mut lock = self.inner.lock().unwrap();
316        let current = lock.stack.last().cloned().unwrap();
317        lock.data.seek(SeekFrom::Start(to)).unwrap();
318        lock.spans
319            .get_mut(&current)
320            .unwrap()
321            .actions
322            .push(Action::Seek(to as usize));
323    }
324}
325
326impl Subscriber for CounterSubscriber {
327    fn register_callsite(&self, _meta: &Metadata<'_>) -> subscriber::Interest {
328        subscriber::Interest::always()
329    }
330
331    fn new_span(&self, new_span: &span::Attributes<'_>) -> Id {
332        let mut lock = self.inner.lock().unwrap();
333
334        let metadata = new_span.metadata();
335        let name = metadata.name();
336        lock.last_id += 1;
337        let id = lock.last_id;
338        let id = Id::from_u64(id);
339
340        lock.spans.insert(id.clone(), ReadSpan::new(name));
341        lock.metadata.insert(id.clone(), metadata);
342        assert_eq!(new_span.parent(), None);
343        assert!(new_span.is_contextual());
344        // TODO set root here if new_span.is_root()?
345        id
346    }
347    fn try_close(&self, _id: Id) -> bool {
348        true
349    }
350    fn current_span(&self) -> Current {
351        let lock = self.inner.lock().unwrap();
352        if let Some(id) = lock.stack.last() {
353            let metadata = lock.metadata[id];
354            Current::new(id.clone(), metadata)
355        } else {
356            Current::none()
357        }
358    }
359
360    fn record_follows_from(&self, _span: &Id, _follows: &Id) {}
361    fn record(&self, _: &Id, _values: &span::Record<'_>) {}
362    fn event(&self, _event: &Event<'_>) {}
363
364    fn enabled(&self, _metadata: &Metadata<'_>) -> bool {
365        true
366    }
367
368    fn enter(&self, span: &Id) {
369        let mut lock = self.inner.lock().unwrap();
370        if let Some(current) = lock.stack.last().cloned() {
371            lock.spans
372                .get_mut(&current)
373                .unwrap()
374                .actions
375                .push(Action::Span(span.clone()));
376        } else {
377            lock.root_span = Some(span.clone());
378        }
379        lock.stack.push(span.clone());
380    }
381    fn exit(&self, span: &Id) {
382        let mut lock = self.inner.lock().unwrap();
383        assert_eq!(&lock.stack.pop().unwrap(), span);
384    }
385}
386
387#[cfg(test)]
388mod test {
389    use std::io::Error;
390
391    use byteorder::{LE, ReadBytesExt};
392    use tracing::instrument;
393
394    use super::*;
395
396    #[instrument(name = "read_nested_stuff", skip_all)]
397    fn read_nested_stuff<R: Read + Seek>(reader: &mut R) -> Result<(), Error> {
398        let _a = reader.read_u32::<LE>()?;
399        Ok(())
400    }
401
402    #[instrument(name = "read_stuff", skip_all)]
403    fn read_stuff<R: Read + Seek>(reader: &mut R) -> Result<(), Error> {
404        let _a = reader.read_u8()?;
405        read_nested_stuff(reader)?;
406        reader.seek(std::io::SeekFrom::Current(1))?;
407        let _c = reader.read_u8()?;
408        reader.seek(std::io::SeekFrom::Current(-1))?;
409        let _c = reader.read_u8()?;
410        Ok(())
411    }
412
413    fn new_reader() -> Cursor<Vec<u8>> {
414        let mut reader = std::io::Cursor::new(vec![
415            1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 16, 17, 18, 19, 20,
416        ]);
417        reader.seek(SeekFrom::Start(2)).unwrap();
418        reader
419    }
420
421    #[test]
422    fn test_trace_read() -> Result<(), Error> {
423        read("trace_read.json", &mut new_reader(), |s| {
424            read_stuff(s)?;
425            read_stuff(s)
426        })?;
427
428        Ok(())
429    }
430
431    #[test]
432    fn test_trace_read_incremental() -> Result<(), Error> {
433        read_incremental("trace_read_incremental.json", &mut new_reader(), |s| {
434            read_stuff(s)?;
435            read_stuff(s)
436        })?;
437
438        Ok(())
439    }
440
441    #[test]
442    fn test_trace_stream() -> Result<(), Error> {
443        let mut s = TraceStream::new("trace_stream.json", new_reader());
444        read_stuff(&mut s)?;
445        read_stuff(&mut s)?;
446
447        Ok(())
448    }
449
450    #[test]
451    fn test_trace_stream_incremental() -> Result<(), Error> {
452        let mut s = TraceStream::new_incremental("trace_stream_incremental.json", new_reader());
453        read_stuff(&mut s)?;
454        read_stuff(&mut s)?;
455
456        Ok(())
457    }
458}
ser_hex/lib.rs

ser_hex/
lib.rs
