From ce502b379b4713fa682034657139cf1099951727 Mon Sep 17 00:00:00 2001
From: yanhe <y1185535289@outlook.com>
Date: Wed, 4 Sep 2024 20:23:40 +0800
Subject: [PATCH] Add try_read_byte method.

---
 src/reader.rs  | 37 +++++++++++++++++++++++++++++++++++++
 src/session.rs |  4 ++++
 2 files changed, 41 insertions(+)

diff --git a/src/reader.rs b/src/reader.rs
index 1f5d8b62..20e87379 100644
--- a/src/reader.rs
+++ b/src/reader.rs
@@ -118,6 +118,7 @@ pub struct Options {
 pub struct NBReader {
     reader: Receiver<result::Result<PipedChar, PipeError>>,
     buffer: String,
+    bytes:Vec<u8>,
     eof: bool,
     timeout: Option<time::Duration>,
 }
@@ -172,6 +173,7 @@ impl NBReader {
         NBReader {
             reader: rx,
             buffer: String::with_capacity(1024),
+            bytes:Vec::with_capacity(1024),
             eof: false,
             timeout: options.timeout_ms.map(time::Duration::from_millis),
         }
@@ -199,6 +201,29 @@ impl NBReader {
         Ok(())
     }
 
+    /// reads all available byte(u8) from the read channel and stores them in self.bytes
+    fn read_into_bytes(&mut self) -> Result<(), Error> {
+        if self.eof {
+            return Ok(());
+        }
+        while let Ok(from_channel) = self.reader.try_recv() {
+            match from_channel {
+                Ok(PipedChar::Char(c)) => self.bytes.push(c),
+                Ok(PipedChar::EOF) => self.eof = true,
+                // this is just from experience, e.g. "sleep 5" returns the other error which
+                // most probably means that there is no stdout stream at all -> send EOF
+                // this only happens on Linux, not on OSX
+                Err(PipeError::IO(ref err)) => {
+                    // For an explanation of why we use `raw_os_error` see:
+                    // https://github.com/zhiburt/ptyprocess/commit/df003c8e3ff326f7d17bc723bc7c27c50495bb62
+                    self.eof = err.raw_os_error() == Some(5)
+                }
+            }
+        }
+        Ok(())
+    }
+
+
     /// Read until needle is found (blocking!) and return tuple with:
     /// 1. yet unread string until and without needle
     /// 2. matched needle
@@ -299,6 +324,18 @@ impl NBReader {
             None
         }
     }
+
+    pub fn try_read_byte(&mut self) -> Option<u8> {
+        // discard eventual errors, EOF will be handled in read_until correctly
+        let _ = self.read_into_bytes();
+        if !self.bytes.is_empty() {
+            let byte=self.bytes.remove(0);
+            return Some(byte)
+        } else {
+            None
+        }
+    }
+
 }
 
 #[cfg(test)]
diff --git a/src/session.rs b/src/session.rs
index 7bfef797..e1de84c0 100644
--- a/src/session.rs
+++ b/src/session.rs
@@ -88,6 +88,10 @@ impl<W: Write> StreamSession<W> {
         self.reader.try_read()
     }
 
+    pub fn try_read_byte(&mut self) -> Option<u8> {
+        self.reader.try_read_byte()
+    }
+
     // wrapper around reader::read_until to give more context for errors
     fn exp(&mut self, needle: &ReadUntil) -> Result<(String, String), Error> {
         self.reader.read_until(needle)