feat: exec mode improvements

Use Bollard library to exec in pure Rust. `--use-cli` cli arg, will then only try to exec into containers using Docker. Only try to exec into a container if the state == Running.
2023-11-17 02:49:30 +00:00
parent 28b0315dd7
commit 0e5ee143b0
15 changed files with 477 additions and 142 deletions
@@ -0,0 +1,319 @@
+use std::{
+    fmt,
+    hash::{Hash, Hasher},
+    io::{Read, Write},
+    sync::{atomic::AtomicBool, Arc},
+};
+
+use bollard::{
+    container,
+    exec::{CreateExecOptions, StartExecOptions, StartExecResults},
+    Docker,
+};
+use crossterm::terminal::enable_raw_mode;
+use futures_util::StreamExt;
+use parking_lot::Mutex;
+use tokio::io::{AsyncReadExt, AsyncWriteExt};
+
+use crate::{
+    app_data::{AppData, ContainerId, State},
+    app_error::AppError,
+    parse_args::CliArgs,
+    ui::{GuiState, Status},
+};
+
+/// TTY location
+const TTY: &str = "/dev/tty";
+
+/// This will be the start of a docker exec emssage if one is unable to actually exec into the container
+const OCI_ERROR: &str = "OCI runtime exec failed";
+
+/// Set the cursor position on the screen to (0,0)
+pub const CURSOR_POS: &str = "\x1B[J\x1B[H";
+
+/// This needs to be written to stdout when exiting the exec mode, else the input handler thread gets confused,
+/// see https://sw.kovidgoyal.net/kitty/keyboard-protocol/#progressive-enhancement
+const KEYBOARD_PROTO: &str = "\x1B[?u\x1B[c";
+
+mod command {
+    pub const PWD: &str = "pwd";
+    pub const DOCKER: &str = "docker";
+    pub const EXEC: &str = "exec";
+    pub const SH: &str = "sh";
+    pub const IT: &str = "-it";
+}
+
+/// Currently known byte output after writing KEYBOARD_PROTO to stdout
+/// valid arm: [91, 63, 54, 49, 59, 54, 59, 55, 59, 50, 50, 59, 50, 51, 59, 50, 52, 59, 50, 56, 59, 51, 50, 59,52, 50] => [?61;6;7;22;23;24;28;32;2
+/// valid x86: [91, 63, 49, 59, 50, 99] => [?1;2c
+/// invalid x86: [91, 63, 49, 59, 48, 99] => [?1;0c
+enum ByteOutput {
+    Arm,
+    X86,
+}
+
+impl ByteOutput {
+    const fn len(&self) -> usize {
+        match self {
+            Self::Arm => 26,
+            Self::X86 => 6,
+        }
+    }
+    const fn last(&self) -> &[u8] {
+        match self {
+            Self::Arm => &[50],
+            Self::X86 => &[99],
+        }
+    }
+}
+
+/// Check the output from tty to see if it matches known sequence.
+/// At the moment we only need to check the length and end digit, as x86 valid and invalid match in these two regards
+fn byte_sequence_valid(bytes: &[u8]) -> bool {
+    [ByteOutput::Arm, ByteOutput::X86]
+        .iter()
+        .any(|i| i.len() == bytes.len() && bytes.ends_with(i.last()))
+}
+
+/// Check if tty is able to be written to, aka not windows
+pub fn tty_readable() -> bool {
+    std::fs::OpenOptions::new()
+        .read(true)
+        .write(false)
+        .open(TTY)
+        .is_ok()
+}
+
+/// Async tty reading, spawned into its own tokio thread
+fn tty(run: Arc<AtomicBool>) -> Option<AsyncTTY> {
+    if tty_readable() {
+        let (tx, rx) = std::sync::mpsc::channel();
+        tokio::spawn(async move {
+            if let Ok(mut f) = tokio::fs::File::open(TTY).await {
+                while run.load(std::sync::atomic::Ordering::SeqCst) {
+                    let mut buf = [0];
+                    if tokio::time::timeout(
+                        std::time::Duration::from_millis(10),
+                        f.read_exact(&mut buf),
+                    )
+                    .await
+                    .is_ok()
+                        && tx.send(buf[0]).is_err()
+                    {
+                        run.store(false, std::sync::atomic::Ordering::SeqCst);
+                    }
+                }
+            }
+        });
+        Some(AsyncTTY { rx })
+    } else {
+        None
+    }
+}
+
+struct AsyncTTY {
+    rx: std::sync::mpsc::Receiver<u8>,
+}
+
+#[derive(Debug, Clone)]
+pub enum ExecMode {
+    // use Bollard Rust library
+    Internal((ContainerId, Arc<Docker>)),
+    // use the external `docker-cli`
+    External(ContainerId),
+}
+
+impl ExecMode {
+    /// Test if we can exec into the selected container, first via the Internal methods, then by the External
+    /// If the container is oxker, it will always return None
+    pub async fn new(app_data: &Arc<Mutex<AppData>>, docker: &Arc<Docker>) -> Option<Self> {
+        let is_oxker = app_data.lock().is_oxker();
+        if is_oxker {
+            return None;
+        }
+
+        let use_cli = app_data.lock().args.use_cli;
+        let container = app_data.lock().get_selected_container_id_state();
+
+        if let Some((id, state)) = container {
+            if state == State::Running {
+                if tty_readable() && !use_cli {
+                    if let Ok(exec) = docker
+                        .create_exec(
+                            id.get(),
+                            CreateExecOptions {
+                                attach_stdout: Some(true),
+                                attach_stderr: Some(true),
+                                cmd: Some(vec![command::PWD]),
+                                ..Default::default()
+                            },
+                        )
+                        .await
+                    {
+                        if let Ok(StartExecResults::Attached { mut output, .. }) =
+                            docker.start_exec(&exec.id, None).await
+                        {
+                            if let Some(Ok(msg)) = output.next().await {
+                                if !msg.to_string().starts_with(OCI_ERROR) {
+                                    return Some(Self::Internal((id.clone(), Arc::clone(docker))));
+                                }
+                            }
+                        }
+                    }
+                }
+
+                if let Ok(output) = std::process::Command::new(command::DOCKER)
+                    .args([command::EXEC, id.get(), command::PWD])
+                    .output()
+                {
+                    if let Ok(output) = String::from_utf8(output.stdout) {
+                        if !output.starts_with(OCI_ERROR) {
+                            return Some(Self::External(id.clone()));
+                        }
+                    }
+                }
+            }
+        }
+        None
+    }
+
+    /// exec into the container using the external docker cli, the result it just piped into oxker
+    fn exec_external(id: &ContainerId) {
+        let mut stdout = std::io::stdout();
+        stdout.write_all(CURSOR_POS.as_bytes()).ok();
+        if let Ok(mut child) = std::process::Command::new(command::DOCKER)
+            .args([command::EXEC, command::IT, id.get(), command::SH])
+            .stdin(std::process::Stdio::inherit())
+            .stdout(std::process::Stdio::inherit())
+            .stderr(std::process::Stdio::inherit())
+            .spawn()
+        {
+            child.wait().ok();
+            if child.kill().is_err() {
+                std::process::exit(1)
+            }
+        }
+    }
+
+    /// Exec into the container via the Bollard library, stdout & stdin on different threads
+    /// Have to deal with strange output once dropped, hence the use of internal_cleanup() method
+    async fn exec_internal(&self, id: &ContainerId, docker: &Arc<Docker>) -> Result<(), AppError> {
+        let run = Arc::new(AtomicBool::new(true));
+
+        if let Ok(exec_result) = docker
+            .create_exec(
+                id.get(),
+                CreateExecOptions {
+                    attach_stdout: Some(true),
+                    attach_stderr: Some(false),
+                    attach_stdin: Some(true),
+                    tty: Some(true),
+                    cmd: Some(vec![command::SH]),
+                    ..Default::default()
+                },
+            )
+            .await
+        {
+            if let Ok(StartExecResults::Attached {
+                mut output,
+                mut input,
+            }) = docker
+                .start_exec(
+                    &exec_result.id,
+                    Some(StartExecOptions {
+                        detach: false,
+                        ..Default::default()
+                    }),
+                )
+                .await
+            {
+                if let Some(async_tty) = tty(Arc::clone(&run)) {
+                    let run_thread = Arc::clone(&run);
+                    tokio::spawn(async move {
+                        enable_raw_mode().ok();
+                        let mut stdout = std::io::stdout();
+                        stdout.write_all(CURSOR_POS.as_bytes()).ok();
+                        stdout.flush().ok();
+
+                        while run_thread.load(std::sync::atomic::Ordering::SeqCst) {
+                            while let Some(Ok(x)) = output.next().await {
+                                stdout.write_all(&x.into_bytes()).ok();
+                                stdout.flush().ok();
+                            }
+                            run_thread.store(false, std::sync::atomic::Ordering::SeqCst);
+                        }
+                    });
+
+                    while let Ok(x) = async_tty.rx.recv() {
+                        input.write(&[x]).await.ok();
+                    }
+
+                    self.internal_cleanup()?;
+                }
+            } else {
+                return Err(AppError::Terminal);
+            }
+        }
+        Ok(())
+    }
+
+    // This is the fix for key pressed not being handled correctly on quit
+    // It writes a special message to the stdout, and then listens out for a valid response
+    // afterwhich it's assumes that we're completely done with TTY
+    fn internal_cleanup(&self) -> Result<(), AppError> {
+        match self {
+            Self::External(_) => Ok(()),
+            Self::Internal(_) => {
+                let waiting = Arc::new(AtomicBool::new(true));
+                let waiting_thread = Arc::clone(&waiting);
+
+                std::thread::spawn(move || {
+                    // At the moment the known max length is 26
+                    let mut bytes = Vec::with_capacity(26);
+                    while waiting_thread.load(std::sync::atomic::Ordering::SeqCst) {
+                        let mut buf = [0];
+                        if let Ok(mut f) = std::fs::File::open(TTY) {
+                            if f.read_exact(&mut buf).is_err() {
+                                waiting_thread.store(false, std::sync::atomic::Ordering::SeqCst);
+                            }
+                            bytes.push(buf[0]);
+                            if byte_sequence_valid(&bytes) {
+                                waiting_thread.store(false, std::sync::atomic::Ordering::SeqCst);
+                            }
+                        };
+                    }
+                });
+
+                let mut stdout = std::io::stdout();
+                stdout.write_all(KEYBOARD_PROTO.as_bytes()).ok();
+                stdout.flush().ok();
+
+                let start = std::time::Instant::now();
+                while waiting.load(std::sync::atomic::Ordering::SeqCst) {
+                    if start.elapsed().as_millis() > 1500 {
+                        waiting.store(false, std::sync::atomic::Ordering::SeqCst);
+                        return Err(AppError::Terminal);
+                    }
+                    std::thread::sleep(std::time::Duration::from_millis(10));
+                }
+                Ok(())
+            }
+        }
+    }
+
+    // RESET TERMINAL BEFROEHAND
+    pub async fn run(
+        &self,
+        app_data: &Arc<Mutex<AppData>>,
+        gui_state: &Arc<Mutex<GuiState>>,
+    ) -> Result<(), AppError> {
+        match self {
+            Self::External(id) => {
+                Self::exec_external(id);
+                Ok(())
+            }
+
+            Self::Internal((id, docker)) => self.exec_internal(id, docker).await,
+        }
+    }
+}