Fix/improve defmt parser around edge cases (#496)

* Slightly refactor defmt framing

* Add tests, fix issues around raw data and partial frames

* Add to changelog

Author: bugadani

CHANGELOG.md

@@ -15,6 +15,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 
+- Fixed printing panic backtraces when using `esp-println` and `defmt` (#496)
+
 ### Changed
 
 ### Removed

espflash/src/cli/monitor/parser/esp_defmt.rs

@@ -8,71 +8,78 @@ use defmt_decoder::{Frame, Table};
 
 use crate::cli::monitor::parser::InputParser;
 
+#[derive(Debug, PartialEq)]
 enum FrameKind<'a> {
-    Defmt(Frame<'a>),
+    Defmt(&'a [u8]),
     Raw(&'a [u8]),
 }
 
 struct FrameDelimiter {
     buffer: Vec<u8>,
-    table: Option<Table>,
     in_frame: bool,
 }
 
 // Framing info added by esp-println
 const FRAME_START: &[u8] = &[0xFF, 0x00];
 const FRAME_END: &[u8] = &[0x00];
 
-fn search(haystack: &[u8], look_for_end: bool) -> Option<(&[u8], usize)> {
-    let needle = if look_for_end { FRAME_END } else { FRAME_START };
-    let start = if look_for_end {
-        // skip leading zeros
-        haystack.iter().position(|&b| b != 0)?
-    } else {
-        0
-    };
-
-    let end = haystack[start..]
-        .windows(needle.len())
-        .position(|window| window == needle)?;
-
-    Some((&haystack[start..][..end], start + end + needle.len()))
-}
-
 impl FrameDelimiter {
-    pub fn feed(&mut self, buffer: &[u8], mut process: impl FnMut(FrameKind<'_>)) {
-        let Some(table) = self.table.as_mut() else {
-            process(FrameKind::Raw(buffer));
-            return;
+    fn new() -> Self {
+        Self {
+            buffer: Vec::new(),
+            in_frame: false,
+        }
+    }
+
+    fn search(haystack: &[u8], look_for_end: bool) -> Option<(&[u8], usize)> {
+        let needle = if look_for_end { FRAME_END } else { FRAME_START };
+        let start = if look_for_end {
+            // skip leading zeros
+            haystack.iter().position(|&b| b != 0)?
+        } else {
+            0
         };
 
-        let mut decoder = table.new_stream_decoder();
+        let end = haystack[start..]
+            .windows(needle.len())
+            .position(|window| window == needle)?;
+
+        Some((&haystack[start..][..end], start + end + needle.len()))
+    }
 
+    pub fn feed(&mut self, buffer: &[u8], mut process: impl FnMut(FrameKind<'_>)) {
         self.buffer.extend_from_slice(buffer);
 
-        while let Some((frame, consumed)) = search(&self.buffer, self.in_frame) {
-            if !self.in_frame {
+        while let Some((frame, consumed)) = Self::search(&self.buffer, self.in_frame) {
+            if self.in_frame {
+                process(FrameKind::Defmt(frame));
+            } else if !frame.is_empty() {
                 process(FrameKind::Raw(frame));
-                self.in_frame = true;
+            }
+            self.in_frame = !self.in_frame;
+
+            self.buffer.drain(..consumed);
+        }
+
+        if !self.in_frame {
+            // If we have a 0xFF byte at the end, we should assume it's the start of a new frame.
+            let consume = if self.buffer.ends_with(&[0xFF]) {
+                &self.buffer[..self.buffer.len() - 1]
             } else {
-                decoder.received(frame);
-                // small reliance on rzcobs internals: we need to feed the terminating zero
-                decoder.received(FRAME_END);
-                if let Ok(frame) = decoder.decode() {
-                    process(FrameKind::Defmt(frame));
-                } else {
-                    log::warn!("Failed to decode defmt frame");
-                }
-                self.in_frame = false;
+                self.buffer.as_slice()
             };
 
-            self.buffer.drain(..consumed);
+            if !consume.is_empty() {
+                process(FrameKind::Raw(consume));
+                self.buffer.drain(..consume.len());
+            }
         }
     }
 }
 
 pub struct EspDefmt {
     delimiter: FrameDelimiter,
+    table: Option<Table>,
 }
 
 impl EspDefmt {
@@ -94,48 +101,187 @@ impl EspDefmt {
 
     pub fn new(elf: Option<&[u8]>) -> Self {
         Self {
-            delimiter: FrameDelimiter {
-                buffer: Vec::new(),
-                table: Self::load_table(elf),
-                in_frame: false,
-            },
+            delimiter: FrameDelimiter::new(),
+            table: Self::load_table(elf),
         }
     }
+
+    fn handle_raw(bytes: &[u8], out: &mut impl Write) {
+        out.write_all(bytes).unwrap();
+    }
+
+    fn handle_defmt(frame: Frame<'_>, out: &mut impl Write) {
+        match frame.level() {
+            Some(level) => {
+                let color = match level {
+                    defmt_parser::Level::Trace => Color::Cyan,
+                    defmt_parser::Level::Debug => Color::Blue,
+                    defmt_parser::Level::Info => Color::Green,
+                    defmt_parser::Level::Warn => Color::Yellow,
+                    defmt_parser::Level::Error => Color::Red,
+                };
+
+                // Print the level before each line.
+                let level = level.as_str().to_uppercase();
+                for line in frame.display_message().to_string().lines() {
+                    out.queue(PrintStyledContent(
+                        format!("[{level}] - {line}\r\n").with(color),
+                    ))
+                    .unwrap();
+                }
+            }
+            None => {
+                out.queue(Print(frame.display_message().to_string()))
+                    .unwrap();
+                out.queue(Print("\r\n")).unwrap();
+            }
+        }
+
+        out.flush().unwrap();
+    }
 }
 
 impl InputParser for EspDefmt {
     fn feed(&mut self, bytes: &[u8], out: &mut impl Write) {
+        let Some(table) = self.table.as_mut() else {
+            Self::handle_raw(bytes, out);
+            return;
+        };
+
+        let mut decoder = table.new_stream_decoder();
+
         self.delimiter.feed(bytes, |frame| match frame {
             FrameKind::Defmt(frame) => {
-                match frame.level() {
-                    Some(level) => {
-                        let color = match level {
-                            defmt_parser::Level::Trace => Color::Cyan,
-                            defmt_parser::Level::Debug => Color::Blue,
-                            defmt_parser::Level::Info => Color::Green,
-                            defmt_parser::Level::Warn => Color::Yellow,
-                            defmt_parser::Level::Error => Color::Red,
-                        };
-
-                        // Print the level before each line.
-                        let level = level.as_str().to_uppercase();
-                        for line in frame.display_message().to_string().lines() {
-                            out.queue(PrintStyledContent(
-                                format!("[{level}] - {line}\r\n").with(color),
-                            ))
-                            .unwrap();
-                        }
-                    }
-                    None => {
-                        out.queue(Print(frame.display_message().to_string()))
-                            .unwrap();
-                        out.queue(Print("\r\n")).unwrap();
-                    }
-                };
+                decoder.received(frame);
+                // small reliance on rzcobs internals: we need to feed the terminating zero
+                decoder.received(FRAME_END);
+
+                if let Ok(frame) = decoder.decode() {
+                    Self::handle_defmt(frame, out);
+                } else {
+                    log::warn!("Failed to decode defmt frame");
+                }
+            }
+            FrameKind::Raw(bytes) => Self::handle_raw(bytes, out),
+        });
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    #[test]
+    fn framing_prints_raw_data_by_default() {
+        let mut parser = FrameDelimiter::new();
+
+        let mut asserted = 0;
+        parser.feed(b"hello", |frame| {
+            assert_eq!(frame, FrameKind::Raw(b"hello"));
+            asserted += 1;
+        });
+        assert_eq!(asserted, 1);
+    }
+
+    #[test]
+    fn start_byte_on_end_is_not_part_of_the_raw_sequence() {
+        let mut parser = FrameDelimiter::new();
+
+        let mut asserted = 0;
+        parser.feed(b"hello\xFF", |frame| {
+            assert_eq!(frame, FrameKind::Raw(b"hello"));
+            asserted += 1;
+        });
+        assert_eq!(asserted, 1);
+    }
+
+    #[test]
+    fn frame_start_on_end_is_not_part_of_the_raw_sequence() {
+        let mut parser = FrameDelimiter::new();
+
+        let mut asserted = 0;
+        parser.feed(b"hello\xFF\x00", |frame| {
+            assert_eq!(frame, FrameKind::Raw(b"hello"));
+            asserted += 1;
+        });
+        assert_eq!(asserted, 1);
+    }
+
+    #[test]
+    fn process_data_after_frame() {
+        let mut parser = FrameDelimiter::new();
 
-                out.flush().unwrap();
+        let mut asserted = 0;
+        parser.feed(b"\xFF\x00frame data\x00hello", |frame| {
+            match asserted {
+                0 => assert_eq!(frame, FrameKind::Defmt(b"frame data")),
+                1 => assert_eq!(frame, FrameKind::Raw(b"hello")),
+                _ => panic!("Too many frames"),
             }
-            FrameKind::Raw(bytes) => out.write_all(bytes).unwrap(),
+            asserted += 1;
+        });
+        assert_eq!(asserted, 2);
+    }
+
+    #[test]
+    fn can_concatenate_partial_defmt_frames() {
+        let mut parser = FrameDelimiter::new();
+
+        let mut asserted = 0;
+        parser.feed(b"\xFF\x00frame", |_| {
+            panic!("Should not have a frame yet");
+        });
+        parser.feed(b" data\x00\xFF", |frame| {
+            assert_eq!(frame, FrameKind::Defmt(b"frame data"));
+            asserted += 1;
+        });
+        parser.feed(b"\x00second frame", |_| {
+            panic!("Should not have a frame yet");
         });
+        parser.feed(b"\x00last part", |frame| {
+            match asserted {
+                1 => assert_eq!(frame, FrameKind::Defmt(b"second frame")),
+                2 => assert_eq!(frame, FrameKind::Raw(b"last part")),
+                _ => panic!("Too many frames"),
+            }
+            asserted += 1;
+        });
+        assert_eq!(asserted, 3);
+    }
+
+    #[test]
+    fn defmt_frames_back_to_back() {
+        let mut parser = FrameDelimiter::new();
+
+        let mut asserted = 0;
+        parser.feed(b"\xFF\x00frame data1\x00\xFF\x00frame data2\x00", |frame| {
+            match asserted {
+                0 => assert_eq!(frame, FrameKind::Defmt(b"frame data1")),
+                1 => assert_eq!(frame, FrameKind::Defmt(b"frame data2")),
+                _ => panic!("Too many frames"),
+            }
+            asserted += 1;
+        });
+        assert_eq!(asserted, 2);
+    }
+
+    #[test]
+    fn output_includes_ff_and_0_bytes() {
+        let mut parser = FrameDelimiter::new();
+
+        let mut asserted = 0;
+        parser.feed(
+            b"some message\xFF with parts of\0 a defmt \0\xFF frame delimiter",
+            |frame| {
+                assert_eq!(
+                    frame,
+                    FrameKind::Raw(
+                        b"some message\xFF with parts of\0 a defmt \0\xFF frame delimiter"
+                    )
+                );
+                asserted += 1;
+            },
+        );
+        assert_eq!(asserted, 1);
     }
 }