Added ReceiverHandler to logger

2 files changed, 143 insertions, 13 deletions

diff --git a/embassy-usb-logger/src/lib.rs b/embassy-usb-logger/src/lib.rs
index 11188b4ef..29c102f10 100644
--- a/embassy-usb-logger/src/lib.rs
+++ b/embassy-usb-logger/src/lib.rs
@@ -3,6 +3,7 @@
 #![warn(missing_docs)]
 
 use core::fmt::Write as _;
+use core::future::Future;
 
 use embassy_futures::join::join;
 use embassy_sync::pipe::Pipe;
@@ -13,6 +14,25 @@ use log::{Metadata, Record};
 
 type CS = embassy_sync::blocking_mutex::raw::CriticalSectionRawMutex;
 
+/// A trait that can be implemented and then passed to the
+pub trait ReceiverHandler {
+    /// Data comes in from the serial port with each command and runs this function
+    fn handle_data(&self, data: &[u8]) -> impl Future<Output = ()> + Send;
+
+    /// Create a new instance of the Handler
+    fn new() -> Self;
+}
+
+/// Use this Handler if you don't wish to use any handler
+pub struct DummyHandler;
+
+impl ReceiverHandler for DummyHandler {
+    async fn handle_data(&self, _data: &[u8]) {}
+    fn new() -> Self {
+        Self {}
+    }
+}
+
 /// The logger state containing buffers that must live as long as the USB peripheral.
 pub struct LoggerState<'d> {
     state: State<'d>,
@@ -39,17 +59,19 @@ impl<'d> LoggerState<'d> {
 pub const MAX_PACKET_SIZE: u8 = 64;
 
 /// The logger handle, which contains a pipe with configurable size for buffering log messages.
-pub struct UsbLogger<const N: usize> {
+pub struct UsbLogger<const N: usize, T: ReceiverHandler + Send + Sync> {
     buffer: Pipe<CS, N>,
     custom_style: Option<fn(&Record, &mut Writer<'_, N>) -> ()>,
+    recieve_handler: Option<T>,
 }
 
-impl<const N: usize> UsbLogger<N> {
+impl<const N: usize, T: ReceiverHandler + Send + Sync> UsbLogger<N, T> {
     /// Create a new logger instance.
     pub const fn new() -> Self {
         Self {
             buffer: Pipe::new(),
             custom_style: None,
+            recieve_handler: None,
         }
     }
 
@@ -58,9 +80,15 @@ impl<const N: usize> UsbLogger<N> {
         Self {
             buffer: Pipe::new(),
             custom_style: Some(custom_style),
+            recieve_handler: None,
         }
     }
 
+    /// Add a command handler to the logger
+    pub fn with_handler(&mut self, handler: T) {
+        self.recieve_handler = Some(handler);
+    }
+
     /// Run the USB logger using the state and USB driver. Never returns.
     pub async fn run<'d, D>(&'d self, state: &'d mut LoggerState<'d>, driver: D) -> !
     where
@@ -118,15 +146,22 @@ impl<const N: usize> UsbLogger<N> {
                 }
             }
         };
-        let discard_fut = async {
-            let mut discard_buf: [u8; MAX_PACKET_SIZE as usize] = [0; MAX_PACKET_SIZE as usize];
+        let reciever_fut = async {
+            let mut reciever_buf: [u8; MAX_PACKET_SIZE as usize] = [0; MAX_PACKET_SIZE as usize];
             receiver.wait_connection().await;
             loop {
-                let _ = receiver.read_packet(&mut discard_buf).await;
+                let n = receiver.read_packet(&mut reciever_buf).await.unwrap();
+                match &self.recieve_handler {
+                    Some(handler) => {
+                        let data = &reciever_buf[..n];
+                        handler.handle_data(data).await;
+                    }
+                    None => (),
+                }
             }
         };
 
-        join(log_fut, discard_fut).await;
+        join(log_fut, reciever_fut).await;
     }
 
     /// Creates the futures needed for the logger from a given class
@@ -142,7 +177,7 @@ impl<const N: usize> UsbLogger<N> {
     }
 }
 
-impl<const N: usize> log::Log for UsbLogger<N> {
+impl<const N: usize, T: ReceiverHandler + Send + Sync> log::Log for UsbLogger<N, T> {
     fn enabled(&self, _metadata: &Metadata) -> bool {
         true
     }
@@ -182,7 +217,7 @@ impl<'d, const N: usize> core::fmt::Write for Writer<'d, N> {
 
 /// Initialize and run the USB serial logger, never returns.
 ///
-/// Arguments specify the buffer size, log level and the USB driver, respectively.
+/// Arguments specify the buffer size, log level and the USB driver, respectively. You can optionally add a RecieverHandler.
 ///
 /// # Usage
 ///
@@ -196,17 +231,27 @@ impl<'d, const N: usize> core::fmt::Write for Writer<'d, N> {
 #[macro_export]
 macro_rules! run {
     ( $x:expr, $l:expr, $p:ident ) => {
-        static LOGGER: ::embassy_usb_logger::UsbLogger<$x> = ::embassy_usb_logger::UsbLogger::new();
+        static LOGGER: ::embassy_usb_logger::UsbLogger<$x, ::embassy_usb_logger::DummyHandler> =
+            ::embassy_usb_logger::UsbLogger::new();
         unsafe {
             let _ = ::log::set_logger_racy(&LOGGER).map(|()| log::set_max_level_racy($l));
         }
         let _ = LOGGER.run(&mut ::embassy_usb_logger::LoggerState::new(), $p).await;
     };
+
+    ( $x:expr, $l:expr, $p:ident, $h:ty ) => {
+        unsafe {
+            static mut LOGGER: ::embassy_usb_logger::UsbLogger<$x, $h> = ::embassy_usb_logger::UsbLogger::new();
+            LOGGER.with_handler(<$h>::new());
+            let _ = ::log::set_logger_racy(&LOGGER).map(|()| log::set_max_level_racy($l));
+            let _ = LOGGER.run(&mut ::embassy_usb_logger::LoggerState::new(), $p).await;
+        }
+    };
 }
 
 /// Initialize the USB serial logger from a serial class and return the future to run it.
 ///
-/// Arguments specify the buffer size, log level and the serial class, respectively.
+/// Arguments specify the buffer size, log level and the serial class, respectively. You can optionally add a RecieverHandler.
 ///
 /// # Usage
 ///
@@ -220,19 +265,29 @@ macro_rules! run {
 #[macro_export]
 macro_rules! with_class {
     ( $x:expr, $l:expr, $p:ident ) => {{
-        static LOGGER: ::embassy_usb_logger::UsbLogger<$x> = ::embassy_usb_logger::UsbLogger::new();
+        static LOGGER: ::embassy_usb_logger::UsbLogger<$x, ::embassy_usb_logger::DummyHandler> =
+            ::embassy_usb_logger::UsbLogger::new();
         unsafe {
             let _ = ::log::set_logger_racy(&LOGGER).map(|()| log::set_max_level_racy($l));
         }
         LOGGER.create_future_from_class($p)
     }};
+
+    ( $x:expr, $l:expr, $p:ident, $h:ty ) => {{
+        unsafe {
+            static mut LOGGER: ::embassy_usb_logger::UsbLogger<$x, $h> = ::embassy_usb_logger::UsbLogger::new();
+            LOGGER.with_handler(<$h>::new());
+            let _ = ::log::set_logger_racy(&LOGGER).map(|()| log::set_max_level_racy($l));
+            LOGGER.create_future_from_class($p)
+        }
+    }};
 }
 
 /// Initialize the USB serial logger from a serial class and return the future to run it.
 /// This version of the macro allows for a custom style function to be passed in.
 /// The custom style function will be called for each log record and is responsible for writing the log message to the buffer.
 ///
-/// Arguments specify the buffer size, log level, the serial class and the custom style function, respectively.
+/// Arguments specify the buffer size, log level, the serial class and the custom style function, respectively. You can optionally add a RecieverHandler.
 ///
 /// # Usage
 ///
@@ -250,10 +305,21 @@ macro_rules! with_class {
 #[macro_export]
 macro_rules! with_custom_style {
     ( $x:expr, $l:expr, $p:ident, $s:expr ) => {{
-        static LOGGER: ::embassy_usb_logger::UsbLogger<$x> = ::embassy_usb_logger::UsbLogger::with_custom_style($s);
+        static LOGGER: ::embassy_usb_logger::UsbLogger<$x, ::embassy_usb_logger::DummyHandler> =
+            ::embassy_usb_logger::UsbLogger::with_custom_style($s);
         unsafe {
             let _ = ::log::set_logger_racy(&LOGGER).map(|()| log::set_max_level_racy($l));
         }
         LOGGER.create_future_from_class($p)
     }};
+
+    ( $x:expr, $l:expr, $p:ident, $s:expr, $h:ty ) => {{
+        unsafe {
+            static mut LOGGER: ::embassy_usb_logger::UsbLogger<$x, $h> =
+                ::embassy_usb_logger::UsbLogger::with_custom_style($s);
+            LOGGER.with_handler(<$h>::new());
+            let _ = ::log::set_logger_racy(&LOGGER).map(|()| log::set_max_level_racy($l));
+            LOGGER.create_future_from_class($p)
+        }
+    }};
 }
diff --git a/examples/rp/src/bin/usb_serial_with_handler.rs b/examples/rp/src/bin/usb_serial_with_handler.rs
new file mode 100644
index 000000000..a9e65be70
--- /dev/null
+++ b/examples/rp/src/bin/usb_serial_with_handler.rs
@@ -0,0 +1,64 @@
+//! This example shows how to use USB (Universal Serial Bus) in the RP2040 chip.
+//!
+//! This creates the possibility to send log::info/warn/error/debug! to USB serial port.
+
+#![no_std]
+#![no_main]
+
+use core::str;
+
+use embassy_executor::Spawner;
+use embassy_rp::bind_interrupts;
+use embassy_rp::peripherals::USB;
+use embassy_rp::rom_data::reset_to_usb_boot;
+use embassy_rp::usb::{Driver, InterruptHandler};
+use embassy_time::Timer;
+use embassy_usb_logger::ReceiverHandler;
+use {defmt_rtt as _, panic_probe as _};
+
+bind_interrupts!(struct Irqs {
+    USBCTRL_IRQ => InterruptHandler<USB>;
+});
+
+struct Handler;
+
+impl ReceiverHandler for Handler {
+    async fn handle_data(&self, data: &[u8]) {
+        if let Ok(data) = str::from_utf8(data) {
+            let data = data.trim();
+
+            // If you are using elf2uf2-term with the '-t' flag, then when closing the serial monitor,
+            // this will automatically put the pico into boot mode
+            if data == "q" || data == "elf2uf2-term" {
+                reset_to_usb_boot(0, 0); // Restart the chip
+            } else if data.eq_ignore_ascii_case("hello") {
+                log::info!("World!");
+            } else {
+                log::info!("Recieved: {:?}", data);
+            }
+        }
+    }
+
+    fn new() -> Self {
+        Self
+    }
+}
+
+#[embassy_executor::task]
+async fn logger_task(driver: Driver<'static, USB>) {
+    embassy_usb_logger::run!(1024, log::LevelFilter::Info, driver, Handler);
+}
+
+#[embassy_executor::main]
+async fn main(spawner: Spawner) {
+    let p = embassy_rp::init(Default::default());
+    let driver = Driver::new(p.USB, Irqs);
+    spawner.spawn(logger_task(driver)).unwrap();
+
+    let mut counter = 0;
+    loop {
+        counter += 1;
+        log::info!("Tick {}", counter);
+        Timer::after_secs(1).await;
+    }
+}