keep/src/compression_engine/program.rs

use anyhow::{Context, Result, anyhow};
use log::*;
use std::env;
use std::fs;
use std::fs::File;
use std::io::{Read, Write};
use std::os::unix::fs::PermissionsExt;
use std::path::PathBuf;
use std::process::{Child, Command, Stdio};

use crate::compression_engine::CompressionEngine;

pub struct ProgramReader {
    process: Child,
    stdout: Option<std::process::ChildStdout>,
}

impl Read for ProgramReader {
    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
        self.stdout.as_mut().unwrap().read(buf)
    }
}

impl Drop for ProgramReader {
    fn drop(&mut self) {
        // Ensure the process is waited on to prevent zombie processes
        let _ = self.process.wait();
    }
}

pub struct ProgramWriter {
    process: Child,
    stdin: Option<std::process::ChildStdin>,
}

impl Write for ProgramWriter {
    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
        self.stdin.as_mut().unwrap().write(buf)
    }

    fn flush(&mut self) -> std::io::Result<()> {
        self.stdin.as_mut().unwrap().flush()
    }
}

impl Drop for ProgramWriter {
    fn drop(&mut self) {
        // Close stdin to signal EOF to the child process
        drop(self.stdin.take());
        // Ensure the process is waited on to prevent zombie processes
        let _ = self.process.wait();
    }
}

#[derive(Debug, Eq, PartialEq, Clone)]
pub struct CompressionEngineProgram {
    pub program: String,
    pub compress: Vec<String>,
    pub decompress: Vec<String>,
    pub supported: bool,
}

impl CompressionEngineProgram {
    pub fn new(
        program: &str,
        compress: Vec<&str>,
        decompress: Vec<&str>,
    ) -> CompressionEngineProgram {
        let program_path = get_program_path(program);
        let supported = program_path.is_ok();

        CompressionEngineProgram {
            program: program_path.unwrap_or(program.to_string()),
            compress: compress.iter().map(|s| s.to_string()).collect(),
            decompress: decompress.iter().map(|s| s.to_string()).collect(),
            supported,
        }
    }
}

fn get_program_path(program: &str) -> Result<String> {
    debug!("COMPRESSION: Looking for executable: {}", program);
    if let Ok(path) = env::var("PATH") {
        for p in path.split(':') {
            let p_str = format!("{}/{}", p, program);
            let stat = fs::metadata(p_str.clone());
            if let Ok(stat) = stat {
                let md = stat;
                let permissions = md.permissions();
                if md.is_file() && permissions.mode() & 0o111 != 0 {
                    return Ok(p_str);
                }
            }
        }
    }
    Err(anyhow!("Unable to find binary {} in PATH", program))
}

impl CompressionEngine for CompressionEngineProgram {
    fn is_supported(&self) -> bool {
        self.supported
    }

    fn open(&self, file_path: PathBuf) -> Result<Box<dyn Read>> {
        debug!("COMPRESSION: Opening {:?} using {:?}", file_path, *self);

        let program = self.program.clone();
        let args = self.decompress.clone();

        debug!(
            "COMPRESSION: Executing command: {:?} {:?} reading from {:?}",
            program, args, file_path
        );

        let file = File::open(file_path).context("Unable to open file for reading")?;

        let mut process = Command::new(program.clone())
            .args(args.clone())
            .stdin(file)
            .stdout(Stdio::piped())
            .spawn()
            .context(anyhow!(
                "Unable to spawn child process: {:?} {:?}",
                program,
                args
            ))?;

        let stdout = process.stdout.take().ok_or_else(|| {
            anyhow!("Failed to capture stdout from child process")
        })?;

        Ok(Box::new(ProgramReader {
            process,
            stdout: Some(stdout),
        }))
    }

    fn create(&self, file_path: PathBuf) -> Result<Box<dyn Write>> {
        debug!("COMPRESSION: Writing to {:?} using {:?}", file_path, *self);

        let program = self.program.clone();
        let args = self.compress.clone();

        debug!(
            "COMPRESSION: Executing command: {:?} {:?} writing to {:?}",
            program, args, file_path
        );

        let file = File::create(file_path).context("Unable to open file for writing")?;

        let mut process = Command::new(program.clone())
            .args(args.clone())
            .stdin(Stdio::piped())
            .stdout(file)
            .spawn()
            .context(anyhow!(
                "Problem spawning child process: {:?} {:?}",
                program,
                args
            ))?;

        let stdin = process.stdin.take().ok_or_else(|| {
            anyhow!("Failed to capture stdin from child process")
        })?;

        Ok(Box::new(ProgramWriter {
            process,
            stdin: Some(stdin),
        }))
    }
}