rsmp4decrypt 0.2.0

use crate::{DecryptionKey, Error};
use std::{
    collections::hash_map::DefaultHasher,
    ffi::OsString,
    fs,
    hash::{Hash, Hasher},
    path::{Path, PathBuf},
    process::{Command, ExitStatus},
    sync::{Mutex, OnceLock},
};

const WORKER_BYTES: &[u8] = include_bytes!(env!("RSMP4DECRYPT_WORKER_EMBED"));
const WORKER_ERROR_PREFIX: &str = "RSMP4DECRYPT_ERROR\t";

static WORKER_INSTALL_LOCK: Mutex<()> = Mutex::new(());
static WORKER_BINARY_PATH: OnceLock<PathBuf> = OnceLock::new();

/// A process-isolated decryptor for file-based workloads.
///
/// Each call spawns a dedicated helper process, which allows multiple decryptions
/// to run in parallel without sharing Bento4 state inside the parent process.
#[derive(Clone, Debug)]
pub struct IsolatedMp4Decryptor {
    keys: Vec<DecryptionKey>,
}

impl IsolatedMp4Decryptor {
    pub(crate) fn new(keys: Vec<DecryptionKey>) -> Self {
        Self { keys }
    }

    /// Decrypts an MP4 file by delegating the work to a separate helper process.
    pub fn decrypt_file<I, O, F>(
        &self,
        input_path: I,
        output_path: O,
        fragments_info_path: Option<F>,
    ) -> Result<(), Error>
    where
        I: AsRef<Path>,
        O: AsRef<Path>,
        F: AsRef<Path>,
    {
        let mut command = Command::new(worker_binary_path()?);
        command.args(worker_args(
            input_path.as_ref(),
            output_path.as_ref(),
            fragments_info_path.as_ref().map(AsRef::as_ref),
            &self.keys,
        ));

        let output = command.output()?;
        if output.status.success() {
            return Ok(());
        }

        Err(parse_worker_failure(&output.status, &output.stderr))
    }

    /// Decrypts an MP4 file asynchronously by delegating the work to a helper process.
    #[cfg(feature = "tokio")]
    #[cfg_attr(docsrs, doc(cfg(feature = "tokio")))]
    pub async fn decrypt_file_async<I, O, F>(
        &self,
        input_path: I,
        output_path: O,
        fragments_info_path: Option<F>,
    ) -> Result<(), Error>
    where
        I: AsRef<Path>,
        O: AsRef<Path>,
        F: AsRef<Path>,
    {
        let mut command = tokio::process::Command::new(worker_binary_path()?);
        command.args(worker_args(
            input_path.as_ref(),
            output_path.as_ref(),
            fragments_info_path.as_ref().map(AsRef::as_ref),
            &self.keys,
        ));

        let output = command.output().await?;
        if output.status.success() {
            return Ok(());
        }

        Err(parse_worker_failure(&output.status, &output.stderr))
    }
}

fn worker_args(
    input_path: &Path,
    output_path: &Path,
    fragments_info_path: Option<&Path>,
    keys: &[DecryptionKey],
) -> Vec<OsString> {
    let mut args = Vec::with_capacity(keys.len() * 2 + 4);
    for key in keys {
        args.push(OsString::from("--key"));
        args.push(OsString::from(key.to_spec()));
    }

    if let Some(path) = fragments_info_path {
        args.push(OsString::from("--fragments-info"));
        args.push(path.as_os_str().to_owned());
    }

    args.push(input_path.as_os_str().to_owned());
    args.push(output_path.as_os_str().to_owned());
    args
}

fn parse_worker_failure(status: &ExitStatus, stderr: &[u8]) -> Error {
    let stderr = String::from_utf8_lossy(stderr).trim().to_owned();
    if let Some(error) = parse_bento4_worker_error(&stderr) {
        return error;
    }

    let stderr = if stderr.is_empty() {
        "worker exited without error output".to_owned()
    } else {
        stderr
    };

    Error::WorkerFailed {
        status: describe_exit_status(status),
        stderr,
    }
}

fn parse_bento4_worker_error(stderr: &str) -> Option<Error> {
    let line = stderr
        .lines()
        .find(|line| line.starts_with(WORKER_ERROR_PREFIX))?;
    let payload = line.strip_prefix(WORKER_ERROR_PREFIX)?;
    let mut parts = payload.splitn(3, '\t');
    let operation = normalize_operation(parts.next()?);
    let code = parts.next()?.parse::<i32>().ok()?;
    let name = parts.next()?.to_owned();

    Some(Error::DecryptionFailed {
        operation,
        code,
        name,
    })
}

fn normalize_operation(operation: &str) -> &'static str {
    match operation {
        "opening input media" => "opening input media",
        "opening output media" => "opening output media",
        "opening fragments info media" => "opening fragments info media",
        "decrypting media" => "decrypting media",
        "finalizing decrypted output" => "finalizing decrypted output",
        _ => "processing media",
    }
}

fn describe_exit_status(status: &ExitStatus) -> String {
    match status.code() {
        Some(code) => code.to_string(),
        None => "terminated by signal".to_owned(),
    }
}

fn worker_binary_path() -> Result<&'static Path, Error> {
    if let Some(path) = WORKER_BINARY_PATH.get() {
        return Ok(path.as_path());
    }

    let _guard = WORKER_INSTALL_LOCK.lock().expect("poisoned worker lock");
    if let Some(path) = WORKER_BINARY_PATH.get() {
        return Ok(path.as_path());
    }

    let path = install_worker_binary()?;
    Ok(WORKER_BINARY_PATH.get_or_init(|| path).as_path())
}

fn install_worker_binary() -> Result<PathBuf, Error> {
    let directory = worker_directory();
    fs::create_dir_all(&directory)?;

    let path = directory.join(worker_file_name());
    let needs_write = match fs::read(&path) {
        Ok(existing) => existing != WORKER_BYTES,
        Err(_) => true,
    };

    if needs_write {
        fs::write(&path, WORKER_BYTES)?;
        make_executable(&path)?;
    }

    Ok(path)
}

fn worker_directory() -> PathBuf {
    let mut hasher = DefaultHasher::new();
    WORKER_BYTES.hash(&mut hasher);

    std::env::temp_dir().join("rsmp4decrypt").join(format!(
        "{}-{:016x}",
        env!("CARGO_PKG_VERSION"),
        hasher.finish()
    ))
}

fn worker_file_name() -> &'static str {
    if cfg!(windows) {
        "rsmp4decrypt-worker.exe"
    } else {
        "rsmp4decrypt-worker"
    }
}

#[cfg(unix)]
fn make_executable(path: &Path) -> Result<(), Error> {
    use std::os::unix::fs::PermissionsExt;

    let permissions = PermissionsExt::from_mode(0o700);
    fs::set_permissions(path, permissions)?;
    Ok(())
}

#[cfg(not(unix))]
fn make_executable(_path: &Path) -> Result<(), Error> {
    Ok(())
}