Fix stdin and stderr handling during docker exec

1 files changed, 21 insertions, 9 deletions

diff --git a/src/tty.rs b/src/tty.rs
index 446ab26..509c468 100644
--- a/src/tty.rs
+++ b/src/tty.rs
@@ -1,26 +1,38 @@
 use std::io::Read;
+use std::io::Cursor;
+use byteorder::{BigEndian, ReadBytesExt};
+
 pub struct Tty {
     pub stdout: String,
     pub stderr: String,
 }
 
+// https://docs.docker.com/engine/api/v1.26/#operation/ContainerAttach
 impl Tty {
     pub fn new(mut stream: Box<Read>) -> Tty {
         let mut stdout: Vec<u8> = vec![];
         let mut stderr: Vec<u8> = vec![];
         loop {
+            // 8 byte header [ STREAM_TYPE, 0, 0, 0, SIZE1, SIZE2, SIZE3, SIZE4 ]
             let mut header = [0; 8];
-            match stream.read(&mut header) {
-                Ok(0) => break,
+            match stream.read_exact(&mut header) {
                 Ok(_) => {
-                    let mut body: Vec<u8> = vec![0; header[7] as usize];
-                    if let Ok(_) = stream.read(&mut body) {
-                        if header[0] == 1 {
-                            stdout.append(&mut body);
-                        }
-                        if header[0] == 2 {
-                            stderr.append(&mut body);
+                    let payload_size: Vec<u8> = header[4..8].to_vec();
+                    let mut buffer = vec![0; Cursor::new(&payload_size).read_u32::<BigEndian>().unwrap() as usize];
+                    match stream.read_exact(&mut buffer) {
+                        Ok(_) => {
+                            match header[0] {
+                                // stdin, unhandled
+                                0 => break,
+                                // stdout
+                                1 => stdout.append(&mut buffer),
+                                // stderr
+                                2 => stderr.append(&mut buffer),
+                                //unhandled
+                                _ => break,
+                            }
                         }
+                        Err(_) => break,
                     };
                 }
                 Err(_) => break,