From bd4c7eff1a50abc8e9ba99d2dcb5381f0b310800 Mon Sep 17 00:00:00 2001
From: Gary Guo <gary.guo@lowrisc.org>
Date: Thu, 7 Dec 2023 13:46:56 +0000
Subject: [PATCH] [opentitanlib] make `uart_read` return at most a line at a
 time

This is a hack to provide temporary solution to address some of our
flaky tests, where `UartConsole::wait_for` over-consumes data.

Signed-off-by: Gary Guo <gary.guo@lowrisc.org>
---
 sw/host/opentitanlib/src/uart/console.rs      | 23 ++++++++++++++++++-
 sw/host/tests/chip/rv_dm/src/ndm_reset_req.rs |  1 +
 2 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/sw/host/opentitanlib/src/uart/console.rs b/sw/host/opentitanlib/src/uart/console.rs
index b03b73387cee56..325e6d97c68b97 100644
--- a/sw/host/opentitanlib/src/uart/console.rs
+++ b/sw/host/opentitanlib/src/uart/console.rs
@@ -86,6 +86,27 @@ impl UartConsole {
             return Ok(ExitStatus::ExitSuccess);
         }
 
+        // HACK(nbdd0121): do a nonblocking read because the UART buffer may still have data in it.
+        // If we wait for mio event now, we might be blocking forever.
+        while self.uart_read(device, Duration::from_millis(0), &mut stdout)? {
+            if self
+                .exit_success
+                .as_ref()
+                .map(|rx| rx.is_match(&self.buffer))
+                == Some(true)
+            {
+                return Ok(ExitStatus::ExitSuccess);
+            }
+            if self
+                .exit_failure
+                .as_ref()
+                .map(|rx| rx.is_match(&self.buffer))
+                == Some(true)
+            {
+                return Ok(ExitStatus::ExitFailure);
+            }
+        }
+
         let mut poll = Poll::new()?;
         let transport_help_token = Self::get_next_token();
         let nonblocking_help = device.nonblocking_help()?;
@@ -176,7 +197,7 @@ impl UartConsole {
     where
         T: ConsoleDevice + ?Sized,
     {
-        let mut buf = [0u8; 256];
+        let mut buf = [0u8; 1];
         let len = device.console_read(&mut buf, timeout)?;
         if len == 0 {
             return Ok(false);
diff --git a/sw/host/tests/chip/rv_dm/src/ndm_reset_req.rs b/sw/host/tests/chip/rv_dm/src/ndm_reset_req.rs
index 6b3d86a5024b39..f6bbe2ac46b343 100644
--- a/sw/host/tests/chip/rv_dm/src/ndm_reset_req.rs
+++ b/sw/host/tests/chip/rv_dm/src/ndm_reset_req.rs
@@ -31,6 +31,7 @@ fn test_ndm_reset_req(opts: &Opts, transport: &TransportWrapper) -> Result<()> {
     // Enable console and wait for the message.
     let uart = transport.uart("console")?;
     uart.set_flow_control(true)?;
+    let _ = UartConsole::wait_for(&*uart, r"Running [^\r\n]*", opts.timeout)?;
 
     log::info!("Waiting for \"wait for ndm reset\" message");
     let _ = UartConsole::wait_for(&*uart, "wait for ndm reset", opts.timeout)?;