ensure read_buf doesn't overflow the buffer

2 files changed, 36 insertions, 3 deletions

diff --git a/src/pty.rs b/src/pty.rs
index 87fb311..a63cc38 100644
--- a/src/pty.rs
+++ b/src/pty.rs
@@ -76,7 +76,7 @@ impl tokio::io::AsyncRead for Pty {
                 std::task::Poll::Ready(guard) => guard,
                 std::task::Poll::Pending => return std::task::Poll::Pending,
             }?;
-            let mut b = [0u8; 4096];
+            let mut b = vec![0u8; buf.capacity()];
             match guard.try_io(|inner| (&inner.get_ref().0).read(&mut b)) {
                 Ok(Ok(bytes)) => {
                     // XXX this is safe, but not particularly efficient
@@ -154,7 +154,7 @@ impl<'a> tokio::io::AsyncRead for ReadPty<'a> {
                 std::task::Poll::Ready(guard) => guard,
                 std::task::Poll::Pending => return std::task::Poll::Pending,
             }?;
-            let mut b = [0u8; 4096];
+            let mut b = vec![0u8; buf.capacity()];
             match guard.try_io(|inner| (&inner.get_ref().0).read(&mut b)) {
                 Ok(Ok(bytes)) => {
                     // XXX this is safe, but not particularly efficient
@@ -267,7 +267,7 @@ impl tokio::io::AsyncRead for OwnedReadPty {
                 std::task::Poll::Ready(guard) => guard,
                 std::task::Poll::Pending => return std::task::Poll::Pending,
             }?;
-            let mut b = [0u8; 4096];
+            let mut b = vec![0u8; buf.capacity()];
             match guard.try_io(|inner| (&inner.get_ref().0).read(&mut b)) {
                 Ok(Ok(bytes)) => {
                     // XXX this is safe, but not particularly efficient
diff --git a/tests/basic.rs b/tests/basic.rs
index 6f276e6..ed75e60 100644
--- a/tests/basic.rs
+++ b/tests/basic.rs
@@ -51,3 +51,36 @@ fn test_cat_async() {
         });
     assert_eq!(status.code().unwrap(), 0);
 }
+
+#[cfg(feature = "async")]
+#[test]
+fn test_yes_async() {
+    use tokio::io::AsyncReadExt as _;
+
+    tokio::runtime::Builder::new_multi_thread()
+        .enable_all()
+        .build()
+        .unwrap()
+        .block_on(async {
+            let mut pty = pty_process::Pty::new().unwrap();
+            let pts = pty.pts().unwrap();
+            pty.resize(pty_process::Size::new(24, 80)).unwrap();
+            let mut child =
+                pty_process::Command::new("yes").spawn(&pts).unwrap();
+
+            let mut buf = [0u8; 3];
+
+            let bytes = pty.read_buf(&mut &mut buf[..]).await.unwrap();
+            assert_eq!(&buf[..bytes], b"y\r\n");
+
+            let (mut pty_r, _pty_w) = pty.split();
+            let bytes = pty_r.read_buf(&mut &mut buf[..]).await.unwrap();
+            assert_eq!(&buf[..bytes], b"y\r\n");
+
+            let (mut pty_r, _pty_w) = pty.into_split();
+            let bytes = pty_r.read_buf(&mut &mut buf[..]).await.unwrap();
+            assert_eq!(&buf[..bytes], b"y\r\n");
+
+            child.kill().await.unwrap()
+        });
+}