ya_packet_trace/
lib.rs

1pub enum WriteTarget {
2    Log,
3    Write(Box<dyn std::io::Write + Send>),
4}
5
6pub fn try_extract_from_ip_frame(frame: impl AsRef<[u8]>) -> Option<Vec<u8>> {
7    let frame: &[u8] = frame.as_ref();
8
9    if frame.len() > 55 && frame[12..14] == [0x08, 0x00] {
10        Some(frame[54..].to_owned())
11    } else {
12        None
13    }
14}
15
16#[cfg(feature = "enable")]
17static WRITE_TARGET: std::sync::Mutex<WriteTarget> = std::sync::Mutex::new(WriteTarget::Log);
18
19#[allow(unused_variables)]
20pub fn set_write_target(wt: WriteTarget) {
21    #[cfg(feature = "enable")]
22    {
23        *WRITE_TARGET.lock().unwrap() = wt;
24    }
25}
26
27/// Logs location and hash of provided data
28#[cfg(feature = "enable")]
29#[macro_export]
30macro_rules! packet_trace {
31    ($location:expr, $payload:block) => {{
32        $crate::helpers::do_write($location, $payload);
33    }};
34}
35
36/// Logs location and hash of provided data
37#[cfg(not(feature = "enable"))]
38#[macro_export]
39macro_rules! packet_trace {
40    ($location:expr, $payload:block) => {};
41}
42
43/// Logs location and hash of provided data
44#[cfg(feature = "enable")]
45#[macro_export]
46macro_rules! packet_trace_maybe {
47    ($location:expr, $maybe_payload:block) => {{
48        if let Some(payload) = $maybe_payload {
49            $crate::packet_trace!($location, { payload })
50        }
51    }};
52}
53
54/// Logs location and hash of provided data
55#[cfg(not(feature = "enable"))]
56#[macro_export]
57macro_rules! packet_trace_maybe {
58    ($location:expr, $maybe_payload:block) => {};
59}
60
61/// Date format used for timestamps
62pub const DATE_FORMAT_STR: &str = "%Y-%m-%dT%H:%M:%S%.6f%z";
63
64/// Macro internals
65///
66/// While this module must be public due to the way Rust
67/// expands declarative macros, no guarantees are made regarding
68/// this module.
69#[cfg(feature = "enable")]
70pub mod helpers {
71    pub fn do_write(location: impl std::fmt::Display, payload: impl AsRef<[u8]>) {
72        use crate::{WriteTarget, WRITE_TARGET};
73
74        let sz = payload.as_ref().len();
75        let hash = do_hash(payload);
76        let ts = ts();
77
78        match &mut *WRITE_TARGET.lock().unwrap() {
79            WriteTarget::Log => {
80                log::trace!(target: "packet-trace", "{},{:016x},{},{}", location, hash, ts, sz);
81            }
82            WriteTarget::Write(w) => {
83                writeln!(w, "{},{:016x},{},{}", location, hash, ts, sz).unwrap();
84            }
85        }
86    }
87
88    pub fn do_hash(data: impl AsRef<[u8]>) -> u64 {
89        use std::hash::Hasher;
90
91        let mut hasher = fxhash::FxHasher64::default();
92        hasher.write(data.as_ref());
93
94        hasher.finish()
95    }
96
97    pub fn ts() -> String {
98        chrono::Utc::now()
99            .format(crate::DATE_FORMAT_STR)
100            .to_string()
101    }
102}
103
104#[cfg(test)]
105mod test {
106    use log::LevelFilter;
107    use once_cell::sync::OnceCell;
108    use serial_test::serial;
109    use std::fmt::Write;
110    use std::sync::atomic::{AtomicBool, Ordering::SeqCst};
111    use std::sync::{Arc, Mutex};
112
113    #[cfg(feature = "enable")]
114    use regex::Regex;
115
116    #[cfg(feature = "enable")]
117    use crate::DATE_FORMAT_STR;
118
119    struct StringLog(Arc<Mutex<String>>);
120
121    static LOGGER: OnceCell<StringLog> = OnceCell::new();
122
123    impl StringLog {
124        fn new() -> Self {
125            StringLog(Arc::new(Mutex::new(String::new())))
126        }
127
128        fn global() -> &'static Self {
129            let first_init = AtomicBool::new(false);
130            let result = LOGGER.get_or_init(|| {
131                first_init.store(true, SeqCst);
132                Self::new()
133            });
134
135            if first_init.load(SeqCst) {
136                log::set_logger(StringLog::global()).unwrap();
137                log::set_max_level(LevelFilter::Trace);
138            }
139
140            result
141        }
142
143        fn get_string() -> String {
144            Self::global().0.lock().unwrap().clone()
145        }
146
147        fn clear() {
148            Self::global().0.lock().unwrap().clear();
149        }
150    }
151
152    impl log::Log for StringLog {
153        fn enabled(&self, _metadata: &log::Metadata) -> bool {
154            true
155        }
156
157        fn log(&self, record: &log::Record) {
158            if record.target() == "packet-trace" {
159                let mut buf = self.0.lock().unwrap();
160                writeln!(&mut buf, "{}", record.args()).unwrap()
161            }
162        }
163
164        fn flush(&self) {}
165    }
166
167    impl std::io::Write for &StringLog {
168        fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
169            let text = std::str::from_utf8(buf).unwrap();
170            let mut string = self.0.lock().unwrap();
171            string.push_str(text);
172
173            Ok(buf.len())
174        }
175
176        fn flush(&mut self) -> std::io::Result<()> {
177            Ok(())
178        }
179    }
180
181    #[test]
182    fn test_invocation_compiles() {
183        if false {
184            packet_trace!("test-1", { &[1, 2, 3] });
185        }
186    }
187
188    #[cfg(feature = "enable")]
189    #[test]
190    #[serial]
191    pub fn test_date() {
192        StringLog::clear();
193
194        packet_trace!("test-date", { &[1, 2, 3] });
195        let output = StringLog::get_string();
196        let date = &output["test-date,0123456789abcdef,".len()..output.len() - ",3\n".len()];
197
198        assert!(chrono::DateTime::parse_from_str(&date, DATE_FORMAT_STR).is_ok());
199    }
200
201    #[cfg(feature = "enable")]
202    #[test]
203    #[serial]
204    pub fn test_hash() {
205        StringLog::clear();
206
207        packet_trace!("test-foo", { &[1, 2, 3] });
208        let expected = Regex::new(r#"test-foo,[0-9A-Fa-f]{16}.*\n"#).unwrap();
209
210        assert!(expected.is_match(&StringLog::get_string()));
211    }
212
213    #[cfg(feature = "enable")]
214    #[test]
215    #[serial]
216    pub fn test_twice() {
217        {
218            StringLog::clear();
219
220            packet_trace!("test-foo", { &[1, 2, 3] });
221            let expected = Regex::new(r#"test-foo,[0-9A-Fa-f]{16}.*\n"#).unwrap();
222
223            assert!(expected.is_match(&StringLog::get_string()));
224        }
225
226        {
227            StringLog::clear();
228
229            packet_trace!("test-bar", { &[1, 2, 3] });
230            let expected = Regex::new(r#"test-bar,[0-9A-Fa-f]{16}.*\n"#).unwrap();
231
232            assert!(expected.is_match(&StringLog::get_string()));
233        }
234    }
235
236    #[cfg(feature = "enable")]
237    #[test]
238    #[serial]
239    pub fn test_seq_3() {
240        StringLog::clear();
241
242        packet_trace!("test-foo", { &[1, 2, 3] });
243        packet_trace!("test-bar", { b"test data" });
244        packet_trace!("test-baz", { vec![0u8, 12, 13, 22] });
245        let expected = Regex::new(&format!(
246            "{}{}{}",
247            r#"test-foo,[0-9A-Fa-f]{16}.*\n"#,
248            r#"test-bar,[0-9A-Fa-f]{16}.*\n"#,
249            r#"test-baz,[0-9A-Fa-f]{16}.*\n"#,
250        ))
251        .unwrap();
252
253        assert!(expected.is_match(&StringLog::get_string()));
254    }
255
256    #[cfg(feature = "enable")]
257    #[test]
258    #[serial]
259    pub fn test_custom_write_target() {
260        use crate::{set_write_target, WriteTarget};
261
262        StringLog::clear();
263
264        set_write_target(WriteTarget::Write(Box::new(StringLog::global())));
265        packet_trace!("test-foo", { &[1, 2, 3] });
266        set_write_target(WriteTarget::Log);
267
268        let output = StringLog::get_string();
269
270        let expected = Regex::new(r#"test-foo,[0-9A-Fa-f]{16}.*\n"#).unwrap();
271        assert!(expected.is_match(&output));
272
273        let date = &output["test-date,0123456789abcdef,".len()..output.len() - ",3\n".len()];
274        assert!(chrono::DateTime::parse_from_str(&date, DATE_FORMAT_STR).is_ok());
275    }
276
277    #[cfg(not(feature = "enable"))]
278    #[test]
279    #[serial]
280    pub fn test_disable() {
281        StringLog::clear();
282
283        packet_trace!("test-foo", { &[1, 2, 3] });
284        packet_trace!("test-bar", { b"test data" });
285        packet_trace!("test-baz", { vec![0u8, 12, 13, 22] });
286        let expected = "";
287
288        assert_eq!(StringLog::get_string(), expected);
289    }
290}
ya_packet_trace/lib.rs

ya_packet_trace/
lib.rs
